millipede_core/autoscale/
aimd.rs

1use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
2
3struct UpdateGuard<'a>(&'a AtomicBool);
4
5impl Drop for UpdateGuard<'_> {
6    fn drop(&mut self) {
7        self.0.store(false, Ordering::Release);
8    }
9}
10
11/// An atomic additive-increase, multiplicative-decrease concurrency controller.
12pub struct AimdController {
13    min: usize,
14    max: usize,
15    increase_after_successes: usize,
16    decrease_factor: f32,
17    desired: AtomicUsize,
18    success_streak: AtomicUsize,
19    updating: AtomicBool,
20}
21
22impl AimdController {
23    /// Creates a controller with normalized bounds and tuning values.
24    pub fn new(
25        min_concurrency: usize,
26        max_concurrency: usize,
27        initial_concurrency: usize,
28        increase_after_successes: usize,
29        decrease_factor: f32,
30    ) -> Self {
31        let min = min_concurrency.max(1);
32        let max = max_concurrency.max(min);
33        let decrease_factor = if decrease_factor > 0.0 && decrease_factor <= 1.0 {
34            decrease_factor
35        } else {
36            0.5
37        };
38        Self {
39            min,
40            max,
41            increase_after_successes: increase_after_successes.max(1),
42            decrease_factor,
43            desired: AtomicUsize::new(initial_concurrency.clamp(min, max)),
44            success_streak: AtomicUsize::new(0),
45            updating: AtomicBool::new(false),
46        }
47    }
48
49    /// Returns the currently desired concurrency.
50    pub fn desired_concurrency(&self) -> usize {
51        self.desired.load(Ordering::Acquire)
52    }
53
54    /// Records a successful attempt, additively increasing after a sustained streak.
55    pub fn record_success(&self) {
56        let _update = self.begin_update();
57        let previous = self
58            .success_streak
59            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |streak| {
60                let next = streak.saturating_add(1);
61                Some(if next >= self.increase_after_successes {
62                    next - self.increase_after_successes
63                } else {
64                    next
65                })
66            })
67            .expect("success streak update always produces a value");
68        if previous >= self.increase_after_successes - 1 {
69            let _ = self
70                .desired
71                .fetch_update(Ordering::AcqRel, Ordering::Acquire, |desired| {
72                    Some(desired.saturating_add(1).min(self.max))
73                });
74        }
75    }
76
77    /// Records a setback, clearing the streak and multiplicatively decreasing concurrency.
78    pub fn record_setback(&self) {
79        let _update = self.begin_update();
80        self.success_streak.store(0, Ordering::Release);
81        let _ = self
82            .desired
83            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |desired| {
84                Some(
85                    ((desired as f32 * self.decrease_factor).round() as usize)
86                        .clamp(self.min, self.max),
87                )
88            });
89    }
90
91    fn begin_update(&self) -> UpdateGuard<'_> {
92        while self
93            .updating
94            .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
95            .is_err()
96        {
97            std::hint::spin_loop();
98        }
99        UpdateGuard(&self.updating)
100    }
101}
millipede_core/autoscale/aimd.rs

millipede_core/autoscale/
aimd.rs
