datanode/heartbeat/handler/
downgrade_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_meta::instruction::{
16    DowngradeRegion, DowngradeRegionReply, DowngradeRegionsReply, InstructionError,
17    InstructionReply,
18};
19use common_telemetry::tracing::info;
20use common_telemetry::{error, warn};
21use futures::future::join_all;
22use store_api::region_engine::{SetRegionRoleStateResponse, SettableRegionRoleState};
23use store_api::region_request::{RegionFlushRequest, RegionRequest};
24use store_api::storage::RegionId;
25
26use crate::heartbeat::handler::{HandlerContext, InstructionHandler};
27use crate::heartbeat::task_tracker::WaitResult;
28
29#[derive(Debug, Clone, Copy, Default)]
30pub struct DowngradeRegionsHandler;
31
32impl DowngradeRegionsHandler {
33    async fn handle_downgrade_region(
34        ctx: &HandlerContext,
35        DowngradeRegion {
36            region_id,
37            flush_timeout,
38        }: DowngradeRegion,
39    ) -> DowngradeRegionReply {
40        let Some(writable) = ctx.region_server.is_region_leader(region_id) else {
41            warn!("Region: {region_id} is not found");
42            return DowngradeRegionReply {
43                region_id,
44                last_entry_id: None,
45                metadata_last_entry_id: None,
46                exists: false,
47                error: None,
48            };
49        };
50
51        let region_server_moved = ctx.region_server.clone();
52
53        // Ignores flush request
54        if !writable {
55            warn!(
56                "Region: {region_id} is not writable, flush_timeout: {:?}",
57                flush_timeout
58            );
59            return ctx.downgrade_to_follower_gracefully(region_id).await;
60        }
61
62        // If flush_timeout is not set, directly convert region to follower.
63        let Some(flush_timeout) = flush_timeout else {
64            return ctx.downgrade_to_follower_gracefully(region_id).await;
65        };
66
67        // Sets region to downgrading,
68        // the downgrading region will reject all write requests.
69        // However, the downgrading region will still accept read, flush requests.
70        match ctx
71            .region_server
72            .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader)
73            .await
74        {
75            Ok(SetRegionRoleStateResponse::Success { .. }) => {}
76            Ok(SetRegionRoleStateResponse::NotFound) => {
77                warn!("Region: {region_id} is not found");
78                return DowngradeRegionReply {
79                    region_id,
80                    last_entry_id: None,
81                    metadata_last_entry_id: None,
82                    exists: false,
83                    error: None,
84                };
85            }
86            Ok(SetRegionRoleStateResponse::InvalidTransition(err)) => {
87                error!(err; "Failed to convert region to downgrading leader - invalid transition");
88                return DowngradeRegionReply {
89                    region_id,
90                    last_entry_id: None,
91                    metadata_last_entry_id: None,
92                    exists: true,
93                    error: Some(InstructionError::from_error(&err)),
94                };
95            }
96            Err(err) => {
97                error!(err; "Failed to convert region to downgrading leader");
98                return DowngradeRegionReply {
99                    region_id,
100                    last_entry_id: None,
101                    metadata_last_entry_id: None,
102                    exists: true,
103                    error: Some(InstructionError::from_error(&err)),
104                };
105            }
106        }
107
108        let register_result = ctx
109            .downgrade_tasks
110            .try_register(
111                region_id,
112                Box::pin(async move {
113                    info!("Flush region: {region_id} before converting region to follower");
114                    region_server_moved
115                        .handle_request(
116                            region_id,
117                            RegionRequest::Flush(RegionFlushRequest::default()),
118                        )
119                        .await?;
120
121                    Ok(())
122                }),
123            )
124            .await;
125
126        if register_result.is_busy() {
127            warn!("Another flush task is running for the region: {region_id}");
128        }
129
130        let mut watcher = register_result.into_watcher();
131        let result = ctx.downgrade_tasks.wait(&mut watcher, flush_timeout).await;
132
133        match result {
134            WaitResult::Timeout => DowngradeRegionReply {
135                region_id,
136                last_entry_id: None,
137                metadata_last_entry_id: None,
138                exists: true,
139                error: Some(InstructionError::legacy_internal_retryable(format!(
140                    "Flush region timeout, region: {region_id}, timeout: {:?}",
141                    flush_timeout
142                ))),
143            },
144            WaitResult::Finish(Ok(_)) => ctx.downgrade_to_follower_gracefully(region_id).await,
145            WaitResult::Finish(Err(err)) => DowngradeRegionReply {
146                region_id,
147                last_entry_id: None,
148                metadata_last_entry_id: None,
149                exists: true,
150                error: Some(InstructionError::from_error(&err)),
151            },
152        }
153    }
154}
155
156#[async_trait::async_trait]
157impl InstructionHandler for DowngradeRegionsHandler {
158    type Instruction = Vec<DowngradeRegion>;
159
160    async fn handle(
161        &self,
162        ctx: &HandlerContext,
163        downgrade_regions: Self::Instruction,
164    ) -> Option<InstructionReply> {
165        let futures = downgrade_regions
166            .into_iter()
167            .map(|downgrade_region| Self::handle_downgrade_region(ctx, downgrade_region));
168        // Join all futures; parallelism is governed by the underlying flush scheduler.
169        let results = join_all(futures).await;
170
171        Some(InstructionReply::DowngradeRegions(
172            DowngradeRegionsReply::new(results),
173        ))
174    }
175}
176
177impl HandlerContext {
178    async fn downgrade_to_follower_gracefully(&self, region_id: RegionId) -> DowngradeRegionReply {
179        match self
180            .region_server
181            .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower)
182            .await
183        {
184            Ok(SetRegionRoleStateResponse::Success(success)) => DowngradeRegionReply {
185                region_id,
186                last_entry_id: success.last_entry_id(),
187                metadata_last_entry_id: success.metadata_last_entry_id(),
188                exists: true,
189                error: None,
190            },
191            Ok(SetRegionRoleStateResponse::NotFound) => {
192                warn!("Region: {region_id} is not found");
193                DowngradeRegionReply {
194                    region_id,
195                    last_entry_id: None,
196                    metadata_last_entry_id: None,
197                    exists: false,
198                    error: None,
199                }
200            }
201            Ok(SetRegionRoleStateResponse::InvalidTransition(err)) => {
202                error!(err; "Failed to convert region to follower - invalid transition");
203                DowngradeRegionReply {
204                    region_id,
205                    last_entry_id: None,
206                    metadata_last_entry_id: None,
207                    exists: true,
208                    error: Some(InstructionError::from_error(&err)),
209                }
210            }
211            Err(err) => {
212                error!(err; "Failed to convert region to {}", SettableRegionRoleState::Follower);
213                DowngradeRegionReply {
214                    region_id,
215                    last_entry_id: None,
216                    metadata_last_entry_id: None,
217                    exists: true,
218                    error: Some(InstructionError::from_error(&err)),
219                }
220            }
221        }
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use std::assert_matches;
228    use std::sync::Arc;
229    use std::time::Duration;
230
231    use common_meta::heartbeat::handler::{HandleControl, HeartbeatResponseHandler};
232    use common_meta::heartbeat::mailbox::MessageMeta;
233    use common_meta::instruction::{DowngradeRegion, Instruction};
234    use common_meta::kv_backend::memory::MemoryKvBackend;
235    use mito2::config::MitoConfig;
236    use mito2::engine::MITO_ENGINE_NAME;
237    use mito2::test_util::{CreateRequestBuilder, TestEnv};
238    use store_api::region_engine::{
239        RegionEngine, RegionRole, SetRegionRoleStateResponse, SetRegionRoleStateSuccess,
240    };
241    use store_api::region_request::RegionRequest;
242    use store_api::storage::RegionId;
243    use tokio::time::Instant;
244
245    use crate::error;
246    use crate::heartbeat::handler::downgrade_region::DowngradeRegionsHandler;
247    use crate::heartbeat::handler::tests::HeartbeatResponseTestEnv;
248    use crate::heartbeat::handler::{
249        HandlerContext, InstructionHandler, RegionHeartbeatResponseHandler,
250    };
251    use crate::tests::{MockRegionEngine, mock_region_server};
252
253    #[tokio::test]
254    async fn test_region_not_exist() {
255        let mut mock_region_server = mock_region_server();
256        let (mock_engine, _) = MockRegionEngine::new(MITO_ENGINE_NAME);
257        mock_region_server.register_engine(mock_engine);
258        let kv_backend = Arc::new(MemoryKvBackend::new());
259        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
260        let region_id = RegionId::new(1024, 1);
261        let waits = vec![None, Some(Duration::from_millis(100u64))];
262
263        for flush_timeout in waits {
264            let reply = DowngradeRegionsHandler
265                .handle(
266                    &handler_context,
267                    vec![DowngradeRegion {
268                        region_id,
269                        flush_timeout,
270                    }],
271                )
272                .await;
273
274            let reply = &reply.unwrap().expect_downgrade_regions_reply()[0];
275            assert!(!reply.exists);
276            assert!(reply.error.is_none());
277            assert!(reply.last_entry_id.is_none());
278        }
279    }
280
281    #[tokio::test]
282    async fn test_region_readonly() {
283        let mock_region_server = mock_region_server();
284        let region_id = RegionId::new(1024, 1);
285        let (mock_engine, _) =
286            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
287                region_engine.mock_role = Some(Some(RegionRole::Follower));
288                region_engine.handle_request_mock_fn = Some(Box::new(|_, req| {
289                    if let RegionRequest::Flush(_) = req {
290                        // Should be unreachable.
291                        unreachable!();
292                    };
293
294                    Ok(0)
295                }));
296                region_engine.handle_set_readonly_gracefully_mock_fn = Some(Box::new(|_| {
297                    Ok(SetRegionRoleStateResponse::success(
298                        SetRegionRoleStateSuccess::mito(1024),
299                    ))
300                }))
301            });
302        mock_region_server.register_test_region(region_id, mock_engine);
303        let kv_backend = Arc::new(MemoryKvBackend::new());
304        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
305
306        let waits = vec![None, Some(Duration::from_millis(100u64))];
307        for flush_timeout in waits {
308            let reply = DowngradeRegionsHandler
309                .handle(
310                    &handler_context,
311                    vec![DowngradeRegion {
312                        region_id,
313                        flush_timeout,
314                    }],
315                )
316                .await;
317
318            let reply = &reply.unwrap().expect_downgrade_regions_reply()[0];
319            assert!(reply.exists);
320            assert!(reply.error.is_none());
321            assert_eq!(reply.last_entry_id.unwrap(), 1024);
322        }
323    }
324
325    #[tokio::test]
326    async fn test_region_flush_timeout() {
327        let mock_region_server = mock_region_server();
328        let region_id = RegionId::new(1024, 1);
329        let (mock_engine, _) =
330            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
331                region_engine.mock_role = Some(Some(RegionRole::Leader));
332                region_engine.handle_request_delay = Some(Duration::from_secs(100));
333                region_engine.handle_set_readonly_gracefully_mock_fn = Some(Box::new(|_| {
334                    Ok(SetRegionRoleStateResponse::success(
335                        SetRegionRoleStateSuccess::mito(1024),
336                    ))
337                }))
338            });
339        mock_region_server.register_test_region(region_id, mock_engine);
340        let kv_backend = Arc::new(MemoryKvBackend::new());
341        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
342
343        let flush_timeout = Duration::from_millis(100);
344        let reply = DowngradeRegionsHandler
345            .handle(
346                &handler_context,
347                vec![DowngradeRegion {
348                    region_id,
349                    flush_timeout: Some(flush_timeout),
350                }],
351            )
352            .await;
353
354        let reply = &reply.unwrap().expect_downgrade_regions_reply()[0];
355        assert!(reply.exists);
356        assert!(reply.error.as_ref().unwrap().message.contains("timeout"));
357        assert!(reply.last_entry_id.is_none());
358    }
359
360    #[tokio::test]
361    async fn test_region_flush_timeout_and_retry() {
362        let mock_region_server = mock_region_server();
363        let region_id = RegionId::new(1024, 1);
364        let (mock_engine, _) =
365            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
366                region_engine.mock_role = Some(Some(RegionRole::Leader));
367                region_engine.handle_request_delay = Some(Duration::from_millis(300));
368                region_engine.handle_set_readonly_gracefully_mock_fn = Some(Box::new(|_| {
369                    Ok(SetRegionRoleStateResponse::success(
370                        SetRegionRoleStateSuccess::mito(1024),
371                    ))
372                }))
373            });
374        mock_region_server.register_test_region(region_id, mock_engine);
375        let kv_backend = Arc::new(MemoryKvBackend::new());
376        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
377
378        let waits = vec![
379            Some(Duration::from_millis(100u64)),
380            Some(Duration::from_millis(100u64)),
381        ];
382
383        for flush_timeout in waits {
384            let reply = DowngradeRegionsHandler
385                .handle(
386                    &handler_context,
387                    vec![DowngradeRegion {
388                        region_id,
389                        flush_timeout,
390                    }],
391                )
392                .await;
393
394            let reply = &reply.unwrap().expect_downgrade_regions_reply()[0];
395            assert!(reply.exists);
396            assert!(reply.error.as_ref().unwrap().message.contains("timeout"));
397            assert!(reply.last_entry_id.is_none());
398        }
399        let timer = Instant::now();
400        let reply = DowngradeRegionsHandler
401            .handle(
402                &handler_context,
403                vec![DowngradeRegion {
404                    region_id,
405                    flush_timeout: Some(Duration::from_millis(500)),
406                }],
407            )
408            .await;
409        // Must less than 300 ms.
410        assert!(timer.elapsed().as_millis() < 300);
411
412        let reply = &reply.unwrap().expect_downgrade_regions_reply()[0];
413        assert!(reply.exists);
414        assert!(reply.error.is_none());
415        assert_eq!(reply.last_entry_id.unwrap(), 1024);
416    }
417
418    #[tokio::test]
419    async fn test_region_flush_timeout_and_retry_error() {
420        let mock_region_server = mock_region_server();
421        let region_id = RegionId::new(1024, 1);
422        let (mock_engine, _) =
423            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
424                region_engine.mock_role = Some(Some(RegionRole::Leader));
425                region_engine.handle_request_delay = Some(Duration::from_millis(300));
426                region_engine.handle_request_mock_fn = Some(Box::new(|_, _| {
427                    error::UnexpectedSnafu {
428                        violated: "mock flush failed",
429                    }
430                    .fail()
431                }));
432                region_engine.handle_set_readonly_gracefully_mock_fn = Some(Box::new(|_| {
433                    Ok(SetRegionRoleStateResponse::success(
434                        SetRegionRoleStateSuccess::mito(1024),
435                    ))
436                }))
437            });
438        mock_region_server.register_test_region(region_id, mock_engine);
439        let kv_backend = Arc::new(MemoryKvBackend::new());
440        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
441
442        let waits = vec![
443            Some(Duration::from_millis(100u64)),
444            Some(Duration::from_millis(100u64)),
445        ];
446
447        for flush_timeout in waits {
448            let reply = DowngradeRegionsHandler
449                .handle(
450                    &handler_context,
451                    vec![DowngradeRegion {
452                        region_id,
453                        flush_timeout,
454                    }],
455                )
456                .await;
457            let reply = &reply.unwrap().expect_downgrade_regions_reply()[0];
458            assert!(reply.exists);
459            assert!(reply.error.as_ref().unwrap().message.contains("timeout"));
460            assert!(reply.last_entry_id.is_none());
461        }
462        let timer = Instant::now();
463        let reply = DowngradeRegionsHandler
464            .handle(
465                &handler_context,
466                vec![DowngradeRegion {
467                    region_id,
468                    flush_timeout: Some(Duration::from_millis(500)),
469                }],
470            )
471            .await;
472        // Must less than 300 ms.
473        assert!(timer.elapsed().as_millis() < 300);
474        let reply = &reply.unwrap().expect_downgrade_regions_reply()[0];
475        assert!(reply.exists);
476        assert!(
477            reply
478                .error
479                .as_ref()
480                .unwrap()
481                .message
482                .contains("flush failed")
483        );
484        assert!(reply.last_entry_id.is_none());
485    }
486
487    #[tokio::test]
488    async fn test_set_region_readonly_not_found() {
489        let mock_region_server = mock_region_server();
490        let region_id = RegionId::new(1024, 1);
491        let (mock_engine, _) =
492            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
493                region_engine.mock_role = Some(Some(RegionRole::Leader));
494                region_engine.handle_set_readonly_gracefully_mock_fn =
495                    Some(Box::new(|_| Ok(SetRegionRoleStateResponse::NotFound)));
496            });
497        mock_region_server.register_test_region(region_id, mock_engine);
498        let kv_backend = Arc::new(MemoryKvBackend::new());
499        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
500        let reply = DowngradeRegionsHandler
501            .handle(
502                &handler_context,
503                vec![DowngradeRegion {
504                    region_id,
505                    flush_timeout: None,
506                }],
507            )
508            .await;
509        let reply = &reply.unwrap().expect_downgrade_regions_reply()[0];
510        assert!(!reply.exists);
511        assert!(reply.error.is_none());
512        assert!(reply.last_entry_id.is_none());
513    }
514
515    #[tokio::test]
516    async fn test_set_region_readonly_error() {
517        let mock_region_server = mock_region_server();
518        let region_id = RegionId::new(1024, 1);
519        let (mock_engine, _) =
520            MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| {
521                region_engine.mock_role = Some(Some(RegionRole::Leader));
522                region_engine.handle_set_readonly_gracefully_mock_fn = Some(Box::new(|_| {
523                    error::UnexpectedSnafu {
524                        violated: "Failed to set region to readonly",
525                    }
526                    .fail()
527                }));
528            });
529        mock_region_server.register_test_region(region_id, mock_engine);
530        let kv_backend = Arc::new(MemoryKvBackend::new());
531        let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend);
532        let reply = DowngradeRegionsHandler
533            .handle(
534                &handler_context,
535                vec![DowngradeRegion {
536                    region_id,
537                    flush_timeout: None,
538                }],
539            )
540            .await;
541        let reply = &reply.unwrap().expect_downgrade_regions_reply()[0];
542        assert!(reply.exists);
543        assert!(
544            reply
545                .error
546                .as_ref()
547                .unwrap()
548                .message
549                .contains("Failed to set region to readonly")
550        );
551        assert!(reply.last_entry_id.is_none());
552    }
553
554    #[tokio::test]
555    async fn test_downgrade_regions() {
556        common_telemetry::init_default_ut_logging();
557
558        let mut region_server = mock_region_server();
559        let kv_backend = Arc::new(MemoryKvBackend::new());
560        let heartbeat_handler =
561            RegionHeartbeatResponseHandler::new(region_server.clone(), kv_backend);
562        let mut engine_env = TestEnv::with_prefix("downgrade-regions").await;
563        let engine = engine_env.create_engine(MitoConfig::default()).await;
564        region_server.register_engine(Arc::new(engine.clone()));
565        let region_id = RegionId::new(1024, 1);
566        let region_id1 = RegionId::new(1024, 2);
567        let builder = CreateRequestBuilder::new();
568        let create_req = builder.build();
569        region_server
570            .handle_request(region_id, RegionRequest::Create(create_req))
571            .await
572            .unwrap();
573        let create_req1 = builder.build();
574        region_server
575            .handle_request(region_id1, RegionRequest::Create(create_req1))
576            .await
577            .unwrap();
578        let meta = MessageMeta::new_test(1, "test", "dn-1", "meta-0");
579        let instruction = Instruction::DowngradeRegions(vec![
580            DowngradeRegion {
581                region_id,
582                flush_timeout: Some(Duration::from_secs(1)),
583            },
584            DowngradeRegion {
585                region_id: region_id1,
586                flush_timeout: Some(Duration::from_secs(1)),
587            },
588        ]);
589        let mut heartbeat_env = HeartbeatResponseTestEnv::new();
590        let mut ctx = heartbeat_env.create_handler_ctx((meta, Default::default(), instruction));
591        let control = heartbeat_handler.handle(&mut ctx).await.unwrap();
592        assert_matches!(control, HandleControl::Continue);
593
594        let (_, reply) = heartbeat_env.receiver.recv().await.unwrap();
595        let reply = reply.expect_downgrade_regions_reply();
596        assert_eq!(reply[0].region_id, region_id);
597        assert!(reply[0].exists);
598        assert!(reply[0].error.is_none());
599        assert_eq!(reply[0].last_entry_id, Some(0));
600        assert_eq!(reply[1].region_id, region_id1);
601        assert!(reply[1].exists);
602        assert!(reply[1].error.is_none());
603        assert_eq!(reply[1].last_entry_id, Some(0));
604
605        assert_eq!(engine.role(region_id).unwrap(), RegionRole::Follower);
606        assert_eq!(engine.role(region_id1).unwrap(), RegionRole::Follower);
607    }
608}
datanode/heartbeat/handler/downgrade_region.rs

datanode/heartbeat/handler/
downgrade_region.rs
