r402_server/settlement/
tracker.rs

1//! In-flight counter for background settlement tasks.
2
3use std::sync::Arc;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::time::Duration;
6
7use tokio::sync::Notify;
8
9/// Shared in-flight counter for background settlement tasks.
10///
11/// Cloning the tracker is cheap and shares state. Drain at shutdown via
12/// [`Self::wait_for_drain`].
13#[derive(Clone, Debug)]
14pub struct BackgroundSettlementTracker {
15    inner: Arc<TrackerInner>,
16}
17
18#[derive(Debug)]
19struct TrackerInner {
20    in_flight: AtomicUsize,
21    drained: Notify,
22}
23
24impl Default for BackgroundSettlementTracker {
25    fn default() -> Self {
26        Self::new()
27    }
28}
29
30impl BackgroundSettlementTracker {
31    /// Constructs a tracker with zero in-flight tasks.
32    #[must_use]
33    pub fn new() -> Self {
34        Self {
35            inner: Arc::new(TrackerInner {
36                in_flight: AtomicUsize::new(0),
37                drained: Notify::new(),
38            }),
39        }
40    }
41
42    /// Returns the current approximate number of in-flight settlement tasks.
43    #[must_use]
44    pub fn in_flight(&self) -> usize {
45        self.inner.in_flight.load(Ordering::SeqCst)
46    }
47
48    /// Increments the in-flight counter and returns a guard that decrements it on drop.
49    pub(crate) fn start(&self) -> SettlementInFlightGuard {
50        let _previous = self.inner.in_flight.fetch_add(1, Ordering::SeqCst);
51        SettlementInFlightGuard {
52            inner: Arc::clone(&self.inner),
53        }
54    }
55
56    /// Awaits the in-flight count to reach zero, bounded by `timeout`.
57    ///
58    /// # Errors
59    ///
60    /// Returns the count of in-flight tasks when the timeout elapses
61    /// before the drain completes.
62    pub async fn wait_for_drain(&self, timeout: Duration) -> Result<(), usize> {
63        if self.in_flight() == 0 {
64            return Ok(());
65        }
66        let deadline = tokio::time::Instant::now() + timeout;
67        loop {
68            let notified = self.inner.drained.notified();
69            tokio::pin!(notified);
70            tokio::select! {
71                () = &mut notified => {}
72                () = tokio::time::sleep_until(deadline) => {
73                    let remaining = self.in_flight();
74                    return if remaining == 0 { Ok(()) } else { Err(remaining) };
75                }
76            }
77            if self.in_flight() == 0 {
78                return Ok(());
79            }
80        }
81    }
82}
83
84/// Drop-guard returned by [`BackgroundSettlementTracker::start`].
85#[derive(Debug)]
86pub(crate) struct SettlementInFlightGuard {
87    inner: Arc<TrackerInner>,
88}
89
90impl Drop for SettlementInFlightGuard {
91    fn drop(&mut self) {
92        let previous = self.inner.in_flight.fetch_sub(1, Ordering::SeqCst);
93        if previous == 1 {
94            self.inner.drained.notify_waiters();
95        }
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[tokio::test]
104    async fn empty_tracker_drains_immediately() {
105        let tracker = BackgroundSettlementTracker::new();
106        assert_eq!(tracker.in_flight(), 0);
107        tracker.wait_for_drain(Duration::ZERO).await.unwrap();
108    }
109
110    #[tokio::test]
111    async fn drain_waits_for_guard_drop() {
112        let tracker = BackgroundSettlementTracker::new();
113        let guard = tracker.start();
114        assert_eq!(tracker.in_flight(), 1);
115
116        let tracker_clone = tracker.clone();
117        let drop_task = tokio::spawn(async move {
118            tokio::time::sleep(Duration::from_millis(10)).await;
119            drop(guard);
120            assert_eq!(tracker_clone.in_flight(), 0);
121        });
122
123        tracker
124            .wait_for_drain(Duration::from_secs(1))
125            .await
126            .expect("drain should complete after the guard drops");
127        drop_task.await.unwrap();
128    }
129
130    #[tokio::test]
131    async fn drain_times_out_when_guards_outlive_deadline() {
132        let tracker = BackgroundSettlementTracker::new();
133        let _guard = tracker.start();
134
135        let result = tracker.wait_for_drain(Duration::from_millis(20)).await;
136        assert_eq!(result, Err(1), "deadline elapses with the guard alive");
137    }
138
139    #[tokio::test]
140    async fn nested_guards_decrement_in_order() {
141        let tracker = BackgroundSettlementTracker::new();
142        let g1 = tracker.start();
143        let g2 = tracker.start();
144        let g3 = tracker.start();
145        assert_eq!(tracker.in_flight(), 3);
146        drop(g2);
147        assert_eq!(tracker.in_flight(), 2);
148        drop(g1);
149        assert_eq!(tracker.in_flight(), 1);
150        drop(g3);
151        assert_eq!(tracker.in_flight(), 0);
152    }
153}
r402_server/settlement/tracker.rs

r402_server/settlement/
tracker.rs
