sqlx_core/ext/
async_stream.rs

1//! A minimalist clone of the `async-stream` crate in 100% safe code, without proc macros.
2//!
3//! This was created initially to get around some weird compiler errors we were getting with
4//! `async-stream`, and now it'd just be more work to replace.
5
6use std::future::{self, Future};
7use std::pin::Pin;
8use std::sync::{Arc, Mutex};
9use std::task::{Context, Poll};
10
11use futures_core::future::BoxFuture;
12use futures_core::stream::Stream;
13use futures_core::FusedFuture;
14use futures_util::future::Fuse;
15use futures_util::FutureExt;
16
17use crate::error::Error;
18
19pub struct TryAsyncStream<'a, T> {
20    yielder: Yielder<T>,
21    future: Fuse<BoxFuture<'a, Result<(), Error>>>,
22}
23
24impl<'a, T> TryAsyncStream<'a, T> {
25    pub fn new<F, Fut>(f: F) -> Self
26    where
27        F: FnOnce(Yielder<T>) -> Fut + Send,
28        Fut: 'a + Future<Output = Result<(), Error>> + Send,
29        T: 'a + Send,
30    {
31        let yielder = Yielder::new();
32
33        let future = f(yielder.duplicate()).boxed().fuse();
34
35        Self { future, yielder }
36    }
37}
38
39pub struct Yielder<T> {
40    // This mutex should never have any contention in normal operation.
41    // We're just using it because `Rc<Cell<Option<T>>>` would not be `Send`.
42    value: Arc<Mutex<Option<T>>>,
43}
44
45impl<T> Yielder<T> {
46    fn new() -> Self {
47        Yielder {
48            value: Arc::new(Mutex::new(None)),
49        }
50    }
51
52    // Don't want to expose a `Clone` impl
53    fn duplicate(&self) -> Self {
54        Yielder {
55            value: self.value.clone(),
56        }
57    }
58
59    /// NOTE: may deadlock the task if called from outside the future passed to `TryAsyncStream`.
60    pub async fn r#yield(&self, val: T) {
61        let replaced = self
62            .value
63            .lock()
64            .expect("BUG: panicked while holding a lock")
65            .replace(val);
66
67        debug_assert!(
68            replaced.is_none(),
69            "BUG: previously yielded value not taken"
70        );
71
72        let mut yielded = false;
73
74        // Allows the generating future to suspend its execution without changing the task priority,
75        // which would happen with `tokio::task::yield_now()`.
76        //
77        // Note that because this has no way to schedule a wakeup, this could deadlock the task
78        // if called in the wrong place.
79        future::poll_fn(|_cx| {
80            if !yielded {
81                yielded = true;
82                Poll::Pending
83            } else {
84                Poll::Ready(())
85            }
86        })
87        .await
88    }
89
90    fn take(&self) -> Option<T> {
91        self.value
92            .lock()
93            .expect("BUG: panicked while holding a lock")
94            .take()
95    }
96}
97
98impl<T> Stream for TryAsyncStream<'_, T> {
99    type Item = Result<T, Error>;
100
101    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
102        if self.future.is_terminated() {
103            return Poll::Ready(None);
104        }
105
106        match self.future.poll_unpin(cx) {
107            Poll::Ready(Ok(())) => {
108                // Future returned without yielding another value,
109                // or else it would have returned `Pending` instead.
110                Poll::Ready(None)
111            }
112            Poll::Ready(Err(e)) => Poll::Ready(Some(Err(e))),
113            Poll::Pending => self
114                .yielder
115                .take()
116                .map_or(Poll::Pending, |val| Poll::Ready(Some(Ok(val)))),
117        }
118    }
119}
120
121#[macro_export]
122macro_rules! try_stream {
123    ($($block:tt)*) => {
124        $crate::ext::async_stream::TryAsyncStream::new(move |yielder| ::tracing::Instrument::in_current_span(async move {
125            // Anti-footgun: effectively pins `yielder` to this future to prevent any accidental
126            // move to another task, which could deadlock.
127            let yielder = &yielder;
128
129            macro_rules! r#yield {
130                ($v:expr) => {{
131                    yielder.r#yield($v).await;
132                }}
133            }
134
135            $($block)*
136        }))
137    }
138}
sqlx_core/ext/async_stream.rs

sqlx_core/ext/
async_stream.rs
