Skip to main content

tor_rtmock/
task.rs

1//! Executor for running tests with mocked environment
2//!
3//! See [`MockExecutor`]
4
5use std::any::Any;
6use std::cell::Cell;
7use std::collections::VecDeque;
8use std::fmt::{self, Debug, Display};
9use std::future::Future;
10use std::io::{self, Write as _};
11use std::iter;
12use std::panic::{AssertUnwindSafe, catch_unwind, panic_any};
13use std::pin::{Pin, pin};
14use std::sync::{Arc, Mutex, MutexGuard, Weak};
15use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
16
17use futures::FutureExt as _;
18use futures::pin_mut;
19use futures::task::{FutureObj, Spawn, SpawnError};
20
21use assert_matches::assert_matches;
22use educe::Educe;
23use itertools::Either::{self, *};
24use itertools::{chain, izip};
25use slotmap_careful::DenseSlotMap;
26use std::backtrace::Backtrace;
27use strum::EnumIter;
28
29// NB: when using traced_test, the trace! and error! output here is generally suppressed
30// in tests of other crates.  To see it, you can write something like this
31// (in the dev-dependencies of the crate whose tests you're running):
32//    tracing-test = { version = "0.2.4", features = ["no-env-filter"] }
33use tracing::{error, trace};
34
35use oneshot_fused_workaround::{self as oneshot, Canceled};
36use tor_error::error_report;
37use tor_rtcompat::{Blocking, ToplevelBlockOn};
38
39use Poll::*;
40use TaskState::*;
41
42/// Type-erased future, one for each of our (normal) tasks
43type TaskFuture = FutureObj<'static, ()>;
44
45/// Future for the argument to `block_on`, which is handled specially
46type MainFuture<'m> = Pin<&'m mut dyn Future<Output = ()>>;
47
48//---------- principal data structures ----------
49
50/// Executor for running tests with mocked environment
51///
52/// For test cases which don't actually wait for anything in the real world.
53///
54/// This is the executor.
55/// It implements [`Spawn`] and [`ToplevelBlockOn`]
56///
57/// It will usually be used as part of a `MockRuntime`.
58///
59/// To run futures, call [`ToplevelBlockOn::block_on`]
60///
61/// # Restricted environment
62///
63/// Tests run with this executor must not attempt to block
64/// on anything "outside":
65/// every future that anything awaits must (eventually) be woken directly
66/// *by some other task* in the same test case.
67///
68/// (By directly we mean that the [`Waker::wake`] call is made
69/// by that waking future, before that future itself awaits anything.)
70///
71/// # Panics
72///
73/// The executor will panic
74/// if the toplevel future (passed to `block_on`)
75/// doesn't complete (without externally blocking),
76/// but instead waits for something.
77///
78/// The executor will malfunction or panic if reentered.
79/// (Eg, if `block_on` is reentered.)
80#[derive(Clone, Default, Educe)]
81#[educe(Debug)]
82pub struct MockExecutor {
83    /// Mutable state
84    #[educe(Debug(ignore))]
85    shared: Arc<Shared>,
86}
87
88/// Shared state and ancillary information
89///
90/// This is always within an `Arc`.
91#[derive(Default)]
92struct Shared {
93    /// Shared state
94    data: Mutex<Data>,
95    /// Condition variable for thread scheduling
96    ///
97    /// Signaled when [`Data.thread_to_run`](struct.Data.html#structfield.thread_to_run)
98    /// is modified.
99    thread_condvar: std::sync::Condvar,
100}
101
102/// Task id, module to hide `Ti` alias
103mod task_id {
104    slotmap_careful::new_key_type! {
105        /// Task ID, usually called `TaskId`
106        ///
107        /// Short name in special `task_id` module so that [`Debug`] is nice
108        pub(super) struct Ti;
109    }
110}
111use task_id::Ti as TaskId;
112
113/// Executor's state
114///
115/// ### Task state machine
116///
117/// A task is created in `tasks`, `Awake`, so also in `awake`.
118///
119/// When we poll it, we take it out of `awake` and set it to `Asleep`,
120/// and then call `poll()`.
121/// Any time after that, it can be made `Awake` again (and put back onto `awake`)
122/// by the waker ([`ActualWaker`], wrapped in [`Waker`]).
123///
124/// The task's future is of course also present here in this data structure.
125/// However, during poll we must release the lock,
126/// so we cannot borrow the future from `Data`.
127/// Instead, we move it out.  So `Task.fut` is an `Option`.
128///
129/// ### "Main" task - the argument to `block_on`
130///
131/// The signature of `BlockOn::block_on` accepts a non-`'static` future
132/// (and a non-`Send`/`Sync` one).
133///
134/// So we cannot store that future in `Data` because `Data` is `'static`.
135/// Instead, this main task future is passed as an argument down the call stack.
136/// In the data structure we simply store a placeholder, `TaskFutureInfo::Main`.
137#[derive(Educe, derive_more::Debug)]
138#[educe(Default)]
139struct Data {
140    /// Tasks
141    ///
142    /// Includes tasks spawned with `spawn`,
143    /// and also the future passed to `block_on`.
144    #[debug("{:?}", DebugTasks(self, || tasks.keys()))]
145    tasks: DenseSlotMap<TaskId, Task>,
146
147    /// `awake` lists precisely: tasks that are `Awake`, plus maybe stale `TaskId`s
148    ///
149    /// Tasks are pushed onto the *back* when woken,
150    /// so back is the most recently woken.
151    #[debug("{:?}", DebugTasks(self, || awake.iter().cloned()))]
152    awake: VecDeque<TaskId>,
153
154    /// If a future from `progress_until_stalled` exists
155    progressing_until_stalled: Option<ProgressingUntilStalled>,
156
157    /// Scheduling policy
158    scheduling: SchedulingPolicy,
159
160    /// (Sub)thread we want to run now
161    ///
162    /// At any one time only one thread is meant to be running.
163    /// Other threads are blocked in condvar wait, waiting for this to change.
164    ///
165    /// **Modified only** within
166    /// [`thread_context_switch_send_instruction_to_run`](Shared::thread_context_switch_send_instruction_to_run),
167    /// which takes responsibility for preserving the following **invariants**:
168    ///
169    ///  1. no-one but the named thread is allowed to modify this field.
170    ///  2. after modifying this field, signal `thread_condvar`
171    #[educe(Default(expression = "ThreadDescriptor::Executor"))]
172    thread_to_run: ThreadDescriptor,
173}
174
175/// How we should schedule?
176#[derive(Debug, Clone, Default, EnumIter)]
177#[non_exhaustive]
178pub enum SchedulingPolicy {
179    /// Task *most* recently woken is run
180    ///
181    /// This is the default.
182    ///
183    /// It will expose starvation bugs if a task never sleeps.
184    /// (Which is a good thing in tests.)
185    #[default]
186    Stack,
187    /// Task *least* recently woken is run.
188    Queue,
189}
190
191/// Record of a single task
192///
193/// Tracks a spawned task, or the main task (the argument to `block_on`).
194///
195/// Stored in [`Data`]`.tasks`.
196struct Task {
197    /// For debugging output
198    desc: String,
199    /// Has this been woken via a waker?  (And is it in `Data.awake`?)
200    ///
201    /// **Set to `Awake` only by [`Task::set_awake`]**,
202    /// preserving the invariant that
203    /// every `Awake` task is in [`Data.awake`](struct.Data.html#structfield.awake).
204    state: TaskState,
205    /// The actual future (or a placeholder for it)
206    ///
207    /// May be `None` briefly in the executor main loop, because we've
208    /// temporarily moved it out so we can poll it,
209    /// or if this is a Subthread task which is currently running sync code
210    /// (in which case we're blocked in the executor waiting to be
211    /// woken up by [`thread_context_switch`](Shared::thread_context_switch).
212    ///
213    /// Note that the `None` can be observed outside the main loop, because
214    /// the main loop unlocks while it polls, so other (non-main-loop) code
215    /// might see it.
216    fut: Option<TaskFutureInfo>,
217}
218
219/// A future as stored in our record of a [`Task`]
220#[derive(Educe)]
221#[educe(Debug)]
222enum TaskFutureInfo {
223    /// The [`Future`].  All is normal.
224    Normal(#[educe(Debug(ignore))] TaskFuture),
225    /// The future isn't here because this task is the main future for `block_on`
226    Main,
227    /// This task is actually a [`Subthread`](MockExecutor::subthread_spawn)
228    ///
229    /// Instead of polling it, we'll switch to it with
230    /// [`thread_context_switch`](Shared::thread_context_switch).
231    Subthread,
232}
233
234/// State of a task - do we think it needs to be polled?
235///
236/// Stored in [`Task`]`.state`.
237#[derive(Debug)]
238enum TaskState {
239    /// Awake - needs to be polled
240    ///
241    /// Established by [`waker.wake()`](Waker::wake)
242    Awake,
243    /// Asleep - does *not* need to be polled
244    ///
245    /// Established each time just before we call the future's [`poll`](Future::poll)
246    Asleep(Vec<SleepLocation>),
247}
248
249/// Actual implementor of `Wake` for use in a `Waker`
250///
251/// Futures (eg, channels from [`futures`]) will use this to wake a task
252/// when it should be polled.
253///
254/// This type must not be `Cloned` with the `Data` lock held.
255/// Consequently, a `Waker` mustn't either.
256struct ActualWaker {
257    /// Executor state
258    ///
259    /// The Waker mustn't to hold a strong reference to the executor,
260    /// since typically a task holds a future that holds a Waker,
261    /// and the executor holds the task - so that would be a cycle.
262    data: Weak<Shared>,
263
264    /// Which task this is
265    id: TaskId,
266}
267
268/// State used for an in-progress call to
269/// [`progress_until_stalled`][`MockExecutor::progress_until_stalled`]
270///
271/// If present in [`Data`], an (async) call to `progress_until_stalled`
272/// is in progress.
273///
274/// The future from `progress_until_stalled`, [`ProgressUntilStalledFuture`]
275/// is a normal-ish future.
276/// It can be polled in the normal way.
277/// When it is polled, it looks here, in `finished`, to see if it's `Ready`.
278///
279/// The future is made ready, and woken (via `waker`),
280/// by bespoke code in the task executor loop.
281///
282/// When `ProgressUntilStalledFuture` (maybe completes and) is dropped,
283/// its `Drop` impl is used to remove this from `Data.progressing_until_stalled`.
284#[derive(Debug)]
285struct ProgressingUntilStalled {
286    /// Have we, in fact, stalled?
287    ///
288    /// Made `Ready` by special code in the executor loop
289    finished: Poll<()>,
290
291    /// Waker
292    ///
293    /// Signalled by special code in the executor loop
294    waker: Option<Waker>,
295}
296
297/// Future from
298/// [`progress_until_stalled`][`MockExecutor::progress_until_stalled`]
299///
300/// See [`ProgressingUntilStalled`] for an overview of this aspect of the contraption.
301///
302/// Existence of this struct implies `Data.progressing_until_stalled` is `Some`.
303/// There can only be one at a time.
304#[derive(Educe)]
305#[educe(Debug)]
306struct ProgressUntilStalledFuture {
307    /// Executor's state; this future's state is in `.progressing_until_stalled`
308    #[educe(Debug(ignore))]
309    shared: Arc<Shared>,
310}
311
312/// Identifies a thread we know about - the executor thread, or a Subthread
313///
314/// Not related to `std::thread::ThreadId`.
315///
316/// See [`spawn_subthread`](MockExecutor::subthread_spawn) for definition of a Subthread.
317///
318/// This being a thread-local and not scoped by which `MockExecutor` we're talking about
319/// means that we can't cope if there are multiple `MockExecutor`s involved in the same thread.
320/// That's OK (and documented).
321#[derive(Copy, Clone, Eq, PartialEq, derive_more::Debug)]
322enum ThreadDescriptor {
323    /// Foreign - neither the (running) executor, nor a Subthread
324    #[debug("FOREIGN")]
325    Foreign,
326    /// The executor.
327    #[debug("Exe")]
328    Executor,
329    /// This task, which is a Subthread.
330    #[debug("{_0:?}")]
331    Subthread(TaskId),
332}
333
334/// Marker indicating that this task is a Subthread, not an async task.
335///
336/// See [`spawn_subthread`](MockExecutor::subthread_spawn) for definition of a Subthread.
337#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
338struct IsSubthread;
339
340/// [`Shared::subthread_yield`] should set our task awake before switching to the executor
341#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
342struct SetAwake;
343
344thread_local! {
345    /// Identifies this thread.
346    pub static THREAD_DESCRIPTOR: Cell<ThreadDescriptor> = const {
347        Cell::new(ThreadDescriptor::Foreign)
348    };
349}
350
351//---------- creation ----------
352
353impl MockExecutor {
354    /// Make a `MockExecutor` with default parameters
355    pub fn new() -> Self {
356        Self::default()
357    }
358
359    /// Make a `MockExecutor` with a specific `SchedulingPolicy`
360    pub fn with_scheduling(scheduling: SchedulingPolicy) -> Self {
361        Data {
362            scheduling,
363            ..Default::default()
364        }
365        .into()
366    }
367}
368
369impl From<Data> for MockExecutor {
370    fn from(data: Data) -> MockExecutor {
371        let shared = Shared {
372            data: Mutex::new(data),
373            thread_condvar: std::sync::Condvar::new(),
374        };
375        MockExecutor {
376            shared: Arc::new(shared),
377        }
378    }
379}
380
381//---------- spawning ----------
382
383impl MockExecutor {
384    /// Spawn a task and return something to identify it
385    ///
386    /// `desc` should `Display` as some kind of short string (ideally without spaces)
387    /// and will be used in the `Debug` impl and trace log messages from `MockExecutor`.
388    ///
389    /// The returned value is an opaque task identifier which is very cheap to clone
390    /// and which can be used by the caller in debug logging,
391    /// if it's desired to correlate with the debug output from `MockExecutor`.
392    /// Most callers will want to ignore it.
393    ///
394    /// This method is infallible.  (The `MockExecutor` cannot be shut down.)
395    pub fn spawn_identified(
396        &self,
397        desc: impl Display,
398        fut: impl Future<Output = ()> + Send + 'static,
399    ) -> impl Debug + Clone + Send + 'static {
400        self.spawn_internal(desc.to_string(), FutureObj::from(Box::new(fut)))
401    }
402
403    /// Spawn a task and return its output for further usage
404    ///
405    /// `desc` should `Display` as some kind of short string (ideally without spaces)
406    /// and will be used in the `Debug` impl and trace log messages from `MockExecutor`.
407    pub fn spawn_join<T: Debug + Send + 'static>(
408        &self,
409        desc: impl Display,
410        fut: impl Future<Output = T> + Send + 'static,
411    ) -> impl Future<Output = T> {
412        let (tx, rx) = oneshot::channel();
413        self.spawn_identified(desc, async move {
414            let res = fut.await;
415            tx.send(res)
416                .expect("Failed to send future's output, did future panic?");
417        });
418        rx.map(|m| m.expect("Failed to receive future's output"))
419    }
420
421    /// Spawn a task and return its `TaskId`
422    ///
423    /// Convenience method for use by `spawn_identified` and `spawn_obj`.
424    /// The future passed to `block_on` is not handled here.
425    fn spawn_internal(&self, desc: String, fut: TaskFuture) -> TaskId {
426        let mut data = self.shared.lock();
427        data.insert_task(desc, TaskFutureInfo::Normal(fut))
428    }
429}
430
431impl Data {
432    /// Insert a task given its `TaskFutureInfo` and return its `TaskId`.
433    fn insert_task(&mut self, desc: String, fut: TaskFutureInfo) -> TaskId {
434        let state = Awake;
435        let id = self.tasks.insert(Task {
436            state,
437            desc,
438            fut: Some(fut),
439        });
440        self.awake.push_back(id);
441        trace!("MockExecutor spawned {:?}={:?}", id, self.tasks[id]);
442        id
443    }
444}
445
446impl Spawn for MockExecutor {
447    fn spawn_obj(&self, future: TaskFuture) -> Result<(), SpawnError> {
448        self.spawn_internal("spawn_obj".into(), future);
449        Ok(())
450    }
451}
452
453impl Blocking for MockExecutor {
454    type ThreadHandle<T: Send + 'static> = Pin<Box<dyn Future<Output = T>>>;
455
456    fn spawn_blocking<F, T>(&self, f: F) -> Self::ThreadHandle<T>
457    where
458        F: FnOnce() -> T + Send + 'static,
459        T: Send + 'static,
460    {
461        assert_matches!(
462            THREAD_DESCRIPTOR.get(),
463            ThreadDescriptor::Executor | ThreadDescriptor::Subthread(_),
464            "MockExecutor::spawn_blocking_io only allowed from future or subthread, being run by this executor"
465        );
466        Box::pin(
467            self.subthread_spawn("spawn_blocking", f)
468                .map(|x| x.expect("Error in spawn_blocking subthread.")),
469        )
470    }
471
472    fn reenter_block_on<F>(&self, future: F) -> F::Output
473    where
474        F: Future,
475        F::Output: Send + 'static,
476    {
477        self.subthread_block_on_future(future)
478    }
479}
480
481//---------- block_on ----------
482
483impl ToplevelBlockOn for MockExecutor {
484    fn block_on<F>(&self, input_fut: F) -> F::Output
485    where
486        F: Future,
487    {
488        let mut value: Option<F::Output> = None;
489
490        // Box this just so that we can conveniently control precisely when it's dropped.
491        // (We could do this with Option and Pin::set but that seems clumsier.)
492        let mut input_fut = Box::pin(input_fut);
493
494        let run_store_fut = {
495            let value = &mut value;
496            let input_fut = &mut input_fut;
497            async {
498                trace!("MockExecutor block_on future...");
499                let t = input_fut.await;
500                trace!("MockExecutor block_on future returned...");
501                *value = Some(t);
502                trace!("MockExecutor block_on future exiting.");
503            }
504        };
505
506        {
507            pin_mut!(run_store_fut);
508
509            let main_id = self
510                .shared
511                .lock()
512                .insert_task("main".into(), TaskFutureInfo::Main);
513            trace!("MockExecutor {main_id:?} is task for block_on");
514            self.execute_to_completion(run_store_fut);
515        }
516
517        #[allow(clippy::let_and_return)] // clarity
518        let value = value.take().unwrap_or_else(|| {
519            // eprintln can be captured by libtest, but the debug_dump goes to io::stderr.
520            // use the latter, so that the debug dump is prefixed by this message.
521            let _: io::Result<()> = writeln!(io::stderr(), "all futures blocked, crashing...");
522            // write to tracing too, so the tracing log is clear about when we crashed
523            error!("all futures blocked, crashing...");
524
525            // Sequencing here is subtle.
526            //
527            // We should do the dump before dropping the input future, because the input
528            // future is likely to own things that, when dropped, wake up other tasks,
529            // rendering the dump inaccurate.
530            //
531            // But also, dropping the input future may well drop a ProgressUntilStalledFuture
532            // which then reenters us.  More generally, we mustn't call user code
533            // with the lock held.
534            //
535            // And, we mustn't panic with the data lock held.
536            //
537            // If value was Some, then this closure is dropped without being called,
538            // which drops the future after it has yielded the value, which is correct.
539            {
540                let mut data = self.shared.lock();
541                data.debug_dump();
542            }
543            drop(input_fut);
544
545            panic!(
546                r"
547all futures blocked. waiting for the real world? or deadlocked (waiting for each other) ?
548"
549            );
550        });
551
552        value
553    }
554}
555
556//---------- execution - core implementation ----------
557
558impl MockExecutor {
559    /// Keep polling tasks until nothing more can be done
560    ///
561    /// Ie, stop when `awake` is empty and `progressing_until_stalled` is `None`.
562    fn execute_to_completion(&self, mut main_fut: MainFuture) {
563        trace!("MockExecutor execute_to_completion...");
564        loop {
565            self.execute_until_first_stall(main_fut.as_mut());
566
567            // Handle `progressing_until_stalled`
568            let pus_waker = {
569                let mut data = self.shared.lock();
570                let pus = &mut data.progressing_until_stalled;
571                trace!("MockExecutor execute_to_completion PUS={:?}", &pus);
572                let Some(pus) = pus else {
573                    // No progressing_until_stalled, we're actually done.
574                    break;
575                };
576                assert_eq!(
577                    pus.finished, Pending,
578                    "ProgressingUntilStalled finished twice?!"
579                );
580                pus.finished = Ready(());
581
582                // Release the lock temporarily so that ActualWaker::clone doesn't deadlock
583                let waker = pus
584                    .waker
585                    .take()
586                    .expect("ProgressUntilStalledFuture not ever polled!");
587                drop(data);
588                let waker_copy = waker.clone();
589                let mut data = self.shared.lock();
590
591                let pus = &mut data.progressing_until_stalled;
592                if let Some(double) = pus
593                    .as_mut()
594                    .expect("progressing_until_stalled updated under our feet!")
595                    .waker
596                    .replace(waker)
597                {
598                    panic!("double progressing_until_stalled.waker! {double:?}");
599                }
600
601                waker_copy
602            };
603            pus_waker.wake();
604        }
605        trace!("MockExecutor execute_to_completion done");
606    }
607
608    /// Keep polling tasks until `awake` is empty
609    ///
610    /// (Ignores `progressing_until_stalled` - so if one is active,
611    /// will return when all other tasks have blocked.)
612    ///
613    /// # Panics
614    ///
615    /// Might malfunction or panic if called reentrantly
616    fn execute_until_first_stall(&self, main_fut: MainFuture) {
617        trace!("MockExecutor execute_until_first_stall ...");
618
619        assert_eq!(
620            THREAD_DESCRIPTOR.get(),
621            ThreadDescriptor::Foreign,
622            "MockExecutor executor re-entered"
623        );
624        THREAD_DESCRIPTOR.set(ThreadDescriptor::Executor);
625
626        let r = catch_unwind(AssertUnwindSafe(|| self.executor_main_loop(main_fut)));
627
628        THREAD_DESCRIPTOR.set(ThreadDescriptor::Foreign);
629
630        match r {
631            Ok(()) => trace!("MockExecutor execute_until_first_stall done."),
632            Err(e) => {
633                trace!("MockExecutor executor, or async task, panicked!");
634                panic_any(e)
635            }
636        }
637    }
638
639    /// Keep polling tasks until `awake` is empty (inner, executor main loop)
640    ///
641    /// This is only called from [`MockExecutor::execute_until_first_stall`],
642    /// so it could also be called `execute_until_first_stall_inner`.
643    fn executor_main_loop(&self, mut main_fut: MainFuture) {
644        'outer: loop {
645            // Take a `Awake` task off `awake` and make it `Asleep`
646            let (id, mut fut) = 'inner: loop {
647                let mut data = self.shared.lock();
648                let Some(id) = data.schedule() else {
649                    break 'outer;
650                };
651                let Some(task) = data.tasks.get_mut(id) else {
652                    trace!("MockExecutor {id:?} vanished");
653                    continue;
654                };
655                task.state = Asleep(vec![]);
656                let fut = task.fut.take().expect("future missing from task!");
657                break 'inner (id, fut);
658            };
659
660            // Poll the selected task
661            trace!("MockExecutor {id:?} polling...");
662            let waker = ActualWaker::make_waker(&self.shared, id);
663            let mut cx = Context::from_waker(&waker);
664            let r: Either<Poll<()>, IsSubthread> = match &mut fut {
665                TaskFutureInfo::Normal(fut) => Left(fut.poll_unpin(&mut cx)),
666                TaskFutureInfo::Main => Left(main_fut.as_mut().poll(&mut cx)),
667                TaskFutureInfo::Subthread => Right(IsSubthread),
668            };
669
670            // Deal with the returned `Poll`
671            let _fut_drop_late;
672            {
673                let mut data = self.shared.lock();
674                let task = data
675                    .tasks
676                    .get_mut(id)
677                    .expect("task vanished while we were polling it");
678
679                match r {
680                    Left(Pending) => {
681                        trace!("MockExecutor {id:?} -> Pending");
682                        if task.fut.is_some() {
683                            panic!("task reinserted while we polled it?!");
684                        }
685                        // The task might have been woken *by its own poll method*.
686                        // That's why we set it to `Asleep` *earlier* rather than here.
687                        // All we need to do is put the future back.
688                        task.fut = Some(fut);
689                    }
690                    Left(Ready(())) => {
691                        trace!("MockExecutor {id:?} -> Ready");
692                        // Oh, it finished!
693                        // It might be in `awake`, but that's allowed to contain stale tasks,
694                        // so we *don't* need to scan that list and remove it.
695                        data.tasks.remove(id);
696                        // It is important that we don't drop `fut` until we have released
697                        // the data lock, since it is an external type and might try to reenter
698                        // us (eg by calling spawn).  If we do that here, we risk deadlock.
699                        // So, move `fut` to a variable with scope outside the block with `data`.
700                        _fut_drop_late = fut;
701                    }
702                    Right(IsSubthread) => {
703                        trace!("MockExecutor {id:?} -> Ready, waking Subthread");
704                        // Task is a subthread, which has called thread_context_switch
705                        // to switch to us.  We "poll" it by switching back.
706
707                        // Put back `TFI::Subthread`, which was moved out temporarily, above.
708                        task.fut = Some(fut);
709
710                        self.shared.thread_context_switch(
711                            data,
712                            ThreadDescriptor::Executor,
713                            ThreadDescriptor::Subthread(id),
714                        );
715
716                        // Now, if the Subthread still exists, that's because it's switched
717                        // back to us, and is waiting in subthread_block_on_future again.
718                        // Or it might have ended, in which case it's not in `tasks` any more.
719                        // In any case we can go back to scheduling futures.
720                    }
721                }
722            }
723        }
724    }
725}
726
727impl Data {
728    /// Return the next task to run
729    ///
730    /// The task is removed from `awake`, but **`state` is not set to `Asleep`**.
731    /// The caller must restore the invariant!
732    fn schedule(&mut self) -> Option<TaskId> {
733        use SchedulingPolicy as SP;
734        match self.scheduling {
735            SP::Stack => self.awake.pop_back(),
736            SP::Queue => self.awake.pop_front(),
737        }
738    }
739}
740
741impl ActualWaker {
742    /// Obtain a strong reference to the executor's data
743    fn upgrade_data(&self) -> Option<Arc<Shared>> {
744        self.data.upgrade()
745    }
746
747    /// Wake the task corresponding to this `ActualWaker`
748    ///
749    /// This is like `<Self as std::task::Wake>::wake()` but takes `&self`, not `Arc`
750    fn wake(&self) {
751        let Some(data) = self.upgrade_data() else {
752            // The executor is gone!  Don't try to wake.
753            return;
754        };
755        let mut data = data.lock();
756        let data = &mut *data;
757        trace!("MockExecutor {:?} wake", &self.id);
758        let Some(task) = data.tasks.get_mut(self.id) else {
759            return;
760        };
761        task.set_awake(self.id, &mut data.awake);
762    }
763
764    /// Create and return a `Waker` for task `id`
765    fn make_waker(shared: &Arc<Shared>, id: TaskId) -> Waker {
766        ActualWaker {
767            data: Arc::downgrade(shared),
768            id,
769        }
770        .new_waker()
771    }
772}
773
774//---------- "progress until stalled" functionality ----------
775
776impl MockExecutor {
777    /// Run tasks in the current executor until every other task is waiting
778    ///
779    /// # Panics
780    ///
781    /// Might malfunction or panic if more than one such call is running at once.
782    ///
783    /// (Ie, you must `.await` or drop the returned `Future`
784    /// before calling this method again.)
785    ///
786    /// Must be called and awaited within a future being run by `self`.
787    pub fn progress_until_stalled(&self) -> impl Future<Output = ()> + use<> {
788        let mut data = self.shared.lock();
789        assert!(
790            data.progressing_until_stalled.is_none(),
791            "progress_until_stalled called more than once"
792        );
793        trace!("MockExecutor progress_until_stalled...");
794        data.progressing_until_stalled = Some(ProgressingUntilStalled {
795            finished: Pending,
796            waker: None,
797        });
798        ProgressUntilStalledFuture {
799            shared: self.shared.clone(),
800        }
801    }
802}
803
804impl Future for ProgressUntilStalledFuture {
805    type Output = ();
806
807    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<()> {
808        let waker = cx.waker().clone();
809        let mut data = self.shared.lock();
810        let pus = data.progressing_until_stalled.as_mut();
811        trace!("MockExecutor progress_until_stalled polling... {:?}", &pus);
812        let pus = pus.expect("ProgressingUntilStalled missing");
813        pus.waker = Some(waker);
814        pus.finished
815    }
816}
817
818impl Drop for ProgressUntilStalledFuture {
819    fn drop(&mut self) {
820        self.shared.lock().progressing_until_stalled = None;
821    }
822}
823
824//---------- (sub)threads ----------
825
826impl MockExecutor {
827    /// Spawn a "Subthread", for processing in a sync context
828    ///
829    /// `call` will be run on a separate thread, called a "Subthread".
830    ///
831    /// But it will **not run simultaneously** with the executor,
832    /// nor with other Subthreads.
833    /// So Subthreads are somewhat like coroutines.
834    ///
835    /// `call` must be capable of making progress without waiting for any other Subthreads.
836    /// `call` may wait for async futures, using
837    /// [`subthread_block_on_future`](MockExecutor::subthread_block_on_future).
838    ///
839    /// Subthreads may be used for cpubound activity,
840    /// or synchronous IO (such as large volumes of disk activity),
841    /// provided that the synchronous code will reliably make progress,
842    /// without waiting (directly or indirectly) for any async task or Subthread -
843    /// except via `subthread_block_on_future`.
844    ///
845    /// # Subthreads vs raw `std::thread` threads
846    ///
847    /// Programs using `MockExecutor` may use `std::thread` threads directly.
848    /// However, this is not recommended.  There are severe limitations:
849    ///
850    ///  * Only a Subthread can re-enter the async context from sync code:
851    ///    this must be done with
852    ///    using [`subthread_block_on_future`](MockExecutor::subthread_block_on_future).
853    ///    (Re-entering the executor with
854    ///    [`block_on`](tor_rtcompat::ToplevelBlockOn::block_on)
855    ///    is not allowed.)
856    ///  * If async tasks want to suspend waiting for synchronous code,
857    ///    the synchronous code must run on a Subthread.
858    ///    This allows the `MockExecutor` to know when
859    ///    that synchronous code is still making progress.
860    ///    (This is needed for
861    ///    [`progress_until_stalled`](MockExecutor::progress_until_stalled)
862    ///    and the facilities which use it, such as
863    ///    [`MockRuntime::advance_until_stalled`](crate::MockRuntime::advance_until_stalled).)
864    ///  * Subthreads never run in parallel -
865    ///    they only run as scheduled deterministically by the `MockExecutor`.
866    ///    So using Subthreads eliminates a source of test nonndeterminism.
867    ///    (Execution order is still varied due to explicitly varying the scheduling policy.)
868    ///
869    /// # Panics, abuse, and malfunctions
870    ///
871    /// If `call` panics and unwinds, `spawn_subthread` yields `Err`.
872    /// The application code should to do something about it if this happens,
873    /// typically, logging errors, tearing things down, or failing a test case.
874    ///
875    /// If the executor doesn't run, the subthread will not run either, and will remain stuck.
876    /// (So, typically, if the thread supposed to run the executor panics,
877    /// for example because a future or the executor itself panics,
878    /// all the subthreads will become stuck - effectively, they'll be leaked.)
879    ///
880    /// `spawn_subthread` panics if OS thread spawning fails.
881    /// (Like `std::thread::spawn()` does.)
882    ///
883    /// `MockExecutor`s will malfunction or panic if
884    /// any executor invocation method (eg `block_on`) is called on a Subthread.
885    pub fn subthread_spawn<T: Send + 'static>(
886        &self,
887        desc: impl Display,
888        call: impl FnOnce() -> T + Send + 'static,
889    ) -> impl Future<Output = Result<T, Box<dyn Any + Send>>> + Unpin + Send + Sync + 'static {
890        let desc = desc.to_string();
891        let (output_tx, output_rx) = oneshot::channel();
892
893        // NB: we don't know which thread we're on!
894        // In principle we might be on another Subthread.
895        // So we can't context switch here.  That would be very confusing.
896        //
897        // Instead, we prepare the new Subthread as follows:
898        //   - There is a task in the executor
899        //   - The task is ready to be polled, whenever the executor decides to
900        //   - The thread starts running right away, but immediately waits until it is scheduled
901        // See `subthread_entrypoint`.
902
903        {
904            let mut data = self.shared.lock();
905            let id = data.insert_task(desc.clone(), TaskFutureInfo::Subthread);
906
907            let _: std::thread::JoinHandle<()> = std::thread::Builder::new()
908                .name(desc)
909                .spawn({
910                    let shared = self.shared.clone();
911                    move || shared.subthread_entrypoint(id, call, output_tx)
912                })
913                .expect("spawn failed");
914        }
915
916        output_rx.map(|r| {
917            r.unwrap_or_else(|_: Canceled| panic!("Subthread cancelled but should be impossible!"))
918        })
919    }
920
921    /// Call an async `Future` from a Subthread
922    ///
923    /// Blocks the Subthread, and arranges to run async tasks,
924    /// including `fut`, until `fut` completes.
925    ///
926    /// `fut` is polled on the executor thread, not on the Subthread.
927    /// (We may change that in the future, allowing passing a non-`Send` future.)
928    ///
929    /// # Panics, abuse, and malfunctions
930    ///
931    /// `subthread_block_on_future` will malfunction or panic
932    /// if called on a thread that isn't a Subthread from the same `MockExecutor`
933    /// (ie a thread made with [`spawn_subthread`](MockExecutor::subthread_spawn)).
934    ///
935    /// If `fut` itself panics, the executor will panic.
936    ///
937    /// If the executor isn't running, `subthread_block_on_future` will hang indefinitely.
938    /// See `spawn_subthread`.
939    pub fn subthread_block_on_future<T: Send + 'static>(&self, fut: impl Future<Output = T>) -> T {
940        let id = match THREAD_DESCRIPTOR.get() {
941            ThreadDescriptor::Subthread(id) => id,
942            ThreadDescriptor::Executor => {
943                panic!("subthread_block_on_future called from MockExecutor thread (async task?)")
944            }
945            ThreadDescriptor::Foreign => panic!(
946                "subthread_block_on_future called on foreign thread (not spawned with spawn_subthread)"
947            ),
948        };
949        trace!("MockExecutor thread {id:?}, subthread_block_on_future...");
950        let mut fut = pin!(fut);
951
952        // We yield once before the first poll, and once after Ready, to shake up the
953        // execution order a bit, depending on the scheduling policy.
954        let yield_ = |set_awake| self.shared.subthread_yield(id, set_awake);
955        yield_(Some(SetAwake));
956
957        let ret = loop {
958            // Poll the provided future
959            trace!("MockExecutor thread {id:?}, s.t._block_on_future polling...");
960            let waker = ActualWaker::make_waker(&self.shared, id);
961            let mut cx = Context::from_waker(&waker);
962            let r: Poll<T> = fut.as_mut().poll(&mut cx);
963
964            if let Ready(r) = r {
965                trace!("MockExecutor thread {id:?}, s.t._block_on_future poll -> Ready");
966                break r;
967            }
968
969            // Pending.  Switch back to the executor thread.
970            // When the future becomes ready, the Waker will be woken, waking the task,
971            // so that the executor will "poll" us again.
972            trace!("MockExecutor thread {id:?}, s.t._block_on_future poll -> Pending");
973
974            yield_(None);
975        };
976
977        yield_(Some(SetAwake));
978
979        trace!("MockExecutor thread {id:?}, subthread_block_on_future complete.");
980
981        ret
982    }
983}
984
985impl Shared {
986    /// Main entrypoint function for a Subthread
987    ///
988    /// Entered on a new `std::thread` thread created by
989    /// [`subthread_spawn`](MockExecutor::subthread_spawn).
990    ///
991    /// When `call` completes, sends its returned value `T` to `output_tx`.
992    fn subthread_entrypoint<T: Send + 'static>(
993        self: Arc<Self>,
994        id: TaskId,
995        call: impl FnOnce() -> T + Send + 'static,
996        output_tx: oneshot::Sender<Result<T, Box<dyn Any + Send>>>,
997    ) {
998        THREAD_DESCRIPTOR.set(ThreadDescriptor::Subthread(id));
999        trace!("MockExecutor thread {id:?}, entrypoint");
1000
1001        // We start out Awake, but we wait for the executor to tell us to run.
1002        // This will be done the first time the task is "polled".
1003        {
1004            let data = self.lock();
1005            self.thread_context_switch_waitfor_instruction_to_run(
1006                data,
1007                ThreadDescriptor::Subthread(id),
1008            );
1009        }
1010
1011        trace!("MockExecutor thread {id:?}, entering user code");
1012
1013        // Run the user's actual thread function.
1014        // This will typically reenter us via subthread_block_on_future.
1015        let ret = catch_unwind(AssertUnwindSafe(call));
1016
1017        trace!("MockExecutor thread {id:?}, completed user code");
1018
1019        // This makes the return value from subthread_spawn ready.
1020        // It will be polled by the executor in due course, presumably.
1021
1022        output_tx.send(ret).unwrap_or_else(
1023            #[allow(clippy::unnecessary_lazy_evaluations)]
1024            |_| {}, // receiver dropped, maybe executor dropped or something?
1025        );
1026
1027        {
1028            let mut data = self.lock();
1029
1030            // Never poll this task again (so never schedule this thread)
1031            let _: Task = data.tasks.remove(id).expect("Subthread task vanished!");
1032
1033            // Tell the executor it is scheduled now.
1034            // We carry on exiting, in parallel (holding the data lock).
1035            self.thread_context_switch_send_instruction_to_run(
1036                &mut data,
1037                ThreadDescriptor::Subthread(id),
1038                ThreadDescriptor::Executor,
1039            );
1040        }
1041    }
1042
1043    /// Yield back to the executor from a subthread
1044    ///
1045    /// Checks that things are in order
1046    /// (in particular, that this task is in the data structure as a subhtread)
1047    /// and switches to the executor thread.
1048    ///
1049    /// The caller must arrange that the task gets woken.
1050    ///
1051    /// With [`SetAwake`], sets our task awake, so that we'll be polled
1052    /// again as soon as we get to the top of the executor's queue.
1053    /// Otherwise, we'll be reentered after someone wakes a [`Waker`] for the task.
1054    fn subthread_yield(&self, us: TaskId, set_awake: Option<SetAwake>) {
1055        let mut data = self.lock();
1056        {
1057            let data = &mut *data;
1058            let task = data.tasks.get_mut(us).expect("Subthread task vanished!");
1059            match &task.fut {
1060                Some(TaskFutureInfo::Subthread) => {}
1061                other => panic!("subthread_block_on_future but TFI {other:?}"),
1062            };
1063            if let Some(SetAwake) = set_awake {
1064                task.set_awake(us, &mut data.awake);
1065            }
1066        }
1067        self.thread_context_switch(
1068            data,
1069            ThreadDescriptor::Subthread(us),
1070            ThreadDescriptor::Executor,
1071        );
1072    }
1073
1074    /// Switch from (sub)thread `us` to (sub)thread `them`
1075    ///
1076    /// Returns when someone calls `thread_context_switch(.., us)`.
1077    fn thread_context_switch(
1078        &self,
1079        mut data: MutexGuard<Data>,
1080        us: ThreadDescriptor,
1081        them: ThreadDescriptor,
1082    ) {
1083        trace!("MockExecutor thread {us:?}, switching to {them:?}");
1084        self.thread_context_switch_send_instruction_to_run(&mut data, us, them);
1085        self.thread_context_switch_waitfor_instruction_to_run(data, us);
1086    }
1087
1088    /// Instruct the (sub)thread `them` to run
1089    ///
1090    /// Update `thread_to_run`, which will wake up `them`'s
1091    /// call to `thread_context_switch_waitfor_instruction_to_run`.
1092    ///
1093    /// Must be called from (sub)thread `us`.
1094    /// Part of `thread_context_switch`, not normally called directly.
1095    fn thread_context_switch_send_instruction_to_run(
1096        &self,
1097        data: &mut MutexGuard<Data>,
1098        us: ThreadDescriptor,
1099        them: ThreadDescriptor,
1100    ) {
1101        assert_eq!(data.thread_to_run, us);
1102        data.thread_to_run = them;
1103        self.thread_condvar.notify_all();
1104    }
1105
1106    /// Await an instruction for this thread, `us`, to run
1107    ///
1108    /// Waits for `thread_to_run` to be `us`,
1109    /// waiting for `thread_condvar` as necessary.
1110    ///
1111    /// Part of `thread_context_switch`, not normally called directly.
1112    fn thread_context_switch_waitfor_instruction_to_run(
1113        &self,
1114        data: MutexGuard<Data>,
1115        us: ThreadDescriptor,
1116    ) {
1117        #[allow(let_underscore_lock)]
1118        let _: MutexGuard<_> = self
1119            .thread_condvar
1120            .wait_while(data, |data| {
1121                let live = data.thread_to_run;
1122                let resume = live == us;
1123                if resume {
1124                    trace!("MockExecutor thread {us:?}, resuming");
1125                } else {
1126                    trace!("MockExecutor thread {us:?}, waiting for {live:?}");
1127                }
1128                // We're in `.wait_while`, not `.wait_until`.  Confusing.
1129                !resume
1130            })
1131            .expect("data lock poisoned");
1132    }
1133}
1134
1135//---------- ancillary and convenience functions ----------
1136
1137/// Trait to let us assert at compile time that something is nicely `Sync` etc.
1138#[allow(dead_code)] // yes, we don't *use* anything from this trait
1139trait EnsureSyncSend: Sync + Send + 'static {}
1140impl EnsureSyncSend for ActualWaker {}
1141impl EnsureSyncSend for MockExecutor {}
1142
1143impl MockExecutor {
1144    /// Return the number of tasks running in this executor
1145    ///
1146    /// One possible use is for a test case to check that task(s)
1147    /// that ought to have exited, have indeed done so.
1148    ///
1149    /// In the usual case, the answer will be at least 1,
1150    /// because it counts the future passed to
1151    /// [`block_on`](MockExecutor::block_on)
1152    /// (perhaps via [`MockRuntime::test_with_various`](crate::MockRuntime::test_with_various)).
1153    pub fn n_tasks(&self) -> usize {
1154        self.shared.lock().tasks.len()
1155    }
1156}
1157
1158impl Shared {
1159    /// Lock and obtain the guard
1160    ///
1161    /// Convenience method which panics on poison
1162    fn lock(&self) -> MutexGuard<Data> {
1163        self.data.lock().expect("data lock poisoned")
1164    }
1165}
1166
1167impl Task {
1168    /// Set task `id` to `Awake` and arrange that it will be polled.
1169    fn set_awake(&mut self, id: TaskId, data_awake: &mut VecDeque<TaskId>) {
1170        match self.state {
1171            Awake => {}
1172            Asleep(_) => {
1173                self.state = Awake;
1174                data_awake.push_back(id);
1175            }
1176        }
1177    }
1178}
1179
1180//---------- ActualWaker as RawWaker ----------
1181
1182/// Using [`ActualWaker`] in a [`RawWaker`]
1183///
1184/// We need to make a
1185/// [`Waker`] (the safe, type-erased, waker, used by actual futures)
1186/// which contains an
1187/// [`ActualWaker`] (our actual waker implementation, also safe).
1188///
1189/// `std` offers `Waker::from<Arc<impl Wake>>`.
1190/// But we want a bespoke `Clone` implementation, so we don't want to use `Arc`.
1191///
1192/// So instead, we implement the `RawWaker` API in terms of `ActualWaker`.
1193/// We keep the `ActualWaker` in a `Box`, and actually `clone` it (and the `Box`).
1194///
1195/// SAFETY
1196///
1197///  * The data pointer is `Box::<ActualWaker>::into_raw()`
1198///  * We share these when we clone
1199///  * No-one is allowed `&mut ActualWaker` unless there are no other clones
1200///  * So we may make references `&ActualWaker`
1201impl ActualWaker {
1202    /// Wrap up an [`ActualWaker`] as a type-erased [`Waker`] for passing to futures etc.
1203    fn new_waker(self) -> Waker {
1204        unsafe { Waker::from_raw(self.raw_new()) }
1205    }
1206
1207    /// Helper: wrap up an [`ActualWaker`] as a [`RawWaker`].
1208    fn raw_new(self) -> RawWaker {
1209        let self_: Box<ActualWaker> = self.into();
1210        let self_: *mut ActualWaker = Box::into_raw(self_);
1211        let self_: *const () = self_ as _;
1212        RawWaker::new(self_, &RAW_WAKER_VTABLE)
1213    }
1214
1215    /// Implementation of [`RawWakerVTable`]'s `clone`
1216    unsafe fn raw_clone(self_: *const ()) -> RawWaker {
1217        unsafe {
1218            let self_: *const ActualWaker = self_ as _;
1219            let self_: &ActualWaker = self_.as_ref().unwrap_unchecked();
1220            let copy: ActualWaker = self_.clone();
1221            copy.raw_new()
1222        }
1223    }
1224
1225    /// Implementation of [`RawWakerVTable`]'s `wake`
1226    unsafe fn raw_wake(self_: *const ()) {
1227        unsafe {
1228            Self::raw_wake_by_ref(self_);
1229            Self::raw_drop(self_);
1230        }
1231    }
1232
1233    /// Implementation of [`RawWakerVTable`]'s `wake_ref_by`
1234    unsafe fn raw_wake_by_ref(self_: *const ()) {
1235        unsafe {
1236            let self_: *const ActualWaker = self_ as _;
1237            let self_: &ActualWaker = self_.as_ref().unwrap_unchecked();
1238            self_.wake();
1239        }
1240    }
1241
1242    /// Implementation of [`RawWakerVTable`]'s `drop`
1243    unsafe fn raw_drop(self_: *const ()) {
1244        unsafe {
1245            let self_: *mut ActualWaker = self_ as _;
1246            let self_: Box<ActualWaker> = Box::from_raw(self_);
1247            drop(self_);
1248        }
1249    }
1250}
1251
1252/// vtable for `Box<ActualWaker>` as `RawWaker`
1253//
1254// This ought to be in the impl block above, but
1255//   "associated `static` items are not allowed"
1256static RAW_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
1257    ActualWaker::raw_clone,
1258    ActualWaker::raw_wake,
1259    ActualWaker::raw_wake_by_ref,
1260    ActualWaker::raw_drop,
1261);
1262
1263//---------- Sleep location tracking and dumping ----------
1264
1265/// We record "where a future went to sleep" as (just) a backtrace
1266///
1267/// This type alias allows us to mock `Backtrace` for miri.
1268/// (It also insulates from future choices about sleep location representation.0
1269#[cfg(not(miri))]
1270type SleepLocation = Backtrace;
1271
1272impl Data {
1273    /// Dump tasks and their sleep location backtraces
1274    fn dump_backtraces(&self, f: &mut fmt::Formatter) -> fmt::Result {
1275        for (id, task) in self.tasks.iter() {
1276            let prefix = |f: &mut fmt::Formatter| write!(f, "{id:?}={task:?}: ");
1277            match &task.state {
1278                Awake => {
1279                    prefix(f)?;
1280                    writeln!(f, "awake")?;
1281                }
1282                Asleep(locs) => {
1283                    let n = locs.len();
1284                    for (i, loc) in locs.iter().enumerate() {
1285                        prefix(f)?;
1286                        writeln!(f, "asleep, backtrace {i}/{n}:\n{loc}",)?;
1287                    }
1288                    if n == 0 {
1289                        prefix(f)?;
1290                        writeln!(f, "asleep, no backtraces, Waker never cloned, stuck!",)?;
1291                    }
1292                }
1293            }
1294        }
1295        writeln!(
1296            f,
1297            "\nNote: there might be spurious traces, see docs for MockExecutor::debug_dump\n"
1298        )?;
1299        Ok(())
1300    }
1301}
1302
1303/// Track sleep locations via `<Waker as Clone>`.
1304///
1305/// See [`MockExecutor::debug_dump`] for the explanation.
1306impl Clone for ActualWaker {
1307    fn clone(&self) -> Self {
1308        let id = self.id;
1309
1310        if let Some(data) = self.upgrade_data() {
1311            // If the executor is gone, there is nothing to adjust
1312            let mut data = data.lock();
1313            if let Some(task) = data.tasks.get_mut(self.id) {
1314                match &mut task.state {
1315                    Awake => trace!("MockExecutor cloned waker for awake task {id:?}"),
1316                    Asleep(locs) => locs.push(SleepLocation::force_capture()),
1317                }
1318            } else {
1319                trace!("MockExecutor cloned waker for dead task {id:?}");
1320            }
1321        }
1322
1323        ActualWaker {
1324            data: self.data.clone(),
1325            id,
1326        }
1327    }
1328}
1329
1330//---------- API for full debug dump ----------
1331
1332/// Debugging dump of a `MockExecutor`'s state
1333///
1334/// Returned by [`MockExecutor::as_debug_dump`]
1335//
1336// Existence implies backtraces have been resolved
1337//
1338// We use `Either` so that we can also use this internally when we have &mut Data.
1339pub struct DebugDump<'a>(Either<&'a Data, MutexGuard<'a, Data>>);
1340
1341impl MockExecutor {
1342    /// Dump the executor's state including backtraces of waiting tasks, to stderr
1343    ///
1344    /// This is considerably more extensive than simply
1345    /// `MockExecutor as Debug`.
1346    ///
1347    /// (This is a convenience method, which wraps
1348    /// [`MockExecutor::as_debug_dump()`].
1349    ///
1350    /// ### Backtrace salience (possible spurious traces)
1351    ///
1352    /// **Summary**
1353    ///
1354    /// The technique used to capture backtraces when futures sleep is not 100% exact.
1355    /// It will usually show all the actual sleeping sites,
1356    /// but it might also show other backtraces which were part of
1357    /// the implementation of some complex relevant future.
1358    ///
1359    /// **Details**
1360    ///
1361    /// When a future's implementation wants to sleep,
1362    /// it needs to record the [`Waker`] (from the [`Context`])
1363    /// so that the "other end" can call `.wake()` on it later,
1364    /// when the future should be woken.
1365    ///
1366    /// Since `Context.waker()` gives `&Waker`, borrowed from the `Context`,
1367    /// the future must clone the `Waker`,
1368    /// and it must do so in within the `poll()` call.
1369    ///
1370    /// A future which is waiting in a `select!` will typically
1371    /// show multiple traces, one for each branch.
1372    /// But,
1373    /// if a future sleeps on one thing, and then when polled again later,
1374    /// sleeps on something different, without waking up in between,
1375    /// both backtrace locations will be shown.
1376    /// And,
1377    /// a complicated future contraption *might* clone the `Waker` more times.
1378    /// So not every backtrace will necessarily be informative.
1379    ///
1380    /// ### Panics
1381    ///
1382    /// Panics on write errors.
1383    pub fn debug_dump(&self) {
1384        self.as_debug_dump().to_stderr();
1385    }
1386
1387    /// Dump the executor's state including backtraces of waiting tasks
1388    ///
1389    /// This is considerably more extensive than simply
1390    /// `MockExecutor as Debug`.
1391    ///
1392    /// Returns an object for formatting with [`Debug`].
1393    /// To simply print the dump to stderr (eg in a test),
1394    /// use [`.debug_dump()`](MockExecutor::debug_dump).
1395    ///
1396    /// **Backtrace salience (possible spurious traces)** -
1397    /// see [`.debug_dump()`](MockExecutor::debug_dump).
1398    pub fn as_debug_dump(&self) -> DebugDump {
1399        let data = self.shared.lock();
1400        DebugDump(Right(data))
1401    }
1402}
1403
1404impl Data {
1405    /// Convenience function: dump including backtraces, to stderr
1406    fn debug_dump(&mut self) {
1407        DebugDump(Left(self)).to_stderr();
1408    }
1409}
1410
1411impl DebugDump<'_> {
1412    /// Convenience function: dump tasks and backtraces to stderr
1413    #[allow(clippy::wrong_self_convention)] // "to_stderr" doesn't mean "convert to stderr"
1414    fn to_stderr(self) {
1415        write!(io::stderr().lock(), "{:?}", self)
1416            .unwrap_or_else(|e| error_report!(e, "failed to write debug dump to stderr"));
1417    }
1418}
1419
1420//---------- bespoke Debug impls ----------
1421
1422impl Debug for DebugDump<'_> {
1423    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1424        let self_: &Data = &self.0;
1425
1426        writeln!(f, "MockExecutor state:\n{self_:#?}")?;
1427        writeln!(f, "MockExecutor task dump:")?;
1428        self_.dump_backtraces(f)?;
1429
1430        Ok(())
1431    }
1432}
1433
1434// See `impl Debug for Data` for notes on the output
1435impl Debug for Task {
1436    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1437        let Task { desc, state, fut } = self;
1438        write!(f, "{:?}", desc)?;
1439        write!(f, "=")?;
1440        match fut {
1441            None => write!(f, "P")?,
1442            Some(TaskFutureInfo::Normal(_)) => write!(f, "f")?,
1443            Some(TaskFutureInfo::Main) => write!(f, "m")?,
1444            Some(TaskFutureInfo::Subthread) => write!(f, "T")?,
1445        }
1446        match state {
1447            Awake => write!(f, "W")?,
1448            Asleep(locs) => write!(f, "s{}", locs.len())?,
1449        };
1450        Ok(())
1451    }
1452}
1453
1454/// Helper: `Debug`s as a list of tasks, given the `Data` for lookups and a list of the ids
1455///
1456/// `Task`s in `Data` are printed as `Ti(ID)"SPEC"=FLAGS"`.
1457///
1458/// `FLAGS` are:
1459///
1460///  * `T`: this task is for a Subthread (from subthread_spawn).
1461///  * `P`: this task is being polled (its `TaskFutureInfo` is absent)
1462///  * `f`: this is a normal task with a future and its future is present in `Data`
1463///  * `m`: this is the main task from `block_on`
1464///
1465///  * `W`: the task is awake
1466///  * `s<n>`: the task is asleep, and `<n>` is the number of recorded sleeping locations
1467//
1468// We do it this way because the naive dump from derive is very expansive
1469// and makes it impossible to see the wood for the trees.
1470// This very compact representation it easier to find a task of interest in the output.
1471//
1472// This is implemented in `impl Debug for Task`.
1473//
1474//
1475// rustc doesn't think automatically-derived Debug impls count for whether a thing is used.
1476// This has caused quite some fallout.  https://github.com/rust-lang/rust/pull/85200
1477// I think derive_more emits #[automatically_derived], so that even though we use this
1478// in our Debug impl, that construction is unused.
1479#[allow(dead_code)]
1480struct DebugTasks<'d, F>(&'d Data, F);
1481
1482// See `impl Debug for Data` for notes on the output
1483impl<F, I> Debug for DebugTasks<'_, F>
1484where
1485    F: Fn() -> I,
1486    I: Iterator<Item = TaskId>,
1487{
1488    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1489        let DebugTasks(data, ids) = self;
1490        for (id, delim) in izip!(ids(), chain!(iter::once(""), iter::repeat(" ")),) {
1491            write!(f, "{delim}{id:?}")?;
1492            match data.tasks.get(id) {
1493                None => write!(f, "-")?,
1494                Some(task) => write!(f, "={task:?}")?,
1495            }
1496        }
1497        Ok(())
1498    }
1499}
1500
1501/// Mock `Backtrace` for miri
1502///
1503/// See also the not-miri `type SleepLocation`, alias above.
1504#[cfg(miri)]
1505mod miri_sleep_location {
1506    #[derive(Debug, derive_more::Display)]
1507    #[display("<SleepLocation>")]
1508    pub(super) struct SleepLocation {}
1509
1510    impl SleepLocation {
1511        pub(super) fn force_capture() -> Self {
1512            SleepLocation {}
1513        }
1514    }
1515}
1516#[cfg(miri)]
1517use miri_sleep_location::SleepLocation;
1518
1519#[cfg(test)]
1520mod test {
1521    // @@ begin test lint list maintained by maint/add_warning @@
1522    #![allow(clippy::bool_assert_comparison)]
1523    #![allow(clippy::clone_on_copy)]
1524    #![allow(clippy::dbg_macro)]
1525    #![allow(clippy::mixed_attributes_style)]
1526    #![allow(clippy::print_stderr)]
1527    #![allow(clippy::print_stdout)]
1528    #![allow(clippy::single_char_pattern)]
1529    #![allow(clippy::unwrap_used)]
1530    #![allow(clippy::unchecked_time_subtraction)]
1531    #![allow(clippy::useless_vec)]
1532    #![allow(clippy::needless_pass_by_value)]
1533    #![allow(clippy::string_slice)] // See arti#2571
1534    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
1535    use super::*;
1536    use futures::channel::mpsc;
1537    use futures::{SinkExt as _, StreamExt as _};
1538    use strum::IntoEnumIterator;
1539    use tracing::info;
1540
1541    #[cfg(not(miri))] // trace! asks for the time, which miri doesn't support
1542    use tracing_test::traced_test;
1543
1544    fn various_mock_executors() -> impl Iterator<Item = MockExecutor> {
1545        // This duplicates the part of the logic in MockRuntime::test_with_various which
1546        // relates to MockExecutor, because we don't have a MockRuntime::builder.
1547        // The only parameter to MockExecutor is its scheduling policy, so this seems fine.
1548        SchedulingPolicy::iter().map(|scheduling| {
1549            eprintln!("===== MockExecutor::with_scheduling({scheduling:?}) =====");
1550            MockExecutor::with_scheduling(scheduling)
1551        })
1552    }
1553
1554    #[cfg_attr(not(miri), traced_test)]
1555    #[test]
1556    fn simple() {
1557        let runtime = MockExecutor::default();
1558        let val = runtime.block_on(async { 42 });
1559        assert_eq!(val, 42);
1560    }
1561
1562    #[cfg_attr(not(miri), traced_test)]
1563    #[test]
1564    fn stall() {
1565        let runtime = MockExecutor::default();
1566
1567        runtime.block_on({
1568            let runtime = runtime.clone();
1569            async move {
1570                const N: usize = 3;
1571                let (mut txs, mut rxs): (Vec<_>, Vec<_>) =
1572                    (0..N).map(|_| mpsc::channel::<usize>(5)).unzip();
1573
1574                let mut rx_n = rxs.pop().unwrap();
1575
1576                for (i, mut rx) in rxs.into_iter().enumerate() {
1577                    runtime.spawn_identified(i, {
1578                        let mut txs = txs.clone();
1579                        async move {
1580                            loop {
1581                                eprintln!("task {i} rx...");
1582                                let v = rx.next().await.unwrap();
1583                                let nv = v + 1;
1584                                eprintln!("task {i} rx {v}, tx {nv}");
1585                                let v = nv;
1586                                txs[v].send(v).await.unwrap();
1587                            }
1588                        }
1589                    });
1590                }
1591
1592                dbg!();
1593                #[allow(deprecated)] // TODO(#2386)
1594                let _: mpsc::TryRecvError = rx_n.try_next().unwrap_err();
1595
1596                dbg!();
1597                runtime.progress_until_stalled().await;
1598
1599                dbg!();
1600                #[allow(deprecated)] // TODO(#2386)
1601                let _: mpsc::TryRecvError = rx_n.try_next().unwrap_err();
1602
1603                dbg!();
1604                txs[0].send(0).await.unwrap();
1605
1606                dbg!();
1607                runtime.progress_until_stalled().await;
1608
1609                dbg!();
1610                let r = rx_n.next().await;
1611                assert_eq!(r, Some(N - 1));
1612
1613                dbg!();
1614                #[allow(deprecated)] // TODO(#2386)
1615                let _: mpsc::TryRecvError = rx_n.try_next().unwrap_err();
1616
1617                runtime.spawn_identified("tx", {
1618                    let txs = txs.clone();
1619                    async {
1620                        eprintln!("sending task...");
1621                        for (i, mut tx) in txs.into_iter().enumerate() {
1622                            eprintln!("sending 0 to {i}...");
1623                            tx.send(0).await.unwrap();
1624                        }
1625                        eprintln!("sending task done");
1626                    }
1627                });
1628
1629                runtime.debug_dump();
1630
1631                for i in 0..txs.len() {
1632                    eprintln!("main {i} wait stall...");
1633                    runtime.progress_until_stalled().await;
1634                    eprintln!("main {i} rx wait...");
1635                    let r = rx_n.next().await;
1636                    eprintln!("main {i} rx = {r:?}");
1637                    assert!(r == Some(0) || r == Some(N - 1));
1638                }
1639
1640                eprintln!("finishing...");
1641                runtime.progress_until_stalled().await;
1642                eprintln!("finished.");
1643            }
1644        });
1645    }
1646
1647    #[cfg_attr(not(miri), traced_test)]
1648    #[test]
1649    fn spawn_blocking() {
1650        let runtime = MockExecutor::default();
1651
1652        runtime.block_on({
1653            let runtime = runtime.clone();
1654            async move {
1655                let thr_1 = runtime.spawn_blocking(|| 42);
1656                let thr_2 = runtime.spawn_blocking(|| 99);
1657
1658                assert_eq!(thr_2.await, 99);
1659                assert_eq!(thr_1.await, 42);
1660            }
1661        });
1662    }
1663
1664    #[cfg_attr(not(miri), traced_test)]
1665    #[test]
1666    fn drop_reentrancy() {
1667        // Check that dropping a completed task future is done *outside* the data lock.
1668        // Involves a contrived future whose Drop impl reenters the executor.
1669        //
1670        // If `_fut_drop_late = fut` in execute_until_first_stall (the main loop)
1671        // is replaced with `drop(fut)` (dropping the future at the wrong moment),
1672        // we do indeed get deadlock, so this test case is working.
1673
1674        struct ReentersOnDrop {
1675            runtime: MockExecutor,
1676        }
1677        impl Future for ReentersOnDrop {
1678            type Output = ();
1679            fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<()> {
1680                Poll::Ready(())
1681            }
1682        }
1683        impl Drop for ReentersOnDrop {
1684            fn drop(&mut self) {
1685                self.runtime
1686                    .spawn_identified("dummy", futures::future::ready(()));
1687            }
1688        }
1689
1690        for runtime in various_mock_executors() {
1691            runtime.block_on(async {
1692                runtime.spawn_identified("trapper", {
1693                    let runtime = runtime.clone();
1694                    ReentersOnDrop { runtime }
1695                });
1696            });
1697        }
1698    }
1699
1700    #[cfg_attr(not(miri), traced_test)]
1701    #[test]
1702    fn subthread_oneshot() {
1703        for runtime in various_mock_executors() {
1704            runtime.block_on(async {
1705                let (tx, rx) = oneshot::channel();
1706                info!("spawning subthread");
1707                let thr = runtime.subthread_spawn("thr1", {
1708                    let runtime = runtime.clone();
1709                    move || {
1710                        info!("subthread_block_on_future...");
1711                        let i = runtime.subthread_block_on_future(rx).unwrap();
1712                        info!("subthread_block_on_future => {i}");
1713                        i + 1
1714                    }
1715                });
1716                info!("main task sending");
1717                tx.send(12).unwrap();
1718                info!("main task sent");
1719                let r = thr.await.unwrap();
1720                info!("main task thr => {r}");
1721                assert_eq!(r, 13);
1722            });
1723        }
1724    }
1725
1726    #[cfg_attr(not(miri), traced_test)]
1727    #[test]
1728    fn subthread_pingpong() {
1729        for runtime in various_mock_executors() {
1730            runtime.block_on(async {
1731                let (mut i_tx, mut i_rx) = mpsc::channel(1);
1732                let (mut o_tx, mut o_rx) = mpsc::channel(1);
1733                info!("spawning subthread");
1734                let thr = runtime.subthread_spawn("thr", {
1735                    let runtime = runtime.clone();
1736                    move || {
1737                        while let Some(i) = {
1738                            info!("thread receiving ...");
1739                            runtime.subthread_block_on_future(i_rx.next())
1740                        } {
1741                            let o = i + 12;
1742                            info!("thread received {i}, sending {o}");
1743                            runtime.subthread_block_on_future(o_tx.send(o)).unwrap();
1744                            info!("thread sent {o}");
1745                        }
1746                        info!("thread exiting");
1747                        42
1748                    }
1749                });
1750                for i in 0..2 {
1751                    info!("main task sending {i}");
1752                    i_tx.send(i).await.unwrap();
1753                    info!("main task sent {i}");
1754                    let o = o_rx.next().await.unwrap();
1755                    info!("main task recv => {o}");
1756                    assert_eq!(o, i + 12);
1757                }
1758                info!("main task dropping sender");
1759                drop(i_tx);
1760                info!("main task awaiting thread");
1761                let r = thr.await.unwrap();
1762                info!("main task complete");
1763                assert_eq!(r, 42);
1764            });
1765        }
1766    }
1767}