Skip to main content

ps_promise/methods/shared/promise/implementations/
future.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    sync::PoisonError,
5    task::{
6        Context,
7        Poll::{self, Pending, Ready},
8        Waker,
9    },
10};
11
12use crate::PromiseRejection;
13
14use super::super::super::constants::LIVELOCK_MAX_SELF_POLLS;
15use super::super::methods::poll_step::PollStep;
16use super::super::SharedPromise;
17
18impl<T, E> Future for SharedPromise<T, E>
19where
20    T: Clone + Send + 'static,
21    E: PromiseRejection + Clone,
22{
23    type Output = Result<T, E>;
24
25    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
26        let mut counter = 0;
27
28        loop {
29            match self.poll_step(cx) {
30                PollStep::Pending => return Pending,
31                PollStep::Rejected(err) => return Ready(Err(err)),
32                PollStep::Resolved(value) => return Ready(Ok(value)),
33                PollStep::Consumed => return Ready(Err(E::already_consumed())),
34                PollStep::ReEnter => continue,
35                PollStep::Woke => counter += 1,
36            }
37
38            if counter >= LIVELOCK_MAX_SELF_POLLS {
39                let wakers: Vec<Waker> = self
40                    .state
41                    .wakers
42                    .lock()
43                    .unwrap_or_else(PoisonError::into_inner)
44                    .values()
45                    .cloned()
46                    .collect();
47
48                for waker in wakers {
49                    waker.wake();
50                }
51
52                return Pending;
53            }
54        }
55    }
56}
57
58#[cfg(test)]
59#[allow(clippy::expect_used)]
60mod tests {
61    use std::{
62        future::Future,
63        pin::Pin,
64        sync::{
65            atomic::{AtomicUsize, Ordering},
66            Arc, Mutex,
67        },
68        task::{Context, Poll, Waker},
69    };
70
71    use crate::{Promise, PromiseRejection, SharedPromise, TaskFailure};
72
73    use super::LIVELOCK_MAX_SELF_POLLS;
74
75    #[derive(Debug, Clone, PartialEq)]
76    enum E {
77        AlreadyConsumed,
78        Fail,
79        TaskFailed,
80    }
81
82    impl PromiseRejection for E {
83        fn already_consumed() -> Self {
84            Self::AlreadyConsumed
85        }
86
87        fn task_failed(_: TaskFailure) -> Self {
88            Self::TaskFailed
89        }
90    }
91
92    fn cx() -> Context<'static> {
93        Context::from_waker(Waker::noop())
94    }
95
96    fn poll<F: std::future::Future + Unpin>(future: &mut F) -> std::task::Poll<F::Output> {
97        std::pin::Pin::new(future).poll(&mut cx())
98    }
99
100    #[test]
101    fn every_clone_observes_the_result() {
102        let shared = Promise::<i32, E>::lazy(async { Ok(42) }).shared();
103        let mut first = shared.clone();
104        let mut second = shared;
105
106        assert_eq!(poll(&mut first), std::task::Poll::Ready(Ok(42)));
107        assert_eq!(poll(&mut second), std::task::Poll::Ready(Ok(42)));
108    }
109
110    #[test]
111    fn inner_promise_runs_once() {
112        let runs = Arc::new(AtomicUsize::new(0));
113        let counter = runs.clone();
114
115        let shared = Promise::<i32, E>::lazy(async move {
116            counter.fetch_add(1, Ordering::Relaxed);
117
118            Ok(7)
119        })
120        .shared();
121
122        let mut first = shared.clone();
123        let mut second = shared;
124
125        assert_eq!(poll(&mut first), std::task::Poll::Ready(Ok(7)));
126        assert_eq!(poll(&mut second), std::task::Poll::Ready(Ok(7)));
127        assert_eq!(runs.load(Ordering::Relaxed), 1);
128    }
129
130    #[test]
131    fn repeated_polls_keep_returning_the_result() {
132        let mut shared = Promise::<i32, E>::lazy(async { Ok(9) }).shared();
133
134        assert_eq!(poll(&mut shared), std::task::Poll::Ready(Ok(9)));
135        assert_eq!(poll(&mut shared), std::task::Poll::Ready(Ok(9)));
136    }
137
138    #[test]
139    fn every_clone_observes_the_rejection() {
140        let shared = Promise::<i32, E>::lazy(async { Err(E::Fail) }).shared();
141        let mut first = shared.clone();
142        let mut second = shared;
143
144        assert_eq!(poll(&mut first), std::task::Poll::Ready(Err(E::Fail)));
145        assert_eq!(poll(&mut second), std::task::Poll::Ready(Err(E::Fail)));
146    }
147
148    #[test]
149    fn pending_then_settled_for_late_consumers() {
150        let (promise, resolve, _reject) = Promise::<i32, E>::with_resolvers();
151
152        let shared = promise.shared();
153        let mut early = shared.clone();
154        let mut late = shared;
155
156        assert!(poll(&mut early).is_pending());
157
158        resolve.resolve(5);
159
160        assert_eq!(poll(&mut late), std::task::Poll::Ready(Ok(5)));
161        assert_eq!(poll(&mut early), std::task::Poll::Ready(Ok(5)));
162    }
163
164    /// Inner future that self-wakes and returns `Pending` on every poll until it
165    /// has been polled `limit` times, then resolves. A correct driver must not
166    /// poll it `limit` times within a single `SharedPromise::poll`: the inline
167    /// re-entry loop is expected to yield to the executor instead of re-entering
168    /// without bound. The terminal resolve is an escape hatch so that a spinning
169    /// driver terminates and fails the assertion rather than hanging the test.
170    struct SpinProbe {
171        polls: Arc<AtomicUsize>,
172        limit: usize,
173    }
174
175    impl Future for SpinProbe {
176        type Output = Result<i32, E>;
177
178        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
179            let polled = self.polls.fetch_add(1, Ordering::Relaxed) + 1;
180
181            if polled >= self.limit {
182                return Poll::Ready(Ok(0));
183            }
184
185            cx.waker().wake_by_ref();
186
187            Poll::Pending
188        }
189    }
190
191    #[test]
192    fn inner_future_is_not_spin_polled_unboundedly() {
193        const LIMIT: usize = 10_000;
194
195        let polls = Arc::new(AtomicUsize::new(0));
196
197        let mut shared = Promise::<i32, E>::lazy(SpinProbe {
198            polls: polls.clone(),
199            limit: LIMIT,
200        })
201        .shared();
202
203        // A single poll must not re-drive the inner future `LIMIT` times: the
204        // driver is expected to yield to the executor rather than spin-poll a
205        // self-waking-pending future inline without bound.
206        let _ = poll(&mut shared);
207
208        let polled = polls.load(Ordering::Relaxed);
209
210        assert!(
211            polled < LIMIT,
212            "inner future was polled {polled} times in one SharedPromise::poll; \
213             the inline re-entry loop must yield instead of spin-polling"
214        );
215    }
216
217    fn poll_with<F: Future + Unpin>(future: &mut F, waker: &Waker) -> Poll<F::Output> {
218        Pin::new(future).poll(&mut Context::from_waker(waker))
219    }
220
221    struct CountingWaker {
222        count: AtomicUsize,
223    }
224
225    impl std::task::Wake for CountingWaker {
226        fn wake(self: Arc<Self>) {
227            self.count.fetch_add(1, Ordering::Relaxed);
228        }
229
230        fn wake_by_ref(self: &Arc<Self>) {
231            self.count.fetch_add(1, Ordering::Relaxed);
232        }
233    }
234
235    #[test]
236    fn resolved_shared_polls_ready_ok() {
237        let mut shared: SharedPromise<i32, E> = Promise::resolve(42).shared();
238
239        assert_eq!(poll(&mut shared), Poll::Ready(Ok(42)));
240    }
241
242    #[test]
243    fn rejected_shared_polls_ready_err() {
244        let inner: Promise<i32, E> = Promise::lazy(async { Err(E::Fail) });
245
246        let mut shared = inner.shared();
247
248        assert_eq!(poll(&mut shared), Poll::Ready(Err(E::Fail)));
249    }
250
251    #[test]
252    fn inner_runs_exactly_once_across_clones() {
253        let runs = Arc::new(AtomicUsize::new(0));
254
255        let runs_inner = Arc::clone(&runs);
256        let inner: Promise<i32, E> = Promise::lazy(async move {
257            runs_inner.fetch_add(1, Ordering::SeqCst);
258            Ok(7)
259        });
260
261        let mut a = inner.shared();
262        let mut b = a.clone();
263        let mut c = a.clone();
264
265        assert_eq!(poll(&mut a), Poll::Ready(Ok(7)));
266        assert_eq!(poll(&mut b), Poll::Ready(Ok(7)));
267        assert_eq!(poll(&mut c), Poll::Ready(Ok(7)));
268
269        assert_eq!(runs.load(Ordering::SeqCst), 1);
270    }
271
272    #[test]
273    fn every_clone_observes_same_result() {
274        let inner: Promise<i32, E> = Promise::lazy(async { Ok(123) });
275
276        let mut a = inner.shared();
277        let mut b = a.clone();
278        let mut c = b.clone();
279
280        assert_eq!(poll(&mut a), Poll::Ready(Ok(123)));
281        assert_eq!(poll(&mut b), Poll::Ready(Ok(123)));
282        assert_eq!(poll(&mut c), Poll::Ready(Ok(123)));
283    }
284
285    #[test]
286    fn repeated_polls_are_idempotent() {
287        let mut shared: SharedPromise<i32, E> = Promise::resolve(9).shared();
288
289        assert_eq!(poll(&mut shared), Poll::Ready(Ok(9)));
290        assert_eq!(poll(&mut shared), Poll::Ready(Ok(9)));
291        assert_eq!(poll(&mut shared), Poll::Ready(Ok(9)));
292    }
293
294    #[test]
295    fn pending_clone_then_resolves() {
296        let (promise, resolve, _reject) = Promise::<i32, E>::with_resolvers();
297
298        let mut shared = promise.shared();
299        let mut clone = shared.clone();
300
301        assert_eq!(poll(&mut shared), Poll::Pending);
302        assert_eq!(poll(&mut clone), Poll::Pending);
303
304        resolve.resolve(55);
305
306        assert_eq!(poll(&mut shared), Poll::Ready(Ok(55)));
307        assert_eq!(poll(&mut clone), Poll::Ready(Ok(55)));
308    }
309
310    #[test]
311    fn pending_consumer_is_woken_on_resolve() {
312        let (promise, resolve, _reject) = Promise::<i32, E>::with_resolvers();
313
314        let mut shared = promise.shared();
315
316        let counter = Arc::new(CountingWaker {
317            count: AtomicUsize::new(0),
318        });
319        let waker = Waker::from(Arc::clone(&counter));
320
321        assert_eq!(poll_with(&mut shared, &waker), Poll::Pending);
322        assert_eq!(counter.count.load(Ordering::SeqCst), 0);
323
324        resolve.resolve(88);
325
326        assert!(counter.count.load(Ordering::SeqCst) >= 1);
327
328        assert_eq!(poll_with(&mut shared, &waker), Poll::Ready(Ok(88)));
329    }
330
331    #[test]
332    fn late_consumer_observes_settled_result() {
333        let inner: Promise<i32, E> = Promise::lazy(async { Ok(321) });
334
335        let mut shared = inner.shared();
336
337        let late = shared.clone();
338
339        assert_eq!(poll(&mut shared), Poll::Ready(Ok(321)));
340
341        let mut late = late;
342
343        assert_eq!(poll(&mut late), Poll::Ready(Ok(321)));
344    }
345
346    type ParkedConsumer = Arc<Mutex<Option<(SharedPromise<i32, E>, Waker)>>>;
347
348    /// A slow self-waking inner future: it returns `Pending` and re-arms its own
349    /// waker on every poll until it has been polled `LIVELOCK_MAX_SELF_POLLS`
350    /// times, which trips the driver's inline self-poll bound. On the poll that
351    /// trips the bound it lets a late consumer park into the just-drained waker
352    /// queue, then keeps returning `Pending` (it would resolve on a later poll
353    /// that the dropped driver never issues).
354    struct SlowStrandProbe {
355        parked: ParkedConsumer,
356        polls: usize,
357    }
358
359    impl Future for SlowStrandProbe {
360        type Output = Result<i32, E>;
361
362        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
363            let this = self.get_mut();
364
365            this.polls += 1;
366
367            cx.waker().wake_by_ref();
368
369            if this.polls == LIVELOCK_MAX_SELF_POLLS {
370                let mut parked = this.parked.lock().expect("parked consumer");
371
372                if let Some((clone, waker)) = parked.as_mut() {
373                    let mut parked_cx = Context::from_waker(waker);
374
375                    assert!(Pin::new(clone).poll(&mut parked_cx).is_pending());
376                }
377
378                drop(parked);
379            }
380
381            Poll::Pending
382        }
383    }
384
385    #[test]
386    fn slow_self_waker_does_not_strand_a_late_consumer_when_driver_is_dropped() {
387        let parked: ParkedConsumer = Arc::new(Mutex::new(None));
388
389        let shared = Promise::<i32, E>::lazy(SlowStrandProbe {
390            parked: parked.clone(),
391            polls: 0,
392        })
393        .shared();
394
395        let late_waker = Arc::new(CountingWaker {
396            count: AtomicUsize::new(0),
397        });
398
399        {
400            let mut slot = parked.lock().expect("register late consumer");
401
402            *slot = Some((shared.clone(), Waker::from(late_waker.clone())));
403        }
404
405        let driver_waker = Arc::new(CountingWaker {
406            count: AtomicUsize::new(0),
407        });
408        let mut driver = shared;
409
410        assert!(poll_with(&mut driver, &Waker::from(driver_waker)).is_pending());
411
412        drop(driver);
413
414        assert!(
415            late_waker.count.load(Ordering::Relaxed) >= 1,
416            "late consumer was stranded: parked, owed a wake, never woken"
417        );
418    }
419}