datanode/heartbeat/handler/
flush_region.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::time::Instant;
16
17use common_meta::instruction::{
18    FlushErrorStrategy, FlushRegionReply, FlushRegions, FlushStrategy, InstructionError,
19    InstructionReply,
20};
21use common_telemetry::{debug, error, warn};
22use store_api::region_request::{RegionFlushReason, RegionFlushRequest, RegionRequest};
23use store_api::storage::RegionId;
24
25use crate::error::{self, RegionNotFoundSnafu, RegionNotReadySnafu, Result};
26use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
27
28pub struct FlushRegionsHandler;
29
30#[async_trait::async_trait]
31impl InstructionHandler for FlushRegionsHandler {
32    type Instruction = FlushRegions;
33
34    async fn handle(
35        &self,
36        ctx: &HandlerContext,
37        flush_regions: FlushRegions,
38    ) -> Option<InstructionReply> {
39        let start_time = Instant::now();
40        let strategy = flush_regions.strategy;
41        let region_ids = flush_regions.region_ids;
42        let error_strategy = flush_regions.error_strategy;
43        let reason = flush_regions.reason;
44
45        let reply = if matches!(strategy, FlushStrategy::Async) {
46            // Asynchronous hint mode: fire-and-forget, no reply expected
47            ctx.handle_flush_hint(region_ids, reason).await;
48            None
49        } else {
50            // Synchronous mode: return reply with results
51            let reply = ctx
52                .handle_flush_sync(region_ids, error_strategy, reason)
53                .await;
54            Some(InstructionReply::FlushRegions(reply))
55        };
56
57        let elapsed = start_time.elapsed();
58        debug!(
59            "FlushRegions strategy: {:?}, elapsed: {:?}, reply: {:?}",
60            strategy, elapsed, reply
61        );
62
63        reply
64    }
65}
66
67impl HandlerContext {
68    /// Performs the actual region flush operation.
69    async fn perform_region_flush(
70        &self,
71        region_id: RegionId,
72        reason: Option<RegionFlushReason>,
73    ) -> Result<()> {
74        let request = RegionRequest::Flush(RegionFlushRequest {
75            reason,
76            ..Default::default()
77        });
78        self.region_server
79            .handle_request(region_id, request)
80            .await?;
81        Ok(())
82    }
83
84    /// Handles asynchronous flush hints (fire-and-forget).
85    async fn handle_flush_hint(
86        &self,
87        region_ids: Vec<RegionId>,
88        reason: Option<RegionFlushReason>,
89    ) {
90        let start_time = Instant::now();
91        for region_id in &region_ids {
92            let result = self.perform_region_flush(*region_id, reason).await;
93            match result {
94                Ok(_) => {}
95                Err(error::Error::RegionNotFound { .. }) => {
96                    warn!(
97                        "Received a flush region hint from meta, but target region: {} is not found.",
98                        region_id
99                    );
100                }
101                Err(err) => {
102                    warn!("Failed to flush region: {}, error: {}", region_id, err);
103                }
104            }
105        }
106        let elapsed = start_time.elapsed();
107        debug!(
108            "Flush regions hint: {:?}, elapsed: {:?}",
109            region_ids, elapsed
110        );
111    }
112
113    /// Handles synchronous flush operations with proper error handling and replies.
114    async fn handle_flush_sync(
115        &self,
116        region_ids: Vec<RegionId>,
117        error_strategy: FlushErrorStrategy,
118        reason: Option<RegionFlushReason>,
119    ) -> FlushRegionReply {
120        let mut results = Vec::with_capacity(region_ids.len());
121
122        for region_id in region_ids {
123            let result = self.flush_single_region_sync(region_id, reason).await;
124
125            match &result {
126                Ok(_) => results.push((region_id, Ok(()))),
127                Err(err) => {
128                    // Log the full error chain here; the reply carries the
129                    // classification and the root cause of the original error.
130                    error!(err; "Failed to flush region, region_id: {}", region_id);
131                    results.push((region_id, Err(InstructionError::from_error(err))));
132
133                    // For fail-fast strategy, abort on first error
134                    if matches!(error_strategy, FlushErrorStrategy::FailFast) {
135                        break;
136                    }
137                }
138            }
139        }
140
141        FlushRegionReply::from_results(results)
142    }
143
144    /// Flushes a single region synchronously with proper error handling.
145    async fn flush_single_region_sync(
146        &self,
147        region_id: RegionId,
148        reason: Option<RegionFlushReason>,
149    ) -> Result<()> {
150        // Check if region is leader and writable
151        let Some(writable) = self.region_server.is_region_leader(region_id) else {
152            return Err(RegionNotFoundSnafu { region_id }.build());
153        };
154
155        if !writable {
156            return Err(RegionNotReadySnafu { region_id }.build());
157        }
158
159        // Register and execute the flush task
160        let region_server_moved = self.region_server.clone();
161        let register_result = self
162            .flush_tasks
163            .try_register(
164                region_id,
165                Box::pin(async move {
166                    region_server_moved
167                        .handle_request(
168                            region_id,
169                            RegionRequest::Flush(RegionFlushRequest {
170                                reason,
171                                ..Default::default()
172                            }),
173                        )
174                        .await?;
175                    Ok(())
176                }),
177            )
178            .await;
179
180        if register_result.is_busy() {
181            warn!("Another flush task is running for the region: {region_id}");
182        }
183
184        let mut watcher = register_result.into_watcher();
185        self.flush_tasks.wait_until_finish(&mut watcher).await
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use std::sync::{Arc, RwLock};
192
193    use common_error::ext::RetryHint;
194    use common_error::status_code::StatusCode;
195    use common_meta::instruction::{FlushErrorStrategy, FlushRegions};
196    use common_meta::kv_backend::memory::MemoryKvBackend;
197    use mito2::engine::MITO_ENGINE_NAME;
198    use store_api::storage::RegionId;
199
200    use super::*;
201    use crate::error::{RegionBusySnafu, UnexpectedSnafu};
202    use crate::tests::{MockRegionEngine, mock_region_server};
203
204    type FlushedRequests = Arc<RwLock<Vec<(RegionId, Option<RegionFlushReason>)>>>;
205
206    #[tokio::test]
207    async fn test_handle_flush_region_hint() {
208        let flushed_requests: FlushedRequests = Arc::new(RwLock::new(Vec::new()));
209
210        let mock_region_server = mock_region_server();
211        let region_ids = (0..16).map(|i| RegionId::new(1024, i)).collect::<Vec<_>>();
212        for region_id in &region_ids {
213            let flushed_requests_ref = flushed_requests.clone();
214            let (mock_engine, _) =
215                MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, move |region_engine| {
216                    region_engine.handle_request_mock_fn =
217                        Some(Box::new(move |region_id, request| {
218                            let RegionRequest::Flush(request) = request else {
219                                panic!("Expected flush request");
220                            };
221                            flushed_requests_ref
222                                .write()
223                                .unwrap()
224                                .push((region_id, request.reason));
225                            Ok(0)
226                        }))
227                });
228            mock_region_server.register_test_region(*region_id, mock_engine);
229        }
230        let kv_backend = Arc::new(MemoryKvBackend::new());
231        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
232
233        // Async hint mode
234        let flush_instruction = FlushRegions::async_batch(region_ids.clone())
235            .with_reason(RegionFlushReason::RemoteWalPrune);
236        let reply = FlushRegionsHandler
237            .handle(&handler_context, flush_instruction)
238            .await;
239        assert!(reply.is_none()); // Hint mode returns no reply
240        let expected = region_ids
241            .iter()
242            .map(|region_id| (*region_id, Some(RegionFlushReason::RemoteWalPrune)))
243            .collect::<Vec<_>>();
244        assert_eq!(*flushed_requests.read().unwrap(), expected);
245
246        // Non-existent regions
247        flushed_requests.write().unwrap().clear();
248        let not_found_region_ids = (0..2).map(|i| RegionId::new(2048, i)).collect::<Vec<_>>();
249        let flush_instruction = FlushRegions::async_batch(not_found_region_ids);
250        let reply = FlushRegionsHandler
251            .handle(&handler_context, flush_instruction)
252            .await;
253        assert!(reply.is_none());
254        assert!(flushed_requests.read().unwrap().is_empty());
255    }
256
257    #[tokio::test]
258    async fn test_handle_flush_region_sync_single() {
259        let flushed_requests: FlushedRequests = Arc::new(RwLock::new(Vec::new()));
260
261        let mock_region_server = mock_region_server();
262        let region_id = RegionId::new(1024, 1);
263
264        let flushed_requests_ref = flushed_requests.clone();
265        let (mock_engine, _) =
266            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, move |region_engine| {
267                region_engine.handle_request_mock_fn = Some(Box::new(move |region_id, request| {
268                    let RegionRequest::Flush(request) = request else {
269                        panic!("Expected flush request");
270                    };
271                    flushed_requests_ref
272                        .write()
273                        .unwrap()
274                        .push((region_id, request.reason));
275                    Ok(0)
276                }))
277            });
278        mock_region_server.register_test_region(region_id, mock_engine);
279        let kv_backend = Arc::new(MemoryKvBackend::new());
280        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
281
282        let flush_instruction =
283            FlushRegions::sync_single(region_id).with_reason(RegionFlushReason::Repartition);
284        let reply = FlushRegionsHandler
285            .handle(&handler_context, flush_instruction)
286            .await;
287        let flush_reply = reply.unwrap().expect_flush_regions_reply();
288        assert!(flush_reply.overall_success);
289        assert_eq!(flush_reply.results.len(), 1);
290        assert_eq!(flush_reply.results[0].0, region_id);
291        assert!(flush_reply.results[0].1.is_ok());
292        assert_eq!(
293            *flushed_requests.read().unwrap(),
294            vec![(region_id, Some(RegionFlushReason::Repartition))]
295        );
296    }
297
298    #[tokio::test]
299    async fn test_handle_flush_region_sync_batch_fail_fast() {
300        let flushed_region_ids: Arc<RwLock<Vec<RegionId>>> = Arc::new(RwLock::new(Vec::new()));
301
302        let mock_region_server = mock_region_server();
303        let region_ids = vec![
304            RegionId::new(1024, 1),
305            RegionId::new(1024, 2),
306            RegionId::new(1024, 3),
307        ];
308
309        // Register only the first region, others will fail
310        let flushed_region_ids_ref = flushed_region_ids.clone();
311        let (mock_engine, _) =
312            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, move |region_engine| {
313                region_engine.handle_request_mock_fn = Some(Box::new(move |region_id, _request| {
314                    flushed_region_ids_ref.write().unwrap().push(region_id);
315                    Ok(0)
316                }))
317            });
318        mock_region_server.register_test_region(region_ids[0], mock_engine);
319        let kv_backend = Arc::new(MemoryKvBackend::new());
320        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
321
322        // Sync batch with fail-fast strategy
323        let flush_instruction =
324            FlushRegions::sync_batch(region_ids.clone(), FlushErrorStrategy::FailFast);
325        let reply = FlushRegionsHandler
326            .handle(&handler_context, flush_instruction)
327            .await;
328        let flush_reply = reply.unwrap().expect_flush_regions_reply();
329        assert!(!flush_reply.overall_success); // Should fail due to non-existent regions
330        // With fail-fast, only process regions until first failure
331        assert!(flush_reply.results.len() <= region_ids.len());
332    }
333
334    #[tokio::test]
335    async fn test_handle_flush_region_sync_batch_try_all() {
336        let flushed_region_ids: Arc<RwLock<Vec<RegionId>>> = Arc::new(RwLock::new(Vec::new()));
337
338        let mock_region_server = mock_region_server();
339        let region_ids = vec![RegionId::new(1024, 1), RegionId::new(1024, 2)];
340
341        // Register only the first region
342        let flushed_region_ids_ref = flushed_region_ids.clone();
343        let (mock_engine, _) =
344            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, move |region_engine| {
345                region_engine.handle_request_mock_fn = Some(Box::new(move |region_id, _request| {
346                    flushed_region_ids_ref.write().unwrap().push(region_id);
347                    Ok(0)
348                }))
349            });
350        mock_region_server.register_test_region(region_ids[0], mock_engine);
351        let kv_backend = Arc::new(MemoryKvBackend::new());
352        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
353
354        // Sync batch with try-all strategy
355        let flush_instruction =
356            FlushRegions::sync_batch(region_ids.clone(), FlushErrorStrategy::TryAll);
357        let reply = FlushRegionsHandler
358            .handle(&handler_context, flush_instruction)
359            .await;
360        let flush_reply = reply.unwrap().expect_flush_regions_reply();
361        assert!(!flush_reply.overall_success); // Should fail due to one non-existent region
362        // With try-all, should process all regions
363        assert_eq!(flush_reply.results.len(), region_ids.len());
364        // First should succeed, second should fail
365        assert!(flush_reply.results[0].1.is_ok());
366        assert!(flush_reply.results[1].1.is_err());
367    }
368
369    /// A task body error used to verify that the reply keeps its classification.
370    #[derive(Clone, Copy, Debug)]
371    enum TaskFlushError {
372        RegionBusy,
373        Unexpected,
374    }
375
376    impl TaskFlushError {
377        fn build(self, region_id: RegionId) -> crate::error::Error {
378            match self {
379                Self::RegionBusy => RegionBusySnafu { region_id }.build(),
380                Self::Unexpected => UnexpectedSnafu {
381                    violated: "flush task boom marker",
382                }
383                .build(),
384            }
385        }
386
387        fn expected(self) -> (StatusCode, RetryHint) {
388            match self {
389                Self::RegionBusy => (StatusCode::RegionBusy, RetryHint::Retryable),
390                Self::Unexpected => (StatusCode::Unexpected, RetryHint::NonRetryable),
391            }
392        }
393
394        fn message_fragment(self, region_id: RegionId) -> String {
395            match self {
396                Self::RegionBusy => format!("Region {region_id} is busy"),
397                Self::Unexpected => "flush task boom marker".to_string(),
398            }
399        }
400    }
401
402    #[tokio::test]
403    async fn test_handle_flush_region_sync_keeps_task_error_classification() {
404        let failing_region = RegionId::new(1024, 1);
405        let succeeding_region = RegionId::new(1024, 2);
406
407        for error in [TaskFlushError::RegionBusy, TaskFlushError::Unexpected] {
408            let flushed_regions: Arc<RwLock<Vec<RegionId>>> = Arc::new(RwLock::new(Vec::new()));
409
410            let mock_region_server = mock_region_server();
411            let flushed_regions_ref = flushed_regions.clone();
412            let (mock_engine, _) =
413                MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, move |region_engine| {
414                    region_engine.handle_request_mock_fn =
415                        Some(Box::new(move |region_id, _request| {
416                            flushed_regions_ref.write().unwrap().push(region_id);
417                            if region_id == failing_region {
418                                Err(error.build(region_id))
419                            } else {
420                                Ok(0)
421                            }
422                        }));
423                });
424            mock_region_server.register_test_region(failing_region, mock_engine.clone());
425            mock_region_server.register_test_region(succeeding_region, mock_engine);
426
427            let kv_backend = Arc::new(MemoryKvBackend::new());
428            let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
429            let instruction = FlushRegions::sync_batch(
430                vec![failing_region, succeeding_region],
431                FlushErrorStrategy::TryAll,
432            );
433            let reply = FlushRegionsHandler
434                .handle(&handler_context, instruction)
435                .await
436                .unwrap()
437                .expect_flush_regions_reply();
438
439            assert!(!reply.overall_success, "{error:?}");
440            assert_eq!(reply.results.len(), 2, "{error:?}");
441            assert_eq!(reply.results[0].0, failing_region, "{error:?}");
442            let reply_error = reply.results[0].1.as_ref().unwrap_err();
443            let (code, retry_hint) = error.expected();
444            assert_eq!(reply_error.code, code, "{error:?}");
445            assert_eq!(reply_error.retry_hint, retry_hint, "{error:?}");
446            assert!(
447                reply_error
448                    .message
449                    .contains(&error.message_fragment(failing_region)),
450                "{error:?}: {}",
451                reply_error.message
452            );
453            assert_eq!(reply.results[1].0, succeeding_region, "{error:?}");
454            assert!(reply.results[1].1.is_ok(), "{error:?}");
455            assert_eq!(
456                *flushed_regions.read().unwrap(),
457                vec![failing_region, succeeding_region],
458                "{error:?}"
459            );
460        }
461    }
462
463    #[test]
464    fn test_flush_regions_display() {
465        let region_id = RegionId::new(1024, 1);
466        let flush_regions = FlushRegions::sync_single(region_id);
467        let display = format!("{}", flush_regions);
468        assert_eq!(
469            display,
470            "FlushRegions(region_ids=[4398046511105(1024, 1)], strategy=Sync, error_strategy=FailFast, reason=None)"
471        );
472    }
473}
datanode/heartbeat/handler/flush_region.rs

datanode/heartbeat/handler/
flush_region.rs
