cocoon_tpm_utils_async/
test.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2023-2025 SUSE LLC
3// Author: Nicolai Stange <nstange@suse.de>
4
5//! Testing [`TestNopSyncTypes`] and a [`Future`] [test
6//! exector](TestAsyncExecutor) implementations.
7
8extern crate alloc;
9use alloc::{boxed::Box, sync};
10
11use crate::{
12    alloc::SyncVec,
13    sync_types::{self, Lock as _},
14};
15use core::{cell, convert, future, marker, ops, pin, sync::atomic, task};
16
17/// Dummy [`Lock`](sync_types::Lock) for testing purposes.
18///
19/// Any attempt to lock an already locked `TestNopLock` will panic.
20pub struct TestNopLock<T: marker::Send> {
21    locked: atomic::AtomicBool,
22    v: cell::UnsafeCell<T>,
23}
24
25impl<T: marker::Send> convert::From<T> for TestNopLock<T> {
26    fn from(value: T) -> Self {
27        Self {
28            locked: atomic::AtomicBool::new(false),
29            v: cell::UnsafeCell::new(value),
30        }
31    }
32}
33
34unsafe impl<T: marker::Send> marker::Send for TestNopLock<T> {}
35unsafe impl<T: marker::Send> marker::Sync for TestNopLock<T> {}
36
37impl<T: marker::Send> sync_types::Lock<T> for TestNopLock<T> {
38    type Guard<'a>
39        = TestNopLockGuard<'a, T>
40    where
41        Self: 'a;
42
43    fn lock(&self) -> Self::Guard<'_> {
44        assert_eq!(
45            self.locked
46                .compare_exchange(false, true, atomic::Ordering::Acquire, atomic::Ordering::Relaxed),
47            Ok(false),
48            "Testing TestNopLocks are not expected to ever be contended."
49        );
50        TestNopLockGuard { lock: self }
51    }
52}
53
54impl<T: marker::Send> sync_types::ConstructibleLock<T> for TestNopLock<T> {
55    fn get_mut(&mut self) -> &mut T {
56        assert!(!self.locked.load(atomic::Ordering::Relaxed));
57        let p = self.v.get();
58        unsafe { &mut *p }
59    }
60}
61
62/// The [locking guard](sync_types::Lock::Guard) associated with
63/// [`TestNopLock`].
64pub struct TestNopLockGuard<'a, T: marker::Send> {
65    lock: &'a TestNopLock<T>,
66}
67
68impl<'a, T: marker::Send> Drop for TestNopLockGuard<'a, T> {
69    fn drop(&mut self) {
70        assert_eq!(
71            self.lock
72                .locked
73                .compare_exchange(true, false, atomic::Ordering::Acquire, atomic::Ordering::Relaxed),
74            Ok(true),
75            "Testing TestNopLock with active lock guard found unlocked."
76        );
77    }
78}
79
80impl<'a, T: marker::Send> ops::Deref for TestNopLockGuard<'a, T> {
81    type Target = T;
82
83    fn deref(&self) -> &Self::Target {
84        let p = self.lock.v.get();
85        // Safety: the very purpose of a Lock is exclusive access, so no aliasing.
86        unsafe { &*p }
87    }
88}
89
90impl<'a, T: marker::Send> ops::DerefMut for TestNopLockGuard<'a, T> {
91    fn deref_mut(&mut self) -> &mut Self::Target {
92        let p = self.lock.v.get();
93        // Safety: the very purpose of a Lock is exclusive access, so no aliasing.
94        unsafe { &mut *p }
95    }
96}
97
98/// Dummy [`RwLock`](sync_types::RwLock) for testing purposes.
99///
100/// Any attempt to lock a read-locked `TestNopLock` for write or vice versa will
101/// panic.
102pub struct TestNopRwLock<T: marker::Send + marker::Sync> {
103    locked: atomic::AtomicIsize,
104    v: cell::UnsafeCell<T>,
105}
106
107impl<T: marker::Send + marker::Sync> convert::From<T> for TestNopRwLock<T> {
108    fn from(value: T) -> Self {
109        Self {
110            locked: atomic::AtomicIsize::new(0),
111            v: cell::UnsafeCell::new(value),
112        }
113    }
114}
115
116unsafe impl<T: marker::Send + marker::Sync> marker::Send for TestNopRwLock<T> {}
117unsafe impl<T: marker::Send + marker::Sync> marker::Sync for TestNopRwLock<T> {}
118
119impl<T: marker::Send + marker::Sync> sync_types::RwLock<T> for TestNopRwLock<T> {
120    type ReadGuard<'a>
121        = TestNopRwLockReadGuard<'a, T>
122    where
123        Self: 'a;
124    type WriteGuard<'a>
125        = TestNopRwLockWriteGuard<'a, T>
126    where
127        Self: 'a;
128
129    fn read(&self) -> Self::ReadGuard<'_> {
130        assert!(
131            self.locked.fetch_add(1, atomic::Ordering::Acquire) >= 0,
132            "Testing TestNopRwLocks are not expected to ever be contended."
133        );
134        TestNopRwLockReadGuard { lock: self }
135    }
136
137    fn write(&self) -> Self::WriteGuard<'_> {
138        assert_eq!(
139            self.locked.fetch_sub(1, atomic::Ordering::Acquire),
140            0,
141            "Testing TestNopRwLocks are not expected to ever be contended."
142        );
143        TestNopRwLockWriteGuard { lock: self }
144    }
145
146    fn get_mut(&mut self) -> &mut T {
147        assert_eq!(self.locked.load(atomic::Ordering::Relaxed), 0);
148        let p = self.v.get();
149        unsafe { &mut *p }
150    }
151}
152
153/// The [read lock guard](sync_types::RwLock::ReadGuard) associated with
154/// [`TestNopRwLock`].
155pub struct TestNopRwLockReadGuard<'a, T: marker::Send + marker::Sync> {
156    lock: &'a TestNopRwLock<T>,
157}
158
159impl<'a, T: marker::Send + marker::Sync> Drop for TestNopRwLockReadGuard<'a, T> {
160    fn drop(&mut self) {
161        assert!(
162            self.lock.locked.fetch_sub(1, atomic::Ordering::Release) > 0,
163            "Testing TestNopRwLock with active read guard found unlocked or write locked."
164        );
165    }
166}
167
168impl<'a, T: marker::Send + marker::Sync> ops::Deref for TestNopRwLockReadGuard<'a, T> {
169    type Target = T;
170
171    fn deref(&self) -> &Self::Target {
172        let p = self.lock.v.get();
173        // Safety: the very purpose of a RwLock is exclusive Writers, so no
174        // aliasing with mutable references.
175        unsafe { &*p }
176    }
177}
178
179/// The [write lock guard](sync_types::RwLock::WriteGuard) associated with
180/// [`TestNopRwLock`].
181pub struct TestNopRwLockWriteGuard<'a, T: marker::Send + marker::Sync> {
182    lock: &'a TestNopRwLock<T>,
183}
184
185impl<'a, T: marker::Send + marker::Sync> Drop for TestNopRwLockWriteGuard<'a, T> {
186    fn drop(&mut self) {
187        assert_eq!(
188            self.lock.locked.fetch_add(1, atomic::Ordering::Release),
189            -1,
190            "Testing TestNopRwLock with active lock write guard found unlocked or read locked."
191        );
192    }
193}
194
195impl<'a, T: marker::Send + marker::Sync> ops::Deref for TestNopRwLockWriteGuard<'a, T> {
196    type Target = T;
197
198    fn deref(&self) -> &Self::Target {
199        let p = self.lock.v.get();
200        // Safety: the very purpose of a RwLock is exclusive Writers, so no
201        // aliasing with mutable references.
202        unsafe { &*p }
203    }
204}
205
206impl<'a, T: marker::Send + marker::Sync> ops::DerefMut for TestNopRwLockWriteGuard<'a, T> {
207    fn deref_mut(&mut self) -> &mut Self::Target {
208        let p = self.lock.v.get();
209        // Safety: the very purpose of a RwLock is exclusive Writers, so no
210        // aliasing with mutable references.
211        unsafe { &mut *p }
212    }
213}
214
215/// Dummy [`SyncTypes`](sync_types::SyncTypes) collection for testing purposes.
216pub struct TestNopSyncTypes;
217
218impl sync_types::SyncTypes for TestNopSyncTypes {
219    type Lock<T: marker::Send> = TestNopLock<T>;
220    type RwLock<T: marker::Send + marker::Sync> = TestNopRwLock<T>;
221    type SyncRcPtrFactory = sync_types::GenericArcFactory;
222}
223
224/// Dyn dispatcher trait to a [`Future`] enqueued at [`TestAsyncExecutor`].
225///
226/// # See also:
227///
228/// * [`QueuedTask`]
229trait QueuedTaskDispatch: marker::Send {
230    /// Poll the wrapped [`Future`].
231    ///
232    /// Return true once the wrapped [`Future`]'s
233    /// [`poll()`](future::Future::poll) returns
234    /// [`Ready`](task::Poll::Ready), `false` otherwise.
235    fn poll_pinned(&mut self, cx: &mut task::Context<'_>) -> bool;
236}
237
238/// Dyn dispatcher to a [`Future`] enqueued at [`TestAsyncExecutor`].
239///
240/// As the individual enqueued [`Future`]s can all have different
241/// [`Output`](future::Future::Output) types, it is not possible to store them
242/// as `dyn` objects and dispatch to their [`poll()`](future::Future::poll)
243/// implementations directly, hence the indirection.
244struct QueuedTask<F: future::Future + Send>
245where
246    F::Output: Send + 'static,
247{
248    /// The queued [`Future`].
249    f: F,
250    /// The result, stored once `f`'s [`poll()`](future::Future::poll) returns
251    /// [`Ready`](task::Poll::Ready).
252    ///
253    /// The `result` is shared with and can get stolen by the task's associated
254    /// [`TestAsyncExecutorTaskWaiter`].
255    result: sync_types::GenericArc<TestNopLock<Option<F::Output>>>,
256}
257
258impl<F: future::Future + Send> QueuedTaskDispatch for QueuedTask<F>
259where
260    F::Output: Send + 'static,
261{
262    fn poll_pinned(&mut self, cx: &mut task::Context<'_>) -> bool {
263        // Safety: always called with actually pinned &mut self, as part of the
264        // contract.
265        let f = unsafe { pin::Pin::new_unchecked(&mut self.f) };
266        match future::Future::poll(f, cx) {
267            task::Poll::Ready(result) => {
268                *self.result.lock() = Some(result);
269                true
270            }
271            task::Poll::Pending => false,
272        }
273    }
274}
275
276/// Runnable status of a [`TaskQueueEntry`] enqueued at [`TestAsyncExecutor`].
277enum TaskStatus {
278    /// The task is blocked, i.e. the associated [`Future`]'s
279    /// [`poll()`](future::Future::poll)
280    /// returned [`Pending`](task::Poll::Pending) and the task has not been
281    /// woken yet.
282    Blocked,
283    /// The task is runnable and can be polled, i.e. it's either been freshly
284    /// enqueued or woken.
285    Runnable,
286}
287
288/// State for top-level [`Future`] entry [enqueued](TestAsyncExecutor::spawn) at
289/// a [`TestAsyncExecutor`].
290struct TaskQueueEntry {
291    /// The task id assigned to the associated enqueued top-level [`Future`]
292    /// by the containing [executor](TestAsyncExecutor).
293    id: u64,
294    /// Runnable status of the top-level [`Future`].
295    status: TaskStatus,
296    /// The enqueued top-level [`Future`].
297    task: Option<pin::Pin<Box<dyn QueuedTaskDispatch>>>,
298    /// [`Waker`](task::Waker) installed on behalf of the
299    /// [`TestAsyncExecutorTaskWaiter`]'s
300    /// [`Future::poll()`](future::Future::poll) implementation, if any.
301    waiter_waker: Option<task::Waker>,
302}
303
304/// A [`Waker`](task::Waker) for the [`TestAsyncExecutor`].
305struct Waker {
306    task_id: u64,
307    executor: sync_types::GenericArc<TestAsyncExecutor>,
308}
309
310impl alloc::task::Wake for Waker {
311    fn wake(self: sync::Arc<Self>) {
312        let executor = &self.executor;
313        let mut tasks = executor.tasks.lock();
314        for t in tasks.iter_mut() {
315            if t.id == self.task_id && matches!(t.status, TaskStatus::Blocked) {
316                t.status = TaskStatus::Runnable
317            }
318        }
319    }
320}
321
322enum TaskWaiterState<T: marker::Send> {
323    Pending {
324        /// Executor the associated top-level [`Future`] had been enqueued to.
325        executor: sync_types::GenericArc<TestAsyncExecutor>,
326        /// The task id assigned to the associated enqueued top-level [`Future`]
327        /// by the `executor`.
328        task_id: u64,
329        /// Pointer to the [`Future::Output`] slot shared with
330        /// [`QueuedTask::result`]
331        result: sync_types::GenericArc<TestNopLock<Option<T>>>,
332    },
333    Done,
334}
335
336/// Waiter to be returned for top level [`Future`]s via
337/// [`TestAsyncExecutor::spawn()`](TestAsyncExecutor::spawn).
338///
339/// A `TestAsyncExecutorTaskWaiter` provides a means to obtain the
340/// [`Output`](future::Future::Output) of the associated enqueued [`Future`]
341/// when [`Ready`](task::Poll::Ready).
342///
343/// Note that a [`TestAsyncExecutorTaskWaiter`] implements [`Future`] itself, so
344/// it can get enqueued to the [`TestAsyncExecutor`] or polled from another
345/// owning [`Future`].
346pub struct TestAsyncExecutorTaskWaiter<T: marker::Send> {
347    state: TaskWaiterState<T>,
348}
349
350impl<T: marker::Send> TestAsyncExecutorTaskWaiter<T> {
351    /// Take the associated enqueued [`Future`]'s
352    /// [`Output`](future::Future::Output) if [`Ready`](task::Poll::Ready).
353    ///
354    /// Returns `None` if the (future::Future::Output) is not
355    /// [`Ready`](task::Poll::Ready) yet or the result has already been
356    /// taken, including through [`Self::poll()`](Self::poll).
357    pub fn take(mut self) -> Option<T> {
358        match &mut self.state {
359            TaskWaiterState::Pending {
360                executor: _,
361                task_id: _,
362                result,
363            } => {
364                let result = result.lock().take();
365                self.state = TaskWaiterState::Done;
366                result
367            }
368            TaskWaiterState::Done => None,
369        }
370    }
371}
372
373impl<T: marker::Send> Drop for TestAsyncExecutorTaskWaiter<T> {
374    fn drop(&mut self) {
375        match &self.state {
376            TaskWaiterState::Pending {
377                executor,
378                task_id,
379                result,
380            } => {
381                if !result.lock().is_some() {
382                    executor.remove_task(*task_id);
383                }
384            }
385            TaskWaiterState::Done => (),
386        }
387    }
388}
389
390impl<T: marker::Send> Unpin for TestAsyncExecutorTaskWaiter<T> {}
391
392impl<T: marker::Send> future::Future for TestAsyncExecutorTaskWaiter<T> {
393    type Output = T;
394
395    fn poll(self: pin::Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
396        let this = self.get_mut();
397        match &this.state {
398            TaskWaiterState::Pending {
399                executor,
400                task_id,
401                result,
402            } => {
403                let mut locked_result = result.lock();
404                if let Some(result) = locked_result.take() {
405                    drop(locked_result);
406                    this.state = TaskWaiterState::Done;
407                    task::Poll::Ready(result)
408                } else {
409                    let mut tasks = executor.tasks.lock();
410                    let task = tasks.iter_mut().find(|task| task.id == *task_id).unwrap();
411                    task.waiter_waker = Some(cx.waker().clone());
412                    task::Poll::Pending
413                }
414            }
415            TaskWaiterState::Done => unreachable!(),
416        }
417    }
418}
419
420/// Single-threaded [`Future`] task executor for testing purposes.
421///
422/// Enqueue top-level [`Future`]s via [`spawn()`](Self::spawn) and
423/// poll all currently enqueued ones to completion via
424/// [`run_to_completion()`](Self::run_to_completion).
425///
426/// Enqueued [`Future`]s may [enqueue](Self::spawn) further ones from their
427/// [`Future::poll()`](future::Future::poll) and poll on the resulting
428/// [`TestAsyncExecutorTaskWaiter`].
429pub struct TestAsyncExecutor {
430    /// All enqueued top-level [`Future`]s.
431    tasks: TestNopLock<SyncVec<TaskQueueEntry>>,
432    /// The task id to assign to the next [`Future`] to get
433    /// [enqueued](Self::spawn).
434    next_id: atomic::AtomicU64,
435}
436
437impl TestAsyncExecutor {
438    /// Create a new [`TestAsyncExecutor`] instance.
439    pub fn new() -> sync_types::GenericArc<Self> {
440        <sync_types::GenericArcFactory as sync_types::SyncRcPtrFactory>::try_new(Self {
441            tasks: TestNopLock::from(SyncVec::new()),
442            next_id: atomic::AtomicU64::new(0),
443        })
444        .unwrap()
445    }
446
447    /// Enqueue a top-level [`Future`] for polling from a subsequent
448    /// [`run_to_completion()`](Self::run_to_completion) invocation.
449    ///
450    /// # See also:
451    /// * [`TestAsyncExecutorTaskWaiter`]
452    /// * [`run_to_completion()`](Self::run_to_completion)
453    pub fn spawn<F: future::Future + Send + 'static>(
454        this: &sync_types::GenericArc<Self>,
455        f: F,
456    ) -> TestAsyncExecutorTaskWaiter<F::Output>
457    where
458        F::Output: Send + 'static,
459    {
460        let id = this.next_id.fetch_add(1, atomic::Ordering::Relaxed);
461
462        let result =
463            <sync_types::GenericArcFactory as sync_types::SyncRcPtrFactory>::try_new(TestNopLock::from(None)).unwrap();
464        let waiter = TestAsyncExecutorTaskWaiter {
465            state: TaskWaiterState::Pending {
466                executor: this.clone(),
467                task_id: id,
468                result: result.clone(),
469            },
470        };
471
472        let task = Box::pin(QueuedTask { f, result }) as pin::Pin<Box<dyn QueuedTaskDispatch>>;
473        let tasks = this.tasks.lock();
474        let (mut tasks, r) = SyncVec::try_reserve_exact(&this.tasks, tasks, 1);
475        r.unwrap();
476        tasks.push(TaskQueueEntry {
477            id,
478            status: TaskStatus::Runnable,
479            task: Some(task),
480            waiter_waker: None,
481        });
482
483        waiter
484    }
485
486    fn remove_task(&self, id: u64) {
487        let mut tasks = self.tasks.lock();
488        if let Some(index) = tasks.iter().position(|task| task.id == id) {
489            // Removing/Dropping the task might drop further TaskWaiter instances held by
490            // that task, which would in turn invoke this function and try to
491            // get self.tasks for writing.
492            let entry = tasks.remove(index);
493            drop(tasks);
494            drop(entry);
495        };
496    }
497
498    /// Poll all currently [enqueued](Self::spawn) [`Future`]s to completion.
499    ///
500    /// A [`Future`] is in "runnable" state right after it has been
501    /// [`enqueued`](Self::spawn) and ceases to once its
502    /// [`Future::poll()`](future::Future::poll) returns a status of
503    /// [`Pending`](task::Poll::Pending). A non-runnable [`Future`] becomes
504    /// runnable again when [woken](task::Waker::wake). There must always be at
505    /// least one runnable [`Future`] left, or the executor will become
506    /// stuck and report the fact via a panic.
507    ///
508    /// `run_to_completion()` polls the runnable enqueued [`Future`] in a
509    /// round-robin fashion, in the order of their enqueueing. That is, if
510    /// futures `[a, b, c]` are enqueued, with `a` non-runnable and `b` and
511    /// `c` runnable, and the cursor is at `a` or `b`, then `b` will get polled
512    /// first, followed by `c`, followed by `a` if that became runnable in the
513    /// course of polling `b` or `c`.
514    pub fn run_to_completion(this: &sync_types::GenericArc<Self>) {
515        let mut last_polled: Option<(usize, u64)> = None;
516        loop {
517            let mut tasks = this.tasks.lock();
518            if tasks.is_empty() {
519                break;
520            }
521
522            // Determine the next task to examine: either the one with the next larger task
523            // id, if any, or wrap around to the beginning.
524            let mut search_begin = match last_polled {
525                Some((last_index, last_task_id)) => {
526                    // The saved index is only an approximate hint, because self.tasks might have
527                    // changed when its lock was released. Search downwards for
528                    // the last entry before index with a task id <= the last
529                    // one, and search upward from there for the one with the next higher id.
530                    let last_index = last_index.min(tasks.len());
531                    let last_before_leq = tasks[..last_index]
532                        .iter()
533                        .rposition(|entry| entry.id <= last_task_id)
534                        .unwrap_or(0);
535                    match tasks
536                        .iter()
537                        .enumerate()
538                        .skip(last_before_leq)
539                        .find(|(_, entry)| entry.id > last_task_id)
540                    {
541                        Some((index, _)) => index,
542                        None => {
543                            // No task with a higher id than the last one. Wrap around.
544                            0
545                        }
546                    }
547                }
548                None => 0,
549            };
550            let index = loop {
551                match tasks
552                    .iter()
553                    .enumerate()
554                    .skip(search_begin)
555                    .find(|(_, entry)| matches!(entry.status, TaskStatus::Runnable))
556                {
557                    Some((index, _)) => break Some(index),
558                    None => {
559                        // Wrap around if the search hasn't started from the beginning already.
560                        if search_begin == 0 {
561                            break None;
562                        }
563                        search_begin = 0;
564                    }
565                }
566            };
567            let index = index.expect("TestAsyncExecutor stuck with no runnable task.");
568
569            let entry = &mut tasks[index];
570            let task_id = entry.id;
571            last_polled = Some((index, task_id));
572            // Temporarily steal the QueueTask for invoking it below with self.tasks[]
573            // unlocked.
574            let mut task = match entry.task.take() {
575                Some(task) => task,
576                None => {
577                    continue;
578                }
579            };
580            // Set the status to blocked now, so that any wake-ups from wakers
581            // won't get missed.
582            entry.status = TaskStatus::Blocked;
583            // Drop the tasks lock for the duration of polling the task --
584            // it might want to spawn more tasks or drop some TaskWaiters.
585            drop(tasks);
586
587            let waker = task::Waker::from(sync::Arc::new(Waker {
588                task_id,
589                executor: this.clone(),
590            }));
591            let mut cx = task::Context::from_waker(&waker);
592            // Safety: poll_pinned() immediately repins it.
593            let done = unsafe { task.as_mut().get_unchecked_mut() }.poll_pinned(&mut cx);
594
595            let task = if done {
596                // Dropping the task might drop further TaskWaiter instances held by
597                // that task, which would invoke Self::remove_task() and try to
598                // get self.tasks for writing. Do it here outside the self.tasks lock.
599                drop(task);
600                None
601            } else {
602                Some(task)
603            };
604
605            let mut tasks = this.tasks.lock();
606            // While the lock had been released, self.tasks[] could potentially have
607            // been mutated. Find the index corresponding to the task_id saved away above.
608            let updated_index = if index < tasks.len() && tasks[index].id == task_id {
609                // Position is unchanged.
610                index
611            } else {
612                match tasks.iter().position(|entry| entry.id == task_id) {
613                    Some(updated_index) => updated_index,
614                    None => {
615                        // The task has gone, presumably because its associated TaskWaiter had
616                        // been dropped.
617                        continue;
618                    }
619                }
620            };
621            last_polled = Some((updated_index, task_id));
622
623            if done {
624                let waiter_waker = tasks[updated_index].waiter_waker.take();
625                tasks.remove(updated_index);
626                if let Some(waiter_waker) = waiter_waker {
627                    drop(tasks);
628                    waiter_waker.wake();
629                }
630            } else {
631                let entry = &mut tasks[updated_index];
632                // Restore the pointer to the QueuedTask, which had temporarily
633                // been taken before the poll() invocation above.
634                entry.task = task;
635            }
636        }
637    }
638}
639
640#[test]
641fn test_test_async_executor_simple() {
642    struct SimpleTask {}
643
644    impl future::Future for SimpleTask {
645        type Output = u32;
646
647        fn poll(self: pin::Pin<&mut Self>, _cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
648            task::Poll::Ready(42)
649        }
650    }
651
652    let executor = TestAsyncExecutor::new();
653    let waiter = TestAsyncExecutor::spawn(&executor, SimpleTask {});
654    TestAsyncExecutor::run_to_completion(&executor);
655    assert_eq!(waiter.take().unwrap(), 42);
656    assert_eq!(sync_types::GenericArc::strong_count(&executor), 1);
657    assert_eq!(sync_types::GenericArc::weak_count(&executor), 0);
658
659    let waiter = TestAsyncExecutor::spawn(&executor, async { async { 42 }.await });
660    TestAsyncExecutor::run_to_completion(&executor);
661    assert_eq!(waiter.take().unwrap(), 42);
662    assert_eq!(sync_types::GenericArc::strong_count(&executor), 1);
663    assert_eq!(sync_types::GenericArc::weak_count(&executor), 0);
664}
665
666#[test]
667fn test_test_async_executor_chained_waiters() {
668    struct SimpleTask {}
669
670    impl future::Future for SimpleTask {
671        type Output = u32;
672
673        fn poll(self: pin::Pin<&mut Self>, _cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
674            task::Poll::Ready(42)
675        }
676    }
677
678    let executor = TestAsyncExecutor::new();
679    let waiter = TestAsyncExecutor::spawn(&executor, SimpleTask {});
680    let waiter = TestAsyncExecutor::spawn(&executor, waiter);
681    let waiter = TestAsyncExecutor::spawn(&executor, waiter);
682    TestAsyncExecutor::run_to_completion(&executor);
683    assert_eq!(waiter.take().unwrap(), 42);
684    assert_eq!(sync_types::GenericArc::strong_count(&executor), 1);
685    assert_eq!(sync_types::GenericArc::weak_count(&executor), 0);
686}
687
688#[test]
689fn test_test_async_executor_recursive_spawning() {
690    use ops::DerefMut as _;
691
692    enum SpawningTask {
693        Init {
694            executor: sync_types::GenericArc<TestAsyncExecutor>,
695            n: u32,
696        },
697        WaitingForSpawn {
698            waiter: TestAsyncExecutorTaskWaiter<u32>,
699        },
700    }
701
702    impl Unpin for SpawningTask {}
703
704    impl future::Future for SpawningTask {
705        type Output = u32;
706
707        fn poll(mut self: pin::Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
708            match self.deref_mut() {
709                Self::Init { executor, n } => {
710                    if *n == 0 {
711                        task::Poll::Ready(0)
712                    } else {
713                        let mut waiter = TestAsyncExecutor::spawn(
714                            executor,
715                            SpawningTask::Init {
716                                executor: executor.clone(),
717                                n: *n - 1,
718                            },
719                        );
720                        match future::Future::poll(pin::Pin::new(&mut waiter), cx) {
721                            task::Poll::Ready(_) => {
722                                // The task associated with the waiter did not have a chance to run
723                                // yet.
724                                unreachable!();
725                            }
726                            task::Poll::Pending => {
727                                *self.deref_mut() = Self::WaitingForSpawn { waiter };
728                                task::Poll::Pending
729                            }
730                        }
731                    }
732                }
733                Self::WaitingForSpawn { waiter } => {
734                    match future::Future::poll(pin::Pin::new(waiter), cx) {
735                        task::Poll::Ready(n) => task::Poll::Ready(n + 1),
736                        task::Poll::Pending => {
737                            // This future's task should have been woken only once
738                            // the waiter has become ready.
739                            unreachable!();
740                        }
741                    }
742                }
743            }
744        }
745    }
746
747    let executor = TestAsyncExecutor::new();
748    let waiter = TestAsyncExecutor::spawn(
749        &executor,
750        SpawningTask::Init {
751            executor: executor.clone(),
752            n: 42,
753        },
754    );
755    TestAsyncExecutor::run_to_completion(&executor);
756    assert_eq!(waiter.take().unwrap(), 42);
757    assert_eq!(sync_types::GenericArc::strong_count(&executor), 1);
758    assert_eq!(sync_types::GenericArc::weak_count(&executor), 0);
759}
760
761#[test]
762fn test_test_async_executor_wake_self() {
763    use ops::Deref as _;
764
765    enum SelfWakingTask {
766        Unpolled,
767        PolledOnce,
768    }
769
770    impl Unpin for SelfWakingTask {}
771
772    impl future::Future for SelfWakingTask {
773        type Output = u32;
774
775        fn poll(mut self: pin::Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
776            match self.deref() {
777                Self::Unpolled => {
778                    cx.waker().wake_by_ref();
779                    *self = Self::PolledOnce;
780                    task::Poll::Pending
781                }
782                Self::PolledOnce => task::Poll::Ready(42),
783            }
784        }
785    }
786
787    let executor = TestAsyncExecutor::new();
788    let waiter = TestAsyncExecutor::spawn(&executor, SelfWakingTask::Unpolled);
789    let waiter = TestAsyncExecutor::spawn(&executor, waiter);
790    TestAsyncExecutor::run_to_completion(&executor);
791    assert_eq!(waiter.take().unwrap(), 42);
792    assert_eq!(sync_types::GenericArc::strong_count(&executor), 1);
793    assert_eq!(sync_types::GenericArc::weak_count(&executor), 0);
794}
cocoon_tpm_utils_async/test.rs

cocoon_tpm_utils_async/
test.rs
