crossflow 0.0.6

/*
 * Copyright (C) 2024 Open Source Robotics Foundation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
*/

use bevy_ecs::{
    hierarchy::ChildOf,
    prelude::{Bundle, Commands, Component, Entity, Event},
};

use std::future::Future;

use crate::{
    AsMapOnce, Cancellable, IntoAsyncMapOnce, IntoBlockingMapOnce, Outcome, Promise, ProvideOnce,
    Sendish, StreamPack, StreamTargetMap, UnusedTarget,
};

mod detach;
pub(crate) use detach::*;

mod finished;
pub(crate) use finished::*;

mod insert;
pub(crate) use insert::*;

mod internal;
pub(crate) use internal::*;

mod map;
pub(crate) use map::*;

mod push;
pub(crate) use push::*;

mod capture;
pub(crate) use capture::*;

mod send_event;
pub(crate) use send_event::*;

mod store;
pub(crate) use store::*;

mod taken;
pub(crate) use taken::*;

/// A series of one or more workflow execution sessions. You can use a [`Series`]
/// to chain the output of one workflow session into the input of another workflow
/// session, and then eventually receive the final output of the whole series.
///
/// You can begin creating a series using [`RequestExt`][crate::RequestExt].
///
/// A series like this can only be a linear sequence---it does not support
/// conditional branching or cycles. If you want a more complex structure than
/// a linear sequence, you will need to spawn a workflow with the structure
/// that you want, and then issue a request into that workflow.
///
/// A series is a one-time-use sequence of sessions. You will have to reconstruct
/// the series each time you want to call it.
///
/// Note that the entire series of sessions will automatically cancel if you
/// drop the final promise of the last session in the series. To prevent that
/// from happening, you can use [`Series::detach`] to lock in the execution of
/// the series (or a subset of the series) no matter what happens downstream.
pub struct Series<'w, 's, 'a, Response, Streams> {
    pub(crate) source: Entity,
    pub(crate) target: Entity,
    pub(crate) commands: &'a mut Commands<'w, 's>,
    pub(crate) _ignore: std::marker::PhantomData<fn(Response, Streams)>,
}

impl<'w, 's, 'a, Response, Streams> Series<'w, 's, 'a, Response, Streams>
where
    Response: 'static + Send + Sync,
    Streams: StreamPack,
{
    /// Keep executing the series up to here even if a downstream dependent gets
    /// dropped. If you continue building the series from this point, then the
    /// later sessions will not be affected by this use of `.detach()` and may
    /// be dropped if its downstream dependent gets dropped.
    ///
    /// Dependency gets dropped in the following situations:
    ///
    /// | Operation                                                 | Drop condition                                        |
    /// |-----------------------------------------------------------|-------------------------------------------------------|
    /// | [`Self::capture`] <br> [`Self::outcome`]                  | The [`Outcome`] is dropped.                           |
    /// | [`Self::store`] <br> [`Self::push`] <br> [`Self::insert`] | The target entity of the operation is despawned.      |
    /// | [`Self::detach`] <br> [`Self::send_event`]                | This will never be dropped                            |
    /// | Using none of the above                                   | The series will immediately be dropped during a flush, so it will never be run at all. <br> This will also push an error into [`UnhandledErrors`](crate::UnhandledErrors). |
    pub fn detach(self) -> Series<'w, 's, 'a, Response, Streams> {
        self.commands.queue(Detach {
            target: self.target,
        });
        self
    }

    /// This is the session ID of the last request so far in the series.
    pub fn session_id(&self) -> Entity {
        self.source
    }

    /// Capture the outcome of the series and all the stream data of the final
    /// provider.
    #[must_use]
    pub fn capture(self) -> Capture<Response, Streams> {
        let target = self.target;
        let mut map = StreamTargetMap::default();
        let stream_receivers = Streams::take_streams(target, &mut map, self.commands);
        self.commands.entity(self.source).insert(map);

        let (outcome, capture_outcome) = Outcome::new();
        self.send_outcome(capture_outcome);

        Capture {
            outcome,
            streams: stream_receivers,
            session: target,
        }
    }

    /// Take the data that comes out of the request, including both the response
    /// and the streams.
    #[must_use]
    #[deprecated(since = "0.0.6", note = "Use .capture() instead")]
    pub fn take(self) -> Recipient<Response, Streams> {
        let (response_sender, response_promise) = Promise::<Response>::new();
        self.commands.queue(AddExecution::new(
            Some(self.source),
            self.target,
            TakenResponse::<Response>::new(response_sender),
        ));
        let mut map = StreamTargetMap::default();
        let stream_receivers = Streams::take_streams(self.target, &mut map, self.commands);
        self.commands.entity(self.source).insert(map);

        Recipient {
            response: response_promise,
            streams: stream_receivers,
            session: self.target,
        }
    }

    /// Capture only the outcome (response) of the series.
    #[must_use]
    pub fn outcome(self) -> Outcome<Response> {
        let (outcome, capture) = Outcome::new();
        self.send_outcome(capture);
        outcome
    }

    /// Take only the response data that comes out of the request.
    #[deprecated(since = "0.0.6", note = "Use .outcome() instead")]
    pub fn take_response(self) -> Promise<Response> {
        let (response_sender, response_promise) = Promise::<Response>::new();
        self.commands.queue(AddExecution::new(
            Some(self.source),
            self.target,
            TakenResponse::<Response>::new(response_sender),
        ));
        response_promise
    }

    /// Pass the outcome of the request to another provider.
    #[must_use]
    pub fn then<P: ProvideOnce<Request = Response>>(
        self,
        provider: P,
    ) -> Series<'w, 's, 'a, P::Response, P::Streams> {
        let source = self.target;
        let target = self
            .commands
            .spawn((Detached::default(), UnusedTarget, SeriesMarker))
            .id();

        // We should automatically delete the previous step in the chain once
        // this one is finished.
        self.commands
            .entity(source)
            .insert((Cancellable::new(cancel_execution), SeriesMarker))
            .remove::<UnusedTarget>()
            .insert(ChildOf(target));
        provider.connect(None, source, target, self.commands);
        Series {
            source,
            target,
            commands: self.commands,
            _ignore: Default::default(),
        }
    }

    /// Apply a one-time callback whose input is the Response of the current
    /// target. The output of the map will become the Response of the returned
    /// target.
    ///
    /// This takes in a regular blocking function, which means all systems will
    /// be blocked from running while the function gets executed.
    #[must_use]
    pub fn map_block<U>(
        self,
        f: impl FnOnce(Response) -> U + 'static + Send + Sync,
    ) -> Series<'w, 's, 'a, U, ()>
    where
        U: 'static + Send + Sync,
    {
        self.then(f.into_blocking_map_once())
    }

    /// Apply a one-time callback whose output is a [`Future`] that will be run
    /// in the [`AsyncComputeTaskPool`][1] (unless the `single_threaded_async`
    /// feature is active). The output of the [`Future`] will be the Response of
    /// the returned Series.
    ///
    /// [1]: bevy_tasks::AsyncComputeTaskPool
    #[must_use]
    pub fn map_async<Task>(
        self,
        f: impl FnOnce(Response) -> Task + 'static + Send + Sync,
    ) -> Series<'w, 's, 'a, Task::Output, ()>
    where
        Task: Future + 'static + Sendish,
        Task::Output: 'static + Send + Sync,
    {
        self.then(f.into_async_map_once())
    }

    /// Apply a one-time map that implements one of
    /// - [`FnOnce(BlockingMap<Request, Streams>) -> Response`](crate::BlockingMap)
    /// - [`FnOnce(AsyncMap<Request, Streams>) -> impl Future<Response>`](crate::AsyncMap)
    ///
    /// If you do not care about providing streams then you can use
    /// [`Self::map_block`] or [`Self::map_async`] instead.
    pub fn map<M, F: AsMapOnce<M>>(
        self,
        f: F,
    ) -> Series<
        'w,
        's,
        'a,
        <F::MapType as ProvideOnce>::Response,
        <F::MapType as ProvideOnce>::Streams,
    >
    where
        F::MapType: ProvideOnce<Request = Response>,
        <F::MapType as ProvideOnce>::Response: 'static + Send + Sync,
        <F::MapType as ProvideOnce>::Streams: StreamPack,
    {
        self.then(f.as_map_once())
    }

    /// Store the response in a [`Storage`] component in the specified entity.
    ///
    /// Each stream will be collected into [`Collection`] components in the
    /// specified entity, one for each stream type. To store the streams in a
    /// different entity, call [`Self::collect_streams`] before this.
    ///
    /// If the entity despawns then the request gets cancelled unless you used
    /// [`Self::detach`] before calling this.
    pub fn store(self, target: Entity) {
        self.commands.queue(AddExecution::new(
            Some(self.source),
            self.target,
            Store::<Response>::new(target),
        ));

        let mut map = StreamTargetMap::default();
        let stream_targets = Streams::collect_streams(self.source, target, &mut map, self.commands);
        self.commands
            .entity(self.source)
            .insert((stream_targets, map));
    }

    /// Collect the stream data into [`Collection<T>`] components in the
    /// specified target, one collection for each stream data type. You must
    /// still decide what to do with the final response data.
    #[must_use]
    pub fn collect_streams(self, target: Entity) -> Series<'w, 's, 'a, Response, ()> {
        let mut map = StreamTargetMap::default();
        let stream_targets = Streams::collect_streams(self.source, target, &mut map, self.commands);
        self.commands
            .entity(self.source)
            .insert((stream_targets, map));

        Series {
            source: self.source,
            target: self.target,
            commands: self.commands,
            _ignore: Default::default(),
        }
    }

    /// Push the response to the back of a [`Collection<T>`] component in an
    /// entity.
    ///
    /// Similar to [`Self::store`] this will also collect streams into this
    /// entity.
    ///
    /// If the entity despawns then the request gets cancelled unless you used
    /// [`Self::detach`] before calling this.
    pub fn push(self, target: Entity) {
        self.commands.queue(AddExecution::new(
            Some(self.source),
            self.target,
            Push::<Response>::new(target, false),
        ));

        let mut map = StreamTargetMap::default();
        let stream_targets = Streams::collect_streams(self.source, target, &mut map, self.commands);
        self.commands
            .entity(self.source)
            .insert((stream_targets, map));
    }

    /// Used internally to implement various ways of capturing an outcome.
    pub(crate) fn send_outcome(self, capture: CaptureOutcome<Response>) {
        self.commands
            .queue(AddExecution::new(Some(self.source), self.target, capture));
    }

    // TODO(@mxgrey): Consider offering ways for users to respond to cancellations.
    // For example, offer an on_cancel method that lets users provide a callback
    // to be triggered when a cancellation happens. Or focus on the terminal end
    // of a series, like offer store_or_else(~), push_or_else(~) etc which accept
    // a callback that will be triggered after a cancellation.
}

impl<'w, 's, 'a, Response, Streams> Series<'w, 's, 'a, Response, Streams>
where
    Response: Bundle,
{
    /// Insert the response as a bundle in the specified entity. Stream data
    /// will be dropped unless you use [`Self::collect_streams`] before this.
    ///
    /// If the entity despawns then the request gets cancelled unless you used
    /// [`Self::detach`] before calling this.
    ///
    /// If the response is not a bundle then you can store it in an entity using
    /// [`Self::store`] or [`Self::push`]. Alternatively you can transform it
    /// into a bundle using [`Self::map_block`] or [`Self::map_async`].
    pub fn insert(self, target: Entity) {
        self.commands.queue(AddExecution::new(
            Some(self.source),
            self.target,
            Insert::<Response>::new(target),
        ));
    }
}

impl<'w, 's, 'a, Response, Streams> Series<'w, 's, 'a, Response, Streams>
where
    Response: Event,
{
    /// Send the response out as an event once it is ready. Stream data will be
    /// dropped unless you use [`Self::collect_streams`] before this.
    ///
    /// Using this will also effectively [detach](Self::detach) the series.
    pub fn send_event(self) {
        self.commands.queue(AddExecution::new(
            Some(self.source),
            self.target,
            SendEvent::<Response>::new(),
        ));
    }
}

/// Contains the final response and streams produced at the end of a series.
pub struct Recipient<Response, Streams: StreamPack> {
    pub response: Promise<Response>,
    pub streams: Streams::StreamReceivers,
    /// The root session ID of the entire series. Every session ID related to
    /// this series is a descendent of this entity.
    ///
    /// This may be counter-intuitive because this is the last session in the
    /// series but structurally we have made it the parent of all the sessions
    /// that will be executed before it. This is done so that despawning behavior
    /// has a more logical relationship with the dependencies between the sessions.
    pub session: Entity,
}

pub struct Capture<Response, Streams: StreamPack> {
    pub outcome: Outcome<Response>,
    pub streams: Streams::StreamReceivers,
    pub session: Entity,
}

/// Used to store a response of a series as a component of an entity.
#[derive(Component)]
pub struct Storage<T> {
    pub data: T,
    pub session: Entity,
}

/// Used to collect responses from multiple series chains into a container
/// attached to an entity.
//
// TODO(@mxgrey): Consider allowing the user to choose the container type.
#[derive(Component)]
pub struct Collection<T> {
    /// The items that have been collected.
    pub items: Vec<Storage<T>>,
}

impl<T> Default for Collection<T> {
    fn default() -> Self {
        Self {
            items: Default::default(),
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::{ContinuousQueueView, prelude::*, testing::*};
    use smallvec::SmallVec;
    use std::{
        sync::{Arc, Mutex},
        time::{Duration, Instant},
    };
    use tokio::sync::mpsc::unbounded_channel;

    #[test]
    fn test_dropped_chain() {
        let mut context = TestingContext::minimal_plugins();

        let (detached_sender, mut detached_receiver) = unbounded_channel();
        let (attached_sender, mut attached_receiver) = unbounded_channel();

        context.command(|commands| {
            let _ = commands
                .request("hello".to_owned(), to_uppercase.into_blocking_map())
                .map_block(move |value| {
                    detached_sender.send(value.clone()).unwrap();
                    value
                })
                .detach()
                .map_block(to_lowercase)
                .map_block(move |value| {
                    attached_sender.send(value.clone()).unwrap();
                    value
                })
                .map_block(to_uppercase);
        });

        context.run(1);
        assert_eq!(detached_receiver.try_recv().unwrap(), "HELLO");
        assert!(attached_receiver.try_recv().is_err());
        assert!(
            context
                .get_unhandled_errors()
                .is_some_and(|e| e.unused_targets.len() == 1)
        );
    }

    #[test]
    fn test_blocking_map() {
        let mut context = TestingContext::minimal_plugins();

        let r = context.resolve_request(String::from("hello"), to_uppercase.into_blocking_map());
        assert_eq!(r, "HELLO");

        let r =
            context.resolve_request(String::from("hello"), to_uppercase.into_blocking_map_once());
        assert_eq!(r, "HELLO");

        let mut outcome = context.command(|commands| {
            commands
                .provide(String::from("hello"))
                .map_block(to_uppercase)
                .outcome()
        });

        context.run_while_pending(&mut outcome);
        context.assert_no_errors();
        assert_eq!(outcome.try_recv().unwrap().unwrap(), "HELLO");

        let mut outcome = context.command(|commands| {
            commands
                .provide(String::from("hello"))
                .map_block(|request| request.to_uppercase())
                .outcome()
        });

        context.run_while_pending(&mut outcome);
        context.assert_no_errors();
        assert_eq!(outcome.try_recv().unwrap().unwrap(), "HELLO");
    }

    #[test]
    fn test_async_map() {
        let mut context = TestingContext::minimal_plugins();

        let request = WaitRequest {
            duration: Duration::from_secs_f64(0.001),
            value: "hello".to_owned(),
        };

        let conditions = FlushConditions::new().with_timeout(Duration::from_secs_f64(5.0));

        let r = context
            .try_resolve_request(request.clone(), wait.into_async_map(), conditions.clone())
            .unwrap();
        assert_eq!(r, "hello");

        let r = context
            .try_resolve_request(
                request.clone(),
                wait.into_async_map_once(),
                conditions.clone(),
            )
            .unwrap();
        assert_eq!(r, "hello");

        let mut outcome =
            context.command(|commands| commands.provide(request.clone()).map_async(wait).outcome());

        assert!(context.run_with_conditions(&mut outcome, conditions.clone()));
        context.assert_no_errors();
        assert_eq!(outcome.try_recv().unwrap().unwrap(), "hello");

        let mut outcome = context.command(|commands| {
            commands
                .provide(request.clone())
                .map_async(|request| {
                    async move {
                        let t = Instant::now();
                        while t.elapsed() < request.duration {
                            // Busy wait
                        }
                        request.value
                    }
                })
                .outcome()
        });

        assert!(context.run_with_conditions(&mut outcome, conditions.clone()));
        context.assert_no_errors();
        assert_eq!(outcome.try_recv().unwrap().unwrap(), "hello");
    }

    #[test]
    fn test_detach() {
        // This is a regression test that covers a bug which existed due to
        // an incorrect handling of detached series when giving input.
        let mut context = TestingContext::minimal_plugins();
        let service = context.spawn_delayed_map(Duration::from_millis(1), |n| *n + 1);

        context.command(|commands| {
            commands.provide(0).then(service).detach();
        });

        let (sender, mut receiver) = tokio::sync::oneshot::channel();
        context.run_with_conditions(&mut receiver, Duration::from_millis(5));
        assert!(
            context.no_unhandled_errors(),
            "Unhandled errors: {:#?}",
            context.get_unhandled_errors(),
        );

        // The promise and sender only exist because run_with_conditions requires
        // them. Moreover we need to make sure that sender does not get dropped
        // prematurely by the compiler, otherwise the promise will have the run
        // exit prematurely. Therefore we call .send(()) here to guarantee the
        // compiler knows to keep it alive until the running is finished.
        //
        // We have observed that using `let (_, mut promise) = ` will cause the
        // sender to drop prematurely, so we don't want to risk that there are
        // other cases where that may happen. It is important for the run to
        // last multiple cycles.
        let _ = sender.send(());
    }

    #[derive(Clone, Debug, PartialEq, Eq, Hash, DeliveryLabel)]
    struct UnitLabel;

    #[derive(Clone, Debug, PartialEq, Eq, Hash, DeliveryLabel)]
    struct StatefulLabel(u64);

    #[test]
    fn test_delivery_instructions() {
        let mut context = TestingContext::minimal_plugins();
        let service = context.spawn_delayed_map_with_viewer(
            Duration::from_secs_f32(0.01),
            |counter: &Arc<Mutex<u64>>| {
                *counter.lock().unwrap() += 1;
            },
            |view: &ContinuousQueueView<_, ()>| {
                assert!(view.len() <= 1);
            },
        );

        verify_delivery_instruction_matrix(service.optional_stream_cast(), &mut context);

        let service = context.spawn_async_delayed_map(
            Duration::from_secs_f32(0.01),
            |counter: Arc<Mutex<u64>>| {
                *counter.lock().unwrap() += 1;
            },
        );

        verify_delivery_instruction_matrix(service, &mut context);

        let async_service = service;
        let service = context.spawn_io_workflow(|scope, builder| {
            builder
                .chain(scope.start)
                .then(async_service)
                .connect(scope.terminate);
        });

        verify_delivery_instruction_matrix(service, &mut context);

        // We don't test blocking services because blocking services are always
        // serial no matter what, so delivery instructions have no effect for them.
    }

    fn verify_delivery_instruction_matrix(
        service: Service<Arc<Mutex<u64>>, ()>,
        context: &mut TestingContext,
    ) {
        // Test for a unit struct
        verify_preemption_matrix(
            service.instruct(UnitLabel),
            service.instruct(UnitLabel.preempt()),
            context,
        );

        // Test for a stateful struct
        verify_preemption_matrix(
            service.instruct(StatefulLabel(5)),
            service.instruct(StatefulLabel(5).preempt()),
            context,
        );

        // Test for a unit struct
        verify_queuing_matrix(service.instruct(UnitLabel), context);

        // Test for a stateful struct
        verify_queuing_matrix(service.instruct(StatefulLabel(7)), context);

        // Test for a unit struct
        verify_ensured_matrix(service, UnitLabel, context);

        // Test for a stateful struct
        verify_ensured_matrix(service, StatefulLabel(2), context);
    }

    fn verify_preemption_matrix(
        queuing_service: ServiceInstructions<Arc<Mutex<u64>>, ()>,
        preempting_service: ServiceInstructions<Arc<Mutex<u64>>, ()>,
        context: &mut TestingContext,
    ) {
        // Test by queuing up a bunch of requests before preempting them all at once.
        verify_preemption(
            1,
            queuing_service.clone(),
            preempting_service.clone(),
            context,
        );
        verify_preemption(
            2,
            queuing_service.clone(),
            preempting_service.clone(),
            context,
        );
        verify_preemption(
            3,
            queuing_service.clone(),
            preempting_service.clone(),
            context,
        );
        verify_preemption(
            4,
            queuing_service.clone(),
            preempting_service.clone(),
            context,
        );

        // Test by repeatedly preempting each request with the next.
        verify_preemption(
            1,
            preempting_service.clone(),
            preempting_service.clone(),
            context,
        );
        verify_preemption(
            2,
            preempting_service.clone(),
            preempting_service.clone(),
            context,
        );
        verify_preemption(
            3,
            preempting_service.clone(),
            preempting_service.clone(),
            context,
        );
        verify_preemption(
            4,
            preempting_service.clone(),
            preempting_service.clone(),
            context,
        );
    }

    fn verify_preemption(
        preemptions: usize,
        preempted_service: ServiceInstructions<Arc<Mutex<u64>>, ()>,
        preempting_service: ServiceInstructions<Arc<Mutex<u64>>, ()>,
        context: &mut TestingContext,
    ) {
        let counter = Arc::new(Mutex::new(0_u64));
        let mut preempted: SmallVec<[Outcome<()>; 16]> = SmallVec::new();
        for _ in 0..preemptions {
            let outcome = context.command(|commands| {
                commands
                    .request(Arc::clone(&counter), preempted_service.clone())
                    .outcome()
            });
            preempted.push(outcome);
        }

        let mut final_outcome = context.command(|commands| {
            commands
                .request(Arc::clone(&counter), preempting_service)
                .outcome()
        });

        for outcome in &mut preempted {
            context.run_with_conditions(outcome, Duration::from_secs(2));
            assert!(outcome.try_recv().unwrap().is_err());
        }

        context.run_with_conditions(&mut final_outcome, Duration::from_secs(2));
        assert!(final_outcome.try_recv().unwrap().is_ok());
        assert_eq!(*counter.lock().unwrap(), 1);
        context.assert_no_errors();
    }

    fn verify_queuing_matrix(
        queuing_service: ServiceInstructions<Arc<Mutex<u64>>, ()>,
        context: &mut TestingContext,
    ) {
        // Test by queuing up a bunch of requests and making sure they all get
        // delivered.
        verify_queuing(2, queuing_service.clone(), context);
        verify_queuing(3, queuing_service.clone(), context);
        verify_queuing(4, queuing_service.clone(), context);
        verify_queuing(5, queuing_service.clone(), context);
    }

    fn verify_queuing(
        queue_size: usize,
        queuing_service: ServiceInstructions<Arc<Mutex<u64>>, ()>,
        context: &mut TestingContext,
    ) {
        let counter = Arc::new(Mutex::new(0_u64));
        let mut queued: SmallVec<[Outcome<()>; 16]> = SmallVec::new();
        for _ in 0..queue_size {
            let outcome = context.command(|commands| {
                commands
                    .request(Arc::clone(&counter), queuing_service.clone())
                    .outcome()
            });
            queued.push(outcome);
        }

        for outcome in &mut queued {
            context.run_with_conditions(outcome, Duration::from_secs(2));
            assert!(outcome.try_recv().unwrap().is_ok());
        }

        assert_eq!(*counter.lock().unwrap(), queue_size as u64);
        context.assert_no_errors();
    }

    fn verify_ensured_matrix<L: DeliveryLabel + Clone>(
        service: Service<Arc<Mutex<u64>>, ()>,
        label: L,
        context: &mut TestingContext,
    ) {
        // Test by queuing up a mix of ensured and unensured requests, and then
        // sending in one that preempts them all. The ensured requests should
        // remain in the queue and execute despite the preempter. The unensured
        // requests should all be cancelled.
        verify_ensured([false, true, false, true], service, label.clone(), context);
        verify_ensured([true, false, false, false], service, label.clone(), context);
        verify_ensured([true, true, false, false], service, label.clone(), context);
        verify_ensured([false, false, true, true], service, label.clone(), context);
        verify_ensured([true, false, false, true], service, label.clone(), context);
        verify_ensured(
            [false, false, false, false],
            service,
            label.clone(),
            context,
        );
        verify_ensured([true, true, true, true], service, label.clone(), context);
    }

    fn verify_ensured<L: DeliveryLabel + Clone>(
        queued: impl IntoIterator<Item = bool>,
        service: Service<Arc<Mutex<u64>>, ()>,
        label: L,
        context: &mut TestingContext,
    ) {
        let counter = Arc::new(Mutex::new(0_u64));
        let mut queued_outcomes: SmallVec<[(Outcome<()>, bool); 16]> = SmallVec::new();
        // This counter starts out at 1 to account for the preempting request.
        let mut expected_count = 1;
        for ensured in queued {
            let srv = if ensured {
                expected_count += 1;
                service.instruct(label.clone().ensure())
            } else {
                service.instruct(label.clone())
            };

            let outcome =
                context.command(|commands| commands.request(Arc::clone(&counter), srv).outcome());

            queued_outcomes.push((outcome, ensured));
        }

        let mut preempter = context.command(|commands| {
            commands
                .request(
                    Arc::clone(&counter),
                    service.instruct(label.clone().preempt()),
                )
                .outcome()
        });

        for (outcome, ensured) in &mut queued_outcomes {
            context.run_with_conditions(outcome, Duration::from_secs(2));
            if *ensured {
                assert!(outcome.try_recv().unwrap().is_ok());
            } else {
                assert!(outcome.try_recv().unwrap().is_err());
            }
        }

        context.run_with_conditions(&mut preempter, Duration::from_secs(2));
        assert!(preempter.is_available());
        assert_eq!(*counter.lock().unwrap(), expected_count);
        context.assert_no_errors();
    }
}