meta_srv/procedure/drain_datanode/
drain_datanode_end.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::{error, info};
19use serde::{Deserialize, Serialize};
20
21use crate::error::{self, Result};
22use crate::procedure::drain_datanode::{Context, State};
23
24/// Finishes the procedure and reports the outcome.
25///
26/// The procedure state carries no output: only its status is observable through the
27/// procedure API, so a drain that left regions behind ends as `Failed` with the failed
28/// regions in the error instead of as `Done` with a report.
29#[derive(Debug, Clone, Serialize, Deserialize)]
30pub(crate) struct DrainDatanodeEnd;
31
32#[async_trait::async_trait]
33#[typetag::serde]
34impl State for DrainDatanodeEnd {
35    async fn next(
36        &mut self,
37        ctx: &mut Context,
38        _procedure_ctx: &ProcedureContext,
39    ) -> Result<(Box<dyn State>, Status)> {
40        let node_id = ctx.persistent_ctx.node_id;
41        let migrated = ctx.volatile_ctx.migrated;
42        let follower_regions = ctx.volatile_ctx.follower_regions().len();
43        let failed_regions = ctx.volatile_ctx.failed_regions.clone();
44
45        if !failed_regions.is_empty() {
46            error!(
47                "Drained datanode {} except {} regions, failed regions: {:?}",
48                node_id,
49                failed_regions.len(),
50                failed_regions
51            );
52            return error::DrainDatanodeAbortSnafu {
53                node_id,
54                failed_regions,
55            }
56            .fail();
57        }
58
59        info!(
60            "Drained datanode {}, migrated leader regions: {}, follower regions left: {}",
61            node_id, migrated, follower_regions
62        );
63
64        Ok((Box::new(DrainDatanodeEnd), Status::done()))
65    }
66
67    fn as_any(&self) -> &dyn Any {
68        self
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use common_procedure::Status;
75    use store_api::storage::RegionId;
76
77    use crate::error;
78    use crate::procedure::drain_datanode::drain_datanode_end::DrainDatanodeEnd;
79    use crate::procedure::drain_datanode::test_util::{TestEnv, new_persistent_context};
80    use crate::procedure::drain_datanode::{DrainDatanodeStrategy, State};
81    use crate::procedure::region_migration::test_util::TestingEnv;
82
83    #[tokio::test]
84    async fn test_drain_datanode_end_transitions_to_done() {
85        let env = TestEnv::new().await;
86
87        let mut ctx = env.context(new_persistent_context(
88            Some(2),
89            DrainDatanodeStrategy::TryAll,
90        ));
91        ctx.volatile_ctx.migrated = 1;
92        let procedure_ctx = TestingEnv::procedure_context();
93
94        let (next, status) = DrainDatanodeEnd
95            .next(&mut ctx, &procedure_ctx)
96            .await
97            .unwrap();
98
99        assert!(next.as_any().downcast_ref::<DrainDatanodeEnd>().is_some());
100        assert!(matches!(status, Status::Done { output: None }));
101    }
102
103    #[tokio::test]
104    async fn test_drain_datanode_end_aborts_when_regions_are_left() {
105        let env = TestEnv::new().await;
106        let region_id = RegionId::new(1024, 1);
107
108        let mut ctx = env.context(new_persistent_context(
109            Some(2),
110            DrainDatanodeStrategy::TryAll,
111        ));
112        ctx.volatile_ctx.failed_regions = vec![region_id];
113        let procedure_ctx = TestingEnv::procedure_context();
114
115        let error = DrainDatanodeEnd
116            .next(&mut ctx, &procedure_ctx)
117            .await
118            .unwrap_err();
119
120        let error::Error::DrainDatanodeAbort { failed_regions, .. } = error else {
121            panic!("expected the drain procedure to abort");
122        };
123        assert_eq!(vec![region_id], failed_regions);
124    }
125}
meta_srv/procedure/drain_datanode/drain_datanode_end.rs

meta_srv/procedure/drain_datanode/
drain_datanode_end.rs
