borsa_core/
stream.rs

1use tokio::task::JoinHandle;
2
3/// Abstraction over a handle that can be queried for completion and aborted.
4pub trait Abortable {
5    /// Abort the underlying task if it is still running.
6    fn abort(&mut self);
7    /// Return `true` if the underlying task has completed.
8    fn is_finished(&self) -> bool;
9}
10
11impl Abortable for JoinHandle<()> {
12    fn abort(&mut self) {
13        // JoinHandle::abort takes &self
14        Self::abort(self);
15    }
16
17    fn is_finished(&self) -> bool {
18        Self::is_finished(self)
19    }
20}
21
22/// Abstraction over a one-shot stop signal.
23pub trait Stoppable {
24    /// Send a best-effort stop signal to request graceful shutdown.
25    fn send(self);
26}
27
28impl Stoppable for tokio::sync::oneshot::Sender<()> {
29    fn send(self) {
30        let _ = Self::send(self, ());
31    }
32}
33
34/// Drop-time logic for stream handles:
35/// - send a best-effort stop signal if present
36/// - abort the task if it hasn't finished yet
37pub fn drop_impl<H, S>(inner: &mut Option<H>, stop_tx: &mut Option<S>)
38where
39    H: Abortable,
40    S: Stoppable,
41{
42    if let Some(tx) = stop_tx.take() {
43        tx.send();
44    }
45    if let Some(mut h) = inner.take()
46        && !h.is_finished()
47    {
48        h.abort();
49    }
50}
51
52/// Minimal stream handle abstraction for long-lived streaming tasks.
53///
54/// Lifecycle contract:
55/// - Prefer calling [`stop`](StreamHandle::stop) to request a graceful shutdown and await completion.
56/// - Call [`abort`](StreamHandle::abort) for immediate, non-graceful termination.
57/// - If dropped without an explicit shutdown, a best-effort stop signal is sent (if available) and
58///   the underlying task is then aborted. The task may not observe the stop signal before abort.
59#[derive(Debug)]
60pub struct StreamHandle {
61    inner: Option<tokio::task::JoinHandle<()>>,
62    stop_tx: Option<tokio::sync::oneshot::Sender<()>>,
63}
64
65impl StreamHandle {
66    /// Create a new `StreamHandle`.
67    ///
68    /// Parameters:
69    /// - `inner`: the spawned task driving the stream.
70    /// - `stop_tx`: a one-shot used to request a graceful stop.
71    ///
72    /// Returns a handle that can be used to stop or abort the stream.
73    #[must_use]
74    pub const fn new(
75        inner: tokio::task::JoinHandle<()>,
76        stop_tx: tokio::sync::oneshot::Sender<()>,
77    ) -> Self {
78        Self {
79            inner: Some(inner),
80            stop_tx: Some(stop_tx),
81        }
82    }
83
84    /// Create a `StreamHandle` that can only abort the task (no graceful stop).
85    ///
86    /// This constructor is intended for connectors that do not support a
87    /// cooperative shutdown signal. Dropping the handle (or calling
88    /// [`abort`](Self::abort)) will force-cancel the underlying task.
89    #[must_use]
90    pub const fn new_abort_only(inner: tokio::task::JoinHandle<()>) -> Self {
91        Self {
92            inner: Some(inner),
93            stop_tx: None,
94        }
95    }
96
97    /// Gracefully stop the underlying stream task and await its completion.
98    ///
99    /// Sends a stop signal if available, then awaits the task. Any errors
100    /// from the task are ignored.
101    pub async fn stop(mut self) {
102        if let Some(tx) = self.stop_tx.take() {
103            let _ = tx.send(());
104        }
105        if let Some(inner) = self.inner.take() {
106            let _ = inner.await;
107        }
108    }
109
110    /// Force-abort the underlying stream task without waiting for completion.
111    ///
112    /// Prefer [`stop`](Self::stop) when possible to allow cleanup.
113    pub fn abort(mut self) {
114        if let Some(inner) = self.inner.take() {
115            inner.abort();
116        }
117    }
118}
119
120impl Drop for StreamHandle {
121    fn drop(&mut self) {
122        crate::stream::drop_impl(&mut self.inner, &mut self.stop_tx);
123    }
124}
borsa_core/stream.rs

borsa_core/
stream.rs
