pulses 0.2.0

//! The [`App`] builder and runtime entry point.

use std::sync::Arc;

use tokio::sync::mpsc;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;

use crate::config::AppConfig;
use crate::core::Broker;
use crate::core::Context;
use crate::core::Handler;
use crate::core::Subscription;
use crate::error::AppError;
use crate::handler_set::HandlerSet;
use crate::handler_set::PoolWiring;
use crate::handler_set::build_routing_table;
use crate::runtime::ack_retrier::AckRetrier;
use crate::runtime::backoff::Backoff;
use crate::runtime::broker_reader::BrokerReader;
use crate::runtime::reclaimer::Reclaimer;
use crate::runtime::router::Router;

/// A configured worker application: a broker, a subscription, a compile-time
/// list of handlers, and runtime tuning.
///
/// Build one with [`App::new`], add handlers with [`App::register`], optionally
/// tune with [`App::with_config`], then drive it with [`App::run`].
///
/// ```no_run
/// # use pulses::{App, Context, Envelope, Handler, HandlerError, Outcome};
/// # use pulses::broker::redis::{RedisBroker, RedisSubscription};
/// # use tokio_util::sync::CancellationToken;
/// struct Worker;
/// impl Handler<RedisBroker> for Worker {
///     const STREAMS: &'static [&'static str] = &["jobs"];
///     async fn handle(&self, _ctx: &Context<RedisBroker>, msg: &Envelope) -> Result<Outcome, HandlerError> {
///         println!("processing {}", msg.id);
///         Ok(Outcome::Ack)
///     }
/// }
/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
/// let broker = RedisBroker::connect("redis://127.0.0.1:6379").await?;
/// let app = App::new(broker, RedisSubscription::for_group("workers", "node-1")).register(Worker);
/// app.run(CancellationToken::new()).await?;
/// # Ok(())
/// # }
/// ```
pub struct App<B, HS>
where
    B: Broker,
    HS: HandlerSet<B>,
{
    broker: B,
    subscription: B::Subscription,
    handlers: HS,
    config: AppConfig,
}

impl<B> App<B, ()>
where
    B: Broker,
{
    /// Start a new application with no handlers yet registered.
    #[must_use]
    pub fn new(broker: B, subscription: B::Subscription) -> Self {
        Self { broker, subscription, handlers: (), config: AppConfig::default() }
    }
}

impl<B, HS> App<B, HS>
where
    B: Broker,
    HS: HandlerSet<B>,
{
    /// Override the runtime configuration.
    #[must_use]
    pub fn with_config(mut self, config: AppConfig) -> Self {
        self.config = config;
        self
    }

    /// Register a handler. The handler's streams are automatically added to the
    /// broker subscription when the app runs.
    #[must_use]
    pub fn register<H>(self, handler: H) -> App<B, (HS, H)>
    where
        H: Handler<B>,
    {
        App {
            broker: self.broker,
            subscription: self.subscription,
            handlers: (self.handlers, handler),
            config: self.config,
        }
    }

    /// Borrow the current configuration.
    pub fn config(&self) -> &AppConfig { &self.config }

    /// Run the application until `cancellation` is triggered.
    ///
    /// Spawns one pool per handler plus the reader, reclaimer, and ack-retrier
    /// actors, then waits. On cancellation the actors stop accepting work and
    /// in-flight handler invocations are drained before returning. If any task
    /// panics, the app shuts the rest down and returns [`AppError::Task`].
    pub async fn run(self, cancellation: CancellationToken) -> Result<(), AppError> {
        self.config.validate()?;

        let routing_table = build_routing_table::<B, HS>();
        let streams: Vec<Arc<str>> = routing_table.keys().cloned().collect();

        let App { broker, mut subscription, handlers, config } = self;
        subscription.set_streams(streams);

        let (ack_retry_tx, ack_retry_rx) = mpsc::channel(config.ack_retry_queue_capacity.max(1));

        let mut tasks: JoinSet<()> = JoinSet::new();
        let mut senders = Vec::new();
        {
            let mut wiring = PoolWiring {
                context: Context::new(broker.clone()),
                config: config.clone(),
                backoff: Arc::new(Backoff::new(config.backoff.clone())),
                ack_retry_tx: ack_retry_tx.clone(),
                cancellation: &cancellation,
                handles: &mut tasks,
                senders: &mut senders,
            };
            handlers.spawn_pools(&mut wiring);
        }

        let router = Router::new(routing_table, senders);

        tasks.spawn(
            AckRetrier::new(
                broker.clone(),
                ack_retry_rx,
                Backoff::new(config.backoff.clone()),
                config.ack_max_attempts,
                cancellation.clone(),
            )
            .run(),
        );
        tasks.spawn(
            BrokerReader::new(
                broker.clone(),
                subscription.clone(),
                router.clone(),
                Backoff::new(config.backoff.clone()),
                config.poll_idle_sleep,
                cancellation.clone(),
            )
            .run(),
        );
        tasks.spawn(Reclaimer::new(broker, subscription, router, config.reclaim_interval, cancellation.clone()).run());

        // Release this scope's sender so the ack retrier can observe closure
        // once all pools have finished.
        drop(ack_retry_tx);

        let mut first_error: Option<AppError> = None;

        loop {
            tokio::select! {
                _ = cancellation.cancelled() => break,
                joined = tasks.join_next() => match joined {
                    // A task exited before shutdown. Ok means a clean stop; an
                    // error means a panic — surface it and bring the rest down.
                    Some(Ok(())) => continue,
                    Some(Err(join_error)) => {
                        tracing::error!(%join_error, "runtime task panicked; shutting down");
                        first_error = Some(AppError::from(join_error));
                        cancellation.cancel();
                        break;
                    }
                    None => break,
                },
            }
        }

        while let Some(joined) = tasks.join_next().await {
            if let Err(join_error) = joined {
                tracing::error!(%join_error, "runtime task failed during shutdown");
                if first_error.is_none() {
                    first_error = Some(AppError::from(join_error));
                }
            }
        }

        match first_error {
            Some(error) => Err(error),
            None => Ok(()),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::sync::atomic::AtomicUsize;
    use std::sync::atomic::Ordering;
    use std::time::Duration;

    use tokio::sync::Mutex;
    use tokio::time::timeout;
    use tokio_util::sync::CancellationToken;

    use super::App;
    use crate::config::AppConfig;
    use crate::core::Context;
    use crate::core::Envelope;
    use crate::core::Handler;
    use crate::core::HandlerError;
    use crate::core::Outcome;
    use crate::test_support::MockBroker;
    use crate::test_support::MockSubscription;
    use crate::test_support::make_envelope;
    use crate::test_support::token;

    #[derive(Clone)]
    struct Recorder {
        seen: Arc<Mutex<Vec<String>>>,
        counter: Arc<AtomicUsize>,
    }

    impl Handler<MockBroker> for Recorder {
        const STREAMS: &'static [&'static str] = &["orders"];
        async fn handle(&self, _ctx: &Context<MockBroker>, msg: &Envelope) -> Result<Outcome, HandlerError> {
            self.seen.lock().await.push(msg.id.to_string());
            self.counter.fetch_add(1, Ordering::SeqCst);
            Ok(Outcome::Ack)
        }
    }

    #[tokio::test]
    async fn processes_then_shuts_down_cleanly() {
        let broker = MockBroker::default();
        broker.enqueue_poll(vec![(make_envelope("orders", "1-1"), token("1-1"))]);

        let seen = Arc::new(Mutex::new(Vec::new()));
        let counter = Arc::new(AtomicUsize::new(0));
        let app = App::new(broker.clone(), MockSubscription::default())
            .with_config(AppConfig { reclaim_interval: Duration::from_millis(50), ..AppConfig::default() })
            .register(Recorder { seen: seen.clone(), counter: counter.clone() });

        let cancellation = CancellationToken::new();
        let run = app.run(cancellation.clone());
        tokio::pin!(run);

        timeout(Duration::from_secs(2), async {
            loop {
                if broker.ack_count() >= 1 {
                    return;
                }
                tokio::select! {
                    _ = tokio::time::sleep(Duration::from_millis(5)) => {}
                    result = &mut run => panic!("app exited early: {result:?}"),
                }
            }
        })
        .await
        .expect("app should process and ack the message");

        assert_eq!(seen.lock().await.clone(), vec!["1-1".to_string()]);
        assert_eq!(broker.ack_count(), 1);

        cancellation.cancel();
        timeout(Duration::from_secs(1), &mut run).await.expect("app should stop").expect("clean exit");
    }

    #[tokio::test]
    async fn rejects_invalid_config() {
        let broker = MockBroker::default();
        let app = App::new(broker, MockSubscription::default())
            .with_config(AppConfig { handler_queue_capacity: 0, ..AppConfig::default() })
            .register(Recorder { seen: Arc::new(Mutex::new(Vec::new())), counter: Arc::new(AtomicUsize::new(0)) });
        let result = app.run(CancellationToken::new()).await;
        assert!(matches!(result, Err(crate::error::AppError::Config(_))));
    }
}