ruststream 0.7.0-rc.10

//! Startup injections for batch handlers: the batch counterpart of [`inject`](super::inject).
//!
//! `#[subscriber("in")] async fn f(batch: &[T], Out(out): Out<P>)`
//! declares the same startup-resolved parameters as a single-message handler: the runtime
//! resolves the [`FromStartup`](super::FromStartup) tuple after the subscription opens and
//! before the first batch, so injected values are live by construction. The batch shape needs
//! its own definition pair because the handler consumes a slice and returns a
//! [`BatchResult`], not a per-message [`HandlerOutcome`](super::HandlerOutcome).

use std::future::Future;

use crate::{BuildBatchContext, IncomingMessage};

use super::batch::{BatchHandler, BatchResult, decode_batch, settle_batch};
use super::context::Context;
use super::dispatch::Workers;
use super::failure::{FailurePolicies, FailurePolicy};
use super::input::{DecodeWith, InputKind};
use super::metadata::{HandlerMetadata, OutgoingMessageMetadata};

/// A batch definition whose handler takes startup-injected parameters.
///
/// Generated by a `&[T]`-taking `#[subscriber]` handler when the signature carries
/// [`Out`](super::Out) parameters; [`Self::Injections`] is their
/// tuple, in declaration order. The metadata surface mirrors [`BatchDef`](super::BatchDef).
pub trait BatchInjectDef: Send + Sync {
    /// The input kind of one batch element (see [`BatchDef::Input`](super::BatchDef::Input)).
    type Input: InputKind;

    /// The subscription source this handler binds to.
    type Source;

    /// The tuple of startup-injected parameters ([`Out`](super::Out), ...).
    type Injections;

    /// The broker's typed subscription-scoped context the batch's handler reads by key (`()` when
    /// the handler names none); see [`BatchDef::Context`](super::BatchDef::Context).
    type Context;

    /// Builds the subscription source (fresh each call).
    fn source(&self) -> Self::Source;

    /// The concurrency policy for this subscriber's dispatch loop (how many batches are in
    /// flight at once).
    fn workers(&self) -> Workers {
        Workers::sequential()
    }

    /// The failure policy for a batch-handler panic and a per-element decode failure.
    fn failure_policies(&self) -> FailurePolicies {
        FailurePolicies::default()
    }

    /// An optional human description (from the handler's doc comment), for `AsyncAPI`.
    fn description(&self) -> Option<&str> {
        None
    }

    /// The element type's serialized JSON Schema, when available.
    fn input_schema(&self) -> Option<String> {
        None
    }

    /// The serialized JSON Schema of the element type's header contract, when one is declared.
    fn headers_schema(&self) -> Option<String> {
        None
    }

    /// The messages this handler publishes, for the `AsyncAPI` `send` operations: every `Out`
    /// slot dictionary entry. The macro fills this in; the default declares nothing. Called
    /// once at registration.
    fn outgoing(&self) -> Vec<OutgoingMessageMetadata> {
        Vec::new()
    }

    /// The element type's [`Message`](crate::MessageInfo) name, when it implements that trait.
    fn message_name(&self) -> Option<&'static str> {
        None
    }

    /// The element type's [`Message`](crate::MessageInfo) description, when it implements that
    /// trait.
    fn message_description(&self) -> Option<&'static str> {
        None
    }
}

/// Runs a [`BatchInjectDef`]'s handler body over an app state of type `S`.
///
/// The same state-generic shape as [`SliceHandler`](super::SliceHandler); see
/// [`PublishingCall`](super::PublishingCall) for the rationale.
pub(crate) trait BatchInjectCall<S>: BatchInjectDef {
    /// Runs the handler body on one decoded batch, with the resolved injections.
    fn call(
        &self,
        batch: &[<Self::Input as InputKind>::Owned],
        injections: &Self::Injections,
        ctx: &mut Context<'_, Self::Context, S>,
    ) -> impl Future<Output = BatchResult> + Send;
}

/// Builds the registration metadata for an injected batch definition mounted under `name`.
pub(crate) fn batch_inject_metadata<D: BatchInjectDef>(name: String, def: &D) -> HandlerMetadata {
    let mut meta = HandlerMetadata::raw(name).with_def_details(
        def.description(),
        def.input_schema(),
        def.headers_schema(),
        def.message_name(),
        def.message_description(),
    );
    meta.input_type = <D::Input as InputKind>::input_label();
    meta.deserialized = <D::Input as InputKind>::DESERIALIZED;
    meta.outgoing = def.outgoing();
    meta
}

/// The batch handler built from a [`BatchInjectDef`] once its injections resolved: decode
/// the batch, run the body with them, settle every delivery.
pub(crate) struct BatchInjectHandler<Def: BatchInjectDef, DecodeCodec> {
    pub(crate) def: Def,
    pub(crate) codec: DecodeCodec,
    pub(crate) injections: Def::Injections,
    pub(crate) decode: FailurePolicy,
}

impl<Def: BatchInjectDef, DecodeCodec> std::fmt::Debug for BatchInjectHandler<Def, DecodeCodec> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BatchInjectHandler").finish_non_exhaustive()
    }
}

impl<Msg, Def, DecodeCodec, State> BatchHandler<Msg, Def::Context, State>
    for BatchInjectHandler<Def, DecodeCodec>
where
    Msg: IncomingMessage,
    Def: BatchInjectCall<State>,
    Def::Input: DecodeWith<DecodeCodec>,
    Def::Injections: Send + Sync,
    Def::Context: BuildBatchContext<Msg> + Send + Sync + 'static,
    DecodeCodec: Send + Sync,
    State: Send + Sync,
{
    type Scratch = Vec<<Def::Input as InputKind>::Owned>;

    async fn handle_batch(
        &self,
        batch: Vec<Msg>,
        values: &mut Self::Scratch,
        ctx: &mut Context<'_, Def::Context, State>,
    ) {
        let subscription = ctx.subscription();
        let accepted = decode_batch::<Msg, Def::Input, DecodeCodec, Def::Context, State>(
            batch,
            values,
            &self.codec,
            self.decode,
            ctx,
        )
        .await;
        if accepted.is_empty() {
            return;
        }
        let delivery = ctx.delivery();
        let result = self.def.call(values, &self.injections, ctx).await;
        values.clear();
        settle_batch(accepted, result, subscription, delivery).await;
    }
}

#[cfg(all(test, feature = "memory", feature = "json"))]
mod tests {
    use std::future::ready;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::{Arc, Mutex};

    use futures::StreamExt;

    use super::super::dispatch::Delivery;
    use super::super::handler::HandlerOutcome;
    use super::super::input::Decoded;
    use super::*;
    use crate::codec::JsonCodec;
    use crate::memory::{ConnectedMemoryBroker, MemoryBroker};
    use crate::testkit::batch::{publish_payloads, pull_batch};
    use crate::{HeaderMap, Name, Subscriber, SubscriptionSource};

    /// A definition whose "injection" is a plain multiplier: the real ones are `Out`
    /// values, but the handler side only cares that the resolved tuple is handed to every call.
    struct Scale {
        calls: Arc<AtomicUsize>,
        seen: Arc<Mutex<Vec<u32>>>,
    }

    impl Scale {
        fn new() -> Self {
            Self {
                calls: Arc::new(AtomicUsize::new(0)),
                seen: Arc::new(Mutex::new(Vec::new())),
            }
        }
    }

    impl BatchInjectDef for Scale {
        type Input = Decoded<u32>;
        type Source = Name;
        type Injections = u32;
        type Context = ();

        fn source(&self) -> Self::Source {
            Name::new("scale")
        }
    }

    // Ignores the app state, so it is generic over it (mounts on any app).
    impl<S: Send + Sync> BatchInjectCall<S> for Scale {
        fn call(
            &self,
            batch: &[u32],
            factor: &u32,
            _ctx: &mut Context<'_, (), S>,
        ) -> impl Future<Output = BatchResult> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            self.seen
                .lock()
                .unwrap()
                .extend(batch.iter().map(|n| n * factor));
            ready(BatchResult::Uniform(HandlerOutcome::ack()))
        }
    }

    #[test]
    fn batch_inject_def_defaults_register_without_documentation() {
        let def = Scale::new();
        // The mount site names the registration after the def's own source.
        let source = def.source();
        let name = SubscriptionSource::<ConnectedMemoryBroker>::name(&source).to_owned();
        let meta = batch_inject_metadata(name, &def);

        assert_eq!(meta.name, "scale");
        assert_eq!(meta.input_type, "u32");
        assert!(meta.description.is_none());
        assert!(meta.payload_schema.is_none());
        assert!(meta.headers_schema.is_none());
        assert!(meta.message_name.is_none());
        assert!(meta.message_description.is_none());
        assert!(meta.outgoing.is_empty());
        assert_eq!(def.workers(), Workers::sequential());
        assert_eq!(def.failure_policies(), FailurePolicies::default());
    }

    #[test]
    fn handler_debug_hides_the_injections() {
        let handler = BatchInjectHandler {
            def: Scale::new(),
            codec: JsonCodec,
            injections: 10,
            decode: FailurePolicy::Drop,
        };
        assert!(format!("{handler:?}").contains("BatchInjectHandler"));
    }

    /// The resolved injections reach every call alongside the decoded batch, and the returned
    /// settlement is applied to the deliveries behind it.
    #[tokio::test]
    async fn resolved_injections_reach_the_body_with_the_batch() {
        let broker = MemoryBroker::new();
        let mut sub = broker.subscribe("scale");
        publish_payloads(&broker, "scale", &[b"1", b"2"]).await;

        let def = Scale::new();
        let seen = Arc::clone(&def.seen);
        let handler = BatchInjectHandler {
            def,
            codec: JsonCodec,
            injections: 10,
            decode: FailurePolicy::Drop,
        };

        let state = ();
        let delivery = Delivery::empty();
        let headers = HeaderMap::new();
        let mut ctx = Context::new("scale", &headers, &state, (), &delivery);
        let batch = pull_batch(&mut sub).await;
        assert_eq!(batch.len(), 2);
        handler.handle_batch(batch, &mut Vec::new(), &mut ctx).await;

        assert_eq!(*seen.lock().unwrap(), [10, 20]);
        // The uniform Ack settled both deliveries, so neither comes back.
        let mut stream = std::pin::pin!(sub.stream());
        assert!(futures::poll!(stream.next()).is_pending());
    }

    /// Every element failing to decode short-circuits: the body is not invoked with an empty
    /// slice, and the dropped elements are not requeued.
    #[tokio::test]
    async fn a_fully_undecodable_batch_never_reaches_the_body() {
        let broker = MemoryBroker::new();
        let mut sub = broker.subscribe("scale");
        publish_payloads(&broker, "scale", &[b"not json", b"also not json"]).await;

        let def = Scale::new();
        let calls = Arc::clone(&def.calls);
        let handler = BatchInjectHandler {
            def,
            codec: JsonCodec,
            injections: 10,
            decode: FailurePolicy::Drop,
        };

        let state = ();
        let delivery = Delivery::empty();
        let headers = HeaderMap::new();
        let mut ctx = Context::new("scale", &headers, &state, (), &delivery);
        let batch = pull_batch(&mut sub).await;
        assert_eq!(batch.len(), 2);
        handler.handle_batch(batch, &mut Vec::new(), &mut ctx).await;

        assert_eq!(calls.load(Ordering::SeqCst), 0);
        let mut stream = std::pin::pin!(sub.stream());
        assert!(futures::poll!(stream.next()).is_pending());
    }
}