meta_srv/procedure/drain_datanode/
migration.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::DatanodeId;
16use common_meta::key::table_route::TableRouteValue;
17use common_meta::peer::Peer;
18use common_procedure::{
19    Context as ProcedureContext, ProcedureContext as SubmissionContext, ProcedureWithId, watcher,
20};
21use common_telemetry::{error, info, warn};
22use snafu::ResultExt;
23use store_api::storage::RegionId;
24
25use crate::error::{self, Result};
26use crate::metrics::METRIC_META_REGION_MIGRATION_FAIL;
27use crate::procedure::drain_datanode::{Context, DrainDatanodeStrategy, utils};
28use crate::procedure::region_migration::DEFAULT_REGION_MIGRATION_TIMEOUT;
29use crate::procedure::region_migration::manager::{
30    RegionMigrationProcedureTask, RegionMigrationTriggerReason,
31};
32
33impl Context {
34    /// Waits for every migration in the batch before applying the drain strategy.
35    ///
36    /// Keep the batch and its results unchanged until all fallible route checks succeed,
37    /// so a metadata retry cannot lose failures or apply progress twice.
38    pub async fn wait_for_inflight_subprocedures(
39        &mut self,
40        procedure_ctx: &ProcedureContext,
41        strategy: DrainDatanodeStrategy,
42    ) -> Result<()> {
43        let node_id = self.persistent_ctx.node_id;
44        let mut migrated = Vec::new();
45        let mut failed = Vec::new();
46        for meta in &self.volatile_ctx.inflight_subprocedures {
47            let receiver = procedure_ctx
48                .provider
49                .procedure_state_receiver(meta.procedure_id)
50                .await
51                .inspect_err(|e| {
52                    error!(e; "Failed to get state receiver for migration sub-procedure {}", meta.procedure_id);
53                })
54                .context(error::QueryProcedureSnafu)?;
55            let Some(mut receiver) = receiver else {
56                warn!(
57                    "Migration sub-procedure {} of region {} is unknown, keep the region pending",
58                    meta.procedure_id, meta.region_id
59                );
60                continue;
61            };
62            match watcher::wait(&mut receiver).await {
63                Ok(_) => migrated.push(meta.region_id),
64                Err(e) => {
65                    error!(e; "Migration sub-procedure {} of region {} failed", meta.procedure_id, meta.region_id);
66                    failed.push(meta.region_id);
67                }
68            }
69        }
70
71        let failed_count = failed.len();
72        let mut failed_regions = Vec::new();
73        for region_id in failed {
74            // A concurrent failover may have achieved the goal despite the child failure.
75            if utils::leads_region(&self.table_metadata_manager, node_id, region_id).await? {
76                failed_regions.push(region_id);
77            } else {
78                migrated.push(region_id);
79            }
80        }
81
82        METRIC_META_REGION_MIGRATION_FAIL.inc_by(failed_count as u64);
83        for region_id in migrated {
84            self.forget_leader_region(region_id);
85            self.volatile_ctx.migrated += 1;
86            info!("Region {} left datanode {}", region_id, node_id);
87        }
88        self.volatile_ctx.inflight_subprocedures.clear();
89        match strategy {
90            DrainDatanodeStrategy::FailFast if !failed_regions.is_empty() => {
91                error::DrainDatanodeAbortSnafu {
92                    node_id,
93                    failed_regions,
94                }
95                .fail()
96            }
97            DrainDatanodeStrategy::TryAll => {
98                for &region_id in &failed_regions {
99                    self.forget_leader_region(region_id);
100                }
101                self.volatile_ctx.failed_regions.extend(failed_regions);
102                Ok(())
103            }
104            DrainDatanodeStrategy::FailFast => Ok(()),
105        }
106    }
107}
108
109impl Context {
110    /// Builds the migration sub-procedure of `region_id`.
111    ///
112    /// Returns `Ok(None)` when the region is no longer led by the drained datanode.
113    /// Otherwise returns the sub-procedure and its target.
114    pub async fn build_migration_procedure(
115        &self,
116        procedure_ctx: &ProcedureContext,
117        region_id: RegionId,
118    ) -> Result<Option<(ProcedureWithId, DatanodeId)>> {
119        let node_id = self.persistent_ctx.node_id;
120        let Some(table_route) = self
121            .table_metadata_manager
122            .table_route_manager()
123            .table_route_storage()
124            .get(region_id.table_id())
125            .await
126            .context(error::TableMetadataManagerSnafu)?
127        else {
128            return Ok(None);
129        };
130        let TableRouteValue::Physical(table_route) = table_route else {
131            return Ok(None);
132        };
133        let Some(region_route) = table_route
134            .region_routes
135            .iter()
136            .find(|region_route| region_route.region.id == region_id)
137        else {
138            return Ok(None);
139        };
140        if region_route.leader_peer.as_ref().map(|peer| peer.id) != Some(node_id) {
141            return Ok(None);
142        }
143
144        let to_peer =
145            utils::select_migration_target(self, region_id, &region_route.follower_peers).await?;
146
147        let submission_context = match procedure_ctx.event_context.clone() {
148            Some(event_context) => SubmissionContext::from_event_context(event_context),
149            None => SubmissionContext::default(),
150        };
151        let target_node_id = to_peer.id;
152        let child = self
153            .region_migration_manager
154            .build_subprocedure(
155                submission_context,
156                RegionMigrationProcedureTask::new(
157                    region_id,
158                    Peer::empty(node_id),
159                    to_peer,
160                    DEFAULT_REGION_MIGRATION_TIMEOUT,
161                    RegionMigrationTriggerReason::Manual,
162                ),
163            )
164            .await?;
165
166        Ok(Some((child, target_node_id)))
167    }
168}
169#[cfg(test)]
170mod tests {
171    use crate::error;
172    use crate::procedure::drain_datanode::test_util::{
173        TestEnv, leader_region, leader_regions, new_persistent_context, seed_plan,
174    };
175    use crate::procedure::drain_datanode::{DrainDatanodeStrategy, SubprocedureMeta};
176    use crate::procedure::test_util::{procedure_context_with_receivers, procedure_state_receiver};
177    use common_error::mock::MockError;
178    use common_error::status_code::StatusCode;
179    use common_procedure::{Error as ProcedureError, ProcedureId, ProcedureState};
180    use std::collections::HashMap;
181    use std::sync::Arc;
182    use store_api::storage::RegionId;
183    use tokio::sync::watch;
184    #[tokio::test]
185    async fn test_build_migration_procedure_checks_route_and_guard() {
186        for leader_id in [1, 2] {
187            let env = TestEnv::new().await;
188            let region_id = RegionId::new(1024, 1);
189            env.create_table(vec![leader_region(region_id, leader_id)])
190                .await;
191            let ctx = env.context(new_persistent_context(
192                Some(2),
193                DrainDatanodeStrategy::FailFast,
194            ));
195            let procedure_ctx =
196                crate::procedure::region_migration::test_util::TestingEnv::procedure_context();
197            let child = ctx
198                .build_migration_procedure(&procedure_ctx, region_id)
199                .await
200                .unwrap();
201            if leader_id == 2 {
202                assert!(child.is_none());
203                assert!(!env.region_migration_manager.tracker().contains(region_id));
204                continue;
205            }
206            let (child, target) = child.unwrap();
207            assert_eq!(2, target);
208            assert!(env.region_migration_manager.tracker().contains(region_id));
209            let error = ctx
210                .build_migration_procedure(&procedure_ctx, region_id)
211                .await
212                .unwrap_err();
213            assert!(matches!(error, error::Error::MigrationRunning { .. }));
214            drop(child);
215            assert!(!env.region_migration_manager.tracker().contains(region_id));
216        }
217    }
218
219    #[tokio::test]
220    async fn test_drain_datanode_waits_for_migrated_batch() {
221        let env = TestEnv::new().await;
222        let region_id = RegionId::new(1024, 1);
223        env.create_table(vec![leader_region(region_id, 1)]).await;
224
225        let mut ctx = env.context(new_persistent_context(
226            Some(2),
227            DrainDatanodeStrategy::FailFast,
228        ));
229        seed_plan(&mut ctx, vec![region_id]);
230        ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta {
231            region_id,
232            procedure_id: ProcedureId::random(),
233        }];
234        let procedure_ctx = procedure_context_with_receivers(HashMap::from([(
235            ctx.volatile_ctx.inflight_subprocedures[0].procedure_id,
236            procedure_state_receiver(ProcedureState::Done { output: None }),
237        )]));
238
239        ctx.wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::FailFast)
240            .await
241            .unwrap();
242
243        assert_eq!(1, ctx.volatile_ctx.migrated);
244        assert!(leader_regions(&ctx).is_empty());
245        assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty());
246    }
247
248    #[tokio::test]
249    async fn test_drain_datanode_wait_failed_child_fail_fast() {
250        let env = TestEnv::new().await;
251        let region_id = RegionId::new(1024, 1);
252        env.create_table(vec![leader_region(region_id, 1)]).await;
253
254        let mut ctx = env.context(new_persistent_context(
255            Some(2),
256            DrainDatanodeStrategy::FailFast,
257        ));
258        seed_plan(&mut ctx, vec![region_id]);
259        ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta {
260            region_id,
261            procedure_id: ProcedureId::random(),
262        }];
263        let procedure_ctx = procedure_context_with_receivers(HashMap::from([(
264            ctx.volatile_ctx.inflight_subprocedures[0].procedure_id,
265            procedure_state_receiver(ProcedureState::failed(Arc::new(ProcedureError::external(
266                MockError::new(StatusCode::Internal),
267            )))),
268        )]));
269
270        let error = ctx
271            .wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::FailFast)
272            .await
273            .unwrap_err();
274
275        // The whole batch is settled before the procedure aborts.
276        assert!(matches!(error, error::Error::DrainDatanodeAbort { .. }));
277        assert_eq!(0, ctx.volatile_ctx.migrated);
278        assert_eq!(vec![region_id], leader_regions(&ctx));
279        assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty());
280    }
281
282    #[tokio::test]
283    async fn test_drain_datanode_wait_failed_child_try_all() {
284        let env = TestEnv::new().await;
285        let region_id = RegionId::new(1024, 1);
286        env.create_table(vec![leader_region(region_id, 1)]).await;
287
288        let mut ctx = env.context(new_persistent_context(
289            Some(2),
290            DrainDatanodeStrategy::FailFast,
291        ));
292        seed_plan(&mut ctx, vec![region_id]);
293        ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta {
294            region_id,
295            procedure_id: ProcedureId::random(),
296        }];
297        let procedure_ctx = procedure_context_with_receivers(HashMap::from([(
298            ctx.volatile_ctx.inflight_subprocedures[0].procedure_id,
299            procedure_state_receiver(ProcedureState::failed(Arc::new(ProcedureError::external(
300                MockError::new(StatusCode::Internal),
301            )))),
302        )]));
303
304        ctx.wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::TryAll)
305            .await
306            .unwrap();
307
308        assert_eq!(vec![region_id], ctx.volatile_ctx.failed_regions);
309        assert_eq!(0, ctx.volatile_ctx.migrated);
310        assert!(leader_regions(&ctx).is_empty());
311    }
312
313    #[tokio::test]
314    async fn test_drain_datanode_wait_failed_child_counts_migrated_when_moved() {
315        let env = TestEnv::new().await;
316        let region_id = RegionId::new(1024, 1);
317        // The route reports another datanode as the leader, which is what a failover that
318        // raced with the migration leaves behind.
319        env.create_table(vec![leader_region(region_id, 2)]).await;
320
321        let mut ctx = env.context(new_persistent_context(
322            Some(2),
323            DrainDatanodeStrategy::FailFast,
324        ));
325        seed_plan(&mut ctx, vec![region_id]);
326        ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta {
327            region_id,
328            procedure_id: ProcedureId::random(),
329        }];
330        let procedure_ctx = procedure_context_with_receivers(HashMap::from([(
331            ctx.volatile_ctx.inflight_subprocedures[0].procedure_id,
332            procedure_state_receiver(ProcedureState::failed(Arc::new(ProcedureError::external(
333                MockError::new(StatusCode::Internal),
334            )))),
335        )]));
336
337        ctx.wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::FailFast)
338            .await
339            .unwrap();
340
341        assert_eq!(1, ctx.volatile_ctx.migrated);
342        assert!(ctx.volatile_ctx.failed_regions.is_empty());
343        assert!(leader_regions(&ctx).is_empty());
344    }
345
346    #[tokio::test]
347    async fn test_drain_datanode_wait_unknown_child_keeps_pending() {
348        let env = TestEnv::new().await;
349        let region_id = RegionId::new(1024, 1);
350        env.create_table(vec![leader_region(region_id, 1)]).await;
351
352        let mut ctx = env.context(new_persistent_context(
353            Some(2),
354            DrainDatanodeStrategy::FailFast,
355        ));
356        seed_plan(&mut ctx, vec![region_id]);
357        ctx.volatile_ctx.inflight_subprocedures = vec![SubprocedureMeta {
358            region_id,
359            procedure_id: ProcedureId::random(),
360        }];
361        let procedure_ctx = procedure_context_with_receivers(HashMap::new());
362
363        ctx.wait_for_inflight_subprocedures(&procedure_ctx, DrainDatanodeStrategy::FailFast)
364            .await
365            .unwrap();
366
367        // The region stays pending and is dispatched again instead of being reported as
368        // failed.
369        assert_eq!(vec![region_id], leader_regions(&ctx));
370        assert_eq!(0, ctx.volatile_ctx.migrated);
371        assert!(ctx.volatile_ctx.failed_regions.is_empty());
372    }
373
374    #[tokio::test]
375    async fn test_drain_datanode_waits_for_entire_batch() {
376        for strategy in [
377            DrainDatanodeStrategy::FailFast,
378            DrainDatanodeStrategy::TryAll,
379        ] {
380            let env = TestEnv::new().await;
381            let failed_region = RegionId::new(1024, 1);
382            let succeeded_region = RegionId::new(1024, 2);
383            env.create_table(vec![
384                leader_region(failed_region, 1),
385                leader_region(succeeded_region, 1),
386            ])
387            .await;
388            let mut ctx = env.context(new_persistent_context(
389                Some(2),
390                DrainDatanodeStrategy::FailFast,
391            ));
392            seed_plan(&mut ctx, vec![failed_region, succeeded_region]);
393            let failed_id = ProcedureId::random();
394            let pending_id = ProcedureId::random();
395            ctx.volatile_ctx.inflight_subprocedures = vec![
396                SubprocedureMeta {
397                    region_id: failed_region,
398                    procedure_id: failed_id,
399                },
400                SubprocedureMeta {
401                    region_id: succeeded_region,
402                    procedure_id: pending_id,
403                },
404            ];
405            let (sender, receiver) = watch::channel(ProcedureState::Running);
406            let procedure_ctx = procedure_context_with_receivers(HashMap::from([
407                (
408                    failed_id,
409                    procedure_state_receiver(ProcedureState::failed(Arc::new(
410                        ProcedureError::external(MockError::new(StatusCode::Internal)),
411                    ))),
412                ),
413                (pending_id, receiver),
414            ]));
415            let mut waiting =
416                Box::pin(ctx.wait_for_inflight_subprocedures(&procedure_ctx, strategy));
417            assert!(futures::poll!(&mut waiting).is_pending());
418            sender.send(ProcedureState::Done { output: None }).unwrap();
419            let result = waiting.await;
420            match strategy {
421                DrainDatanodeStrategy::FailFast => assert!(
422                    matches!(result, Err(error::Error::DrainDatanodeAbort { failed_regions, .. }) if failed_regions == vec![failed_region])
423                ),
424                DrainDatanodeStrategy::TryAll => {
425                    result.unwrap();
426                    assert_eq!(vec![failed_region], ctx.volatile_ctx.failed_regions);
427                }
428            }
429            assert_eq!(1, ctx.volatile_ctx.migrated);
430            assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty());
431        }
432    }
433
434    #[tokio::test]
435    async fn test_drain_datanode_wait_preserves_batch_on_metadata_retry() {
436        use common_meta::key::TableMetadataManager;
437        use common_meta::kv_backend::test_util::MockKvBackendBuilder;
438
439        for strategy in [
440            DrainDatanodeStrategy::FailFast,
441            DrainDatanodeStrategy::TryAll,
442        ] {
443            let env = TestEnv::new().await;
444            let failed_region = RegionId::new(1024, 1);
445            let succeeded_region = RegionId::new(1024, 2);
446            env.create_table(vec![
447                leader_region(failed_region, 1),
448                leader_region(succeeded_region, 1),
449            ])
450            .await;
451            let mut ctx = env.context(new_persistent_context(
452                Some(2),
453                DrainDatanodeStrategy::FailFast,
454            ));
455            seed_plan(&mut ctx, vec![failed_region, succeeded_region]);
456            let failed_id = ProcedureId::random();
457            let succeeded_id = ProcedureId::random();
458            ctx.volatile_ctx.inflight_subprocedures = vec![
459                SubprocedureMeta {
460                    region_id: succeeded_region,
461                    procedure_id: succeeded_id,
462                },
463                SubprocedureMeta {
464                    region_id: failed_region,
465                    procedure_id: failed_id,
466                },
467            ];
468            let procedure_ctx = procedure_context_with_receivers(HashMap::from([
469                (
470                    succeeded_id,
471                    procedure_state_receiver(ProcedureState::Done { output: None }),
472                ),
473                (
474                    failed_id,
475                    procedure_state_receiver(ProcedureState::failed(Arc::new(
476                        ProcedureError::external(MockError::new(StatusCode::Internal)),
477                    ))),
478                ),
479            ]));
480            ctx.table_metadata_manager = Arc::new(TableMetadataManager::new(Arc::new(
481                MockKvBackendBuilder::default()
482                    .range_fn(Arc::new(|_| {
483                        common_meta::error::ElectionNoLeaderSnafu.fail()
484                    }))
485                    .build()
486                    .unwrap(),
487            )));
488            let error = ctx
489                .wait_for_inflight_subprocedures(&procedure_ctx, strategy)
490                .await
491                .unwrap_err();
492            assert!(error.is_retryable());
493            assert_eq!(2, ctx.volatile_ctx.inflight_subprocedures.len());
494            assert_eq!(0, ctx.volatile_ctx.migrated);
495            assert!(ctx.volatile_ctx.failed_regions.is_empty());
496            assert_eq!(vec![failed_region, succeeded_region], leader_regions(&ctx));
497
498            ctx.table_metadata_manager = env.table_metadata_manager().clone();
499            let result = ctx
500                .wait_for_inflight_subprocedures(&procedure_ctx, strategy)
501                .await;
502            match strategy {
503                DrainDatanodeStrategy::FailFast => assert!(
504                    matches!(result, Err(error::Error::DrainDatanodeAbort { failed_regions, .. }) if failed_regions == vec![failed_region])
505                ),
506                DrainDatanodeStrategy::TryAll => {
507                    result.unwrap();
508                    assert_eq!(vec![failed_region], ctx.volatile_ctx.failed_regions);
509                }
510            }
511            assert_eq!(1, ctx.volatile_ctx.migrated);
512            assert!(ctx.volatile_ctx.inflight_subprocedures.is_empty());
513        }
514    }
515}
meta_srv/procedure/drain_datanode/migration.rs

meta_srv/procedure/drain_datanode/
migration.rs
