Mirror of https://github.com/roostorg/osprey
github.com/roostorg/osprey
1use std::{
2 future::Future,
3 pin::Pin,
4 task::{Context, Poll},
5};
6
7use tokio::task::{JoinError, JoinHandle};
8
9mod unbounded_receiver_chunker;
10
11/// A helper struct that will invoke [`JoinHandle::abort`] when dropped.
12#[derive(Debug)]
13pub struct AbortOnDrop<T = ()> {
14 join_handle: Option<JoinHandle<T>>,
15}
16
17impl<T> AbortOnDrop<T> {
18 pub fn new(join_handle: JoinHandle<T>) -> Self {
19 Self {
20 join_handle: Some(join_handle),
21 }
22 }
23
24 /// Disarms the [`AbortOnDrop`], causing it to no longer abort the join handle when dropped.
25 pub fn run_forever(mut self) {
26 self.join_handle = None;
27 }
28
29 pub fn into_inner(mut self) -> JoinHandle<T> {
30 self.join_handle
31 .take()
32 .expect("invariant: the inner join handle should always be present")
33 }
34}
35
36impl<T> Drop for AbortOnDrop<T> {
37 fn drop(&mut self) {
38 if let Some(join_handle) = self.join_handle.take() {
39 join_handle.abort();
40 }
41 }
42}
43
44impl<T> From<JoinHandle<T>> for AbortOnDrop<T> {
45 fn from(join_handle: JoinHandle<T>) -> Self {
46 Self::new(join_handle)
47 }
48}
49
50impl<T> Future for AbortOnDrop<T> {
51 type Output = Result<T, JoinError>;
52
53 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
54 let join_handle = self
55 .join_handle
56 .as_mut()
57 .expect("polling after `AbortOnDrop` already completed");
58
59 Pin::new(join_handle).poll(cx)
60 }
61}