ruststream 0.7.0-rc.11

//! Batch publishing subscribers: consume a batch, publish the derived replies, then ack.
//!
//! Generated by `#[subscriber("in", publish("reply-name"))]` on a `&[T]` handler and mounted with
//! [`BrokerScope::include`](super::BrokerScope::include), chaining
//! [`.out(Reply, ..)`](super::RouterWith::out) to name the reply publish policy (without it, the
//! statement commits with the broker's default policy); on a [`Router`](super::Router) the same
//! chain ends in `.build()`, which is what commits the
//! registration. At startup the wiring pairs into a [`ReplyPublisher`]: a plain one publishes
//! each reply independently, while a `.transactional()` one makes the whole batch's replies
//! visible atomically - the consume-transform-produce pattern.

use std::future::Future;

use bytes::BytesMut;
use serde::Serialize;
use tracing::warn;

#[cfg(feature = "testing")]
use crate::testing::coordinator::{Origin, publishing_to};
use crate::{BuildBatchContext, IncomingMessage, PayloadForm};

use super::batch::{BatchHandler, BatchResult, decode_batch, settle_batch};
use super::context::Context;
use super::dispatch::Workers;
use super::failure::{FailurePolicies, FailurePolicy};
use super::handler::HandlerOutcome;
use super::input::{DecodeWith, InputKind};
use super::metadata::{HandlerMetadata, OutgoingMessageMetadata};
use super::publish::{PublishContext, PublishIdentity, PublishPipeline, ReplyPublisher};

/// A batch subscriber definition that produces replies to publish.
///
/// The batch counterpart of [`PublishingDef`](super::PublishingDef): the handler consumes the
/// whole decoded batch and returns the replies for it, all-or-nothing. Selective per-element
/// outcomes are not available here: a handler needing them publishes manually from a plain batch
/// handler, through a publisher held in the typed application state.
pub trait BatchPublishingDef: Send + Sync {
    /// The input kind of one batch element ([`Decoded<T>`](super::Decoded); see
    /// [`BatchDef::Input`](super::BatchDef::Input)). The handler consumes `&[T]`.
    type Input: InputKind;

    /// The tuple of startup-injected parameters ([`Out`](super::Out), ...; `()` when the
    /// signature carries none), resolved like
    /// [`InjectDef::Injections`](super::InjectDef::Injections).
    type Injections;

    /// The broker's typed subscription-scoped context the batch's handler reads by key (`()`
    /// when the handler names none), exactly as on
    /// [`BatchDef::Context`](super::BatchDef::Context): built once per batch from its first
    /// delivery, so a replying batch body reaches the subscription's handles (a reposition
    /// handle, a consumer-group name) the same way a settling one does.
    type Context;

    /// The reply element type; each entry of the returned `Vec` is encoded and published.
    type Reply;

    /// The subscription source this handler binds to (see
    /// [`SubscriberDef::Source`](super::SubscriberDef::Source)).
    type Source;

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

    /// The name (subject / channel) the replies are published to.
    fn reply_name(&self) -> &str;

    /// The concurrency policy for this subscriber's dispatch loop (how many batches are in
    /// flight at once). The macro fills this in from the `workers(..)` argument; the default is
    /// sequential dispatch.
    fn workers(&self) -> Workers {
        Workers::sequential()
    }

    /// The failure policy for a batch-handler panic and a per-element decode failure. The macro
    /// fills this in from the `on_failure(panic = .., decode = ..)` argument; the default fails
    /// fast on a panic and drops on a 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 it implements [`schemars::JsonSchema`]
    /// and the `asyncapi` feature is on. The macro fills this in; the default omits it.
    fn input_schema(&self) -> Option<String> {
        None
    }

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

    /// The messages this handler publishes, for the `AsyncAPI` `send` operations: the reply
    /// message, plus 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. The
    /// macro fills this in; the default omits it.
    fn message_name(&self) -> Option<&'static str> {
        None
    }

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

/// Runs a [`BatchPublishingDef`]'s handler body over an app state of type `S`.
///
/// Split from [`BatchPublishingDef`] so a handler that ignores the app state is generic over `S`
/// (mounts on any app), while one that reads it via [`Context::state`](super::Context::state)
/// implements this only for its declared `S` - the state match is checked at compile time without
/// pinning a single `State` on the def.
pub(crate) trait BatchPublishingCall<S>: BatchPublishingDef {
    /// Runs the handler body on one decoded batch.
    ///
    /// `Ok(replies)` publishes every reply to [`reply_name`](BatchPublishingDef::reply_name) and
    /// acks the batch; `Err(result)` publishes nothing and settles the batch with `result` -
    /// one uniform outcome, or one outcome per element.
    fn call(
        &self,
        batch: &[<Self::Input as InputKind>::Owned],
        injections: &Self::Injections,
        ctx: &mut Context<'_, Self::Context, S>,
    ) -> impl Future<Output = Result<Vec<Self::Reply>, BatchResult>> + Send;
}

/// Builds the registration metadata for a batch publishing definition mounted under `name`.
pub(crate) fn batch_publishing_metadata<D: BatchPublishingDef>(
    name: String,
    def: &D,
) -> HandlerMetadata {
    let mut meta = HandlerMetadata::raw(name)
        .with_output_type(std::any::type_name::<D::Reply>())
        .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 `BatchHandler` built from a [`BatchPublishingDef`]: decode the batch, run the handler,
/// publish the replies through the [`ReplyPublisher`], then ack the batch.
///
/// Elements that fail to decode are nacked individually and never reach the handler. A failed
/// reply publish (or a failed transaction) retries the whole batch instead of losing replies;
/// with a plain publisher a mid-batch failure can therefore re-publish the earlier replies on
/// redelivery (at-least-once), while a `.transactional()` wiring never leaves them half-visible.
pub(crate) struct BatchPublishingHandler<D: BatchPublishingDef, C, R, PP = PublishIdentity> {
    pub(crate) def: D,
    pub(crate) codec: C,
    pub(crate) publisher: R,
    pub(crate) pipeline: PP,
    pub(crate) injections: D::Injections,
    pub(crate) decode: FailurePolicy,
    /// The broker the reply publisher was paired against, which its replies go to.
    #[cfg(feature = "testing")]
    pub(crate) origin: Origin,
}

impl<D: BatchPublishingDef, C, R, PP> std::fmt::Debug for BatchPublishingHandler<D, C, R, PP> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BatchPublishingHandler")
            .finish_non_exhaustive()
    }
}

impl<M, D, C, R, PP, S> BatchHandler<M, D::Context, S> for BatchPublishingHandler<D, C, R, PP>
where
    M: IncomingMessage,
    D: BatchPublishingCall<S>,
    D::Input: DecodeWith<C>,
    D::Injections: Send + Sync,
    D::Reply: Serialize + Send + Sync,
    D::Context: BuildBatchContext<M> + Send + Sync + 'static,
    C: Send + Sync,
    R: ReplyPublisher<D::Context>,
    PP: PublishPipeline,
    S: Send + Sync,
{
    type Scratch = Vec<<D::Input as InputKind>::Owned>;

    async fn handle_batch(
        &self,
        batch: Vec<M>,
        values: &mut Self::Scratch,
        ctx: &mut Context<'_, D::Context, S>,
    ) {
        let subscription = ctx.subscription();
        let accepted = decode_batch::<M, D::Input, C, D::Context, S>(
            batch,
            values,
            &self.codec,
            self.decode,
            ctx,
        )
        .await;
        if accepted.is_empty() {
            return;
        }
        // The batch the broker delivered is the batch the handler answers for, whole: the size it
        // was built at is the registration's own, so there is nothing left to split here.
        let result = match self.def.call(values, &self.injections, ctx).await {
            Ok(replies) => {
                let name = self.def.reply_name();
                // As on the single-message path: the loop's buffer where the dispatch lent one,
                // and one of this batch's own where it did not.
                let mut own = BytesMut::new();
                let encode = ctx.take_encode_buffer().unwrap_or(&mut own);
                let slot = <R::Payload as PayloadForm>::encode_slot(encode);
                let pubcx = PublishContext::new(ctx.name(), ctx.delivery_headers(), ctx.cx_ref());
                let publish =
                    self.publisher
                        .publish_batch(name, &replies, &self.pipeline, &pubcx, slot);
                #[cfg(feature = "testing")]
                let publish = publishing_to(self.origin, publish);
                match publish.await {
                    Ok(()) => BatchResult::Uniform(HandlerOutcome::ack()),
                    Err(err) => {
                        warn!(
                            target: "ruststream::dispatch",
                            subscription = %subscription,
                            reply = %name,
                            reply_type = std::any::type_name::<D::Reply>(),
                            error = %err,
                            "batch reply publish failed",
                        );
                        BatchResult::Uniform(HandlerOutcome::retry())
                    }
                }
            }
            Err(result) => result,
        };
        values.clear();
        settle_batch(accepted, result, subscription, ctx.delivery()).await;
    }
}

#[cfg(all(test, feature = "memory", feature = "json"))]
mod tests;