fragile/
futures.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4
5use crate::{Fragile, SemiSticky, Sticky};
6
7impl<F: Future> Future for Fragile<F> {
8    type Output = F::Output;
9
10    #[track_caller]
11    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
12        self.get_pin_mut().poll(cx)
13    }
14}
15
16impl<F: Future> Future for Sticky<F> {
17    type Output = F::Output;
18
19    #[track_caller]
20    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
21        // SAFETY: We do not move the wrapper through this mutable reference.
22        let this = unsafe { self.as_mut().get_unchecked_mut() };
23        this.with_mut(|inner| {
24            // SAFETY: `Sticky<F>` is `Unpin` only when `F` is `Unpin`, and the
25            // out-of-line value is dropped in place. Pinning the wrapper therefore
26            // pins the value for the rest of its lifetime.
27            unsafe { Pin::new_unchecked(inner) }.poll(cx)
28        })
29    }
30}
31
32impl<F: Future> Future for SemiSticky<F> {
33    type Output = F::Output;
34
35    #[track_caller]
36    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
37        // SAFETY: We do not move the wrapper through this mutable reference.
38        let this = unsafe { self.as_mut().get_unchecked_mut() };
39        this.with_mut(|inner| {
40            // SAFETY: `SemiSticky<F>` is `Unpin` only when `F` is `Unpin`, and its
41            // storage does not move the value while the wrapper is pinned.
42            unsafe { Pin::new_unchecked(inner) }.poll(cx)
43        })
44    }
45}
46
47#[cfg(feature = "stream")]
48mod stream {
49    use super::*;
50    use futures_core::Stream;
51
52    impl<S: Stream> Stream for Fragile<S> {
53        type Item = S::Item;
54
55        #[track_caller]
56        fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
57            self.get_pin_mut().poll_next(cx)
58        }
59
60        #[inline]
61        fn size_hint(&self) -> (usize, Option<usize>) {
62            match self.try_get() {
63                Ok(x) => x.size_hint(),
64                Err(_) => (0, None),
65            }
66        }
67    }
68
69    impl<S: Stream> Stream for Sticky<S> {
70        type Item = S::Item;
71
72        #[track_caller]
73        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
74            // SAFETY: See the corresponding `Future` implementation.
75            let this = unsafe { self.as_mut().get_unchecked_mut() };
76            this.with_mut(|inner| {
77                // SAFETY: See the corresponding `Future` implementation.
78                unsafe { Pin::new_unchecked(inner) }.poll_next(cx)
79            })
80        }
81
82        #[inline]
83        fn size_hint(&self) -> (usize, Option<usize>) {
84            self.try_with(|inner| inner.size_hint())
85                .unwrap_or((0, None))
86        }
87    }
88
89    impl<S: Stream> Stream for SemiSticky<S> {
90        type Item = S::Item;
91
92        #[track_caller]
93        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
94            // SAFETY: See the corresponding `Future` implementation.
95            let this = unsafe { self.as_mut().get_unchecked_mut() };
96            this.with_mut(|inner| {
97                // SAFETY: See the corresponding `Future` implementation.
98                unsafe { Pin::new_unchecked(inner) }.poll_next(cx)
99            })
100        }
101
102        #[inline]
103        fn size_hint(&self) -> (usize, Option<usize>) {
104            self.try_with(|inner| inner.size_hint())
105                .unwrap_or((0, None))
106        }
107    }
108
109    #[test]
110    fn test_stream() {
111        use futures_executor as executor;
112        use futures_util::{future, stream, StreamExt};
113        let mut w1 = Fragile::new(stream::once(future::ready(42)));
114        let mut w2 = Fragile::new(stream::once(future::ready(42)));
115        assert_eq!(
116            format!("{:?}", executor::block_on(w1.next())),
117            format!("{:?}", executor::block_on(w2.next())),
118        );
119
120        let mut w1 = Sticky::new(stream::once(future::ready(42)));
121        let mut w2 = Sticky::new(stream::once(future::ready(42)));
122        assert_eq!(
123            format!("{:?}", executor::block_on(w1.next())),
124            format!("{:?}", executor::block_on(w2.next())),
125        );
126
127        let mut w1 = SemiSticky::new(stream::once(future::ready(42)));
128        let mut w2 = SemiSticky::new(stream::once(future::ready(42)));
129        assert_eq!(
130            format!("{:?}", executor::block_on(w1.next())),
131            format!("{:?}", executor::block_on(w2.next())),
132        );
133    }
134
135    #[test]
136    fn test_stream_panic() {
137        use futures_executor as executor;
138        use futures_util::{future, stream, StreamExt};
139
140        let mut w = Fragile::new(stream::once(future::ready(42)));
141        let t = std::thread::spawn(move || executor::block_on(w.next()));
142        assert!(t.join().is_err());
143
144        let mut w = Sticky::new(stream::once(future::ready(42)));
145        let t = std::thread::spawn(move || executor::block_on(w.next()));
146        assert!(t.join().is_err());
147
148        let mut w = SemiSticky::new(stream::once(future::ready(42)));
149        let t = std::thread::spawn(move || executor::block_on(w.next()));
150        assert!(t.join().is_err());
151    }
152}
153
154#[cfg(test)]
155struct PendingUntilDrop {
156    address: std::cell::Cell<*const PendingUntilDrop>,
157    registration: Option<std::sync::Arc<std::sync::atomic::AtomicPtr<PendingUntilDrop>>>,
158    _pin: std::marker::PhantomPinned,
159}
160
161#[cfg(test)]
162impl PendingUntilDrop {
163    fn new() -> Self {
164        PendingUntilDrop {
165            address: std::cell::Cell::new(std::ptr::null()),
166            registration: None,
167            _pin: std::marker::PhantomPinned,
168        }
169    }
170
171    fn with_registration(
172        registration: std::sync::Arc<std::sync::atomic::AtomicPtr<PendingUntilDrop>>,
173    ) -> Self {
174        PendingUntilDrop {
175            address: std::cell::Cell::new(std::ptr::null()),
176            registration: Some(registration),
177            _pin: std::marker::PhantomPinned,
178        }
179    }
180}
181
182#[cfg(test)]
183impl Future for PendingUntilDrop {
184    type Output = ();
185
186    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
187        let current = &*self as *const PendingUntilDrop;
188        if self.address.get().is_null() {
189            self.address.set(current);
190        } else {
191            assert_eq!(self.address.get(), current);
192        }
193        if let Some(registration) = &self.registration {
194            registration.store(
195                current as *mut PendingUntilDrop,
196                std::sync::atomic::Ordering::SeqCst,
197            );
198        }
199        Poll::Pending
200    }
201}
202
203#[cfg(all(test, feature = "stream"))]
204impl futures_core::Stream for PendingUntilDrop {
205    type Item = ();
206
207    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
208        Future::poll(self, cx).map(Some)
209    }
210}
211
212#[cfg(test)]
213impl Drop for PendingUntilDrop {
214    fn drop(&mut self) {
215        if !self.address.get().is_null() {
216            assert_eq!(self.address.get(), self as *const PendingUntilDrop);
217        }
218        if let Some(registration) = &self.registration {
219            registration.store(std::ptr::null_mut(), std::sync::atomic::Ordering::SeqCst);
220        }
221    }
222}
223
224#[test]
225fn test_pinned_future_dropped_in_place() {
226    let waker = futures_util::task::noop_waker();
227    let mut context = Context::from_waker(&waker);
228
229    let mut fragile = Box::pin(Fragile::new(PendingUntilDrop::new()));
230    assert!(Future::poll(fragile.as_mut(), &mut context).is_pending());
231    drop(fragile);
232
233    let mut sticky = Box::pin(Sticky::new(PendingUntilDrop::new()));
234    assert!(Future::poll(sticky.as_mut(), &mut context).is_pending());
235    drop(sticky);
236
237    let mut semi_sticky = Box::pin(SemiSticky::new(PendingUntilDrop::new()));
238    assert!(Future::poll(semi_sticky.as_mut(), &mut context).is_pending());
239    drop(semi_sticky);
240}
241
242#[cfg(feature = "stream")]
243#[test]
244fn test_pinned_stream_dropped_in_place() {
245    use futures_core::Stream;
246
247    let waker = futures_util::task::noop_waker();
248    let mut context = Context::from_waker(&waker);
249
250    let mut fragile = Box::pin(Fragile::new(PendingUntilDrop::new()));
251    assert!(Stream::poll_next(fragile.as_mut(), &mut context).is_pending());
252    drop(fragile);
253
254    let mut sticky = Box::pin(Sticky::new(PendingUntilDrop::new()));
255    assert!(Stream::poll_next(sticky.as_mut(), &mut context).is_pending());
256    drop(sticky);
257
258    let mut semi_sticky = Box::pin(SemiSticky::new(PendingUntilDrop::new()));
259    assert!(Stream::poll_next(semi_sticky.as_mut(), &mut context).is_pending());
260    drop(semi_sticky);
261}
262
263#[test]
264fn test_pinned_fragile_drop_on_wrong_thread_keeps_storage_alive() {
265    use std::sync::atomic::{AtomicPtr, Ordering};
266    use std::sync::Arc;
267
268    let registration = Arc::new(AtomicPtr::new(std::ptr::null_mut()));
269    let future = PendingUntilDrop::with_registration(registration.clone());
270    let mut fragile = Box::pin(Fragile::new(future));
271
272    let waker = futures_util::task::noop_waker();
273    let mut context = Context::from_waker(&waker);
274    assert!(Future::poll(fragile.as_mut(), &mut context).is_pending());
275
276    assert!(std::thread::spawn(move || drop(fragile)).join().is_err());
277
278    let registered = registration.load(Ordering::SeqCst);
279    assert!(!registered.is_null());
280    // SAFETY: Polling registered this pointer only after the future was pinned,
281    // and its destructor clears the registration before invalidating storage.
282    // A wrong-thread `Fragile` drop must therefore keep that storage alive.
283    assert_eq!(unsafe { (*registered).address.get() }, registered);
284}
285
286#[test]
287fn test_pinned_sticky_cleanup_at_thread_exit() {
288    use std::sync::atomic::{AtomicPtr, Ordering};
289    use std::sync::Arc;
290
291    fn check<W: Future<Output = ()> + Send + 'static>(wrap: fn(PendingUntilDrop) -> W) {
292        let registration = Arc::new(AtomicPtr::new(std::ptr::null_mut()));
293        let thread_registration = registration.clone();
294        std::thread::spawn(move || {
295            let future = PendingUntilDrop::with_registration(thread_registration.clone());
296            let mut wrapped = Box::pin(wrap(future));
297            let waker = futures_util::task::noop_waker();
298            let mut context = Context::from_waker(&waker);
299            assert!(wrapped.as_mut().poll(&mut context).is_pending());
300            assert!(wrapped.as_mut().poll(&mut context).is_pending());
301            std::thread::spawn(move || drop(wrapped)).join().unwrap();
302            assert!(!thread_registration.load(Ordering::SeqCst).is_null());
303        })
304        .join()
305        .unwrap();
306        assert!(registration.load(Ordering::SeqCst).is_null());
307    }
308
309    check(Sticky::new);
310    check(SemiSticky::new);
311}
312
313#[cfg(test)]
314struct NoDropFuture {
315    address: std::cell::Cell<*const NoDropFuture>,
316    _pin: std::marker::PhantomPinned,
317}
318
319#[cfg(test)]
320impl NoDropFuture {
321    fn new() -> Self {
322        NoDropFuture {
323            address: std::cell::Cell::new(std::ptr::null()),
324            _pin: std::marker::PhantomPinned,
325        }
326    }
327}
328
329#[cfg(test)]
330impl Future for NoDropFuture {
331    type Output = ();
332
333    fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<()> {
334        let address = &*self as *const Self;
335        if self.address.get().is_null() {
336            self.address.set(address);
337        } else {
338            assert_eq!(self.address.get(), address);
339        }
340        Poll::Pending
341    }
342}
343
344#[cfg(all(test, feature = "stream"))]
345impl futures_core::Stream for NoDropFuture {
346    type Item = ();
347
348    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<()>> {
349        self.poll(cx).map(Some)
350    }
351}
352
353#[test]
354fn test_pinned_no_drop_future_on_wrong_thread() {
355    assert!(!std::mem::needs_drop::<NoDropFuture>());
356    let wrappers: [Pin<Box<dyn Future<Output = ()> + Send>>; 3] = [
357        Box::pin(Fragile::new(NoDropFuture::new())),
358        Box::pin(Sticky::new(NoDropFuture::new())),
359        Box::pin(SemiSticky::new(NoDropFuture::new())),
360    ];
361    let waker = futures_util::task::noop_waker();
362    let mut context = Context::from_waker(&waker);
363    for mut wrapped in wrappers {
364        assert!(wrapped.as_mut().poll(&mut context).is_pending());
365        assert!(wrapped.as_mut().poll(&mut context).is_pending());
366        std::thread::spawn(move || drop(wrapped)).join().unwrap();
367    }
368}
369
370#[cfg(feature = "stream")]
371#[test]
372fn test_pinned_no_drop_stream_on_wrong_thread() {
373    let wrappers: [Pin<Box<dyn futures_core::Stream<Item = ()> + Send>>; 3] = [
374        Box::pin(Fragile::new(NoDropFuture::new())),
375        Box::pin(Sticky::new(NoDropFuture::new())),
376        Box::pin(SemiSticky::new(NoDropFuture::new())),
377    ];
378    let waker = futures_util::task::noop_waker();
379    let mut context = Context::from_waker(&waker);
380    for mut wrapped in wrappers {
381        assert!(wrapped.as_mut().poll_next(&mut context).is_pending());
382        assert!(wrapped.as_mut().poll_next(&mut context).is_pending());
383        std::thread::spawn(move || drop(wrapped)).join().unwrap();
384    }
385}
386
387#[test]
388fn test_future() {
389    use futures_executor as executor;
390    use futures_util::future;
391    let w1 = Fragile::new(future::ready(42));
392    let w2 = w1.clone();
393    assert_eq!(
394        format!("{:?}", executor::block_on(w1)),
395        format!("{:?}", executor::block_on(w2)),
396    );
397
398    let w1 = Sticky::new(future::ready(42));
399    let w2 = w1.clone();
400    assert_eq!(
401        format!("{:?}", executor::block_on(w1)),
402        format!("{:?}", executor::block_on(w2)),
403    );
404
405    let w1 = SemiSticky::new(future::ready(42));
406    let w2 = w1.clone();
407    assert_eq!(
408        format!("{:?}", executor::block_on(w1)),
409        format!("{:?}", executor::block_on(w2)),
410    );
411}
412
413#[test]
414fn test_future_panic() {
415    use futures_executor as executor;
416    use futures_util::future;
417    let w = Fragile::new(future::ready(42));
418    let t = std::thread::spawn(move || executor::block_on(w));
419    assert!(t.join().is_err());
420
421    let w = Sticky::new(future::ready(42));
422    let t = std::thread::spawn(move || executor::block_on(w));
423    assert!(t.join().is_err());
424
425    let w = SemiSticky::new(future::ready(42));
426    let t = std::thread::spawn(move || executor::block_on(w));
427    assert!(t.join().is_err());
428}
fragile/futures.rs

fragile/
futures.rs
