ruststream 0.7.0-rc.11

//! The runtime collector routers and scopes mount into: type-erased starters plus metadata.

use std::{fmt, future::Future, num::NonZeroUsize, sync::Arc};

use tokio::task::JoinHandle;

use crate::{BatchSubscriber, Broker, Connected, Subscriber, SubscriptionSource};

use crate::runtime::batch::BatchHandler;
use crate::runtime::dispatch::{
    Workers, spawn_batch_dispatch, spawn_dispatch, spawn_dispatch_workers,
};
use crate::runtime::failure::{DispatchFailure, FailurePolicies};
use crate::runtime::handler::Handler;
use crate::runtime::lifecycle::{BoxError, BoxFuture};
use crate::runtime::metadata::HandlerMetadata;
use crate::runtime::redelivery::{
    CopyPathAddress, RetrySetup, ScopeDelivery, open_mounted_subscriber, open_subscription,
};
use crate::runtime::shutdown::Shutdown;

use super::SourceMessage;

/// A deferred registration: given the broker's connected form (produced by
/// [`Broker::connect`](crate::Broker::connect) at startup), shared state, the scope's publish
/// context and the shutdown signal, it opens the subscription and spawns the dispatch task. The
/// source and handler are captured and type-erased.
pub(crate) type BoundStarter<B, State> = Box<
    dyn FnOnce(
            Arc<Connected<B>>,
            Arc<State>,
            Arc<ScopeDelivery>,
            Shutdown,
        ) -> BoxFuture<'static, Result<JoinHandle<()>, BoxError>>
        + Send,
>;

/// The runtime collector a router mounts into: type-erased starters plus handler metadata.
///
/// `State` is the app's shared state type, threaded so a handler is only mounted on a sink whose state
/// type it matches.
///
/// Created and drained inside the application; a [`RouterDef`](crate::runtime::RouterDef) pushes
/// into it during [`include_router`](crate::runtime::BrokerScope::include_router). You do not
/// construct one directly.
pub struct RouterSink<B: Broker, State = ()> {
    starters: Vec<BoundStarter<B, State>>,
    handlers: Vec<HandlerMetadata>,
}

impl<B: Broker, State> fmt::Debug for RouterSink<B, State> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("RouterSink")
            .field("handlers", &self.handlers.len())
            .finish_non_exhaustive()
    }
}

impl<B: Broker + 'static, State: Send + Sync + 'static> RouterSink<B, State> {
    pub(crate) fn new() -> Self {
        Self {
            starters: Vec::new(),
            handlers: Vec::new(),
        }
    }

    /// Erases an already-created subscriber and its handler into a starter.
    pub(crate) fn push_handle<S, H, Cx>(
        &mut self,
        subscriber: S,
        handler: H,
        meta: HandlerMetadata,
        policies: FailurePolicies,
    ) where
        S: Subscriber + Send + 'static,
        Cx: crate::BuildContext<S::Message> + Send + Sync + 'static,
        H: Handler<S::Message, Cx, State> + 'static,
    {
        let handler = Arc::new(handler);
        let name: Arc<str> = Arc::from(meta.name.as_ref());
        self.starters
            .push(Box::new(move |_connected, state, scope, shutdown| {
                Box::pin(async move {
                    // No descriptor to ask where a copy of a delivery would go, and no mount chain
                    // that can bind one here, so this mount carries no retry path at all.
                    let delivery = open_mounted_subscriber(&scope, &name);
                    let failure = DispatchFailure::new(policies, shutdown.clone());
                    Ok(spawn_dispatch(
                        subscriber, handler, shutdown, name, state, delivery, failure,
                    ))
                })
            }));
        self.handlers.push(meta);
    }

    /// Erases a source and its batch handler into a starter driving
    /// [`BatchSubscriber::batches`]; the subscription opens against the connected broker. `Cx`
    /// is the definition's subscription-scoped batch context, passed explicitly because the
    /// adapter handlers are generic over it.
    // See `spawn_dispatch_workers`: each piece is the registration's own and bundling them would
    // hide that.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn push_subscribe_batch<S, H, Cx>(
        &mut self,
        source: S,
        handler: H,
        meta: HandlerMetadata,
        policies: FailurePolicies,
        workers: Workers,
        batch_size: NonZeroUsize,
        setup: RetrySetup<B, Cx>,
    ) where
        S: SubscriptionSource<Connected<B>> + Send + 'static,
        S::Subscriber: BatchSubscriber + Send + 'static,
        S::Copies: CopyPathAddress<Connected<B>, S>,
        SourceMessage<B, S>: Send + 'static,
        Cx: crate::BuildBatchContext<SourceMessage<B, S>> + Send + Sync + 'static,
        H: BatchHandler<SourceMessage<B, S>, Cx, State> + 'static,
    {
        let handler = Arc::new(handler);
        let name: Arc<str> = Arc::from(meta.name.as_ref());
        self.starters.push(Box::new(
            move |connected: Arc<Connected<B>>, state, scope, shutdown| {
                Box::pin(async move {
                    let (subscriber, delivery) = open_subscription::<B, _, _>(
                        source,
                        connected.as_ref(),
                        &scope,
                        &name,
                        setup,
                    )
                    .await?;
                    let failure = DispatchFailure::new(policies, shutdown.clone());
                    // Turbofish: the adapter handlers are generic over the batch context, so
                    // the pushed definition's own context names it.
                    Ok(spawn_batch_dispatch::<_, _, Cx, _>(
                        subscriber, handler, shutdown, name, state, delivery, failure, workers,
                        batch_size,
                    )?)
                })
            },
        ));
        self.handlers.push(meta);
    }

    /// Erases a source and its handler into a starter dispatching under the `workers` policy;
    /// the subscription opens against the connected broker.
    pub(crate) fn push_subscribe_workers<S, H, Cx>(
        &mut self,
        source: S,
        handler: H,
        meta: HandlerMetadata,
        policies: FailurePolicies,
        workers: Workers,
        setup: RetrySetup<B, Cx>,
    ) where
        S: SubscriptionSource<Connected<B>> + Send + 'static,
        S::Subscriber: Send + 'static,
        S::Copies: CopyPathAddress<Connected<B>, S>,
        SourceMessage<B, S>: Send + Sync + 'static,
        Cx: crate::BuildContext<SourceMessage<B, S>> + Send + Sync + 'static,
        H: Handler<SourceMessage<B, S>, Cx, State> + 'static,
    {
        let handler = Arc::new(handler);
        let name: Arc<str> = Arc::from(meta.name.as_ref());
        self.starters.push(Box::new(
            move |connected: Arc<Connected<B>>, state, scope, shutdown| {
                Box::pin(async move {
                    let (subscriber, delivery) = open_subscription::<B, _, _>(
                        source,
                        connected.as_ref(),
                        &scope,
                        &name,
                        setup,
                    )
                    .await?;
                    let failure = DispatchFailure::new(policies, shutdown.clone());
                    Ok(spawn_dispatch_workers(
                        subscriber, handler, shutdown, name, state, delivery, failure, workers,
                    )?)
                })
            },
        ));
        self.handlers.push(meta);
    }

    /// Pushes a fully custom starter, for the one mount the factory helpers cannot express:
    /// applying a [`BlanketLayer`](crate::runtime::BlanketLayer) inside the startup closure
    /// (its RPITIT return captures the layer borrow, so the applied handler cannot leave a
    /// factory closure).
    pub(crate) fn push_raw(&mut self, starter: BoundStarter<B, State>, meta: HandlerMetadata) {
        self.starters.push(starter);
        self.handlers.push(meta);
    }

    /// The batch counterpart of [`push_injected_workers`](Self::push_injected_workers): the
    /// factory resolves the injections off the opened subscriber, then the loop drives
    /// [`BatchSubscriber::batches`].
    // See `push_subscribe_batch`.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn push_injected_batch<Source, MakeHandler, HandlerFut, NewHandler, HandlerCx>(
        &mut self,
        source: Source,
        make_handler: MakeHandler,
        meta: HandlerMetadata,
        policies: FailurePolicies,
        workers: Workers,
        batch_size: NonZeroUsize,
        setup: RetrySetup<B, HandlerCx>,
    ) where
        Source: SubscriptionSource<Connected<B>> + Send + 'static,
        Source::Subscriber: BatchSubscriber + Send + 'static,
        Source::Copies: CopyPathAddress<Connected<B>, Source>,
        SourceMessage<B, Source>: Send + 'static,
        MakeHandler: FnOnce(Arc<Connected<B>>, Source::Subscriber) -> HandlerFut + Send + 'static,
        HandlerFut: Future<Output = Result<(Source::Subscriber, NewHandler), BoxError>> + Send,
        HandlerCx: crate::BuildBatchContext<SourceMessage<B, Source>> + Send + Sync + 'static,
        NewHandler: BatchHandler<SourceMessage<B, Source>, HandlerCx, State> + 'static,
    {
        let name: Arc<str> = Arc::from(meta.name.as_ref());
        self.starters.push(Box::new(
            move |connected: Arc<Connected<B>>, state, scope, shutdown| {
                Box::pin(async move {
                    let (subscriber, delivery) = open_subscription::<B, _, _>(
                        source,
                        connected.as_ref(),
                        &scope,
                        &name,
                        setup,
                    )
                    .await?;
                    let (subscriber, handler) =
                        make_handler(Arc::clone(&connected), subscriber).await?;
                    let failure = DispatchFailure::new(policies, shutdown.clone());
                    Ok(spawn_batch_dispatch::<_, _, HandlerCx, _>(
                        subscriber,
                        Arc::new(handler),
                        shutdown,
                        name,
                        state,
                        delivery,
                        failure,
                        workers,
                        batch_size,
                    )?)
                })
            },
        ));
        self.handlers.push(meta);
    }

    pub(crate) fn into_parts(self) -> (Vec<BoundStarter<B, State>>, Vec<HandlerMetadata>) {
        (self.starters, self.handlers)
    }
}