oximedia_distributed/
worker_draining.rs

1// Copyright 2024 OxiMedia Project
2// Licensed under the Apache License, Version 2.0
3
4//! Worker draining — graceful shutdown for distributed encoding workers.
5//!
6//! Draining allows a worker to finish its current in-flight tasks without
7//! accepting new ones, then transition to an `Idle` or `Offline` state.
8//! This is essential for rolling upgrades and planned maintenance.
9
10use std::collections::{HashMap, HashSet};
11use std::time::{Duration, Instant};
12use thiserror::Error;
13
14/// Errors produced by the draining subsystem.
15#[derive(Debug, Error, PartialEq, Eq)]
16pub enum DrainError {
17    /// The worker ID is not registered.
18    #[error("worker '{0}' not found")]
19    WorkerNotFound(String),
20
21    /// An operation was attempted on a worker that is already draining.
22    #[error("worker '{0}' is already draining")]
23    AlreadyDraining(String),
24
25    /// An operation was attempted on a worker that has not started draining.
26    #[error("worker '{0}' is not in draining state")]
27    NotDraining(String),
28
29    /// A task was submitted to a worker that is draining.
30    #[error("worker '{0}' is draining and cannot accept new tasks")]
31    WorkerDraining(String),
32}
33
34/// The operational state of a managed worker.
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub enum WorkerState {
37    /// Worker is healthy and accepting tasks.
38    Active,
39    /// Worker is finishing in-flight tasks; no new tasks accepted.
40    Draining,
41    /// Worker has finished all tasks and is ready to be removed.
42    Drained,
43    /// Worker is offline / shut down.
44    Offline,
45}
46
47/// Internal record kept per worker.
48#[derive(Debug)]
49struct WorkerRecord {
50    state: WorkerState,
51    in_flight: HashSet<String>,
52    drain_started: Option<Instant>,
53}
54
55/// Manages graceful draining of workers in a distributed cluster.
56///
57/// Tasks are tracked per worker. When draining is requested the worker
58/// stops accepting new tasks and a [`WorkerState::Drained`] transition
59/// is triggered automatically once all in-flight tasks complete.
60#[derive(Debug, Default)]
61pub struct DrainManager {
62    workers: HashMap<String, WorkerRecord>,
63}
64
65impl DrainManager {
66    /// Create a new drain manager with no workers.
67    pub fn new() -> Self {
68        Self::default()
69    }
70
71    /// Register a new worker in the `Active` state.
72    pub fn register(&mut self, worker_id: impl Into<String>) {
73        self.workers.insert(
74            worker_id.into(),
75            WorkerRecord {
76                state: WorkerState::Active,
77                in_flight: HashSet::new(),
78                drain_started: None,
79            },
80        );
81    }
82
83    /// Deregister a worker.  Returns `true` if it was present.
84    pub fn deregister(&mut self, worker_id: &str) -> bool {
85        self.workers.remove(worker_id).is_some()
86    }
87
88    /// Assign a task to a worker.
89    ///
90    /// Fails with [`DrainError::WorkerDraining`] if the worker is draining or
91    /// drained, and with [`DrainError::WorkerNotFound`] if unknown.
92    pub fn assign_task(
93        &mut self,
94        worker_id: &str,
95        task_id: impl Into<String>,
96    ) -> Result<(), DrainError> {
97        let record = self
98            .workers
99            .get_mut(worker_id)
100            .ok_or_else(|| DrainError::WorkerNotFound(worker_id.to_owned()))?;
101
102        if record.state != WorkerState::Active {
103            return Err(DrainError::WorkerDraining(worker_id.to_owned()));
104        }
105
106        record.in_flight.insert(task_id.into());
107        Ok(())
108    }
109
110    /// Mark a task as complete on a worker.
111    ///
112    /// If the worker is draining and has no remaining in-flight tasks,
113    /// it transitions to [`WorkerState::Drained`] automatically.
114    pub fn complete_task(&mut self, worker_id: &str, task_id: &str) -> Result<(), DrainError> {
115        let record = self
116            .workers
117            .get_mut(worker_id)
118            .ok_or_else(|| DrainError::WorkerNotFound(worker_id.to_owned()))?;
119
120        record.in_flight.remove(task_id);
121
122        // Auto-transition to Drained when drain is in progress and no tasks remain.
123        if record.state == WorkerState::Draining && record.in_flight.is_empty() {
124            record.state = WorkerState::Drained;
125        }
126
127        Ok(())
128    }
129
130    /// Initiate graceful draining of a worker.
131    ///
132    /// Returns [`DrainError::AlreadyDraining`] if already in progress.
133    /// If the worker has no in-flight tasks it immediately becomes `Drained`.
134    pub fn start_drain(&mut self, worker_id: &str) -> Result<(), DrainError> {
135        let record = self
136            .workers
137            .get_mut(worker_id)
138            .ok_or_else(|| DrainError::WorkerNotFound(worker_id.to_owned()))?;
139
140        if record.state == WorkerState::Draining || record.state == WorkerState::Drained {
141            return Err(DrainError::AlreadyDraining(worker_id.to_owned()));
142        }
143
144        record.drain_started = Some(Instant::now());
145
146        if record.in_flight.is_empty() {
147            record.state = WorkerState::Drained;
148        } else {
149            record.state = WorkerState::Draining;
150        }
151
152        Ok(())
153    }
154
155    /// Force-complete draining (e.g., after a timeout), discarding in-flight tasks.
156    pub fn force_drain(&mut self, worker_id: &str) -> Result<Vec<String>, DrainError> {
157        let record = self
158            .workers
159            .get_mut(worker_id)
160            .ok_or_else(|| DrainError::WorkerNotFound(worker_id.to_owned()))?;
161
162        let discarded: Vec<String> = record.in_flight.drain().collect();
163        record.state = WorkerState::Drained;
164        Ok(discarded)
165    }
166
167    /// Return the current state of a worker.
168    pub fn state(&self, worker_id: &str) -> Option<&WorkerState> {
169        self.workers.get(worker_id).map(|r| &r.state)
170    }
171
172    /// Return the number of in-flight tasks on a worker.
173    pub fn in_flight_count(&self, worker_id: &str) -> Option<usize> {
174        self.workers.get(worker_id).map(|r| r.in_flight.len())
175    }
176
177    /// Return the time elapsed since draining started, if applicable.
178    pub fn drain_elapsed(&self, worker_id: &str) -> Option<Duration> {
179        self.workers
180            .get(worker_id)
181            .and_then(|r| r.drain_started.map(|t| t.elapsed()))
182    }
183
184    /// Collect all workers whose drain has exceeded `timeout`.
185    ///
186    /// These workers can then be force-drained to unblock rolling upgrades.
187    pub fn timed_out_drains(&self, timeout: Duration) -> Vec<String> {
188        self.workers
189            .iter()
190            .filter_map(|(id, r)| {
191                if r.state == WorkerState::Draining {
192                    r.drain_started
193                        .filter(|t| t.elapsed() >= timeout)
194                        .map(|_| id.clone())
195                } else {
196                    None
197                }
198            })
199            .collect()
200    }
201
202    /// Return ids of all `Active` workers (candidates to receive new tasks).
203    pub fn active_workers(&self) -> Vec<String> {
204        let mut ids: Vec<String> = self
205            .workers
206            .iter()
207            .filter(|(_, r)| r.state == WorkerState::Active)
208            .map(|(id, _)| id.clone())
209            .collect();
210        ids.sort();
211        ids
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218
219    fn manager_with_two_workers() -> DrainManager {
220        let mut m = DrainManager::new();
221        m.register("w1");
222        m.register("w2");
223        m
224    }
225
226    #[test]
227    fn test_register_and_active_state() {
228        let m = manager_with_two_workers();
229        assert_eq!(m.state("w1"), Some(&WorkerState::Active));
230    }
231
232    #[test]
233    fn test_assign_task_success() {
234        let mut m = manager_with_two_workers();
235        m.assign_task("w1", "t1").unwrap();
236        assert_eq!(m.in_flight_count("w1"), Some(1));
237    }
238
239    #[test]
240    fn test_assign_task_to_draining_fails() {
241        let mut m = manager_with_two_workers();
242        m.assign_task("w1", "t1").unwrap();
243        m.start_drain("w1").unwrap();
244        assert!(matches!(
245            m.assign_task("w1", "t2").unwrap_err(),
246            DrainError::WorkerDraining(_)
247        ));
248    }
249
250    #[test]
251    fn test_drain_with_no_tasks_immediately_drained() {
252        let mut m = manager_with_two_workers();
253        m.start_drain("w1").unwrap();
254        assert_eq!(m.state("w1"), Some(&WorkerState::Drained));
255    }
256
257    #[test]
258    fn test_drain_completes_when_tasks_finish() {
259        let mut m = manager_with_two_workers();
260        m.assign_task("w1", "t1").unwrap();
261        m.assign_task("w1", "t2").unwrap();
262        m.start_drain("w1").unwrap();
263        assert_eq!(m.state("w1"), Some(&WorkerState::Draining));
264        m.complete_task("w1", "t1").unwrap();
265        assert_eq!(m.state("w1"), Some(&WorkerState::Draining));
266        m.complete_task("w1", "t2").unwrap();
267        assert_eq!(m.state("w1"), Some(&WorkerState::Drained));
268    }
269
270    #[test]
271    fn test_already_draining_error() {
272        let mut m = manager_with_two_workers();
273        m.assign_task("w1", "t1").unwrap();
274        m.start_drain("w1").unwrap();
275        assert!(matches!(
276            m.start_drain("w1").unwrap_err(),
277            DrainError::AlreadyDraining(_)
278        ));
279    }
280
281    #[test]
282    fn test_force_drain_discards_tasks() {
283        let mut m = manager_with_two_workers();
284        m.assign_task("w1", "t1").unwrap();
285        m.assign_task("w1", "t2").unwrap();
286        m.start_drain("w1").unwrap();
287        let discarded = m.force_drain("w1").unwrap();
288        assert_eq!(discarded.len(), 2);
289        assert_eq!(m.state("w1"), Some(&WorkerState::Drained));
290        assert_eq!(m.in_flight_count("w1"), Some(0));
291    }
292
293    #[test]
294    fn test_deregister() {
295        let mut m = manager_with_two_workers();
296        assert!(m.deregister("w1"));
297        assert!(m.state("w1").is_none());
298        assert!(!m.deregister("ghost"));
299    }
300
301    #[test]
302    fn test_active_workers_excludes_draining() {
303        let mut m = manager_with_two_workers();
304        m.assign_task("w1", "t1").unwrap();
305        m.start_drain("w1").unwrap();
306        let active = m.active_workers();
307        assert_eq!(active, vec!["w2"]);
308    }
309
310    #[test]
311    fn test_unknown_worker_errors() {
312        let mut m = manager_with_two_workers();
313        assert!(matches!(
314            m.assign_task("ghost", "t").unwrap_err(),
315            DrainError::WorkerNotFound(_)
316        ));
317        assert!(matches!(
318            m.start_drain("ghost").unwrap_err(),
319            DrainError::WorkerNotFound(_)
320        ));
321    }
322}
oximedia_distributed/worker_draining.rs

oximedia_distributed/
worker_draining.rs
