oximedia_workflow/
workflow_throttle.rs

1//! Workflow throttling and rate-limiting for `oximedia-workflow`.
2//!
3//! [`WorkflowThrottler`] controls how many workflows or tasks may execute
4//! concurrently, applying a [`ThrottlePolicy`] and tracking current load via
5//! [`ThrottleState`].
6
7#![allow(dead_code)]
8
9use serde::{Deserialize, Serialize};
10use std::collections::VecDeque;
11
12// ---------------------------------------------------------------------------
13// Throttle policy
14// ---------------------------------------------------------------------------
15
16/// Strategy used to limit concurrent execution.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
18pub enum ThrottlePolicy {
19    /// No throttling; allow unlimited concurrency.
20    Unlimited,
21    /// Fixed maximum number of concurrent executions.
22    FixedMax,
23    /// Token-bucket rate limiter (tokens refill over time).
24    TokenBucket,
25    /// Adaptive: adjusts concurrency based on system load.
26    Adaptive,
27}
28
29impl ThrottlePolicy {
30    /// Returns a short label.
31    #[must_use]
32    pub const fn label(self) -> &'static str {
33        match self {
34            Self::Unlimited => "Unlimited",
35            Self::FixedMax => "Fixed Max",
36            Self::TokenBucket => "Token Bucket",
37            Self::Adaptive => "Adaptive",
38        }
39    }
40
41    /// Returns all variants.
42    #[must_use]
43    pub const fn all() -> &'static [ThrottlePolicy] {
44        &[
45            Self::Unlimited,
46            Self::FixedMax,
47            Self::TokenBucket,
48            Self::Adaptive,
49        ]
50    }
51}
52
53impl std::fmt::Display for ThrottlePolicy {
54    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        f.write_str(self.label())
56    }
57}
58
59// ---------------------------------------------------------------------------
60// Throttle state
61// ---------------------------------------------------------------------------
62
63/// Current load snapshot used by the throttler.
64#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
65pub struct ThrottleState {
66    /// Number of currently running executions.
67    pub active: usize,
68    /// Number of executions waiting in the queue.
69    pub queued: usize,
70    /// Total number of executions admitted since the throttler was created.
71    pub total_admitted: u64,
72    /// Total number of executions rejected / delayed.
73    pub total_rejected: u64,
74}
75
76impl ThrottleState {
77    /// Returns utilisation as a fraction `[0.0, 1.0]` relative to a capacity.
78    #[must_use]
79    #[allow(clippy::cast_precision_loss)]
80    pub fn utilisation(&self, capacity: usize) -> f64 {
81        if capacity == 0 {
82            return 0.0;
83        }
84        self.active as f64 / capacity as f64
85    }
86
87    /// Returns `true` if there is room for at least one more execution
88    /// given the specified capacity.
89    #[must_use]
90    pub fn has_capacity(&self, capacity: usize) -> bool {
91        self.active < capacity
92    }
93}
94
95// ---------------------------------------------------------------------------
96// Workflow throttler
97// ---------------------------------------------------------------------------
98
99/// Controls concurrent workflow / task execution according to a policy.
100#[derive(Debug, Clone)]
101pub struct WorkflowThrottler {
102    policy: ThrottlePolicy,
103    max_concurrent: usize,
104    state: ThrottleState,
105    wait_queue: VecDeque<String>,
106}
107
108impl Default for WorkflowThrottler {
109    fn default() -> Self {
110        Self {
111            policy: ThrottlePolicy::Unlimited,
112            max_concurrent: usize::MAX,
113            state: ThrottleState::default(),
114            wait_queue: VecDeque::new(),
115        }
116    }
117}
118
119impl WorkflowThrottler {
120    /// Creates a throttler with a fixed maximum concurrency.
121    #[must_use]
122    pub fn fixed(max_concurrent: usize) -> Self {
123        Self {
124            policy: ThrottlePolicy::FixedMax,
125            max_concurrent,
126            state: ThrottleState::default(),
127            wait_queue: VecDeque::new(),
128        }
129    }
130
131    /// Creates an unlimited throttler.
132    #[must_use]
133    pub fn unlimited() -> Self {
134        Self::default()
135    }
136
137    /// Returns the current policy.
138    #[must_use]
139    pub fn policy(&self) -> ThrottlePolicy {
140        self.policy
141    }
142
143    /// Returns the maximum concurrency limit.
144    #[must_use]
145    pub fn max_concurrent(&self) -> usize {
146        self.max_concurrent
147    }
148
149    /// Returns a snapshot of the current state.
150    #[must_use]
151    pub fn state(&self) -> &ThrottleState {
152        &self.state
153    }
154
155    /// Attempts to admit a workflow for execution.
156    ///
157    /// Returns `true` if admitted, `false` if queued.
158    pub fn try_admit(&mut self, workflow_id: impl Into<String>) -> bool {
159        let wf_id = workflow_id.into();
160        if self.state.active < self.max_concurrent {
161            self.state.active += 1;
162            self.state.total_admitted += 1;
163            true
164        } else {
165            self.wait_queue.push_back(wf_id);
166            self.state.queued = self.wait_queue.len();
167            self.state.total_rejected += 1;
168            false
169        }
170    }
171
172    /// Signals that a workflow has finished, releasing a concurrency slot.
173    ///
174    /// Returns the next queued workflow ID to admit, if any.
175    pub fn release(&mut self) -> Option<String> {
176        if self.state.active > 0 {
177            self.state.active -= 1;
178        }
179        if let Some(next) = self.wait_queue.pop_front() {
180            self.state.active += 1;
181            self.state.total_admitted += 1;
182            self.state.queued = self.wait_queue.len();
183            Some(next)
184        } else {
185            None
186        }
187    }
188
189    /// Returns the number of workflows currently waiting.
190    #[must_use]
191    pub fn queued_count(&self) -> usize {
192        self.wait_queue.len()
193    }
194
195    /// Returns the current utilisation as a fraction.
196    #[must_use]
197    pub fn utilisation(&self) -> f64 {
198        self.state.utilisation(self.max_concurrent)
199    }
200
201    /// Resets the throttler to its initial state.
202    pub fn reset(&mut self) {
203        self.state = ThrottleState::default();
204        self.wait_queue.clear();
205    }
206}
207
208// ===========================================================================
209// Tests
210// ===========================================================================
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215
216    // -- ThrottlePolicy -----------------------------------------------------
217
218    #[test]
219    fn test_policy_label() {
220        assert_eq!(ThrottlePolicy::Unlimited.label(), "Unlimited");
221        assert_eq!(ThrottlePolicy::FixedMax.label(), "Fixed Max");
222    }
223
224    #[test]
225    fn test_policy_display() {
226        assert_eq!(format!("{}", ThrottlePolicy::TokenBucket), "Token Bucket");
227    }
228
229    #[test]
230    fn test_policy_all() {
231        assert_eq!(ThrottlePolicy::all().len(), 4);
232    }
233
234    // -- ThrottleState ------------------------------------------------------
235
236    #[test]
237    fn test_state_default() {
238        let s = ThrottleState::default();
239        assert_eq!(s.active, 0);
240        assert_eq!(s.queued, 0);
241    }
242
243    #[test]
244    fn test_state_utilisation() {
245        let s = ThrottleState {
246            active: 3,
247            ..Default::default()
248        };
249        assert!((s.utilisation(10) - 0.3).abs() < 1e-6);
250    }
251
252    #[test]
253    fn test_state_utilisation_zero_capacity() {
254        let s = ThrottleState::default();
255        assert_eq!(s.utilisation(0), 0.0);
256    }
257
258    #[test]
259    fn test_state_has_capacity() {
260        let s = ThrottleState {
261            active: 5,
262            ..Default::default()
263        };
264        assert!(s.has_capacity(10));
265        assert!(!s.has_capacity(5));
266    }
267
268    // -- WorkflowThrottler --------------------------------------------------
269
270    #[test]
271    fn test_throttler_unlimited() {
272        let mut t = WorkflowThrottler::unlimited();
273        assert!(t.try_admit("wf-1"));
274        assert!(t.try_admit("wf-2"));
275        assert_eq!(t.state().active, 2);
276    }
277
278    #[test]
279    fn test_throttler_fixed_admit() {
280        let mut t = WorkflowThrottler::fixed(2);
281        assert!(t.try_admit("wf-1"));
282        assert!(t.try_admit("wf-2"));
283        // Third should be queued
284        assert!(!t.try_admit("wf-3"));
285        assert_eq!(t.queued_count(), 1);
286    }
287
288    #[test]
289    fn test_throttler_release_promotes_queued() {
290        let mut t = WorkflowThrottler::fixed(1);
291        assert!(t.try_admit("wf-1"));
292        assert!(!t.try_admit("wf-2"));
293
294        let next = t.release();
295        assert_eq!(next.as_deref(), Some("wf-2"));
296        assert_eq!(t.state().active, 1);
297        assert_eq!(t.queued_count(), 0);
298    }
299
300    #[test]
301    fn test_throttler_release_empty_queue() {
302        let mut t = WorkflowThrottler::fixed(2);
303        assert!(t.try_admit("wf-1"));
304        let next = t.release();
305        assert!(next.is_none());
306        assert_eq!(t.state().active, 0);
307    }
308
309    #[test]
310    fn test_throttler_utilisation() {
311        let mut t = WorkflowThrottler::fixed(4);
312        t.try_admit("wf-1");
313        t.try_admit("wf-2");
314        assert!((t.utilisation() - 0.5).abs() < 1e-6);
315    }
316
317    #[test]
318    fn test_throttler_reset() {
319        let mut t = WorkflowThrottler::fixed(2);
320        t.try_admit("wf-1");
321        t.try_admit("wf-2");
322        t.try_admit("wf-3");
323        t.reset();
324        assert_eq!(t.state().active, 0);
325        assert_eq!(t.queued_count(), 0);
326    }
327
328    #[test]
329    fn test_throttler_total_counters() {
330        let mut t = WorkflowThrottler::fixed(1);
331        t.try_admit("wf-1");
332        t.try_admit("wf-2"); // rejected/queued
333        assert_eq!(t.state().total_admitted, 1);
334        assert_eq!(t.state().total_rejected, 1);
335        t.release(); // promotes wf-2
336        assert_eq!(t.state().total_admitted, 2);
337    }
338}
oximedia_workflow/workflow_throttle.rs

oximedia_workflow/
workflow_throttle.rs
