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