fedimint_core/task/
inner.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::time::{Duration, SystemTime};
4
5use fedimint_core::time::now;
6use fedimint_logging::LOG_TASK;
7use slotmap::SlotMap;
8use tokio::sync::watch;
9use tracing::{debug, error, info, warn};
10
11use super::{TaskGroup, TaskShutdownToken};
12use crate::runtime::{JoinError, JoinHandle};
13use crate::util::FmtCompact as _;
14
15#[derive(Debug)]
16pub struct TaskGroupInner {
17    on_shutdown_tx: watch::Sender<bool>,
18    // It is necessary to keep at least one `Receiver` around,
19    // otherwise shutdown writes are lost.
20    on_shutdown_rx: watch::Receiver<bool>,
21    pub(crate) active_tasks_join_handles:
22        std::sync::Mutex<slotmap::SlotMap<slotmap::DefaultKey, (String, JoinHandle<()>)>>,
23    // using blocking Mutex to avoid `async` in `shutdown` and `add_subgroup`
24    // it's OK as we don't ever need to yield
25    subgroups: std::sync::Mutex<Vec<TaskGroup>>,
26}
27
28impl Default for TaskGroupInner {
29    fn default() -> Self {
30        let (on_shutdown_tx, on_shutdown_rx) = watch::channel(false);
31        Self {
32            on_shutdown_tx,
33            on_shutdown_rx,
34            active_tasks_join_handles: std::sync::Mutex::new(SlotMap::default()),
35            subgroups: std::sync::Mutex::new(vec![]),
36        }
37    }
38}
39
40impl TaskGroupInner {
41    pub fn shutdown(&self) {
42        // Note: set the flag before starting to call shutdown handlers
43        // to avoid confusion.
44        self.on_shutdown_tx
45            .send(true)
46            .expect("We must have on_shutdown_rx around so this never fails");
47
48        let subgroups = self.subgroups.lock().expect("locking failed").clone();
49        for subgroup in subgroups {
50            subgroup.inner.shutdown();
51        }
52    }
53
54    #[inline]
55    pub fn is_shutting_down(&self) -> bool {
56        *self.on_shutdown_tx.borrow()
57    }
58
59    #[inline]
60    pub fn make_shutdown_rx(&self) -> TaskShutdownToken {
61        TaskShutdownToken::new(self.on_shutdown_rx.clone())
62    }
63
64    #[inline]
65    pub fn add_subgroup(&self, tg: TaskGroup) {
66        self.subgroups.lock().expect("locking failed").push(tg);
67    }
68
69    #[inline]
70    pub async fn join_all(&self, deadline: Option<SystemTime>, errors: &mut Vec<JoinError>) {
71        let subgroups = self.subgroups.lock().expect("locking failed").clone();
72        for subgroup in subgroups {
73            info!(target: LOG_TASK, "Waiting for subgroup to finish");
74            subgroup.join_all_inner(deadline, errors).await;
75            info!(target: LOG_TASK, "Subgroup finished");
76        }
77
78        // drop lock early
79        let tasks: Vec<_> = self
80            .active_tasks_join_handles
81            .lock()
82            .expect("Lock failed")
83            .drain()
84            .collect();
85        for (_, (name, join)) in tasks {
86            debug!(target: LOG_TASK, task=%name, "Waiting for task to finish");
87
88            let timeout = deadline.map(|deadline| {
89                deadline
90                    .duration_since(now())
91                    .unwrap_or(Duration::from_millis(10))
92            });
93
94            #[cfg(not(target_family = "wasm"))]
95            let join_future: Pin<Box<dyn Future<Output = _> + Send>> =
96                if let Some(timeout) = timeout {
97                    Box::pin(crate::runtime::timeout(timeout, join))
98                } else {
99                    Box::pin(async { Ok(join.await) })
100                };
101
102            #[cfg(target_family = "wasm")]
103            let join_future: Pin<Box<dyn Future<Output = _>>> = if let Some(timeout) = timeout {
104                Box::pin(crate::runtime::timeout(timeout, join))
105            } else {
106                Box::pin(async { Ok(join.await) })
107            };
108
109            match join_future.await {
110                Ok(Ok(())) => {
111                    debug!(target: LOG_TASK, task=%name, "Task finished");
112                }
113                Ok(Err(err)) => {
114                    error!(target: LOG_TASK, task=%name, err=%err.fmt_compact(), "Task panicked");
115                    errors.push(err);
116                }
117                Err(_) => {
118                    warn!(
119                        target: LOG_TASK, task=%name,
120                        "Timeout waiting for task to shut down"
121                    );
122                }
123            }
124        }
125    }
126}
fedimint_core/task/inner.rs

fedimint_core/task/
inner.rs
