ruststream 0.4.0

//! Batch subscriber definitions: handlers that consume a whole decoded batch per invocation.
//!
//! Generated by `#[subscriber(batch(..))]` and mounted with
//! [`BrokerScope::include_batch`](super::BrokerScope::include_batch); the mount site requires the
//! source's subscriber to implement [`BatchSubscriber`](crate::BatchSubscriber) (natively, or via
//! the [`Buffered`](crate::Buffered) adapter). Mirrors the single-message pipeline:
//! [`BatchDef`] is the [`SubscriberDef`](super::SubscriberDef) counterpart, [`SliceHandler`] the
//! [`Handler`](super::Handler) counterpart, and [`TypedBatch`] the [`Typed`](super::Typed) decode
//! adapter.

use std::{future::Future, marker::PhantomData};

use serde::de::DeserializeOwned;
use tracing::{error, warn};

use crate::IncomingMessage;
use crate::codec::Codec;

use super::context::Context;
use super::dispatch::Workers;
use super::failure::{FailurePolicies, FailurePolicy};
use super::handler::{HandlerResult, Settle};
use super::metadata::HandlerMetadata;

/// The settlement of one dispatched batch.
///
/// Returned by batch handlers (usually through [`IntoBatchResult`]) and applied by the
/// dispatcher to each message's own `ack` / `nack`.
///
/// Per-element [`Settle`]s carry an optional `and_after` continuation each; the uniform form is a
/// bare outcome, since one continuation cannot fan out to every message of the batch.
#[derive(Debug)]
#[non_exhaustive]
pub enum BatchResult {
    /// One outcome settles every message of the batch.
    Uniform(HandlerResult),
    /// Settlement `i` settles slice element `i`, each with its own optional post-settle
    /// continuation. A length mismatch with the dispatched batch is a bug in the handler: the
    /// unmatched remainder is retried (an extra redelivery beats a silently lost message) and the
    /// mismatch is logged.
    PerElement(Vec<Settle>),
}

/// Conversion into a [`BatchResult`], so `#[subscriber(batch(..))]` handlers can return a plain
/// value.
///
/// Implemented for [`BatchResult`] (identity), [`HandlerResult`] / `()` / `Result<(), E>` /
/// `Result<HandlerResult, E>` (one outcome for the whole batch), and a per-element vector
/// (`Vec<Settle>`, or `Vec<HandlerResult>`) where element `i` settles slice element `i`.
pub trait IntoBatchResult {
    /// Converts `self` into the settlement the dispatcher applies.
    fn into_batch_result(self) -> BatchResult;
}

impl IntoBatchResult for BatchResult {
    fn into_batch_result(self) -> BatchResult {
        self
    }
}

impl IntoBatchResult for HandlerResult {
    fn into_batch_result(self) -> BatchResult {
        BatchResult::Uniform(self)
    }
}

impl IntoBatchResult for () {
    fn into_batch_result(self) -> BatchResult {
        BatchResult::Uniform(HandlerResult::Ack)
    }
}

impl<E> IntoBatchResult for Result<(), E> {
    fn into_batch_result(self) -> BatchResult {
        BatchResult::Uniform(match self {
            Ok(()) => HandlerResult::Ack,
            Err(_) => HandlerResult::drop(),
        })
    }
}

impl<E> IntoBatchResult for Result<HandlerResult, E> {
    fn into_batch_result(self) -> BatchResult {
        BatchResult::Uniform(self.unwrap_or_else(|_| HandlerResult::drop()))
    }
}

impl IntoBatchResult for Vec<Settle> {
    fn into_batch_result(self) -> BatchResult {
        BatchResult::PerElement(self)
    }
}

impl IntoBatchResult for Vec<HandlerResult> {
    fn into_batch_result(self) -> BatchResult {
        BatchResult::PerElement(self.into_iter().map(Settle::from).collect())
    }
}

/// A handler invoked with one whole decoded batch.
///
/// The batch parameter is a slice: per-message broker handles stay with the dispatcher, which
/// settles every message of the batch according to the returned [`BatchResult`] - one uniform
/// outcome, or one outcome per element.
///
/// # Examples
///
/// Closures returning any [`IntoBatchResult`] implement `SliceHandler` automatically:
///
/// ```
/// use ruststream::runtime::{Context, HandlerResult, SliceHandler};
///
/// fn assert_slice_handler<T, H: SliceHandler<T>>(_: H) {}
///
/// fn use_closures() {
///     // One outcome for the whole batch.
///     assert_slice_handler::<u32, _>(|batch: &[u32], _ctx: &mut Context| {
///         let _ = batch.len();
///         async { HandlerResult::Ack }
///     });
///     // One outcome per element: entries that are not ready yet retry individually.
///     assert_slice_handler::<u32, _>(|batch: &[u32], _ctx: &mut Context| {
///         let outcomes: Vec<HandlerResult> = batch
///             .iter()
///             .map(|n| {
///                 if *n == 0 {
///                     HandlerResult::retry()
///                 } else {
///                     HandlerResult::Ack
///                 }
///             })
///             .collect();
///         async move { outcomes }
///     });
/// }
/// ```
pub trait SliceHandler<T>: Send + Sync {
    /// Handles one decoded batch, with the per-batch [`Context`].
    fn handle_slice(
        &self,
        batch: &[T],
        ctx: &mut Context,
    ) -> impl Future<Output = BatchResult> + Send;
}

impl<T, F, Fut> SliceHandler<T> for F
where
    F: Fn(&[T], &mut Context) -> Fut + Send + Sync,
    Fut: Future + Send,
    Fut::Output: IntoBatchResult,
{
    fn handle_slice(
        &self,
        batch: &[T],
        ctx: &mut Context,
    ) -> impl Future<Output = BatchResult> + Send {
        // Build the inner future before the async block so the returned future does not hold
        // `&[T]` (which would demand `T: Sync` for it to be `Send`).
        let fut = (self)(batch, ctx);
        async move { fut.await.into_batch_result() }
    }
}

/// A batch handler definition produced by `#[subscriber(batch(..))]`.
///
/// The batch counterpart of [`SubscriberDef`](super::SubscriberDef): same metadata surface, but
/// the handler consumes `&[Input]` and the mount site drives
/// [`BatchSubscriber::batches`](crate::BatchSubscriber::batches) instead of
/// [`Subscriber::stream`](crate::Subscriber::stream).
pub trait BatchDef: Sized {
    /// The decoded element type; the handler consumes `&[Input]`.
    type Input;

    /// The concrete handler type over batches of [`Input`](Self::Input).
    type Handler: SliceHandler<Self::Input>;

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

    /// Consumes the definition, returning the handler.
    fn into_handler(self) -> Self::Handler;
}

/// Builds the registration metadata for a batch definition mounted under `name`.
pub(crate) fn batch_metadata<D: BatchDef>(name: String, def: &D) -> HandlerMetadata {
    HandlerMetadata::typed::<D::Input>(name).with_def_details(
        def.description(),
        def.input_schema(),
        def.message_name(),
        def.message_description(),
    )
}

/// The dispatch-side consumer of one raw batch: decode, run the handler, settle every delivery.
/// The batch counterpart of [`Handler`](super::Handler) at the raw-message level.
pub(crate) trait BatchHandler<M>: Send + Sync {
    /// Consumes one batch of raw deliveries, acknowledging each of them.
    fn handle_batch(&self, batch: Vec<M>, ctx: &mut Context) -> impl Future<Output = ()> + Send;
}

/// Build a [`TypedBatch`] that decodes each element with `codec` into `T` and forwards the batch
/// as `&[T]` to `inner`.
pub(crate) fn typed_batch<M, T, C, H>(codec: C, inner: H) -> TypedBatch<M, T, C, H>
where
    M: IncomingMessage,
    T: DeserializeOwned + Send + Sync,
    C: Codec,
    H: SliceHandler<T>,
{
    TypedBatch {
        codec,
        inner,
        decode: FailurePolicy::Drop,
        _phantom: PhantomData,
    }
}

/// The decode adapter for batches, the [`Typed`](super::Typed) counterpart.
///
/// Each element decodes independently: failures are settled individually per the `decode`
/// [`FailurePolicy`] and never reach the handler; the rest are passed on as one slice. The
/// [`BatchResult`] the handler returns settles the deliveries behind that slice - uniformly, or
/// element by element.
pub struct TypedBatch<M, T, C, H> {
    codec: C,
    inner: H,
    decode: FailurePolicy,
    _phantom: PhantomData<fn(M, T)>,
}

impl<M, T, C, H> TypedBatch<M, T, C, H> {
    /// Sets the policy applied when the codec fails to decode an element.
    #[must_use]
    pub(crate) fn with_decode(mut self, decode: FailurePolicy) -> Self {
        self.decode = decode;
        self
    }
}

impl<M, T, C, H> std::fmt::Debug for TypedBatch<M, T, C, H> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TypedBatch")
            .field("decode", &self.decode)
            .finish_non_exhaustive()
    }
}

impl<M, T, C, H> BatchHandler<M> for TypedBatch<M, T, C, H>
where
    M: IncomingMessage,
    T: DeserializeOwned + Send + Sync,
    C: Codec,
    H: SliceHandler<T>,
{
    async fn handle_batch(&self, batch: Vec<M>, ctx: &mut Context<'_>) {
        let subscription = ctx.name().to_owned();
        let (values, accepted) = decode_batch(batch, &self.codec, self.decode, ctx).await;
        if accepted.is_empty() {
            return;
        }
        let tasks = ctx.tasks().clone();
        match self.inner.handle_slice(&values, ctx).await {
            BatchResult::Uniform(result) => {
                for msg in accepted {
                    settle(msg, result, &subscription).await;
                }
            }
            BatchResult::PerElement(results) => {
                if results.len() != accepted.len() {
                    error!(
                        target: "ruststream::dispatch",
                        subscription = %subscription,
                        expected = accepted.len(),
                        returned = results.len(),
                        "per-element outcome count does not match the batch; \
                         retrying the unmatched remainder",
                    );
                }
                let mut results = results.into_iter();
                for msg in accepted {
                    // An unmatched message gets retried: an extra redelivery beats losing it.
                    let mut result = results
                        .next()
                        .unwrap_or_else(|| HandlerResult::retry().into());
                    let after = result.take_after();
                    settle(msg, result.outcome(), &subscription).await;
                    // The continuation runs after this element is settled, on the tracked set so a
                    // graceful shutdown drains it. At-most-once: a lost or panicking continuation
                    // never redelivers the already-settled message.
                    if let Some(after) = after {
                        tasks.spawn(after);
                    }
                }
            }
        }
    }
}

/// Decodes each element of one raw batch independently: failures are settled per the `decode`
/// [`FailurePolicy`] and never reach the handler; the rest pass through, each decoded value paired
/// with its delivery (`values[i]` decodes `accepted[i]`). A `fail_fast` decode policy tears the
/// service down via `ctx` and drops the offending element so it is not requeued into the failure.
///
/// Takes `&mut Context` (rather than `&Context`) so the future stays `Send`: `Context` carries
/// post-settle hooks (boxed `Send` futures that are not `Sync`), so holding a shared `&Context`
/// across an `.await` would wrongly require `Context: Sync`.
// The `&mut` is for Send-ness, not mutation (only `&self` methods are called), so the lint fires.
#[allow(clippy::needless_pass_by_ref_mut)]
pub(crate) async fn decode_batch<M, T, C>(
    batch: Vec<M>,
    codec: &C,
    decode: FailurePolicy,
    ctx: &mut Context<'_>,
) -> (Vec<T>, Vec<M>)
where
    M: IncomingMessage,
    T: DeserializeOwned,
    C: Codec,
{
    let subscription = ctx.name().to_owned();
    let mut values = Vec::with_capacity(batch.len());
    let mut accepted = Vec::with_capacity(batch.len());
    for msg in batch {
        match codec.decode::<T>(msg.payload()) {
            Ok(value) => {
                values.push(value);
                accepted.push(msg);
            }
            Err(err) => {
                warn!(
                    target: "ruststream::dispatch",
                    subscription = %subscription,
                    message_type = std::any::type_name::<T>(),
                    error = %err,
                    "codec decode failed",
                );
                let outcome = match decode {
                    FailurePolicy::FailFast => {
                        ctx.fail_fast(&format!("batch decode failed: {err}"));
                        HandlerResult::drop()
                    }
                    other => other.settlement().unwrap_or_else(HandlerResult::drop),
                };
                settle(msg, outcome, &subscription).await;
            }
        }
    }
    (values, accepted)
}

/// Applies one settlement to one delivery's own `ack` / `nack`.
pub(crate) async fn settle<M: IncomingMessage>(msg: M, result: HandlerResult, subscription: &str) {
    let ack_result = match result {
        HandlerResult::Ack => msg.ack().await,
        HandlerResult::Nack { requeue } => msg.nack(requeue).await,
        HandlerResult::NackAfter { delay } => msg.nack_after(delay).await,
    };
    if let Err(err) = ack_result {
        warn!(
            target: "ruststream::dispatch",
            subscription = %subscription,
            error = %err,
            "ack / nack failed",
        );
    }
}

#[cfg(all(test, feature = "memory", feature = "json"))]
mod tests {
    use futures::StreamExt;

    use super::super::context::State;
    use super::super::dispatch::Delivery;
    use super::*;
    use crate::codec::JsonCodec;
    use crate::memory::{MemoryBroker, MemoryMessage, MemorySubscriber};
    use crate::{BatchSubscriber, Headers, OutgoingMessage, Publisher, Subscriber};

    async fn publish_numbers(broker: &MemoryBroker, name: &str, numbers: &[u32]) {
        let publisher = broker.publisher();
        for n in numbers {
            publisher
                .publish(OutgoingMessage::new(name, &serde_json::to_vec(n).unwrap()))
                .await
                .unwrap();
        }
    }

    async fn pull_batch(sub: &mut MemorySubscriber) -> Vec<MemoryMessage> {
        let mut stream = std::pin::pin!(sub.batches());
        stream.next().await.unwrap().unwrap()
    }

    #[tokio::test]
    async fn per_element_outcomes_settle_individually() {
        let broker = MemoryBroker::new();
        let mut sub = broker.subscribe("selective");
        publish_numbers(&broker, "selective", &[0, 1, 2]).await;

        // 0 acks, 1 retries, 2 drops: only 1 may come back.
        let handler = typed_batch(JsonCodec, |batch: &[u32], _ctx: &mut Context| {
            let outcomes: Vec<HandlerResult> = batch
                .iter()
                .map(|n| match n {
                    1 => HandlerResult::retry(),
                    2 => HandlerResult::drop(),
                    _ => HandlerResult::Ack,
                })
                .collect();
            async move { outcomes }
        });

        let state = State::default();
        let delivery = Delivery::empty();
        let headers = Headers::new();
        let mut ctx = Context::new("selective", &headers, &state, &delivery);
        let batch = pull_batch(&mut sub).await;
        assert_eq!(batch.len(), 3);
        handler.handle_batch(batch, &mut ctx).await;

        let redelivered = pull_batch(&mut sub).await;
        let payloads: Vec<&[u8]> = redelivered.iter().map(IncomingMessage::payload).collect();
        assert_eq!(payloads, [b"1"]);
        for msg in redelivered {
            msg.ack().await.unwrap();
        }
        let mut stream = std::pin::pin!(sub.stream());
        assert!(futures::poll!(stream.next()).is_pending());
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn per_element_continuations_run_after_settle() {
        use std::sync::Arc;
        use tokio::sync::Notify;
        use tokio_util::task::TaskTracker;

        let broker = MemoryBroker::new();
        let mut sub = broker.subscribe("after-batch");
        publish_numbers(&broker, "after-batch", &[0, 1]).await;

        // Element 0 acks with a continuation; element 1 retries with no continuation.
        let ran = Arc::new(Notify::new());
        let signal = Arc::clone(&ran);
        let handler = typed_batch(JsonCodec, move |batch: &[u32], _ctx: &mut Context| {
            let signal = Arc::clone(&signal);
            let outcomes: Vec<Settle> = batch
                .iter()
                .map(|n| {
                    if *n == 0 {
                        let signal = Arc::clone(&signal);
                        HandlerResult::ack().and_after(async move { signal.notify_one() })
                    } else {
                        HandlerResult::retry().into()
                    }
                })
                .collect();
            async move { outcomes }
        });

        let tasks = TaskTracker::new();
        let state = State::default();
        let delivery = Delivery::with_tasks(tasks.clone());
        let headers = Headers::new();
        let mut ctx = Context::new("after-batch", &headers, &state, &delivery);
        let batch = pull_batch(&mut sub).await;
        handler.handle_batch(batch, &mut ctx).await;

        // The continuation for element 0 runs on the tracked set after settling.
        ran.notified().await;
        tasks.close();
        tasks.wait().await;

        // Element 1 (no continuation) retried and comes back; element 0 is gone.
        let redelivered = pull_batch(&mut sub).await;
        let payloads: Vec<&[u8]> = redelivered.iter().map(IncomingMessage::payload).collect();
        assert_eq!(payloads, [b"1"]);
        for msg in redelivered {
            msg.ack().await.unwrap();
        }
    }

    #[tokio::test]
    async fn unmatched_remainder_is_retried() {
        let broker = MemoryBroker::new();
        let mut sub = broker.subscribe("short");
        publish_numbers(&broker, "short", &[0, 1, 2]).await;

        // A buggy handler returning one outcome for a batch of three: the unmatched two retry.
        let handler = typed_batch(JsonCodec, |_batch: &[u32], _ctx: &mut Context| async {
            vec![HandlerResult::Ack]
        });

        let state = State::default();
        let delivery = Delivery::empty();
        let headers = Headers::new();
        let mut ctx = Context::new("short", &headers, &state, &delivery);
        let batch = pull_batch(&mut sub).await;
        assert_eq!(batch.len(), 3);
        handler.handle_batch(batch, &mut ctx).await;

        let redelivered = pull_batch(&mut sub).await;
        let payloads: Vec<&[u8]> = redelivered.iter().map(IncomingMessage::payload).collect();
        assert_eq!(payloads, [b"1", b"2"]);
        for msg in redelivered {
            msg.ack().await.unwrap();
        }
    }

    // Paused time (current-thread runtime): the per-element delay auto-advances.
    #[tokio::test(start_paused = true)]
    async fn per_element_outcomes_carry_delays() {
        let broker = MemoryBroker::new();
        let mut sub = broker.subscribe("delayed");
        publish_numbers(&broker, "delayed", &[0, 1]).await;

        // 0 acks; 1 retries no sooner than five seconds from now.
        let handler = typed_batch(JsonCodec, |batch: &[u32], _ctx: &mut Context| {
            let outcomes: Vec<HandlerResult> = batch
                .iter()
                .map(|n| match n {
                    1 => HandlerResult::retry_after(std::time::Duration::from_secs(5)),
                    _ => HandlerResult::Ack,
                })
                .collect();
            async move { outcomes }
        });

        let state = State::default();
        let delivery = Delivery::empty();
        let headers = Headers::new();
        let mut ctx = Context::new("delayed", &headers, &state, &delivery);
        let batch = pull_batch(&mut sub).await;
        handler.handle_batch(batch, &mut ctx).await;

        let mut stream = std::pin::pin!(sub.stream());
        assert!(futures::poll!(stream.next()).is_pending());
        tokio::time::advance(std::time::Duration::from_secs(5)).await;
        tokio::task::yield_now().await;

        let redelivered = stream.next().await.unwrap().unwrap();
        assert_eq!(redelivered.payload(), b"1");
        redelivered.ack().await.unwrap();
    }

    #[tokio::test]
    async fn uniform_outcome_settles_the_whole_batch() {
        let broker = MemoryBroker::new();
        let mut sub = broker.subscribe("uniform");
        publish_numbers(&broker, "uniform", &[0, 1]).await;

        let handler = typed_batch(JsonCodec, |_batch: &[u32], _ctx: &mut Context| async {
            HandlerResult::retry()
        });

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

        let redelivered = pull_batch(&mut sub).await;
        assert_eq!(redelivered.len(), 2);
        for msg in redelivered {
            msg.ack().await.unwrap();
        }
    }
}