grommet 0.1.2

//! Request/response on top of a fire-and-forget scheduler.
//!
//! Submission is deliberately one-way: it reports whether work was *accepted*,
//! not what it produced. That is the right primitive for pipelines that reply
//! by writing to a socket, batch their answers, or have no answer at all.
//!
//! When a caller does want the result, wrap the work in a [`Call`] and use
//! [`Router::call`]. The reply channel then travels with the work, which means
//! every path that hands work back: a full mailbox, a downed shard, a passed
//! deadline: hands back the means to answer the caller too.

use crate::clock::Clock;
use crate::router::{Router, SubmitError};
use crate::work::Work;
use grommet_core::ClassId;
use std::error::Error;
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::sync::oneshot;

/// The caller stopped waiting, or the processor finished without answering.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Cancelled;

impl fmt::Display for Cancelled {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("the call was dropped without a response")
    }
}

impl Error for Cancelled {}

/// Why a call did not produce a response.
#[derive(Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum CallError<W> {
    /// Never accepted. The work is returned so the caller can shed it
    /// deliberately.
    Rejected(SubmitError<W>),
    /// Accepted, but the responder was dropped without an answer: the
    /// processor returned without replying, or its future panicked.
    Cancelled,
}

impl<W> fmt::Display for CallError<W> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Rejected(_) => f.write_str("the call was not accepted"),
            Self::Cancelled => Cancelled.fmt(f),
        }
    }
}

impl<W: fmt::Debug> Error for CallError<W> {}

/// Work paired with the channel its answer travels back on.
pub struct Call<W: Work, R> {
    inner: W,
    respond: oneshot::Sender<R>,
}

impl<W: Work, R> Call<W, R> {
    /// Pair `work` with a fresh reply channel, returning the call to submit
    /// and the [`Answer`] to await.
    pub fn new(work: W) -> (Self, Answer<R>) {
        let (respond, receive) = oneshot::channel();
        (Self { inner: work, respond }, Answer(receive))
    }

    pub fn work(&self) -> &W {
        &self.inner
    }

    /// Split into the work and the means to answer it, so a processor can
    /// consume one and hold the other across awaits.
    pub fn into_parts(self) -> (W, Responder<R>) {
        (self.inner, Responder(self.respond))
    }

    /// Answer and discard the work. Handy from
    /// [`Processor::on_expired`](crate::processor::Processor::on_expired).
    pub fn reply(self, response: R) {
        let _ = self.respond.send(response);
    }

    /// Whether the caller has stopped waiting.
    pub fn is_cancelled(&self) -> bool {
        self.respond.is_closed()
    }
}

impl<W: Work, R: Send + 'static> Work for Call<W, R> {
    type Key = W::Key;
    type Id = W::Id;

    fn key(&self) -> Self::Key {
        self.inner.key()
    }

    fn request_id(&self) -> Option<Self::Id> {
        self.inner.request_id()
    }

    fn class(&self) -> ClassId {
        self.inner.class()
    }

    fn time_to_live(&self) -> Option<Duration> {
        self.inner.time_to_live()
    }
}

/// The waiting half of a [`Call`]: a future resolving to the response, or to
/// [`Cancelled`] when the work finished without one.
///
/// Awaiting this is the only thing it does. It exists so that the channel
/// underneath can be replaced without changing what callers wrote.
#[repr(transparent)]
pub struct Answer<R>(oneshot::Receiver<R>);

impl<R> Future for Answer<R> {
    type Output = Result<R, Cancelled>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // The receiver is `Unpin`, so this projection needs no `unsafe`.
        Pin::new(&mut self.get_mut().0).poll(cx).map_err(|_| Cancelled)
    }
}

impl<R> fmt::Debug for Answer<R> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Answer").finish_non_exhaustive()
    }
}

/// The answering half of a [`Call`].
pub struct Responder<R>(oneshot::Sender<R>);

impl<R> Responder<R> {
    /// Answer the caller. A caller that has already given up is not an error.
    pub fn send(self, response: R) {
        let _ = self.0.send(response);
    }

    /// Whether the caller has stopped waiting. Worth checking before starting
    /// expensive work, since nobody will read the result.
    pub fn is_cancelled(&self) -> bool {
        self.0.is_closed()
    }
}

impl<W, R, C, const CLASSES: usize> Router<Call<W, R>, C, CLASSES>
where
    W: Work,
    R: Send + 'static,
    C: Clock,
{
    /// Submit and await the response, waiting on backpressure if the target
    /// shard's mailbox is full.
    pub async fn call(&self, work: W) -> Result<R, CallError<W>> {
        let (call, receive) = Call::new(work);
        self.submit(call)
            .await
            .map_err(|error| CallError::Rejected(error.map(|call| call.inner)))?;
        receive.await.map_err(|Cancelled| CallError::Cancelled)
    }

    /// Submit without waiting, reporting a full mailbox immediately and
    /// returning the response as a separate future.
    ///
    /// Rejection is synchronous and the answer is not, which is exactly the
    /// shape a load-shedding caller needs: it can give up before committing to
    /// an await.
    pub fn try_call(&self, work: W) -> Result<Answer<R>, SubmitError<W>> {
        let (call, receive) = Call::new(work);
        self.try_submit(call).map_err(|error| error.map(|call| call.inner))?;
        Ok(receive)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::clock::ManualClock;
    use crate::mailbox::channel;
    use crate::metrics::ShardStats;
    use crate::processor::Processor;
    use crate::shard::{self, ShardConfig};
    use grommet_core::Disposition;
    use std::sync::Arc;

    #[derive(Debug, PartialEq, Eq)]
    struct Job {
        key: u64,
        ttl: Option<Duration>,
    }

    impl Work for Job {
        type Key = u64;
        type Id = ();
        fn key(&self) -> u64 {
            self.key
        }
        fn class(&self) -> ClassId {
            0
        }
        fn time_to_live(&self) -> Option<Duration> {
            self.ttl
        }
    }

    /// Counts per key and answers with the running total, unless told to stay
    /// silent so the cancellation path can be observed.
    #[derive(Clone, Copy)]
    struct Counter {
        answer: bool,
    }

    impl Processor for Counter {
        type Work = Call<Job, u64>;
        type State = u64;
        type Error = std::convert::Infallible;

        async fn process(
            &self,
            _key: u64,
            state: Option<u64>,
            call: Call<Job, u64>,
        ) -> Result<Disposition<u64>, Self::Error> {
            let total = state.unwrap_or(0) + 1;
            let (_job, responder) = call.into_parts();
            if self.answer {
                responder.send(total);
            }
            Ok(Disposition::Keep(total))
        }

        fn on_expired(&self, _key: u64, call: Call<Job, u64>) {
            call.reply(u64::MAX);
        }
    }

    async fn with_shard<F, Fut>(processor: Counter, mailbox: usize, driver: F)
    where
        F: FnOnce(Arc<Router<Call<Job, u64>, ManualClock, 2>>) -> Fut,
        Fut: Future<Output = ()>,
    {
        let clock = ManualClock::new();
        let (tx, rx) = channel(mailbox);
        let router = Arc::new(Router::new(vec![tx], clock.clone()));
        let stats = Arc::new(ShardStats::<2>::default());
        let mut cfg = ShardConfig::new([4, 4]);
        cfg.tick = Duration::from_millis(1);
        let engine = shard::run(rx, processor, clock, stats, cfg);
        tokio::join!(engine, driver(router));
    }

    #[tokio::test(start_paused = true)]
    async fn a_call_carries_its_response_back_to_the_caller() {
        with_shard(Counter { answer: true }, 8, |router| async move {
            assert_eq!(router.call(Job { key: 3, ttl: None }).await, Ok(1));
            assert_eq!(router.call(Job { key: 3, ttl: None }).await, Ok(2));
            assert_eq!(router.call(Job { key: 4, ttl: None }).await, Ok(1), "state is per key");
        })
        .await;
    }

    #[tokio::test(start_paused = true)]
    async fn a_processor_that_never_answers_cancels_rather_than_hangs() {
        with_shard(Counter { answer: false }, 8, |router| async move {
            assert_eq!(router.call(Job { key: 1, ttl: None }).await, Err(CallError::Cancelled));
        })
        .await;
    }

    #[tokio::test(start_paused = true)]
    async fn shed_work_is_returned_unwrapped_so_the_caller_can_answer_it() {
        let clock = ManualClock::new();
        let (tx, _rx) = channel(1);
        let router = Router::<Call<Job, u64>, ManualClock, 2>::new(vec![tx], clock);

        // Fill the one mailbox slot; nothing is consuming it.
        assert!(router.try_call(Job { key: 1, ttl: None }).is_ok(), "the first call fits");
        let Err(error) = router.try_call(Job { key: 2, ttl: None }) else {
            panic!("a full mailbox must shed");
        };

        assert!(matches!(error, SubmitError::Full(_)));
        assert_eq!(error.into_work().key, 2, "the caller gets its own work back, not a wrapper");
    }

    #[tokio::test(start_paused = true)]
    async fn expired_work_is_answered_from_the_expiry_hook() {
        with_shard(Counter { answer: true }, 8, |router| async move {
            let response = router.call(Job { key: 9, ttl: Some(Duration::ZERO) }).await;
            assert_eq!(response, Ok(u64::MAX), "the deadline path still answers the caller");
        })
        .await;
    }

    #[tokio::test(start_paused = true)]
    async fn a_responder_notices_that_its_caller_gave_up() {
        let (call, receive) = Call::<Job, u64>::new(Job { key: 1, ttl: None });
        assert!(!call.is_cancelled());
        let (_job, responder) = call.into_parts();
        assert!(!responder.is_cancelled());
        drop(receive);
        assert!(responder.is_cancelled(), "an abandoned call is worth detecting before working");
    }

    #[test]
    fn a_call_notices_the_abandonment_before_it_is_split() {
        let (call, receive) = Call::<Job, u64>::new(Job { key: 1, ttl: None });
        assert!(!call.is_cancelled(), "the caller is still waiting");
        drop(receive);
        assert!(
            call.is_cancelled(),
            "shedding an abandoned call is only possible if the wrapper reports it"
        );
    }

    /// Work carrying values that are not the defaults for their types, so a
    /// delegation that answered `Default::default()` instead of forwarding
    /// would be visible.
    #[derive(Debug, PartialEq, Eq)]
    struct Tagged;

    impl Work for Tagged {
        type Key = u64;
        type Id = u32;
        fn key(&self) -> u64 {
            77
        }
        fn request_id(&self) -> Option<u32> {
            Some(5)
        }
        fn class(&self) -> ClassId {
            1
        }
        fn time_to_live(&self) -> Option<Duration> {
            Some(Duration::from_secs(3))
        }
    }

    #[test]
    fn wrapping_work_in_a_call_forwards_every_scheduling_answer_unchanged() {
        let (call, _receive) = Call::<Tagged, u64>::new(Tagged);
        assert_eq!(call.key(), 77);
        assert_eq!(call.request_id(), Some(5));
        assert_eq!(call.class(), 1, "a wrapped item must not be rerouted to another ring");
        assert_eq!(call.time_to_live(), Some(Duration::from_secs(3)));
        assert_eq!(call.work(), &Tagged);
    }

    #[test]
    fn the_failure_types_say_which_of_the_two_things_went_wrong() {
        assert_eq!(Cancelled.to_string(), "the call was dropped without a response");
        assert_eq!(
            CallError::<Job>::Cancelled.to_string(),
            Cancelled.to_string(),
            "a cancelled call reads the same however it is reported"
        );
        assert_eq!(
            CallError::Rejected(SubmitError::Full(Job { key: 1, ttl: None })).to_string(),
            "the call was not accepted",
        );
    }
}