Skip to main content

veilid_tools/
future_race_completion.rs

1//! FutureRaceCompletionPool
2//!
3//! Spawn a future to run to completion in the background and optionally
4//! receive its result through a channel. If the receiver is dropped, the
5//! spawned future still runs to completion (its result is discarded). The
6//! pool tracks lifetime via a read/write lock so a caller can wait for all
7//! in-flight handoffs to drain before shutting down.
8
9use super::*;
10
11use futures_util::future::{select, Either};
12use std::pin::Pin;
13
14/// Outcome of `race_complete_on`
15#[derive(Debug)]
16pub enum FutureRaceCompletion<T, R> {
17    /// The original future finished first
18    Original(T),
19    /// The race future finished first; the original is now running in the pool
20    Race(R),
21}
22
23/// Pool that runs handed-off futures to completion. Cheap to clone.
24#[derive(Debug, Clone)]
25pub struct FutureRaceCompletionPool {
26    /// Read guards held by in-flight futures; shutdown awaits an exclusive write.
27    inner: Arc<AsyncRwLock<()>>,
28    /// Diagnostic name used when spawning pool tasks
29    name: Arc<String>,
30}
31
32impl FutureRaceCompletionPool {
33    /// Creates an empty pool; `name` labels spawned tasks for diagnostics.
34    pub fn new(name: impl Into<String>) -> Self {
35        Self {
36            inner: Arc::new(AsyncRwLock::new(())),
37            name: Arc::new(name.into()),
38        }
39    }
40
41    /// Run `future` to completion in the background; result is dropped.
42    pub fn spawn<F>(&self, future: F)
43    where
44        F: Future<Output = ()> + Send + 'static,
45    {
46        let inner = self.inner.clone();
47        let name = self.name.clone();
48        spawn_detached(name.as_str(), async move {
49            let read_guard = inner.read_arc().await;
50            future.await;
51            drop(read_guard);
52        });
53    }
54
55    /// Run `future` to completion in the background and return a receiver
56    /// that yields its result. Dropping the receiver does not cancel the
57    /// future; the value is just discarded when it eventually arrives.
58    pub fn channel_spawn<F, T>(&self, future: F) -> flume::Receiver<T>
59    where
60        F: Future<Output = T> + Send + 'static,
61        T: Send + 'static,
62    {
63        let (tx, rx) = flume::bounded(1);
64        let inner = self.inner.clone();
65        let name = self.name.clone();
66        spawn_detached(name.as_str(), async move {
67            let read_guard = inner.read_arc().await;
68            let result = future.await;
69            let _ = tx.send(result);
70            drop(read_guard);
71        });
72        rx
73    }
74
75    /// Wait for every in-flight future to complete. Acquiring the write lock
76    /// implies all read guards have been released.
77    pub async fn shutdown(&self) {
78        let _w = self.inner.write_arc().await;
79    }
80}
81
82/// Trait that adds `race_complete_on` to any future
83pub trait FutureRaceCompletionPoolExt: Future + Sized {
84    /// Race `self` against `race_fut`. If `race_fut` finishes first, hand `self`
85    /// off to `pool` to run to completion in the background and return `Race(_)`.
86    /// Otherwise return `Original(_)` with the value `self` produced.
87    fn race_complete_on<R, F>(
88        self,
89        pool: FutureRaceCompletionPool,
90        race_fut: F,
91    ) -> impl Future<Output = FutureRaceCompletion<Self::Output, R>> + Send
92    where
93        Self: Send + 'static,
94        Self::Output: Send + 'static,
95        F: Future<Output = R> + Send,
96        R: Send;
97}
98
99impl<Fut> FutureRaceCompletionPoolExt for Fut
100where
101    Fut: Future,
102{
103    // async fn would drop the explicit `+ Send` bound from the trait signature
104    #[expect(clippy::manual_async_fn)]
105    fn race_complete_on<R, F>(
106        self,
107        pool: FutureRaceCompletionPool,
108        race_fut: F,
109    ) -> impl Future<Output = FutureRaceCompletion<Fut::Output, R>> + Send
110    where
111        Self: Send + 'static,
112        Self::Output: Send + 'static,
113        F: Future<Output = R> + Send,
114        R: Send,
115    {
116        async move {
117            let pinned_self: Pin<Box<Fut>> = Box::pin(self);
118            let pinned_race = Box::pin(race_fut);
119            match select(pinned_self, pinned_race).await {
120                Either::Left((output, _race)) => FutureRaceCompletion::Original(output),
121                Either::Right((race_output, original)) => {
122                    pool.spawn(async move {
123                        let _ = original.await;
124                    });
125                    FutureRaceCompletion::Race(race_output)
126                }
127            }
128        }
129    }
130}