meta_srv/procedure/drain_datanode/
manager.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::HashMap;
16use std::collections::hash_map::Entry;
17use std::sync::{Arc, RwLock};
18
19use common_meta::DatanodeId;
20use common_meta::key::TableMetadataManagerRef;
21use common_procedure::{ProcedureContext, ProcedureId, ProcedureManagerRef, ProcedureWithId};
22use common_telemetry::info;
23use common_time::util::current_time_millis;
24use snafu::ResultExt;
25
26use crate::error::{self, Result};
27use crate::metasrv::SelectorContext;
28use crate::node_scheduling::NodeSchedulingManagerRef;
29use crate::procedure::drain_datanode::{
30    Context as DrainDatanodeContext, DEFAULT_DRAIN_TIMEOUT, DrainDatanodeProcedure,
31    DrainDatanodeProcedureTask, MAX_PARALLELISM, PersistentContext,
32};
33use crate::procedure::region_migration::manager::RegionMigrationManagerRef;
34
35pub type DrainDatanodeManagerRef = Arc<DrainDatanodeManager>;
36
37/// Tracks the running drains, indexed by the drained datanode.
38///
39/// Mirrors `RegionMigrationProcedureTracker`: a drain of a datanode that is already being
40/// drained is rejected at submission instead of waiting for the procedure lock.
41#[derive(Default, Clone)]
42pub struct DrainDatanodeProcedureTracker {
43    running_procedures: Arc<RwLock<HashMap<DatanodeId, DrainDatanodeProcedureTask>>>,
44}
45
46impl DrainDatanodeProcedureTracker {
47    /// Returns the guard of a drain of `task.node_id`, or `None` if that datanode is already
48    /// being drained.
49    pub(crate) fn insert_running_procedure(
50        &self,
51        task: &DrainDatanodeProcedureTask,
52    ) -> Option<DrainDatanodeProcedureGuard> {
53        let mut procedures = self.running_procedures.write().unwrap();
54        match procedures.entry(task.node_id) {
55            Entry::Occupied(_) => None,
56            Entry::Vacant(v) => {
57                v.insert(task.clone());
58                Some(DrainDatanodeProcedureGuard {
59                    node_id: task.node_id,
60                    running_procedures: self.running_procedures.clone(),
61                })
62            }
63        }
64    }
65}
66
67/// Removes the drained datanode from [`DrainDatanodeProcedureTracker`] when dropped.
68pub(crate) struct DrainDatanodeProcedureGuard {
69    node_id: DatanodeId,
70    running_procedures: Arc<RwLock<HashMap<DatanodeId, DrainDatanodeProcedureTask>>>,
71}
72
73impl Drop for DrainDatanodeProcedureGuard {
74    fn drop(&mut self) {
75        self.running_procedures
76            .write()
77            .unwrap()
78            .remove(&self.node_id);
79    }
80}
81
82/// Manages the [`DrainDatanodeProcedure`]s.
83///
84/// It registers the procedure loader and submits drain tasks. Procedures of the same
85/// datanode are serialized by the procedure lock declared in
86/// [`DrainDatanodeProcedure::lock_key`].
87pub struct DrainDatanodeManager {
88    procedure_manager: ProcedureManagerRef,
89    table_metadata_manager: TableMetadataManagerRef,
90    node_scheduling_manager: NodeSchedulingManagerRef,
91    region_migration_manager: RegionMigrationManagerRef,
92    selector_ctx: SelectorContext,
93
94    tracker: DrainDatanodeProcedureTracker,
95}
96
97impl DrainDatanodeManager {
98    /// Returns a new [`DrainDatanodeManager`].
99    pub fn new(
100        procedure_manager: ProcedureManagerRef,
101        table_metadata_manager: TableMetadataManagerRef,
102        node_scheduling_manager: NodeSchedulingManagerRef,
103        region_migration_manager: RegionMigrationManagerRef,
104        selector_ctx: SelectorContext,
105    ) -> Self {
106        Self {
107            procedure_manager,
108            table_metadata_manager,
109            node_scheduling_manager,
110            region_migration_manager,
111            selector_ctx,
112
113            tracker: DrainDatanodeProcedureTracker::default(),
114        }
115    }
116
117    /// Returns the tracker of the running drains.
118    #[cfg(test)]
119    pub(crate) fn tracker(&self) -> DrainDatanodeProcedureTracker {
120        self.tracker.clone()
121    }
122
123    /// Returns the context of a drain run.
124    fn context(&self, persistent_ctx: PersistentContext) -> DrainDatanodeContext {
125        DrainDatanodeContext::new(
126            self.table_metadata_manager.clone(),
127            self.node_scheduling_manager.clone(),
128            self.region_migration_manager.clone(),
129            self.selector_ctx.clone(),
130            persistent_ctx,
131        )
132    }
133
134    /// Submits a [`DrainDatanodeProcedure`] and returns its id.
135    pub async fn submit_procedure(
136        &self,
137        procedure_context: ProcedureContext,
138        task: DrainDatanodeProcedureTask,
139    ) -> Result<ProcedureId> {
140        if task.target_node_id == Some(task.node_id) {
141            return error::InvalidArgumentsSnafu {
142                err_msg: "the target datanode must differ from the drained datanode",
143            }
144            .fail();
145        }
146        if task.parallelism == 0 || task.parallelism > MAX_PARALLELISM {
147            return error::InvalidArgumentsSnafu {
148                err_msg: format!("the parallelism must be in 1..={}", MAX_PARALLELISM),
149            }
150            .fail();
151        }
152        let timeout = task.timeout.unwrap_or(DEFAULT_DRAIN_TIMEOUT);
153        if timeout.is_zero() {
154            return error::InvalidArgumentsSnafu {
155                err_msg: "the timeout must be greater than zero",
156            }
157            .fail();
158        }
159
160        let Some(guard) = self.tracker.insert_running_procedure(&task) else {
161            return error::DrainDatanodeRunningSnafu {
162                node_id: task.node_id,
163            }
164            .fail();
165        };
166
167        let persistent_ctx = PersistentContext {
168            node_id: task.node_id,
169            target_node_id: task.target_node_id,
170            parallelism: task.parallelism,
171            strategy: task.strategy,
172            // The budget starts when the task is submitted, so neither a slow start nor a
173            // restart extends it.
174            deadline_ms: current_time_millis() + timeout.as_millis() as i64,
175        };
176        let procedure = DrainDatanodeProcedure::new(self.context(persistent_ctx), vec![guard]);
177        let procedure_with_id =
178            ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context);
179        let procedure_id = procedure_with_id.id;
180        info!(
181            "Starting drain datanode procedure {procedure_id} for datanode {}, timeout: {:?}",
182            task.node_id, timeout
183        );
184        self.procedure_manager
185            .submit(procedure_with_id)
186            .await
187            .context(error::SubmitProcedureSnafu)?;
188
189        Ok(procedure_id)
190    }
191
192    /// Registers the loader of [`DrainDatanodeProcedure`].
193    pub(crate) fn try_start(&self) -> Result<()> {
194        let table_metadata_manager = self.table_metadata_manager.clone();
195        let node_scheduling_manager = self.node_scheduling_manager.clone();
196        let region_migration_manager = self.region_migration_manager.clone();
197        let selector_ctx = self.selector_ctx.clone();
198        let tracker = self.tracker.clone();
199        self.procedure_manager
200            .register_loader(
201                DrainDatanodeProcedure::TYPE_NAME,
202                Box::new(move |json| {
203                    let table_metadata_manager = table_metadata_manager.clone();
204                    let node_scheduling_manager = node_scheduling_manager.clone();
205                    let region_migration_manager = region_migration_manager.clone();
206                    let selector_ctx = selector_ctx.clone();
207                    let factory = move |persistent_ctx| {
208                        DrainDatanodeContext::new(
209                            table_metadata_manager,
210                            node_scheduling_manager,
211                            region_migration_manager,
212                            selector_ctx,
213                            persistent_ctx,
214                        )
215                    };
216                    DrainDatanodeProcedure::from_json(json, &tracker, factory)
217                        .map(|procedure| Box::new(procedure) as _)
218                }),
219            )
220            .context(error::RegisterProcedureLoaderSnafu {
221                type_name: DrainDatanodeProcedure::TYPE_NAME,
222            })
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use std::time::Duration;
229
230    use common_procedure::Procedure;
231
232    use crate::error;
233    use crate::procedure::drain_datanode::manager::DrainDatanodeProcedureTracker;
234    use crate::procedure::drain_datanode::test_util::{TestEnv, new_persistent_context};
235    use crate::procedure::drain_datanode::{
236        DEFAULT_PARALLELISM, DrainDatanodeProcedure, DrainDatanodeProcedureTask,
237        DrainDatanodeStrategy,
238    };
239
240    #[tokio::test]
241    async fn test_submit_procedure_rejects_invalid_task() {
242        let env = TestEnv::new().await;
243        let manager = env.drain_manager();
244
245        // The pinned target must differ from the drained datanode.
246        let error = manager
247            .submit_procedure(
248                common_procedure::ProcedureContext::default(),
249                DrainDatanodeProcedureTask {
250                    node_id: 1,
251                    target_node_id: Some(1),
252                    parallelism: DEFAULT_PARALLELISM,
253                    strategy: DrainDatanodeStrategy::FailFast,
254                    timeout: None,
255                },
256            )
257            .await
258            .unwrap_err();
259        assert!(matches!(error, error::Error::InvalidArguments { .. }));
260
261        // The parallelism must be in range.
262        for parallelism in [0, crate::procedure::drain_datanode::MAX_PARALLELISM + 1] {
263            let error = manager
264                .submit_procedure(
265                    common_procedure::ProcedureContext::default(),
266                    DrainDatanodeProcedureTask {
267                        node_id: 1,
268                        target_node_id: Some(2),
269                        parallelism,
270                        strategy: DrainDatanodeStrategy::FailFast,
271                        timeout: None,
272                    },
273                )
274                .await
275                .unwrap_err();
276            assert!(matches!(error, error::Error::InvalidArguments { .. }));
277        }
278
279        // The timeout must be greater than zero.
280        let error = manager
281            .submit_procedure(
282                common_procedure::ProcedureContext::default(),
283                DrainDatanodeProcedureTask {
284                    node_id: 1,
285                    target_node_id: Some(2),
286                    parallelism: DEFAULT_PARALLELISM,
287                    strategy: DrainDatanodeStrategy::FailFast,
288                    timeout: Some(Duration::ZERO),
289                },
290            )
291            .await
292            .unwrap_err();
293        assert!(matches!(error, error::Error::InvalidArguments { .. }));
294    }
295
296    #[tokio::test]
297    async fn test_submit_procedure_rejects_a_running_drain() {
298        let env = TestEnv::new().await;
299        let manager = env.drain_manager();
300        let task = DrainDatanodeProcedureTask {
301            node_id: 1,
302            target_node_id: Some(2),
303            parallelism: DEFAULT_PARALLELISM,
304            strategy: DrainDatanodeStrategy::FailFast,
305            timeout: None,
306        };
307
308        // A drain of datanode 1 is running, for example one recovered after a restart.
309        let _guard = manager.tracker().insert_running_procedure(&task).unwrap();
310
311        let error = manager
312            .submit_procedure(common_procedure::ProcedureContext::default(), task)
313            .await
314            .unwrap_err();
315
316        assert!(matches!(error, error::Error::DrainDatanodeRunning { .. }));
317        assert!(!error.is_retryable());
318    }
319
320    #[tokio::test]
321    async fn test_drain_datanode_takes_the_place_back_on_reload() {
322        let env = TestEnv::new().await;
323        let context = env.context(new_persistent_context(
324            Some(2),
325            DrainDatanodeStrategy::FailFast,
326        ));
327        let procedure = DrainDatanodeProcedure::new(context, vec![]);
328        let json = procedure.dump().unwrap();
329        let tracker = DrainDatanodeProcedureTracker::default();
330
331        // Reloading takes the place back, so a second drain of the datanode is rejected.
332        let reloaded = DrainDatanodeProcedure::from_json(&json, &tracker, |persistent_ctx| {
333            env.context(persistent_ctx)
334        })
335        .unwrap();
336        let task = DrainDatanodeProcedureTask {
337            node_id: 1,
338            target_node_id: Some(2),
339            parallelism: DEFAULT_PARALLELISM,
340            strategy: DrainDatanodeStrategy::FailFast,
341            timeout: None,
342        };
343        assert!(tracker.insert_running_procedure(&task).is_none());
344
345        // Dropping the reloaded procedure releases the place.
346        drop(reloaded);
347        assert!(tracker.insert_running_procedure(&task).is_some());
348    }
349}
meta_srv/procedure/drain_datanode/manager.rs

meta_srv/procedure/drain_datanode/
manager.rs
