meta_srv/procedure/drain_datanode/
utils.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::collections::BTreeSet;
16
17use common_meta::DatanodeId;
18use common_meta::key::TableMetadataManagerRef;
19use common_meta::key::table_route::TableRouteValue;
20use common_meta::peer::Peer;
21use common_meta::rpc::router::region_distribution;
22use common_time::util::current_time_millis;
23use futures::TryStreamExt;
24use snafu::ResultExt;
25use store_api::storage::RegionId;
26
27use crate::error::{self, Result};
28use crate::procedure::drain_datanode::{Context, DrainDatanodePlan};
29use crate::selector::common;
30
31/// Builds the drain plan of `node_id`.
32///
33/// Walks the datanode table index and reads the authoritative table routes, so the plan
34/// reflects the routes rather than the index. Each table route is split by role with
35/// [`region_distribution`], which keys the regions by datanode id.
36pub(crate) async fn build_datanode_plan(
37    table_metadata_manager: &TableMetadataManagerRef,
38    node_id: DatanodeId,
39) -> Result<DrainDatanodePlan> {
40    let mut leader_regions = BTreeSet::new();
41    let mut follower_regions = BTreeSet::new();
42
43    let mut tables = table_metadata_manager
44        .datanode_table_manager()
45        .tables(node_id);
46    while let Some(value) = tables
47        .try_next()
48        .await
49        .context(error::TableMetadataManagerSnafu)?
50    {
51        let Some(table_route) = table_metadata_manager
52            .table_route_manager()
53            .table_route_storage()
54            .get(value.table_id)
55            .await
56            .context(error::TableMetadataManagerSnafu)?
57        else {
58            continue;
59        };
60        let TableRouteValue::Physical(table_route) = table_route else {
61            continue;
62        };
63
64        let mut distribution = region_distribution(&table_route.region_routes);
65        let Some(roles) = distribution.remove(&node_id) else {
66            continue;
67        };
68        leader_regions.extend(
69            roles
70                .leader_regions
71                .into_iter()
72                .map(|number| RegionId::new(value.table_id, number)),
73        );
74        follower_regions.extend(
75            roles
76                .follower_regions
77                .into_iter()
78                .map(|number| RegionId::new(value.table_id, number)),
79        );
80    }
81
82    Ok(DrainDatanodePlan {
83        leader_regions: leader_regions.into_iter().collect(),
84        follower_regions: follower_regions.into_iter().collect(),
85    })
86}
87
88/// Returns true if `node_id` still leads `region_id`.
89///
90/// The table route is authoritative, so a region that is not listed anymore is reported as
91/// not led by the datanode.
92pub(crate) async fn leads_region(
93    table_metadata_manager: &TableMetadataManagerRef,
94    node_id: DatanodeId,
95    region_id: RegionId,
96) -> Result<bool> {
97    let Some(table_route) = table_metadata_manager
98        .table_route_manager()
99        .table_route_storage()
100        .get(region_id.table_id())
101        .await
102        .context(error::TableMetadataManagerSnafu)?
103    else {
104        return Ok(false);
105    };
106    let TableRouteValue::Physical(table_route) = table_route else {
107        return Ok(false);
108    };
109    let Some(region_route) = table_route
110        .region_routes
111        .iter()
112        .find(|region_route| region_route.region.id == region_id)
113    else {
114        return Ok(false);
115    };
116
117    Ok(region_route.leader_peer.as_ref().map(|peer| peer.id) == Some(node_id))
118}
119
120/// Returns an error if the drain exceeded its deadline.
121pub(crate) fn ensure_drain_deadline(ctx: &Context) -> Result<()> {
122    if current_time_millis() < ctx.persistent_ctx.deadline_ms {
123        return Ok(());
124    }
125
126    let remaining_leader_regions = ctx
127        .volatile_ctx
128        .plan
129        .as_ref()
130        .map_or(0, |plan| plan.leader_regions.len());
131    error::ExceededDeadlineSnafu {
132        operation: format!(
133            "Drain datanode {}, migrated: {}, remaining leader regions: {}",
134            ctx.persistent_ctx.node_id, ctx.volatile_ctx.migrated, remaining_leader_regions
135        ),
136    }
137    .fail()
138}
139
140/// Selects the migration destination of `region_id` off the drained datanode.
141///
142/// Returns the pinned target when the procedure has one, otherwise the schedulable
143/// datanode with the smallest peer id. Candidates come from the
144/// peer discovery, so the returned peer carries its address and was not marked
145/// unschedulable; the drained datanode and the existing followers are excluded.
146pub(crate) async fn select_migration_target(
147    ctx: &Context,
148    region_id: RegionId,
149    follower_peers: &[Peer],
150) -> Result<Peer> {
151    let node_id = ctx.persistent_ctx.node_id;
152    let target_node_id = ctx.persistent_ctx.target_node_id;
153    let target = common::active_schedulable_datanodes(&ctx.selector_ctx, None)
154        .await?
155        .into_iter()
156        .map(|datanode| datanode.peer)
157        .filter(|peer| peer.id != node_id)
158        .filter(|peer| target_node_id.is_none_or(|target| peer.id == target))
159        .filter(|peer| !follower_peers.iter().any(|follower| follower.id == peer.id))
160        .min_by_key(|peer| peer.id);
161
162    match target {
163        Some(peer) => Ok(peer),
164        None => error::DrainDatanodeAbortSnafu {
165            node_id,
166            failed_regions: vec![region_id],
167        }
168        .fail(),
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use common_meta::peer::Peer;
175    use store_api::storage::RegionId;
176
177    use crate::error;
178    use crate::key::NodeSchedulingRole;
179    use crate::procedure::drain_datanode::test_util::{
180        TestEnv, follower_region, leader_region, new_persistent_context,
181    };
182    use crate::procedure::drain_datanode::{DrainDatanodeStrategy, utils};
183
184    #[tokio::test]
185    async fn test_build_datanode_plan() {
186        let env = TestEnv::new().await;
187        let first = RegionId::new(1024, 1);
188        let second = RegionId::new(1024, 2);
189        env.create_table(vec![leader_region(first, 1), follower_region(second, 2, 1)])
190            .await;
191
192        let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1)
193            .await
194            .unwrap();
195
196        assert_eq!(vec![first], plan.leader_regions);
197        assert_eq!(vec![second], plan.follower_regions);
198    }
199
200    #[tokio::test]
201    async fn test_build_datanode_plan_of_empty_datanode() {
202        let env = TestEnv::new().await;
203
204        let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1)
205            .await
206            .unwrap();
207
208        assert!(plan.leader_regions.is_empty());
209        assert!(plan.follower_regions.is_empty());
210    }
211
212    #[tokio::test]
213    async fn test_build_datanode_plan_reads_the_table_route() {
214        let env = TestEnv::new().await;
215        let region_id = RegionId::new(1024, 1);
216        // The datanode is listed as a follower of the region, so it is not a leader even
217        // though the table index of the datanode contains the table.
218        env.create_table(vec![follower_region(region_id, 2, 1)])
219            .await;
220
221        let plan = utils::build_datanode_plan(env.table_metadata_manager(), 1)
222            .await
223            .unwrap();
224
225        assert!(plan.leader_regions.is_empty());
226        assert_eq!(vec![region_id], plan.follower_regions);
227    }
228
229    #[tokio::test]
230    async fn test_select_migration_target() {
231        let env = TestEnv::new().await;
232        let region_id = RegionId::new(1024, 1);
233
234        let mut ctx = env.context(new_persistent_context(
235            None,
236            DrainDatanodeStrategy::FailFast,
237        ));
238        let peer = utils::select_migration_target(&ctx, region_id, &[])
239            .await
240            .unwrap();
241        assert_eq!(2, peer.id);
242
243        // The existing followers and the drained datanode are excluded.
244        let peer = utils::select_migration_target(&ctx, region_id, &[Peer::empty(2)])
245            .await
246            .unwrap();
247        assert_eq!(3, peer.id);
248        assert_eq!("dn3:3001", peer.addr);
249
250        // A pinned target is resolved to the active peer, so it carries its address.
251        ctx.persistent_ctx.target_node_id = Some(2);
252        let peer = utils::select_migration_target(&ctx, region_id, &[])
253            .await
254            .unwrap();
255        assert_eq!(2, peer.id);
256        assert_eq!("dn2:3001", peer.addr);
257
258        // A target that is already a follower, the drained datanode itself or an
259        // unschedulable datanode is rejected.
260        ctx.persistent_ctx.target_node_id = Some(2);
261        let error = utils::select_migration_target(&ctx, region_id, &[Peer::empty(2)])
262            .await
263            .unwrap_err();
264        assert!(matches!(error, error::Error::DrainDatanodeAbort { .. }));
265
266        ctx.persistent_ctx.target_node_id = Some(1);
267        let error = utils::select_migration_target(&ctx, region_id, &[])
268            .await
269            .unwrap_err();
270        assert!(matches!(error, error::Error::DrainDatanodeAbort { .. }));
271
272        env.node_scheduling_manager
273            .mark_unschedulable(NodeSchedulingRole::Datanode, 3)
274            .await
275            .unwrap();
276        ctx.persistent_ctx.target_node_id = Some(3);
277        let error = utils::select_migration_target(&ctx, region_id, &[])
278            .await
279            .unwrap_err();
280        assert!(matches!(error, error::Error::DrainDatanodeAbort { .. }));
281    }
282}
meta_srv/procedure/drain_datanode/utils.rs

meta_srv/procedure/drain_datanode/
utils.rs
