meta_srv/procedure/drain_datanode/
drain_datanode_start.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::key::NodeSchedulingRole;
23use crate::procedure::drain_datanode::check_drained::CheckDrained;
24use crate::procedure::drain_datanode::{Context, State};
25
26/// Fences the datanode before scanning its regions.
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub(crate) struct DrainDatanodeStart;
29
30#[async_trait::async_trait]
31#[typetag::serde]
32impl State for DrainDatanodeStart {
33    async fn next(
34        &mut self,
35        ctx: &mut Context,
36        _procedure_ctx: &ProcedureContext,
37    ) -> Result<(Box<dyn State>, Status)> {
38        let node_id = ctx.persistent_ctx.node_id;
39        let marked = ctx
40            .node_scheduling_manager
41            .mark_unschedulable(NodeSchedulingRole::Datanode, node_id)
42            .await?;
43        if marked {
44            info!("Marked datanode {} as unschedulable", node_id);
45        }
46
47        // The only persisted status of the procedure: persisting it registers the procedure,
48        // so that a metasrv restart resumes it. The plan and the progress of a run stay in
49        // the volatile context, and a resumed procedure rebuilds the plan by rescanning the
50        // datanode.
51        Ok((Box::new(CheckDrained), Status::executing(true)))
52    }
53
54    fn as_any(&self) -> &dyn Any {
55        self
56    }
57}
58
59#[cfg(test)]
60mod tests {
61    use common_procedure::Status;
62    use store_api::storage::RegionId;
63
64    use crate::key::NodeSchedulingRole;
65    use crate::procedure::drain_datanode::check_drained::CheckDrained;
66    use crate::procedure::drain_datanode::drain_datanode_start::DrainDatanodeStart;
67    use crate::procedure::drain_datanode::test_util::{
68        TestEnv, leader_region, new_persistent_context,
69    };
70    use crate::procedure::drain_datanode::{DrainDatanodeStrategy, State};
71    use crate::procedure::region_migration::test_util::TestingEnv;
72
73    #[tokio::test]
74    async fn test_drain_datanode_start_transitions_to_check_drained() {
75        let env = TestEnv::new().await;
76        env.create_table(vec![leader_region(RegionId::new(1024, 1), 1)])
77            .await;
78
79        let mut ctx = env.context(new_persistent_context(
80            Some(2),
81            DrainDatanodeStrategy::FailFast,
82        ));
83        let procedure_ctx = TestingEnv::procedure_context();
84
85        let (next, status) = DrainDatanodeStart
86            .next(&mut ctx, &procedure_ctx)
87            .await
88            .unwrap();
89
90        assert!(next.as_any().downcast_ref::<CheckDrained>().is_some());
91        // Registering the procedure is the only persisted step of the run.
92        assert!(matches!(status, Status::Executing { persist: true, .. }));
93        assert!(ctx.volatile_ctx.plan.is_none());
94        assert!(
95            !env.node_scheduling_manager
96                .is_schedulable(NodeSchedulingRole::Datanode, 1)
97                .await
98                .unwrap()
99        );
100    }
101}
meta_srv/procedure/drain_datanode/drain_datanode_start.rs

meta_srv/procedure/drain_datanode/
drain_datanode_start.rs
