millipede_core/autoscale/
signals.rs

1use super::{LoadSignal, LoadSnapshot};
2use crate::errors::CrawlError;
3use std::{
4    collections::VecDeque,
5    sync::{Arc, Mutex},
6    time::Duration,
7};
8use tokio::{task::JoinHandle, time::Instant};
9use tokio_util::sync::CancellationToken;
10
11const HISTORY_MAX_AGE: Duration = Duration::from_secs(300);
12const HISTORY_MAX_LEN: usize = 4096;
13const MIN_SAMPLE_INTERVAL: Duration = Duration::from_millis(1);
14
15#[derive(Default)]
16struct SnapshotHistory {
17    snapshots: Mutex<VecDeque<LoadSnapshot>>,
18}
19
20impl SnapshotHistory {
21    fn push(&self, snapshot: LoadSnapshot) {
22        let mut snapshots = self
23            .snapshots
24            .lock()
25            .unwrap_or_else(|error| error.into_inner());
26        snapshots.push_back(snapshot);
27
28        if let Some(cutoff) = snapshot.at.checked_sub(HISTORY_MAX_AGE) {
29            while snapshots.front().is_some_and(|entry| entry.at < cutoff) {
30                snapshots.pop_front();
31            }
32        }
33        while snapshots.len() > HISTORY_MAX_LEN {
34            snapshots.pop_front();
35        }
36    }
37
38    fn sample(&self, window: Duration) -> Vec<LoadSnapshot> {
39        let snapshots = self
40            .snapshots
41            .lock()
42            .unwrap_or_else(|error| error.into_inner());
43        let Some(cutoff) = Instant::now().checked_sub(window) else {
44            return snapshots.iter().copied().collect();
45        };
46        snapshots
47            .iter()
48            .filter(|snapshot| snapshot.at >= cutoff)
49            .copied()
50            .collect()
51    }
52}
53
54/// Options for periodic system CPU load sampling.
55#[derive(Debug, Clone)]
56#[non_exhaustive]
57#[must_use = "CPU load signal options do nothing unless passed to CpuLoadSignal::new"]
58pub struct CpuLoadSignalOptions {
59    /// Maximum used CPU fraction before the signal reports overload.
60    pub max_used_cpu_ratio: f32,
61    /// Interval between CPU usage samples.
62    pub sample_interval: Duration,
63}
64
65impl Default for CpuLoadSignalOptions {
66    fn default() -> Self {
67        Self {
68            max_used_cpu_ratio: 0.95,
69            sample_interval: Duration::from_secs(1),
70        }
71    }
72}
73
74/// Periodically samples aggregate system CPU usage.
75pub struct CpuLoadSignal {
76    options: CpuLoadSignalOptions,
77    history: Arc<SnapshotHistory>,
78    cancel: CancellationToken,
79    task: Mutex<Option<JoinHandle<()>>>,
80}
81
82impl CpuLoadSignal {
83    /// Creates a CPU load signal with the supplied options.
84    pub fn new(options: CpuLoadSignalOptions) -> Self {
85        Self {
86            options,
87            history: Arc::new(SnapshotHistory::default()),
88            cancel: CancellationToken::new(),
89            task: Mutex::new(None),
90        }
91    }
92}
93
94impl Default for CpuLoadSignal {
95    fn default() -> Self {
96        Self::new(CpuLoadSignalOptions::default())
97    }
98}
99
100#[async_trait::async_trait]
101impl LoadSignal for CpuLoadSignal {
102    fn name(&self) -> &str {
103        "cpu"
104    }
105
106    fn overload_threshold(&self) -> f32 {
107        self.options.max_used_cpu_ratio
108    }
109
110    async fn start(&self) -> Result<(), CrawlError> {
111        let mut task = self.task.lock().unwrap_or_else(|error| error.into_inner());
112        if task.is_some() || self.cancel.is_cancelled() {
113            return Ok(());
114        }
115
116        let history = Arc::clone(&self.history);
117        let cancel = self.cancel.child_token();
118        let threshold = self.options.max_used_cpu_ratio;
119        let interval = self
120            .options
121            .sample_interval
122            .max(sysinfo::MINIMUM_CPU_UPDATE_INTERVAL);
123        *task = Some(tokio::spawn(async move {
124            let mut system = sysinfo::System::new();
125            let mut ticker = tokio::time::interval(interval);
126            loop {
127                tokio::select! {
128                    _ = cancel.cancelled() => break,
129                    _ = ticker.tick() => {
130                        system.refresh_cpu_usage();
131                        let usage = system.global_cpu_usage() / 100.0;
132                        history.push(LoadSnapshot {
133                            at: Instant::now(),
134                            overloaded: usage > threshold,
135                        });
136                    }
137                }
138            }
139        }));
140        Ok(())
141    }
142
143    async fn stop(&self) -> Result<(), CrawlError> {
144        self.cancel.cancel();
145        let task = self
146            .task
147            .lock()
148            .unwrap_or_else(|error| error.into_inner())
149            .take();
150        if let Some(task) = task {
151            let _ = task.await;
152        }
153        Ok(())
154    }
155
156    fn sample(&self, window: Duration) -> Vec<LoadSnapshot> {
157        self.history.sample(window)
158    }
159}
160
161/// Options for periodic system memory load sampling.
162#[derive(Debug, Clone)]
163#[non_exhaustive]
164#[must_use = "memory load signal options do nothing unless passed to MemoryLoadSignal::new"]
165pub struct MemoryLoadSignalOptions {
166    /// Maximum used memory fraction before the signal reports overload.
167    pub max_used_memory_ratio: f32,
168    /// Optional byte budget used instead of total system memory.
169    pub memory_bytes: Option<u64>,
170    /// Interval between memory usage samples. Values below 1 ms are clamped to 1 ms.
171    pub sample_interval: Duration,
172}
173
174impl Default for MemoryLoadSignalOptions {
175    fn default() -> Self {
176        Self {
177            max_used_memory_ratio: 0.9,
178            memory_bytes: None,
179            sample_interval: Duration::from_secs(1),
180        }
181    }
182}
183
184/// Periodically samples used system memory against a configurable budget.
185pub struct MemoryLoadSignal {
186    options: MemoryLoadSignalOptions,
187    history: Arc<SnapshotHistory>,
188    cancel: CancellationToken,
189    task: Mutex<Option<JoinHandle<()>>>,
190}
191
192impl MemoryLoadSignal {
193    /// Creates a memory load signal with the supplied options.
194    pub fn new(options: MemoryLoadSignalOptions) -> Self {
195        Self {
196            options,
197            history: Arc::new(SnapshotHistory::default()),
198            cancel: CancellationToken::new(),
199            task: Mutex::new(None),
200        }
201    }
202}
203
204impl Default for MemoryLoadSignal {
205    fn default() -> Self {
206        Self::new(MemoryLoadSignalOptions::default())
207    }
208}
209
210#[async_trait::async_trait]
211impl LoadSignal for MemoryLoadSignal {
212    fn name(&self) -> &str {
213        "memory"
214    }
215
216    fn overload_threshold(&self) -> f32 {
217        self.options.max_used_memory_ratio
218    }
219
220    async fn start(&self) -> Result<(), CrawlError> {
221        let mut task = self.task.lock().unwrap_or_else(|error| error.into_inner());
222        if task.is_some() || self.cancel.is_cancelled() {
223            return Ok(());
224        }
225
226        let history = Arc::clone(&self.history);
227        let cancel = self.cancel.child_token();
228        let threshold = self.options.max_used_memory_ratio;
229        let memory_bytes = self.options.memory_bytes;
230        let sample_interval = self.options.sample_interval.max(MIN_SAMPLE_INTERVAL);
231        *task = Some(tokio::spawn(async move {
232            let mut system = sysinfo::System::new();
233            let mut ticker = tokio::time::interval(sample_interval);
234            loop {
235                tokio::select! {
236                    _ = cancel.cancelled() => break,
237                    _ = ticker.tick() => {
238                        system.refresh_memory();
239                        let budget = memory_bytes.unwrap_or_else(|| system.total_memory());
240                        let overloaded = budget != 0
241                            && system.used_memory() as f64 / budget as f64 > f64::from(threshold);
242                        history.push(LoadSnapshot { at: Instant::now(), overloaded });
243                    }
244                }
245            }
246        }));
247        Ok(())
248    }
249
250    async fn stop(&self) -> Result<(), CrawlError> {
251        self.cancel.cancel();
252        let task = self
253            .task
254            .lock()
255            .unwrap_or_else(|error| error.into_inner())
256            .take();
257        if let Some(task) = task {
258            let _ = task.await;
259        }
260        Ok(())
261    }
262
263    fn sample(&self, window: Duration) -> Vec<LoadSnapshot> {
264        self.history.sample(window)
265    }
266}
267
268/// Options for detecting Tokio executor scheduling lag.
269#[derive(Debug, Clone)]
270#[non_exhaustive]
271#[must_use = "Tokio load signal options do nothing unless passed to TokioRuntimeLoadSignal::new"]
272pub struct TokioRuntimeLoadSignalOptions {
273    /// Maximum scheduling lag before the signal reports overload.
274    pub max_lag: Duration,
275    /// Interval between scheduling-lag probes. Values below 1 ms are clamped to 1 ms.
276    pub sample_interval: Duration,
277}
278
279impl Default for TokioRuntimeLoadSignalOptions {
280    fn default() -> Self {
281        Self {
282            max_lag: Duration::from_millis(50),
283            sample_interval: Duration::from_millis(250),
284        }
285    }
286}
287
288/// Detects Tokio executor load by measuring stable-API timer scheduling lag.
289///
290/// This deliberately uses timer lag instead of the unstable Tokio runtime metrics
291/// sketched in `INTERFACE.md` section 13, so it does not require `tokio_unstable`.
292pub struct TokioRuntimeLoadSignal {
293    options: TokioRuntimeLoadSignalOptions,
294    history: Arc<SnapshotHistory>,
295    cancel: CancellationToken,
296    task: Mutex<Option<JoinHandle<()>>>,
297}
298
299impl TokioRuntimeLoadSignal {
300    /// Creates a Tokio runtime load signal with the supplied options.
301    pub fn new(options: TokioRuntimeLoadSignalOptions) -> Self {
302        Self {
303            options,
304            history: Arc::new(SnapshotHistory::default()),
305            cancel: CancellationToken::new(),
306            task: Mutex::new(None),
307        }
308    }
309}
310
311impl Default for TokioRuntimeLoadSignal {
312    fn default() -> Self {
313        Self::new(TokioRuntimeLoadSignalOptions::default())
314    }
315}
316
317fn lag_overloaded(lag: Duration, max_lag: Duration) -> bool {
318    lag > max_lag
319}
320
321#[async_trait::async_trait]
322impl LoadSignal for TokioRuntimeLoadSignal {
323    fn name(&self) -> &str {
324        "tokio-runtime"
325    }
326
327    /// Returns the lag-to-sampling-interval ratio as informational metadata.
328    ///
329    /// [`SystemStatus`](super::SystemStatus) consumes recorded overload flags
330    /// directly rather than interpreting this value.
331    fn overload_threshold(&self) -> f32 {
332        self.options.max_lag.as_secs_f32()
333            / self
334                .options
335                .sample_interval
336                .max(MIN_SAMPLE_INTERVAL)
337                .as_secs_f32()
338    }
339
340    async fn start(&self) -> Result<(), CrawlError> {
341        let mut task = self.task.lock().unwrap_or_else(|error| error.into_inner());
342        if task.is_some() || self.cancel.is_cancelled() {
343            return Ok(());
344        }
345
346        let history = Arc::clone(&self.history);
347        let cancel = self.cancel.child_token();
348        let max_lag = self.options.max_lag;
349        let sample_interval = self.options.sample_interval.max(MIN_SAMPLE_INTERVAL);
350        *task = Some(tokio::spawn(async move {
351            loop {
352                let target = Instant::now() + sample_interval;
353                tokio::select! {
354                    _ = cancel.cancelled() => break,
355                    _ = tokio::time::sleep_until(target) => {
356                        let lag = Instant::now().saturating_duration_since(target);
357                        history.push(LoadSnapshot {
358                            at: Instant::now(),
359                            overloaded: lag_overloaded(lag, max_lag),
360                        });
361                    }
362                }
363            }
364        }));
365        Ok(())
366    }
367
368    async fn stop(&self) -> Result<(), CrawlError> {
369        self.cancel.cancel();
370        let task = self
371            .task
372            .lock()
373            .unwrap_or_else(|error| error.into_inner())
374            .take();
375        if let Some(task) = task {
376            let _ = task.await;
377        }
378        Ok(())
379    }
380
381    fn sample(&self, window: Duration) -> Vec<LoadSnapshot> {
382        self.history.sample(window)
383    }
384}
385
386/// A manually-fed signal representing downstream client throttling.
387///
388/// Use [`ClientLoadSignal::instrument_storage`] to automatically record successful storage
389/// operations and backend rate-limit errors. Callers can also feed observations directly through
390/// [`ClientLoadSignalHandle`].
391pub struct ClientLoadSignal {
392    history: Arc<SnapshotHistory>,
393}
394
395impl ClientLoadSignal {
396    /// Creates an empty client load signal.
397    pub fn new() -> Self {
398        Self {
399            history: Arc::new(SnapshotHistory::default()),
400        }
401    }
402
403    /// Returns a cloneable handle for recording client health observations.
404    pub fn handle(&self) -> ClientLoadSignalHandle {
405        ClientLoadSignalHandle {
406            history: Arc::clone(&self.history),
407        }
408    }
409
410    /// Wraps a storage client so its successful and rate-limited operations feed this signal.
411    pub fn instrument_storage(
412        &self,
413        client: std::sync::Arc<dyn crate::storage::StorageClient>,
414    ) -> std::sync::Arc<dyn crate::storage::StorageClient> {
415        crate::storage::RateLimitReportingClient::new(client, self.handle())
416    }
417}
418
419impl Default for ClientLoadSignal {
420    fn default() -> Self {
421        Self::new()
422    }
423}
424
425#[async_trait::async_trait]
426impl LoadSignal for ClientLoadSignal {
427    fn name(&self) -> &str {
428        "client"
429    }
430
431    fn overload_threshold(&self) -> f32 {
432        1.0
433    }
434
435    fn sample(&self, window: Duration) -> Vec<LoadSnapshot> {
436        self.history.sample(window)
437    }
438}
439
440/// Cloneable manual observation handle for a [`ClientLoadSignal`].
441#[derive(Clone)]
442pub struct ClientLoadSignalHandle {
443    history: Arc<SnapshotHistory>,
444}
445
446impl ClientLoadSignalHandle {
447    /// Records a rate-limit response as an overloaded observation.
448    pub fn record_rate_limited(&self) {
449        self.history.push(LoadSnapshot {
450            at: Instant::now(),
451            overloaded: true,
452        });
453    }
454
455    /// Records a successful client interaction as a healthy observation.
456    pub fn record_healthy(&self) {
457        self.history.push(LoadSnapshot {
458            at: Instant::now(),
459            overloaded: false,
460        });
461    }
462}
463
464#[cfg(test)]
465mod tests {
466    use super::*;
467
468    #[test]
469    fn lag_overload_uses_strict_boundary() {
470        let max_lag = Duration::from_millis(50);
471        assert!(!lag_overloaded(max_lag, max_lag));
472        assert!(lag_overloaded(max_lag + Duration::from_nanos(1), max_lag));
473    }
474
475    #[tokio::test(start_paused = true)]
476    async fn snapshot_history_caps_length_and_prunes_old_entries() {
477        let history = SnapshotHistory::default();
478        for _ in 0..=HISTORY_MAX_LEN {
479            history.push(LoadSnapshot {
480                at: Instant::now(),
481                overloaded: false,
482            });
483        }
484        assert_eq!(history.sample(Duration::MAX).len(), HISTORY_MAX_LEN);
485
486        tokio::time::advance(HISTORY_MAX_AGE + Duration::from_secs(1)).await;
487        history.push(LoadSnapshot {
488            at: Instant::now(),
489            overloaded: true,
490        });
491        let samples = history.sample(Duration::MAX);
492        assert_eq!(samples.len(), 1);
493        assert!(samples[0].overloaded);
494    }
495}
millipede_core/autoscale/signals.rs

millipede_core/autoscale/
signals.rs
