datanode/heartbeat/handler/
upgrade_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 common_error::ext::{BoxedError, ErrorExt};
16use common_error::status_code::StatusCode;
17use common_meta::instruction::{
18    InstructionError, InstructionReply, UpgradeRegion, UpgradeRegionReply, UpgradeRegionsReply,
19};
20use common_telemetry::{debug, info, warn};
21use store_api::region_request::{RegionCatchupRequest, ReplayCheckpoint};
22use store_api::storage::RegionId;
23
24use crate::error::Result;
25use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
26
27#[derive(Debug, Clone, Copy, Default)]
28pub struct UpgradeRegionsHandler {
29    pub upgrade_region_parallelism: usize,
30}
31
32#[cfg(test)]
33impl UpgradeRegionsHandler {
34    fn new_test() -> UpgradeRegionsHandler {
35        UpgradeRegionsHandler {
36            upgrade_region_parallelism: 8,
37        }
38    }
39}
40
41impl UpgradeRegionsHandler {
42    fn convert_responses_to_replies(
43        responses: Result<Vec<(RegionId, std::result::Result<(), BoxedError>)>>,
44        catchup_regions: &[RegionId],
45    ) -> Vec<UpgradeRegionReply> {
46        match responses {
47            Ok(responses) => responses
48                .into_iter()
49                .map(|(region_id, result)| match result {
50                    Ok(()) => UpgradeRegionReply {
51                        region_id,
52                        ready: true,
53                        exists: true,
54                        error: None,
55                    },
56                    Err(err) => {
57                        if err.status_code() == StatusCode::RegionNotFound {
58                            UpgradeRegionReply {
59                                region_id,
60                                ready: false,
61                                exists: false,
62                                error: Some(InstructionError::from_error(&err)),
63                            }
64                        } else {
65                            UpgradeRegionReply {
66                                region_id,
67                                ready: false,
68                                exists: true,
69                                error: Some(InstructionError::from_error(&err)),
70                            }
71                        }
72                    }
73                })
74                .collect::<Vec<_>>(),
75            Err(err) => catchup_regions
76                .iter()
77                .map(|region_id| UpgradeRegionReply {
78                    region_id: *region_id,
79                    ready: false,
80                    exists: true,
81                    error: Some(InstructionError::from_error(&err)),
82                })
83                .collect::<Vec<_>>(),
84        }
85    }
86}
87
88impl UpgradeRegionsHandler {
89    // Handles upgrade regions instruction.
90    //
91    // Returns batch of upgrade region replies, the order of the replies is not guaranteed.
92    async fn handle_upgrade_regions(
93        &self,
94        ctx: &HandlerContext,
95        upgrade_regions: Vec<UpgradeRegion>,
96    ) -> Vec<UpgradeRegionReply> {
97        let num_upgrade_regions = upgrade_regions.len();
98        let mut replies = Vec::with_capacity(num_upgrade_regions);
99        let mut catchup_requests = Vec::with_capacity(num_upgrade_regions);
100        let mut catchup_regions = Vec::with_capacity(num_upgrade_regions);
101        let mut timeout = None;
102
103        for upgrade_region in upgrade_regions {
104            let Some(writable) = ctx.region_server.is_region_leader(upgrade_region.region_id)
105            else {
106                // Region is not found.
107                debug!("Region {} is not found", upgrade_region.region_id);
108                replies.push(UpgradeRegionReply {
109                    region_id: upgrade_region.region_id,
110                    ready: false,
111                    exists: false,
112                    error: None,
113                });
114                continue;
115            };
116
117            // Ignores the catchup requests for writable regions.
118            if writable {
119                warn!(
120                    "Region {} is writable, ignores the catchup request",
121                    upgrade_region.region_id
122                );
123                replies.push(UpgradeRegionReply {
124                    region_id: upgrade_region.region_id,
125                    ready: true,
126                    exists: true,
127                    error: None,
128                });
129            } else {
130                let UpgradeRegion {
131                    last_entry_id,
132                    metadata_last_entry_id,
133                    replay_entry_id,
134                    metadata_replay_entry_id,
135                    replay_timeout,
136                    ..
137                } = upgrade_region;
138                match timeout {
139                    Some(timeout) => {
140                        debug_assert_eq!(timeout, replay_timeout);
141                    }
142                    None => {
143                        // TODO(weny): required the replay_timeout.
144                        timeout = Some(replay_timeout);
145                    }
146                }
147
148                let checkpoint = match (replay_entry_id, metadata_replay_entry_id) {
149                    (Some(entry_id), metadata_entry_id) => Some(ReplayCheckpoint {
150                        entry_id,
151                        metadata_entry_id,
152                    }),
153                    _ => None,
154                };
155
156                catchup_regions.push(upgrade_region.region_id);
157                catchup_requests.push((
158                    upgrade_region.region_id,
159                    RegionCatchupRequest {
160                        set_writable: true,
161                        entry_id: last_entry_id,
162                        metadata_entry_id: metadata_last_entry_id,
163                        checkpoint,
164                    },
165                ));
166            }
167        }
168
169        let Some(timeout) = timeout else {
170            // No replay timeout, so we don't need to catchup the regions.
171            info!("All regions are writable, no need to catchup");
172            debug_assert_eq!(replies.len(), num_upgrade_regions);
173            return replies;
174        };
175
176        match tokio::time::timeout(
177            timeout,
178            ctx.region_server
179                .handle_batch_catchup_requests(self.upgrade_region_parallelism, catchup_requests),
180        )
181        .await
182        {
183            Ok(responses) => {
184                replies.extend(Self::convert_responses_to_replies(
185                    responses,
186                    &catchup_regions,
187                ));
188            }
189            Err(_) => {
190                replies.extend(catchup_regions.iter().map(|region_id| UpgradeRegionReply {
191                    region_id: *region_id,
192                    ready: false,
193                    exists: true,
194                    error: None,
195                }));
196            }
197        }
198
199        replies
200    }
201}
202
203#[async_trait::async_trait]
204impl InstructionHandler for UpgradeRegionsHandler {
205    type Instruction = Vec<UpgradeRegion>;
206
207    async fn handle(
208        &self,
209        ctx: &HandlerContext,
210        upgrade_regions: Self::Instruction,
211    ) -> Option<InstructionReply> {
212        let replies = self.handle_upgrade_regions(ctx, upgrade_regions).await;
213
214        Some(InstructionReply::UpgradeRegions(UpgradeRegionsReply::new(
215            replies,
216        )))
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use std::sync::Arc;
223    use std::time::Duration;
224
225    use common_meta::instruction::UpgradeRegion;
226    use common_meta::kv_backend::memory::MemoryKvBackend;
227    use mito2::engine::MITO_ENGINE_NAME;
228    use store_api::region_engine::RegionRole;
229    use store_api::storage::RegionId;
230
231    use crate::error;
232    use crate::heartbeat::handler::upgrade_region::UpgradeRegionsHandler;
233    use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
234    use crate::tests::{MockRegionEngine, mock_region_server};
235
236    #[tokio::test]
237    async fn test_region_not_exist() {
238        let mut mock_region_server = mock_region_server();
239        let (mock_engine, _) = MockRegionEngine::new(MITO_ENGINE_NAME);
240        mock_region_server.register_engine(mock_engine);
241        let kv_backend = Arc::new(MemoryKvBackend::new());
242        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
243
244        let region_id = RegionId::new(1024, 1);
245        let region_id2 = RegionId::new(1024, 2);
246        let replay_timeout = Duration::from_millis(100u64);
247        let reply = UpgradeRegionsHandler::new_test()
248            .handle(
249                &handler_context,
250                vec![
251                    UpgradeRegion {
252                        region_id,
253                        replay_timeout,
254                        ..Default::default()
255                    },
256                    UpgradeRegion {
257                        region_id: region_id2,
258                        replay_timeout,
259                        ..Default::default()
260                    },
261                ],
262            )
263            .await;
264
265        let replies = &reply.unwrap().expect_upgrade_regions_reply();
266        assert_eq!(replies[0].region_id, region_id);
267        assert_eq!(replies[1].region_id, region_id2);
268        for reply in replies {
269            assert!(!reply.exists);
270            assert!(reply.error.is_none());
271        }
272    }
273
274    #[tokio::test]
275    async fn test_region_writable() {
276        let mock_region_server = mock_region_server();
277        let region_id = RegionId::new(1024, 1);
278        let region_id2 = RegionId::new(1024, 2);
279
280        let (mock_engine, _) =
281            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
282                region_engine.mock_role = Some(Some(RegionRole::Leader));
283                region_engine.handle_request_mock_fn = Some(Box::new(|_, _| {
284                    // Should be unreachable.
285                    unreachable!();
286                }));
287            });
288        mock_region_server.register_test_region(region_id, mock_engine.clone());
289        mock_region_server.register_test_region(region_id2, mock_engine);
290        let kv_backend = Arc::new(MemoryKvBackend::new());
291        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
292        let replay_timeout = Duration::from_millis(100u64);
293        let reply = UpgradeRegionsHandler::new_test()
294            .handle(
295                &handler_context,
296                vec![
297                    UpgradeRegion {
298                        region_id,
299                        replay_timeout,
300                        ..Default::default()
301                    },
302                    UpgradeRegion {
303                        region_id: region_id2,
304                        replay_timeout,
305                        ..Default::default()
306                    },
307                ],
308            )
309            .await;
310
311        let replies = &reply.unwrap().expect_upgrade_regions_reply();
312        assert_eq!(replies[0].region_id, region_id);
313        assert_eq!(replies[1].region_id, region_id2);
314        for reply in replies {
315            assert!(reply.ready);
316            assert!(reply.exists);
317            assert!(reply.error.is_none());
318        }
319    }
320
321    #[tokio::test]
322    async fn test_legacy_wal_index_hint() {
323        for location_id in [None, Some(42)] {
324            let mock_region_server = mock_region_server();
325            let region_id = RegionId::new(1024, 1);
326            let (mock_engine, _) =
327                MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |engine| {
328                    engine.mock_role = Some(Some(RegionRole::Follower));
329                    engine.handle_request_mock_fn = Some(Box::new(|_, request| {
330                        let store_api::region_request::RegionRequest::Catchup(request) = request
331                        else {
332                            panic!("expected catchup request");
333                        };
334                        assert!(request.set_writable);
335                        assert_eq!(request.entry_id, Some(10));
336                        assert_eq!(request.metadata_entry_id, Some(5));
337                        Ok(0)
338                    }));
339                });
340            mock_region_server.register_test_region(region_id, mock_engine);
341            let ctx =
342                HandlerContext::new_for_test(mock_region_server, Arc::new(MemoryKvBackend::new()));
343            let reply = UpgradeRegionsHandler::new_test()
344                .handle(
345                    &ctx,
346                    vec![UpgradeRegion {
347                        region_id,
348                        last_entry_id: Some(10),
349                        metadata_last_entry_id: Some(5),
350                        location_id,
351                        replay_timeout: Duration::from_secs(1),
352                        ..Default::default()
353                    }],
354                )
355                .await
356                .unwrap();
357            let reply = &reply.expect_upgrade_regions_reply()[0];
358            assert!(reply.ready);
359            assert!(reply.exists);
360            assert!(reply.error.is_none());
361        }
362    }
363
364    #[tokio::test]
365    async fn test_region_not_ready() {
366        let mock_region_server = mock_region_server();
367        let region_id = RegionId::new(1024, 1);
368
369        let (mock_engine, _) =
370            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
371                // Region is not ready.
372                region_engine.mock_role = Some(Some(RegionRole::Follower));
373                region_engine.handle_request_mock_fn = Some(Box::new(|_, _| Ok(0)));
374                // Note: Don't change.
375                region_engine.handle_request_delay = Some(Duration::from_secs(100));
376            });
377        mock_region_server.register_test_region(region_id, mock_engine);
378        let kv_backend = Arc::new(MemoryKvBackend::new());
379        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
380        let replay_timeout = Duration::from_millis(100u64);
381        let reply = UpgradeRegionsHandler::new_test()
382            .handle(
383                &handler_context,
384                vec![UpgradeRegion {
385                    region_id,
386                    replay_timeout,
387                    ..Default::default()
388                }],
389            )
390            .await;
391
392        let reply = &reply.unwrap().expect_upgrade_regions_reply()[0];
393        assert!(!reply.ready);
394        assert!(reply.exists);
395        assert!(reply.error.is_none(), "error: {:?}", reply.error);
396    }
397
398    #[tokio::test]
399    async fn test_region_not_ready_with_retry() {
400        common_telemetry::init_default_ut_logging();
401        let mock_region_server = mock_region_server();
402        let region_id = RegionId::new(1024, 1);
403
404        let (mock_engine, _) =
405            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
406                // Region is not ready.
407                region_engine.mock_role = Some(Some(RegionRole::Follower));
408                region_engine.handle_request_mock_fn = Some(Box::new(|_, _| Ok(0)));
409                region_engine.handle_request_delay = Some(Duration::from_millis(300));
410            });
411        mock_region_server.register_test_region(region_id, mock_engine);
412        let waits = vec![Duration::from_millis(100u64), Duration::from_millis(100u64)];
413        let kv_backend = Arc::new(MemoryKvBackend::new());
414        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
415        for replay_timeout in waits {
416            let reply = UpgradeRegionsHandler::new_test()
417                .handle(
418                    &handler_context,
419                    vec![UpgradeRegion {
420                        region_id,
421                        replay_timeout,
422                        ..Default::default()
423                    }],
424                )
425                .await;
426
427            let reply = &reply.unwrap().expect_upgrade_regions_reply()[0];
428            assert!(!reply.ready);
429            assert!(reply.exists);
430            assert!(reply.error.is_none(), "error: {:?}", reply.error);
431        }
432
433        let reply = UpgradeRegionsHandler::new_test()
434            .handle(
435                &handler_context,
436                vec![UpgradeRegion {
437                    region_id,
438                    replay_timeout: Duration::from_millis(500),
439                    ..Default::default()
440                }],
441            )
442            .await;
443        let reply = &reply.unwrap().expect_upgrade_regions_reply()[0];
444        assert!(reply.ready);
445        assert!(reply.exists);
446        assert!(reply.error.is_none(), "error: {:?}", reply.error);
447    }
448
449    #[tokio::test]
450    async fn test_region_error() {
451        common_telemetry::init_default_ut_logging();
452        let region_id = RegionId::new(1024, 1);
453
454        // With the default (zero) `replay_timeout`, the handler times out before the
455        // mock handle's 100ms delay elapses: it didn't wait for handle returns; it
456        // had no idea about the error.
457        let region_server = mock_region_server();
458        let (mock_engine, _) =
459            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
460                // Region is not ready.
461                region_engine.mock_role = Some(Some(RegionRole::Follower));
462                region_engine.handle_request_mock_fn = Some(Box::new(|_, _| {
463                    error::UnexpectedSnafu {
464                        violated: "mock_error".to_string(),
465                    }
466                    .fail()
467                }));
468                // Note: Don't change.
469                region_engine.handle_request_delay = Some(Duration::from_millis(100));
470            });
471        region_server.register_test_region(region_id, mock_engine);
472        let kv_backend = Arc::new(MemoryKvBackend::new());
473        let handler_context = HandlerContext::new_for_test(region_server, kv_backend);
474        let reply = UpgradeRegionsHandler::new_test()
475            .handle(
476                &handler_context,
477                vec![UpgradeRegion {
478                    region_id,
479                    ..Default::default()
480                }],
481            )
482            .await;
483
484        let reply = &reply.unwrap().expect_upgrade_regions_reply()[0];
485        assert!(!reply.ready);
486        assert!(reply.exists);
487        assert!(reply.error.is_none());
488
489        // With a 200ms `replay_timeout` and a mock handle that returns the error
490        // immediately (no delay), the handler waits for the handle and propagates
491        // its error deterministically: the catchup future completes on its first
492        // poll, so the `replay_timeout` can never fire first. This avoids racing
493        // real wall-clock time (a 100ms delay vs 200ms timeout is flaky on
494        // Windows CI).
495        let mock_region_server = mock_region_server();
496        let (mock_engine, _) =
497            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
498                // Region is not ready.
499                region_engine.mock_role = Some(Some(RegionRole::Follower));
500                region_engine.handle_request_mock_fn = Some(Box::new(|_, _| {
501                    error::UnexpectedSnafu {
502                        violated: "mock_error".to_string(),
503                    }
504                    .fail()
505                }));
506            });
507        mock_region_server.register_test_region(region_id, mock_engine);
508        let kv_backend = Arc::new(MemoryKvBackend::new());
509        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
510        let reply = UpgradeRegionsHandler::new_test()
511            .handle(
512                &handler_context,
513                vec![UpgradeRegion {
514                    region_id,
515                    replay_timeout: Duration::from_millis(200),
516                    ..Default::default()
517                }],
518            )
519            .await;
520
521        let reply = &reply.unwrap().expect_upgrade_regions_reply()[0];
522        assert!(!reply.ready);
523        assert!(reply.exists);
524        assert!(reply.error.is_some());
525        assert!(reply.error.as_ref().unwrap().message.contains("mock_error"));
526    }
527}
datanode/heartbeat/handler/upgrade_region.rs

datanode/heartbeat/handler/
upgrade_region.rs
