monoloop_loop/transaction/lifecycle/
task_spawner.rs1use 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
11pub(crate) struct SpawnRequest {
13 pub class: TaskClass,
14 pub future: BoxFuture,
15 pub reply: oneshot::Sender<TaskId>,
16}
17
18#[derive(Clone, Debug)]
20pub struct TransactionTaskSpawner {
21 tx: mpsc::Sender<SpawnRequest>,
22}
23
24pub enum SpawnReject {
26 Busy {
28 future: BoxFuture,
30 },
31 Rejected {
33 future: BoxFuture,
35 },
36 Orphaned,
40}
41
42impl TransactionTaskSpawner {
43 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 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 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 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 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(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 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}