Skip to main content

monoloop_loop/transaction/lifecycle/
task_spawner.rs

1//! Worker-facing spawn proxy into the supervisor-owned [`TaskSupervisor`] (v2 §7.3 / §16).
2
3use super::task_supervisor::{TaskClass, TaskId};
4use std::future::Future;
5use std::pin::Pin;
6use tokio::sync::mpsc::error::TrySendError;
7use tokio::sync::{mpsc, oneshot};
8
9type BoxFuture = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
10
11/// Request to register+spawn a future on the supervisor's TaskSupervisor.
12pub(crate) struct SpawnRequest {
13    pub class: TaskClass,
14    pub future: BoxFuture,
15    pub reply: oneshot::Sender<TaskId>,
16}
17
18/// Cloneable handle workers use to spawn owned tasks without holding JoinHandles.
19#[derive(Clone, Debug)]
20pub struct TransactionTaskSpawner {
21    tx: mpsc::Sender<SpawnRequest>,
22}
23
24/// Why [`TransactionTaskSpawner::spawn`] rejected or could not confirm ownership.
25pub enum SpawnReject {
26    /// Mailbox full before accept; caller still owns the future (drive or drop).
27    Busy {
28        /// Unspawned future.
29        future: BoxFuture,
30    },
31    /// Channel closed before accept; caller still owns the future.
32    Rejected {
33        /// Unspawned future.
34        future: BoxFuture,
35    },
36    /// `try_send` succeeded but TaskId reply was lost (e.g. shutdown drained the
37    /// request without spawning, or reply send failed). The future is **not**
38    /// returned — caller MUST NOT drive a substitute (Law 23 / 25). Fail closed.
39    Orphaned,
40}
41
42impl TransactionTaskSpawner {
43    /// Create a spawner and the receiver drained by the supervisor loop.
44    pub(crate) fn channel(capacity: usize) -> (Self, mpsc::Receiver<SpawnRequest>) {
45        let (tx, rx) = mpsc::channel(capacity.max(1));
46        (Self { tx }, rx)
47    }
48
49    /// Register then spawn `future` under `class`.
50    ///
51    /// Uses `try_send` so workers never block forever on a full mailbox while the
52    /// supervisor is in `abort_and_drain`. On Busy/Rejected before accept, the
53    /// boxed future is returned so the caller can drive cleanup inline. On
54    /// [`SpawnReject::Orphaned`], the future is gone from the caller — fail closed.
55    pub async fn spawn<F>(&self, class: TaskClass, future: F) -> Result<TaskId, SpawnReject>
56    where
57        F: Future<Output = ()> + Send + 'static,
58    {
59        self.spawn_boxed(class, Box::pin(future)).await
60    }
61
62    /// Same as [`Self::spawn`] with an already-boxed future (Busy retry).
63    pub async fn spawn_boxed(
64        &self,
65        class: TaskClass,
66        future: BoxFuture,
67    ) -> Result<TaskId, SpawnReject> {
68        let (reply_tx, reply_rx) = oneshot::channel();
69        let req = SpawnRequest {
70            class,
71            future,
72            reply: reply_tx,
73        };
74        match self.tx.try_send(req) {
75            Ok(()) => match reply_rx.await {
76                Ok(id) => Ok(id),
77                // Accepted into mailbox; do not invent a dummy future (Law 23/25).
78                Err(_) => Err(SpawnReject::Orphaned),
79            },
80            Err(TrySendError::Full(req)) => Err(SpawnReject::Busy { future: req.future }),
81            Err(TrySendError::Closed(req)) => Err(SpawnReject::Rejected { future: req.future }),
82        }
83    }
84
85    /// Sync try-spawn for `ToolRuntime::start` (no await). On success the supervisor
86    /// owns the future; `TaskId` reply may be dropped (caller does not need it).
87    pub fn try_spawn_owned<F>(&self, class: TaskClass, future: F) -> Result<(), SpawnReject>
88    where
89        F: Future<Output = ()> + Send + 'static,
90    {
91        let (reply_tx, reply_rx) = oneshot::channel();
92        // Drop receiver: we only need ownership registration, not the TaskId.
93        drop(reply_rx);
94        let req = SpawnRequest {
95            class,
96            future: Box::pin(future),
97            reply: reply_tx,
98        };
99        match self.tx.try_send(req) {
100            Ok(()) => Ok(()),
101            Err(TrySendError::Full(req)) => Err(SpawnReject::Busy { future: req.future }),
102            Err(TrySendError::Closed(req)) => Err(SpawnReject::Rejected { future: req.future }),
103        }
104    }
105
106    /// Prefer supervisor ownership: bounded Busy retries, then return the last reject.
107    ///
108    /// Does not drive the future inline — caller decides fail-closed vs last-resort join.
109    pub async fn spawn_with_busy_retry<F, C>(
110        &self,
111        class: TaskClass,
112        future: F,
113        max_retries: u32,
114        mut is_cancelled: C,
115    ) -> Result<TaskId, SpawnReject>
116    where
117        F: Future<Output = ()> + Send + 'static,
118        C: FnMut() -> bool,
119    {
120        let mut future: BoxFuture = Box::pin(future);
121        let mut attempt = 0u32;
122        loop {
123            if is_cancelled() {
124                return Err(SpawnReject::Rejected { future });
125            }
126            match self.spawn_boxed(class.clone(), future).await {
127                Ok(id) => return Ok(id),
128                Err(SpawnReject::Busy { future: f }) => {
129                    future = f;
130                    if attempt >= max_retries {
131                        return Err(SpawnReject::Busy { future });
132                    }
133                    attempt = attempt.saturating_add(1);
134                    tokio::task::yield_now().await;
135                    tokio::time::sleep(std::time::Duration::from_millis(1)).await;
136                }
137                Err(other) => return Err(other),
138            }
139        }
140    }
141}