Skip to main content

futu_gateway_core/bridge/
background_runtime.rs

1//! Background task runtime boundary for [`super::GatewayBridge`].
2//!
3//! Bridge startup wires several long-running or delayed background workers.
4//! Keeping their `JoinHandle`s here gives tests and embedded runtimes a clear
5//! owner boundary; dropping the runtime aborts any still-running tasks instead
6//! of leaving detached work behind.
7
8use std::sync::{Arc, Weak};
9
10use parking_lot::Mutex;
11
12const PRUNE_MIN_TRACKED_HANDLES: usize = 128;
13const PRUNE_EVERY_TRACKS: usize = 64;
14
15#[derive(Clone, Default)]
16pub struct BackgroundRuntime {
17    inner: Arc<BackgroundRuntimeInner>,
18}
19
20#[derive(Clone)]
21pub struct BackgroundTaskTracker {
22    inner: Weak<BackgroundRuntimeInner>,
23}
24
25#[derive(Default)]
26struct BackgroundRuntimeInner {
27    handles: Mutex<TrackedBackgroundHandles>,
28}
29
30#[derive(Default)]
31struct TrackedBackgroundHandles {
32    handles: Vec<tokio::task::JoinHandle<()>>,
33    tracks_since_prune: usize,
34}
35
36impl BackgroundRuntime {
37    pub fn new() -> Self {
38        Self::default()
39    }
40
41    pub fn track(&self, handle: tokio::task::JoinHandle<()>) {
42        self.inner.track(handle);
43    }
44
45    pub fn tracker(&self) -> BackgroundTaskTracker {
46        BackgroundTaskTracker {
47            inner: Arc::downgrade(&self.inner),
48        }
49    }
50
51    pub fn tracked_len(&self) -> usize {
52        let mut handles = self.inner.handles.lock();
53        handles.prune_finished();
54        handles.handles.len()
55    }
56
57    pub fn abort_all(&self) -> usize {
58        let mut handles = self.inner.handles.lock();
59        handles.prune_finished();
60        let count = handles.handles.len();
61        for handle in handles.handles.drain(..) {
62            handle.abort();
63        }
64        count
65    }
66}
67
68impl BackgroundTaskTracker {
69    pub fn track(&self, handle: tokio::task::JoinHandle<()>) {
70        if let Some(inner) = self.inner.upgrade() {
71            inner.track(handle);
72        } else {
73            handle.abort();
74        }
75    }
76}
77
78impl BackgroundRuntimeInner {
79    fn track(&self, handle: tokio::task::JoinHandle<()>) {
80        let mut handles = self.handles.lock();
81        handles.push(handle);
82    }
83}
84
85impl Drop for BackgroundRuntimeInner {
86    fn drop(&mut self) {
87        for handle in self.handles.get_mut().handles.drain(..) {
88            handle.abort();
89        }
90    }
91}
92
93impl TrackedBackgroundHandles {
94    fn push(&mut self, handle: tokio::task::JoinHandle<()>) {
95        self.handles.push(handle);
96        self.tracks_since_prune = self.tracks_since_prune.saturating_add(1);
97        if self.handles.len() >= PRUNE_MIN_TRACKED_HANDLES
98            && self.tracks_since_prune >= PRUNE_EVERY_TRACKS
99        {
100            self.prune_finished();
101        }
102    }
103
104    fn prune_finished(&mut self) {
105        self.handles.retain(|handle| !handle.is_finished());
106        self.tracks_since_prune = 0;
107    }
108}
109
110#[cfg(test)]
111mod tests;