ruststream 0.3.1

//! 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::handler::HandlerResult;
use super::metadata::HandlerMetadata;
use super::typed::DecodeFailure;

/// 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`.
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum BatchResult {
    /// One outcome settles every message of the batch.
    Uniform(HandlerResult),
    /// Outcome `i` settles slice element `i`. 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<HandlerResult>),
}

/// 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, with the same conventions as
/// [`IntoHandlerResult`](super::IntoHandlerResult)), and `Vec<HandlerResult>` (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<HandlerResult> {
    fn into_batch_result(self) -> BatchResult {
        BatchResult::PerElement(self)
    }
}

/// 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()
    }

    /// 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,
        on_decode_failure: DecodeFailure::default(),
        _phantom: PhantomData,
    }
}

/// The decode adapter for batches, the [`Typed`](super::Typed) counterpart.
///
/// Each element decodes independently: failures are settled individually per the
/// [`DecodeFailure`] policy 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,
    on_decode_failure: DecodeFailure,
    _phantom: PhantomData<fn(M, T)>,
}

impl<M, T, C, H> TypedBatch<M, T, C, H> {
    /// Override the behaviour when the codec fails to decode an element.
    #[must_use]
    #[allow(dead_code)] // The mount paths use the default; kept for parity with `Typed`.
    pub(crate) fn on_decode_failure(mut self, mode: DecodeFailure) -> Self {
        self.on_decode_failure = mode;
        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("on_decode_failure", &self.on_decode_failure)
            .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 (values, accepted) = decode_batch(batch, &self.codec, self.on_decode_failure).await;
        if accepted.is_empty() {
            return;
        }
        match self.inner.handle_slice(&values, ctx).await {
            BatchResult::Uniform(result) => {
                for msg in accepted {
                    settle(msg, result).await;
                }
            }
            BatchResult::PerElement(results) => {
                if results.len() != accepted.len() {
                    error!(
                        target: "ruststream::dispatch",
                        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 result = results.next().unwrap_or_else(HandlerResult::retry);
                    settle(msg, result).await;
                }
            }
        }
    }
}

/// Decodes each element of one raw batch independently: failures are nacked per
/// `on_decode_failure` and never reach the handler; the rest pass through, each decoded value
/// paired with its delivery (`values[i]` decodes `accepted[i]`).
pub(crate) async fn decode_batch<M, T, C>(
    batch: Vec<M>,
    codec: &C,
    on_decode_failure: DecodeFailure,
) -> (Vec<T>, Vec<M>)
where
    M: IncomingMessage,
    T: DeserializeOwned,
    C: Codec,
{
    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",
                    error = %err,
                    "codec decode failed",
                );
                let requeue = matches!(on_decode_failure, DecodeFailure::Requeue);
                if let Err(err) = msg.nack(requeue).await {
                    warn!(target: "ruststream::dispatch", error = %err, "nack failed");
                }
            }
        }
    }
    (values, accepted)
}

/// Applies one settlement to one delivery's own `ack` / `nack`.
pub(crate) async fn settle<M: IncomingMessage>(msg: M, result: HandlerResult) {
    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", 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]
    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();
        }
    }
}