clickhouse_arrow/
spawn.rs

1//! A wrapper around a `JoinHandle` that handles panics.
2//!
3//! Taken from [Datafusion](https://github.com/apache/datafusion/blob/ca0b760af6137c0dbec8b07daa5f48e262420cb5/datafusion/common-runtime/src/common.rs)
4use std::future::Future;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7
8use tokio::task::{JoinError, JoinHandle};
9
10/// Helper that  provides a simple API to spawn a single task and join it.
11/// Provides guarantees of aborting on `Drop` to keep it cancel-safe.
12/// Note that if the task was spawned with `spawn_blocking`, it will only be
13/// aborted if it hasn't started yet.
14///
15/// Technically, it's just a wrapper of a `JoinHandle` overriding drop.
16#[derive(Debug)]
17pub struct SpawnedTask<R> {
18    inner: JoinHandle<R>,
19}
20
21impl<R: 'static> SpawnedTask<R> {
22    pub fn spawn<T>(task: T) -> Self
23    where
24        T: Future<Output = R>,
25        T: Send + 'static,
26        R: Send,
27    {
28        // Ok to use spawn here as SpawnedTask handles aborting/cancelling the task on Drop
29        #[allow(clippy::disallowed_methods)]
30        let inner = tokio::task::spawn(task);
31        Self { inner }
32    }
33
34    pub fn spawn_blocking<T>(task: T) -> Self
35    where
36        T: FnOnce() -> R,
37        T: Send + 'static,
38        R: Send,
39    {
40        // Ok to use spawn_blocking here as SpawnedTask handles aborting/cancelling the task on Drop
41        #[allow(clippy::disallowed_methods)]
42        let inner = tokio::task::spawn_blocking(task);
43        Self { inner }
44    }
45
46    /// Joins the task, returning the result of join (`Result<R, JoinError>`).
47    /// Same as awaiting the spawned task, but left for backwards compatibility.
48    ///
49    /// # Errors
50    /// Returns an error if the underlying task cannot be polled.
51    pub async fn join(self) -> Result<R, JoinError> { self.await }
52
53    /// Joins the task and unwinds the panic if it happens.
54    ///
55    /// # Errors
56    /// Returns an error if the underlying task cannot be polled, the task panicked, or was
57    /// cancelled.
58    pub async fn join_unwind(self) -> Result<R, JoinError> {
59        self.await.map_err(|e| {
60            // `JoinError` can be caused either by panic or cancellation. We have to handle panics:
61            if e.is_panic() {
62                std::panic::resume_unwind(e.into_panic());
63            } else {
64                // Cancellation may be caused by two reasons:
65                // 1. Abort is called, but since we consumed `self`, it's not our case (`JoinHandle`
66                //    not accessible outside).
67                // 2. The runtime is shutting down.
68                tracing::warn!("SpawnedTask was polled during shutdown");
69                e
70            }
71        })
72    }
73}
74
75impl<R> Future for SpawnedTask<R> {
76    type Output = Result<R, JoinError>;
77
78    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
79        Pin::new(&mut self.inner).poll(cx)
80    }
81}
82
83impl<R> Drop for SpawnedTask<R> {
84    fn drop(&mut self) { self.inner.abort(); }
85}
86
87#[cfg(test)]
88mod tests {
89    use std::future::{Pending, pending};
90
91    use tokio::runtime::Runtime;
92    use tokio::sync::oneshot;
93
94    use super::*;
95
96    #[tokio::test]
97    async fn runtime_shutdown() {
98        let rt = Runtime::new().unwrap();
99        #[allow(clippy::async_yields_async)]
100        let task = rt
101            .spawn(async {
102                SpawnedTask::spawn(async {
103                    let fut: Pending<()> = pending();
104                    fut.await;
105                    unreachable!("should never return");
106                })
107            })
108            .await
109            .unwrap();
110
111        // caller shutdown their DF runtime (e.g. timeout, error in caller, etc)
112        rt.shutdown_background();
113
114        // race condition
115        // poll occurs during shutdown (buffered stream poll calls, etc)
116        assert!(matches!(
117            task.join_unwind().await,
118            Err(e) if e.is_cancelled()
119        ));
120    }
121
122    #[tokio::test]
123    #[should_panic(expected = "foo")]
124    async fn panic_resume() {
125        // this should panic w/o an `unwrap`
126        let _ = SpawnedTask::spawn(async { panic!("foo") }).join_unwind().await.ok();
127    }
128
129    #[tokio::test]
130    async fn cancel_not_started_task() {
131        let (sender, receiver) = oneshot::channel::<i32>();
132        let task = SpawnedTask::spawn(async {
133            // Shouldn't be reached.
134            sender.send(42).unwrap();
135        });
136
137        drop(task);
138
139        // If the task was cancelled, the sender was also dropped,
140        // and awaiting the receiver should result in an error.
141        assert!(receiver.await.is_err());
142    }
143
144    #[tokio::test]
145    async fn cancel_ongoing_task() {
146        let (sender, mut receiver) = tokio::sync::mpsc::channel(1);
147        let task = SpawnedTask::spawn(async move {
148            sender.send(1).await.unwrap();
149            // This line will never be reached because the channel has a buffer
150            // of 1.
151            sender.send(2).await.unwrap();
152        });
153        // Let the task start.
154        assert_eq!(receiver.recv().await.unwrap(), 1);
155        drop(task);
156
157        // The sender was dropped so we receive `None`.
158        assert!(receiver.recv().await.is_none());
159    }
160}
clickhouse_arrow/spawn.rs

clickhouse_arrow/
spawn.rs
