klukai_types/
spawn.rs

1//! Keep track of how many futures of a certain kind are running (in an [AtomicUsize])
2
3use std::pin::Pin;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::task::{Context, Poll};
6use std::time::Duration;
7
8use futures::Future;
9use pin_project_lite::pin_project;
10use tracing::{info, trace};
11
12/// Global counter for [spawn_counted] and [spawn_counted_w_handle]
13pub static PENDING_HANDLES: AtomicUsize = AtomicUsize::new(0);
14
15/// Spawn `fut` as a [CountedFut] (increments/decrements an [AtomicUsize])
16#[track_caller]
17pub fn spawn_counted<F>(fut: F) -> tokio::task::JoinHandle<F::Output>
18where
19    F: Future + Send + 'static,
20    F::Output: Send,
21{
22    tokio::spawn(CountedFut::new(fut, &PENDING_HANDLES))
23}
24
25/// Spawn `fut` as a [CountedFut] (increments/decrements an [AtomicUsize])
26/// on the given [tokio::runtime::Handle].
27#[track_caller]
28pub fn spawn_counted_w_handle<F>(
29    fut: F,
30    h: &tokio::runtime::Handle,
31) -> tokio::task::JoinHandle<F::Output>
32where
33    F: Future + Send + 'static,
34    F::Output: Send,
35{
36    h.spawn(CountedFut::new(fut, &PENDING_HANDLES))
37}
38
39/// Spawn blocking `fut` as a [CountedFut] (increments/decrements an [AtomicUsize])
40/// on the given [tokio::runtime::Handle].
41#[track_caller]
42pub fn spawn_blocking_counted_w_handle<F, R>(
43    func: F,
44    h: &tokio::runtime::Handle,
45) -> tokio::task::JoinHandle<R>
46where
47    F: FnOnce() -> R + Send + 'static,
48    R: Send + 'static,
49{
50    h.spawn_blocking(move || {
51        let _guard = CountedGuard::new(&PENDING_HANDLES);
52        func()
53    })
54}
55
56pub struct CountedGuard {
57    pendings: &'static AtomicUsize,
58}
59
60impl CountedGuard {
61    pub fn new(pendings: &'static AtomicUsize) -> Self {
62        pendings.fetch_add(1, Ordering::SeqCst);
63        Self { pendings }
64    }
65}
66
67impl Drop for CountedGuard {
68    fn drop(&mut self) {
69        let count = self.pendings.fetch_sub(1, Ordering::SeqCst);
70        trace!("dropping counted guard, count: {}", count - 1);
71    }
72}
73
74pin_project! {
75    /// An [AtomicUsize] is incremented when building this future, and decremented
76    /// again when the future completes.
77    #[project = CountedProj]
78    pub struct CountedFut<F> {
79        #[pin]
80        fut: F,
81        pendings: &'static AtomicUsize,
82    }
83
84    impl<F> PinnedDrop for CountedFut<F> {
85        fn drop(this: Pin<&mut Self>) {
86            let count = this.pendings.fetch_sub(1, Ordering::SeqCst);
87            trace!("dropping counted future, count: {}", count - 1);
88        }
89    }
90}
91
92impl<F> CountedFut<F> {
93    /// Create a new [CountedFut], immediately incrementing `pendings`
94    pub fn new(fut: F, pendings: &'static AtomicUsize) -> Self {
95        pendings.fetch_add(1, Ordering::SeqCst);
96        Self { fut, pendings }
97    }
98}
99
100impl<F> Future for CountedFut<F>
101where
102    F: Future,
103{
104    type Output = F::Output;
105
106    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
107        trace!("polling counted future");
108        let this = self.project();
109        this.fut.poll(cx)
110    }
111}
112
113/// Waits for [PENDING_HANDLES] to reach zero. All counted futures must be
114/// [Tripwire]-aware, and the tripwire must've been tripped, otherwise this
115/// is just going to sleep for 600 seconds.
116pub async fn wait_for_all_pending_handles() {
117    let mut rounds = 0;
118
119    for _ in 0..600 {
120        match PENDING_HANDLES.load(Ordering::SeqCst) {
121            0 => {
122                info!("All spawned futures done!");
123                break;
124            }
125            n => {
126                rounds += 1;
127                if rounds % 10 == 0 {
128                    info!("Waiting on {n} spawned futures before exiting")
129                }
130            }
131        }
132        tokio::time::sleep(Duration::from_millis(100)).await;
133    }
134}
klukai_types/spawn.rs

klukai_types/
spawn.rs
