pulses 0.2.0

A robust, high-performance background job processing library for Rust.
Documentation
//! Continuously polls the broker for new messages and routes them to handler
//! pools. Self-healing: transient poll failures back off and retry rather than
//! terminating the actor.

use std::time::Duration;

use tokio_util::sync::CancellationToken;

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

pub(crate) struct BrokerReader<B: Broker> {
    broker: B,
    subscription: B::Subscription,
    router: Router<B::AckToken>,
    backoff: Backoff,
    poll_idle_sleep: Duration,
    cancellation: CancellationToken,
}

impl<B: Broker> BrokerReader<B> {
    pub(crate) fn new(
        broker: B, subscription: B::Subscription, router: Router<B::AckToken>, backoff: Backoff,
        poll_idle_sleep: Duration, cancellation: CancellationToken,
    ) -> Self {
        Self { broker, subscription, router, backoff, poll_idle_sleep, cancellation }
    }

    pub(crate) async fn run(self) {
        let mut attempt: u32 = 1;

        loop {
            let poll_result = tokio::select! {
                _ = self.cancellation.cancelled() => return,
                result = self.broker.poll(&self.subscription) => result,
            };

            match poll_result {
                Ok(batch) => {
                    attempt = 1;
                    let was_empty = batch.is_empty();

                    for (envelope, token) in batch {
                        if let Err(error) = self.router.route(envelope, token, &self.cancellation).await {
                            tracing::error!(%error, "failed to route polled message");
                        }
                        if self.cancellation.is_cancelled() {
                            return;
                        }
                    }

                    // Guard against a hot loop if the broker is configured for
                    // non-blocking reads.
                    if was_empty {
                        tokio::select! {
                            _ = self.cancellation.cancelled() => return,
                            _ = tokio::time::sleep(self.poll_idle_sleep) => {}
                        }
                    }
                }
                Err(error) => {
                    tracing::warn!(%error, "broker poll failed; backing off");
                    let delay = self.backoff.delay(attempt);
                    attempt = attempt.saturating_add(1);
                    tokio::select! {
                        _ = self.cancellation.cancelled() => return,
                        _ = tokio::time::sleep(delay) => {}
                    }
                }
            }
        }
    }
}

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

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

    use super::BrokerReader;
    use crate::config::BackoffConfig;
    use crate::handler_set::RoutingTable;
    use crate::runtime::backoff::Backoff;
    use crate::runtime::router::Router;
    use crate::test_support::MockBroker;
    use crate::test_support::make_envelope;
    use crate::test_support::token;

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

        let (tx, mut rx) = mpsc::channel(8);
        let mut table = RoutingTable::new();
        table.insert(Arc::from("orders"), 1 << 0);
        let router = Router::new(table, vec![tx]);
        let cancellation = CancellationToken::new();

        let reader = BrokerReader::new(
            broker,
            crate::test_support::MockSubscription::default(),
            router,
            Backoff::new(BackoffConfig::default()),
            Duration::from_millis(5),
            cancellation.clone(),
        );
        let handle = tokio::spawn(reader.run());

        let task = timeout(Duration::from_secs(1), rx.recv()).await.unwrap().unwrap();
        assert_eq!(task.envelope.id.as_ref(), "1-1");

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