pulses 0.2.0

//! Retries acknowledgements that failed in a handler pool, with bounded
//! attempts so a permanently-failing token cannot wedge the queue.

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

use crate::core::Broker;
use crate::runtime::backoff::Backoff;

pub(crate) struct AckRetrier<B: Broker> {
    broker: B,
    input_rx: mpsc::Receiver<B::AckToken>,
    backoff: Backoff,
    max_attempts: u32,
    cancellation: CancellationToken,
}

impl<B: Broker> AckRetrier<B> {
    pub(crate) fn new(
        broker: B, input_rx: mpsc::Receiver<B::AckToken>, backoff: Backoff, max_attempts: u32,
        cancellation: CancellationToken,
    ) -> Self {
        Self { broker, input_rx, backoff, max_attempts: max_attempts.max(1), cancellation }
    }

    /// Process tokens until the input channel closes.
    ///
    /// The channel closes only after every handler pool (and the [`App`]'s own
    /// sender) has been dropped at shutdown, so this naturally drains any acks
    /// queued during the pools' final drain. Cancellation is *not* a stop
    /// signal here — it only tells [`retry_one`](Self::retry_one) to stop
    /// scheduling further retries — which avoids discarding queued ack tokens
    /// (and the duplicate redeliveries that would cause).
    pub(crate) async fn run(mut self) {
        while let Some(token) = self.input_rx.recv().await {
            self.retry_one(token).await;
        }
    }

    async fn retry_one(&self, mut token: B::AckToken) {
        let mut attempt: u32 = 1;
        loop {
            // The ack itself is always allowed to complete (never abandoned
            // mid-flight on cancellation), so a token already being acked is
            // not lost during shutdown.
            let retry_token = token.clone();
            match self.broker.ack(token).await {
                Ok(()) => return,
                Err(error) => {
                    if attempt >= self.max_attempts {
                        tracing::error!(attempt, %error, "ack retry exhausted; giving up (message may be redelivered)");
                        return;
                    }
                    if self.cancellation.is_cancelled() {
                        tracing::warn!("shutting down; abandoning further ack retries");
                        return;
                    }
                }
            }

            let delay = self.backoff.delay(attempt);
            attempt = attempt.saturating_add(1);
            token = retry_token;

            tokio::select! {
                _ = self.cancellation.cancelled() => return,
                _ = tokio::time::sleep(delay) => {}
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

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

    use super::AckRetrier;
    use crate::config::BackoffConfig;
    use crate::runtime::backoff::Backoff;
    use crate::test_support::MockBroker;
    use crate::test_support::token;

    fn fast_backoff() -> Backoff {
        Backoff::new(BackoffConfig {
            initial: Duration::from_millis(1),
            max: Duration::from_millis(2),
            factor: 2.0,
            jitter: 0.0,
        })
    }

    #[tokio::test]
    async fn retries_until_success() {
        let broker = MockBroker::default();
        broker.fail_next_acks(2);
        let (tx, rx) = mpsc::channel(8);
        let cancellation = CancellationToken::new();
        let retrier = AckRetrier::new(broker.clone(), rx, fast_backoff(), 8, cancellation.clone());
        let handle = tokio::spawn(retrier.run());

        tx.send(token("1-1")).await.unwrap();

        timeout(Duration::from_secs(1), async {
            while broker.ack_count() < 1 {
                tokio::time::sleep(Duration::from_millis(1)).await;
            }
        })
        .await
        .expect("ack should eventually succeed");

        // First two attempts failed, third succeeded.
        assert_eq!(broker.ack_attempts(), 3);

        cancellation.cancel();
        drop(tx);
        timeout(Duration::from_secs(1), handle).await.unwrap().unwrap();
    }

    #[tokio::test]
    async fn gives_up_after_max_attempts() {
        let broker = MockBroker::default();
        broker.fail_next_acks(100);
        let (tx, rx) = mpsc::channel(8);
        let cancellation = CancellationToken::new();
        let retrier = AckRetrier::new(broker.clone(), rx, fast_backoff(), 3, cancellation.clone());
        let handle = tokio::spawn(retrier.run());

        tx.send(token("1-1")).await.unwrap();

        timeout(Duration::from_secs(1), async {
            while broker.ack_attempts() < 3 {
                tokio::time::sleep(Duration::from_millis(1)).await;
            }
        })
        .await
        .expect("retrier should attempt up to the cap");

        // Give it a moment to ensure it does not exceed the cap.
        tokio::time::sleep(Duration::from_millis(20)).await;
        assert_eq!(broker.ack_attempts(), 3);

        cancellation.cancel();
        drop(tx);
        timeout(Duration::from_secs(1), handle).await.unwrap().unwrap();
    }
}