future_local_storage/
future.rs

1//! Future types.
2
3use std::{
4    future::Future,
5    pin::Pin,
6    task::{Context, Poll},
7};
8
9use pin_project::{pin_project, pinned_drop};
10
11use crate::{imp::FutureLocalKey, FutureLocalStorage};
12
13impl<F: Future> FutureLocalStorage for F {
14    fn with_scope<T, S>(self, scope: &'static S, value: T) -> ScopedFutureWithValue<T, Self>
15    where
16        T: Send,
17        S: AsRef<FutureLocalKey<T>>,
18    {
19        let scope = scope.as_ref();
20        ScopedFutureWithValue {
21            inner: self,
22            scope,
23            value: Some(value),
24        }
25    }
26}
27
28/// A [`Future`] that sets a value `T` of a future local for the future `F` during its execution.
29/// Unlike the [`ScopedFutureWithValue`] this future discards the future local value.
30#[pin_project]
31#[derive(Debug)]
32pub struct ScopedFuture<T, F>(#[pin] ScopedFutureWithValue<T, F>)
33where
34    T: Send + 'static,
35    F: Future;
36
37impl<T, F> Future for ScopedFuture<T, F>
38where
39    T: Send,
40    F: Future,
41{
42    type Output = F::Output;
43
44    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
45        self.project().0.poll(cx).map(|(_value, result)| result)
46    }
47}
48
49impl<T, F> ScopedFutureWithValue<T, F>
50where
51    T: Send,
52    F: Future,
53{
54    /// Discards the future local value from the future output.
55    pub fn discard_value(self) -> ScopedFuture<T, F> {
56        ScopedFuture(self)
57    }
58}
59
60/// A [`Future`] that sets a value `T` of a future local for the future `F` during its execution.
61///
62/// This future also returns a future local value after execution.
63#[pin_project(PinnedDrop)]
64#[derive(Debug)]
65pub struct ScopedFutureWithValue<T, F>
66where
67    T: Send + 'static,
68    F: Future,
69{
70    // TODO Implement manually drop to provide scope access to the future Drop.
71    #[pin]
72    inner: F,
73    scope: &'static FutureLocalKey<T>,
74    value: Option<T>,
75}
76
77#[pinned_drop]
78impl<T, F> PinnedDrop for ScopedFutureWithValue<T, F>
79where
80    F: Future,
81    T: Send + 'static,
82{
83    fn drop(self: Pin<&mut Self>) {}
84}
85
86impl<T, F> Future for ScopedFutureWithValue<T, F>
87where
88    T: Send,
89    F: Future,
90{
91    type Output = (T, F::Output);
92
93    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
94        let this = self.project();
95        // Swap in future local key.
96        FutureLocalKey::swap(this.scope, this.value);
97        // Poll the underlying future.
98        let result = this.inner.poll(cx);
99        // Swap future local key back.
100        FutureLocalKey::swap(this.scope, this.value);
101
102        let result = std::task::ready!(result);
103        // Take the scoped value to return it back to the future caller.
104        let value = this.value.take().unwrap();
105        Poll::Ready((value, result))
106    }
107}
108
109impl<T, F> From<ScopedFutureWithValue<T, F>> for ScopedFuture<T, F>
110where
111    T: Send,
112    F: Future,
113{
114    fn from(value: ScopedFutureWithValue<T, F>) -> Self {
115        Self(value)
116    }
117}
future_local_storage/future.rs

future_local_storage/
future.rs
