meta_srv/procedure/drain_datanode/
check_drained.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::any::Any;
16
17use common_procedure::{Context as ProcedureContext, Status};
18use common_telemetry::info;
19use serde::{Deserialize, Serialize};
20
21use crate::error::Result;
22use crate::procedure::drain_datanode::dispatch::Dispatch;
23use crate::procedure::drain_datanode::drain_datanode_end::DrainDatanodeEnd;
24use crate::procedure::drain_datanode::{Context, State, utils};
25
26/// Rescans the datanode for regions that still need a migration procedure.
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub(crate) struct CheckDrained;
29
30#[async_trait::async_trait]
31#[typetag::serde]
32impl State for CheckDrained {
33    async fn next(
34        &mut self,
35        ctx: &mut Context,
36        _procedure_ctx: &ProcedureContext,
37    ) -> Result<(Box<dyn State>, Status)> {
38        utils::ensure_drain_deadline(ctx)?;
39
40        let node_id = ctx.persistent_ctx.node_id;
41        // Reconcile against the table routes after settling all batches of this round.
42        ctx.rebuild_plan().await?;
43        // Regions given up by `DrainDatanodeStrategy::TryAll` are still led by the datanode;
44        // they must not trigger another rescan round.
45        // Migrations registered by recovery or another caller own their lifecycle. They
46        // need neither a new child nor a restored parent link, even if still in progress.
47        let tracker = ctx.region_migration_manager.tracker();
48        let volatile_ctx = &mut ctx.volatile_ctx;
49        if let Some(plan) = volatile_ctx.plan.as_mut() {
50            plan.leader_regions.retain(|region_id| {
51                !volatile_ctx.failed_regions.contains(region_id) && !tracker.contains(*region_id)
52            });
53        }
54
55        let migrated = ctx.volatile_ctx.migrated;
56        let rescan_round = ctx.volatile_ctx.rescan_round;
57        let remaining_leader_regions = ctx.volatile_ctx.leader_regions().len();
58        if remaining_leader_regions == 0 {
59            let follower_regions = ctx.volatile_ctx.follower_regions().len();
60            info!(
61                "No more migrations to dispatch for datanode {}, migrated regions: {}, follower regions: {}",
62                node_id, migrated, follower_regions
63            );
64            return Ok((Box::new(DrainDatanodeEnd), Status::executing(false)));
65        }
66
67        info!(
68            "Datanode {} still leads {} regions, rescan round: {}",
69            node_id, remaining_leader_regions, rescan_round
70        );
71
72        Ok((Box::new(Dispatch), Status::executing(false)))
73    }
74
75    fn as_any(&self) -> &dyn Any {
76        self
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use std::time::Duration;
83
84    use common_meta::peer::Peer;
85    use common_procedure::Status;
86    use store_api::storage::RegionId;
87
88    use crate::error;
89    use crate::procedure::drain_datanode::check_drained::CheckDrained;
90    use crate::procedure::drain_datanode::dispatch::Dispatch;
91    use crate::procedure::drain_datanode::drain_datanode_end::DrainDatanodeEnd;
92    use crate::procedure::drain_datanode::test_util::{
93        TestEnv, leader_region, leader_regions, new_persistent_context,
94    };
95    use crate::procedure::drain_datanode::{DrainDatanodeStrategy, MAX_RESCAN_ROUNDS, State};
96    use crate::procedure::region_migration::manager::{
97        RegionMigrationProcedureTask, RegionMigrationTriggerReason,
98    };
99    use crate::procedure::region_migration::test_util::TestingEnv;
100
101    #[tokio::test]
102    async fn test_drain_datanode_check_drained_excludes_running_migrations() {
103        for has_pending_region in [false, true] {
104            let env = TestEnv::new().await;
105            let running = RegionId::new(1024, 1);
106            let pending = RegionId::new(1024, 2);
107            let failed = RegionId::new(1024, 3);
108            let mut routes = vec![leader_region(running, 1), leader_region(failed, 1)];
109            if has_pending_region {
110                routes.push(leader_region(pending, 1));
111            }
112            env.create_table(routes).await;
113            let guard = env
114                .region_migration_manager
115                .tracker()
116                .insert_running_procedure(&RegionMigrationProcedureTask::new(
117                    running,
118                    Peer::empty(1),
119                    Peer::empty(2),
120                    Duration::from_secs(10),
121                    RegionMigrationTriggerReason::Manual,
122                ))
123                .unwrap();
124            let mut ctx = env.context(new_persistent_context(
125                Some(2),
126                DrainDatanodeStrategy::TryAll,
127            ));
128            ctx.volatile_ctx.failed_regions.push(failed);
129            let procedure_ctx = TestingEnv::procedure_context();
130
131            let (next, status) = CheckDrained.next(&mut ctx, &procedure_ctx).await.unwrap();
132            assert!(matches!(status, Status::Executing { persist: false, .. }));
133            assert_eq!(vec![failed], ctx.volatile_ctx.failed_regions);
134            assert_eq!(0, ctx.volatile_ctx.migrated);
135            if has_pending_region {
136                assert!(next.as_any().is::<Dispatch>());
137                assert_eq!(vec![pending], leader_regions(&ctx));
138            } else {
139                assert!(next.as_any().is::<DrainDatanodeEnd>());
140                assert!(leader_regions(&ctx).is_empty());
141            }
142
143            // Once the migration releases its guard, an unchanged route needs a new child.
144            drop(guard);
145            let (next, _) = CheckDrained.next(&mut ctx, &procedure_ctx).await.unwrap();
146            assert!(next.as_any().is::<Dispatch>());
147            let expected = if has_pending_region {
148                vec![running, pending]
149            } else {
150                vec![running]
151            };
152            assert_eq!(expected, leader_regions(&ctx));
153        }
154    }
155
156    #[tokio::test]
157    async fn test_drain_datanode_check_drained_transitions_to_end() {
158        let env = TestEnv::new().await;
159
160        let mut ctx = env.context(new_persistent_context(
161            Some(2),
162            DrainDatanodeStrategy::FailFast,
163        ));
164        let procedure_ctx = TestingEnv::procedure_context();
165
166        let (next, status) = CheckDrained.next(&mut ctx, &procedure_ctx).await.unwrap();
167
168        assert!(next.as_any().downcast_ref::<DrainDatanodeEnd>().is_some());
169        assert!(matches!(status, Status::Executing { persist: false, .. }));
170    }
171
172    #[tokio::test]
173    async fn test_drain_datanode_check_drained_rescan_requeues() {
174        let env = TestEnv::new().await;
175        let region_id = RegionId::new(1024, 1);
176        env.create_table(vec![leader_region(region_id, 1)]).await;
177
178        // The plan of the procedure is stale: a region still on the datanode is not listed.
179        let mut ctx = env.context(new_persistent_context(
180            None,
181            DrainDatanodeStrategy::FailFast,
182        ));
183        let procedure_ctx = TestingEnv::procedure_context();
184
185        let (next, status) = CheckDrained.next(&mut ctx, &procedure_ctx).await.unwrap();
186
187        assert!(matches!(status, Status::Executing { .. }));
188        assert!(next.as_any().downcast_ref::<Dispatch>().is_some());
189        assert_eq!(vec![region_id], leader_regions(&ctx));
190        assert_eq!(1, ctx.volatile_ctx.rescan_round);
191    }
192
193    #[tokio::test]
194    async fn test_drain_datanode_check_drained_exceeded_rounds() {
195        let env = TestEnv::new().await;
196        let region_id = RegionId::new(1024, 1);
197        env.create_table(vec![leader_region(region_id, 1)]).await;
198
199        let mut ctx = env.context(new_persistent_context(
200            None,
201            DrainDatanodeStrategy::FailFast,
202        ));
203        ctx.volatile_ctx.rescan_round = MAX_RESCAN_ROUNDS;
204        let procedure_ctx = TestingEnv::procedure_context();
205
206        let error = CheckDrained
207            .next(&mut ctx, &procedure_ctx)
208            .await
209            .unwrap_err();
210
211        assert!(matches!(
212            error,
213            error::Error::DrainDatanodeRescanExceeded { .. }
214        ));
215        assert!(!error.is_retryable());
216    }
217
218    #[tokio::test]
219    async fn test_drain_datanode_check_drained_ignores_failed_regions() {
220        let env = TestEnv::new().await;
221        let region_id = RegionId::new(1024, 1);
222        env.create_table(vec![leader_region(region_id, 1)]).await;
223
224        let mut ctx = env.context(new_persistent_context(None, DrainDatanodeStrategy::TryAll));
225        ctx.volatile_ctx.failed_regions = vec![region_id];
226        let procedure_ctx = TestingEnv::procedure_context();
227
228        let (next, status) = CheckDrained.next(&mut ctx, &procedure_ctx).await.unwrap();
229
230        assert!(matches!(status, Status::Executing { .. }));
231        assert!(next.as_any().downcast_ref::<DrainDatanodeEnd>().is_some());
232    }
233}
meta_srv/procedure/drain_datanode/check_drained.rs

meta_srv/procedure/drain_datanode/
check_drained.rs
