use std::fmt;
use std::future::{Future, poll_fn};
use std::hash::{DefaultHasher, Hash, Hasher};
use std::num::NonZeroUsize;
use std::panic::AssertUnwindSafe;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context as TaskContext, Poll};
use std::time::Duration;
use bytes::{Bytes, BytesMut};
use futures::{FutureExt, Stream};
#[cfg(test)]
use tokio::runtime::Handle;
use tokio::task::{JoinError, JoinHandle, JoinSet};
use tokio_util::task::TaskTracker;
use tracing::{debug, error, warn};
use crate::{
AckError, BatchSubscriber, HeaderMap, IncomingMessage, OutgoingMessage, RetryDeclaration, Str,
Subscriber,
};
use super::batch::{BatchHandler, TakenBatch};
use super::context::{Context, FromDelivery};
use super::failure::{DispatchFailure, FailurePolicy, panic_reason};
use super::handler::{Handler, HandlerResult};
use super::main_runtime::MainRuntime;
#[cfg(all(test, feature = "poll-diagnostics"))]
use super::poll_diagnostics::PollDiagnostics;
#[cfg(feature = "poll-diagnostics")]
use super::poll_diagnostics::PollSampler;
use super::publish::PublishContext;
#[cfg(all(test, feature = "memory"))]
use super::redelivery::ErasedRetryPublisher;
use super::redelivery::{DeferredRetry, ScopeDelivery};
use super::shutdown::Shutdown;
#[cfg(feature = "testing")]
use crate::testing::coordinator::{
Coordinator, Delivered, HarnessScope, Record, TestHooks, in_harness_scope,
};
pub(crate) use threads::StartThreadError;
use threads::Threads;
pub const RETRY_COUNT_HEADER: &str = "x-ruststream-retry-count";
fn current_retry_count(headers: &HeaderMap) -> u64 {
headers
.get_str(RETRY_COUNT_HEADER)
.and_then(|v| v.parse().ok())
.unwrap_or(0)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Workers {
count: NonZeroUsize,
by_key: bool,
placement: Placement,
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Placement {
Runtime,
Threads,
}
impl Workers {
#[must_use]
pub const fn sequential() -> Self {
Self {
count: NonZeroUsize::MIN,
by_key: false,
placement: Placement::Runtime,
}
}
#[must_use]
pub const fn pool(count: NonZeroUsize) -> Self {
Self {
count,
by_key: false,
placement: Placement::Runtime,
}
}
#[must_use]
pub const fn keyed(count: NonZeroUsize) -> Self {
Self {
count,
by_key: true,
placement: Placement::Runtime,
}
}
#[must_use]
pub const fn threads(count: NonZeroUsize) -> Self {
Self {
count,
by_key: false,
placement: Placement::Threads,
}
}
#[must_use]
pub const fn threads_keyed(count: NonZeroUsize) -> Self {
Self {
count,
by_key: true,
placement: Placement::Threads,
}
}
#[must_use]
pub const fn count(&self) -> NonZeroUsize {
self.count
}
#[must_use]
pub const fn by_key(&self) -> bool {
self.by_key
}
#[must_use]
pub const fn placement(&self) -> Placement {
self.placement
}
pub(crate) const fn is_sequential(&self) -> bool {
self.count.get() == 1 && matches!(self.placement, Placement::Runtime)
}
#[cfg(feature = "testing")]
pub(crate) const fn on_runtime(self) -> Self {
Self {
placement: Placement::Runtime,
..self
}
}
}
impl Default for Workers {
fn default() -> Self {
Self::sequential()
}
}
pub(crate) struct Delivery<C = ()> {
pub(crate) retry: Option<DeferredRetry<C>>,
pub(crate) declaration: RetryDeclaration,
pub(crate) tasks: TaskTracker,
pub(crate) runtime: MainRuntime,
pub(crate) thread_tasks: Option<TaskTracker>,
#[cfg(feature = "testing")]
pub(crate) hooks: Arc<TestHooks>,
#[cfg(feature = "testing")]
pub(crate) scope_id: usize,
#[cfg(feature = "testing")]
pub(crate) subscription: usize,
#[cfg(feature = "poll-diagnostics")]
pub(crate) poll: PollSampler,
}
#[cfg(feature = "poll-diagnostics")]
macro_rules! observed {
($delivery:expr, $future:expr) => {
$delivery.poll.observe($future)
};
}
#[cfg(not(feature = "poll-diagnostics"))]
macro_rules! observed {
($delivery:expr, $future:expr) => {
$future
};
}
impl<C> Delivery<C> {
#[cfg_attr(not(feature = "poll-diagnostics"), allow(unused_variables))]
pub(crate) fn for_subscription(
scope: &ScopeDelivery,
subscription: &str,
retry: Option<DeferredRetry<C>>,
declaration: RetryDeclaration,
) -> Self {
Self {
retry,
declaration,
tasks: scope.tasks().clone(),
runtime: MainRuntime::current(),
thread_tasks: None,
#[cfg(feature = "testing")]
hooks: Arc::clone(scope.hooks()),
#[cfg(feature = "testing")]
scope_id: scope.scope_id(),
#[cfg(feature = "testing")]
subscription: scope.subscription(),
#[cfg(feature = "poll-diagnostics")]
poll: scope.poll().register(subscription),
}
}
#[cfg(test)]
pub(crate) fn detached(retry: Option<DeferredRetry<C>>, tasks: TaskTracker) -> Self {
Self {
retry,
declaration: RetryDeclaration::new(),
tasks,
runtime: MainRuntime::new(test_runtime()),
thread_tasks: None,
#[cfg(feature = "testing")]
hooks: Arc::new(TestHooks::detached()),
#[cfg(feature = "testing")]
scope_id: 0,
#[cfg(feature = "testing")]
subscription: 0,
#[cfg(feature = "poll-diagnostics")]
poll: PollDiagnostics::new().register("detached"),
}
}
#[cfg(all(test, feature = "memory"))]
pub(crate) fn deferring_to(
publisher: Arc<dyn ErasedRetryPublisher<C>>,
destination: &str,
tasks: TaskTracker,
) -> Self {
Self::detached(
Some(DeferredRetry {
publisher,
destination: Some(Arc::from(destination)),
}),
tasks,
)
}
#[cfg(all(test, feature = "memory"))]
pub(crate) fn declaring(mut self, declaration: RetryDeclaration) -> Self {
self.declaration = declaration;
self
}
#[cfg(test)]
pub(crate) fn empty() -> Self {
Self::with_tasks(TaskTracker::new())
}
#[cfg(test)]
pub(crate) fn with_tasks(tasks: TaskTracker) -> Self {
Self::detached(None, tasks)
}
}
impl<C> Delivery<C> {
pub(crate) fn spawn_after<Work>(&self, work: Work)
where
Work: Future<Output = ()> + Send + 'static,
{
match &self.thread_tasks {
None => drop(self.tasks.spawn(work)),
Some(own) => drop(self.tasks.spawn(own.track_future(work))),
}
}
fn on_thread(&self, own: TaskTracker) -> Self {
Self {
retry: self.retry.as_ref().map(|retry| DeferredRetry {
publisher: Arc::clone(&retry.publisher),
destination: retry.destination.clone(),
}),
declaration: self.declaration.clone(),
tasks: self.tasks.clone(),
runtime: self.runtime.clone(),
thread_tasks: Some(own),
#[cfg(feature = "testing")]
hooks: Arc::clone(&self.hooks),
#[cfg(feature = "testing")]
scope_id: self.scope_id,
#[cfg(feature = "testing")]
subscription: self.subscription,
#[cfg(feature = "poll-diagnostics")]
poll: self.poll.for_thread(),
}
}
}
#[cfg(test)]
fn test_runtime() -> Handle {
static IDLE: std::sync::LazyLock<tokio::runtime::Runtime> = std::sync::LazyLock::new(|| {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("a current-thread runtime builds")
});
Handle::try_current().unwrap_or_else(|_| IDLE.handle().clone())
}
impl<C> fmt::Debug for Delivery<C> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Delivery")
.field(
"retry_destination",
&self.retry.as_ref().map(|retry| &retry.destination),
)
.field("declaration", &self.declaration)
.field("pending_continuations", &self.tasks.len())
.finish_non_exhaustive()
}
}
pub(crate) struct Slot<M>(Option<M>);
const SETTLED: &str = "the delivery has already been settled";
impl<M> Slot<M> {
pub(crate) const fn new(msg: M) -> Self {
Self(Some(msg))
}
const fn empty() -> Self {
Self(None)
}
fn fill(&mut self, msg: M) {
self.0 = Some(msg);
}
pub(crate) fn get(&self) -> &M {
self.0.as_ref().expect(SETTLED)
}
pub(crate) fn take(&mut self) -> M {
self.0.take().expect(SETTLED)
}
}
pub(crate) fn spawn_dispatch<S, H, C, St>(
mut subscriber: S,
handler: Arc<H>,
shutdown: Shutdown,
name: Arc<str>,
state: Arc<St>,
delivery: Arc<Delivery<C>>,
failure: DispatchFailure,
) -> JoinHandle<()>
where
S: Subscriber + Send + 'static,
H: Handler<S::Message, C, St> + 'static,
C: crate::BuildContext<S::Message> + Send + Sync + 'static,
St: Send + Sync + 'static,
{
tokio::spawn(async move {
let mut stream = std::pin::pin!(subscriber.stream());
let mut cancelled = std::pin::pin!(shutdown.cancelled());
let mut encode = BytesMut::new();
let mut slot = Slot::empty();
loop {
match turn_into(&shutdown, stream.as_mut(), cancelled.as_mut(), &mut slot).await {
Turn::Delivery(Ok(())) => {
let handling = std::pin::pin!(dispatch(
&*handler,
&mut slot,
&mut encode,
&name,
&state,
&delivery,
&failure
));
handling.await;
}
Turn::Delivery(Err(err)) => {
error!(
target: "ruststream::dispatch",
error = %err,
"subscriber stream error",
);
}
Turn::Ended => {
debug!(
target: "ruststream::dispatch",
subscriber = %name,
"subscriber stream ended",
);
break;
}
Turn::Shutdown => break,
}
}
})
}
fn turn_into<'s, St, Message, Failure, Wait>(
shutdown: &Shutdown,
stream: Pin<&'s mut St>,
cancelled: Pin<&'s mut Wait>,
slot: &'s mut Slot<Message>,
) -> TurnInto<'s, St, Message, Wait>
where
St: Stream<Item = Result<Message, Failure>>,
Wait: Future<Output = ()>,
{
TurnInto {
shutdown: shutdown.is_cancelled(),
stream,
cancelled,
slot,
}
}
struct TurnInto<'s, St, Message, Wait> {
shutdown: bool,
stream: Pin<&'s mut St>,
cancelled: Pin<&'s mut Wait>,
slot: &'s mut Slot<Message>,
}
impl<St, Message, Failure, Wait> Future for TurnInto<'_, St, Message, Wait>
where
St: Stream<Item = Result<Message, Failure>>,
Wait: Future<Output = ()>,
{
type Output = Turn<Result<(), Failure>>;
#[inline]
fn poll(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
if this.shutdown {
return Poll::Ready(Turn::Shutdown);
}
match this.stream.as_mut().poll_next(cx) {
Poll::Ready(Some(Ok(msg))) => {
this.slot.fill(msg);
Poll::Ready(Turn::Delivery(Ok(())))
}
Poll::Ready(Some(Err(err))) => Poll::Ready(Turn::Delivery(Err(err))),
Poll::Ready(None) => Poll::Ready(Turn::Ended),
Poll::Pending => this.cancelled.as_mut().poll(cx).map(|()| Turn::Shutdown),
}
}
}
enum Turn<T> {
Delivery(T),
Ended,
Shutdown,
}
async fn reap<Wait, Done>(
tasks: &mut JoinSet<Done>,
mut cancelled: Pin<&mut Wait>,
) -> Turn<Result<Done, JoinError>>
where
Wait: Future<Output = ()>,
Done: 'static,
{
poll_fn(|cx| match tasks.poll_join_next(cx) {
Poll::Ready(Some(joined)) => Poll::Ready(Turn::Delivery(joined)),
Poll::Ready(None) => Poll::Ready(Turn::Ended),
Poll::Pending => cancelled.as_mut().poll(cx).map(|()| Turn::Shutdown),
})
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn spawn_dispatch_workers<S, H, C, St>(
subscriber: S,
handler: Arc<H>,
shutdown: Shutdown,
name: Arc<str>,
state: Arc<St>,
delivery: Arc<Delivery<C>>,
failure: DispatchFailure,
workers: Workers,
) -> Result<JoinHandle<()>, StartThreadError>
where
S: Subscriber + Send + 'static,
S::Message: Send + Sync + 'static,
H: Handler<S::Message, C, St> + 'static,
C: crate::BuildContext<S::Message> + Send + Sync + 'static,
St: Send + Sync + 'static,
{
#[cfg(feature = "poll-diagnostics")]
if matches!(workers.placement, Placement::Threads) {
delivery.poll.on_dedicated_threads();
}
#[cfg(feature = "testing")]
let workers = if delivery.hooks.coordinator().is_some() {
workers.on_runtime()
} else {
workers
};
if workers.is_sequential() {
return Ok(spawn_dispatch(
subscriber, handler, shutdown, name, state, delivery, failure,
));
}
if matches!(workers.placement, Placement::Threads) {
return spawn_dispatch_threads(
subscriber, handler, shutdown, name, state, delivery, failure, workers,
);
}
Ok(pool::spawn_dispatch_pool(
subscriber, handler, shutdown, name, state, delivery, failure, workers,
))
}
#[allow(clippy::too_many_arguments)]
fn spawn_dispatch_threads<S, H, C, St>(
mut subscriber: S,
handler: Arc<H>,
shutdown: Shutdown,
name: Arc<str>,
state: Arc<St>,
delivery: Arc<Delivery<C>>,
failure: DispatchFailure,
workers: Workers,
) -> Result<JoinHandle<()>, StartThreadError>
where
S: Subscriber + Send + 'static,
S::Message: Send + Sync + 'static,
H: Handler<S::Message, C, St> + 'static,
C: crate::BuildContext<S::Message> + Send + Sync + 'static,
St: Send + Sync + 'static,
{
let shared = Arc::new(pool::Shared {
handler,
name,
state,
delivery,
failure,
});
let threads = Threads::start(&shared.name, workers.count.get(), workers.by_key, |own| {
let shared = Arc::clone(&shared);
let delivery = shared.delivery.on_thread(own);
let mut encode = BytesMut::new();
async move |msg: S::Message| {
let mut slot = Slot::new(msg);
dispatch(
&*shared.handler,
&mut slot,
&mut encode,
&shared.name,
&shared.state,
&delivery,
&shared.failure,
)
.await;
}
})?;
Ok(tokio::spawn(async move {
threads
.feed(
subscriber.stream(),
&shared.name,
&shutdown,
<S::Message as IncomingMessage>::partition_key,
)
.await;
}))
}
fn lane_of(key: &[u8], lanes: usize) -> usize {
let mut hasher = DefaultHasher::new();
key.hash(&mut hasher);
#[allow(clippy::cast_possible_truncation)]
{
(hasher.finish() % lanes as u64) as usize
}
}
fn log_worker_exit<Done>(joined: Result<Done, JoinError>) -> Option<Done> {
match joined {
Ok(done) => Some(done),
Err(err) => {
error!(target: "ruststream::dispatch", error = %err, "worker task failed");
None
}
}
}
#[allow(clippy::too_many_arguments)] pub(crate) fn spawn_batch_dispatch<S, H, C, St>(
mut subscriber: S,
handler: Arc<H>,
shutdown: Shutdown,
name: Arc<str>,
state: Arc<St>,
delivery: Arc<Delivery<C>>,
failure: DispatchFailure,
workers: Workers,
batch_size: NonZeroUsize,
) -> Result<JoinHandle<()>, StartThreadError>
where
S: BatchSubscriber + Send + 'static,
S::Message: Send + 'static,
S::Batch: 'static,
H: BatchHandler<S::Batch, C, St> + 'static,
C: Send + Sync + 'static,
St: Send + Sync + 'static,
{
#[cfg(feature = "poll-diagnostics")]
if matches!(workers.placement, Placement::Threads) {
delivery.poll.on_dedicated_threads();
}
#[cfg(feature = "testing")]
let workers = if delivery.hooks.coordinator().is_some() {
workers.on_runtime()
} else {
workers
};
if matches!(workers.placement, Placement::Threads) {
return spawn_batch_threads(
subscriber, handler, shutdown, name, state, delivery, failure, workers, batch_size,
);
}
Ok(tokio::spawn(async move {
let mut stream = std::pin::pin!(subscriber.batches(batch_size));
let mut cancelled = std::pin::pin!(shutdown.cancelled());
let mut tasks = JoinSet::new();
let mut scratch = <H as BatchHandler<S::Batch, C, St>>::Scratch::default();
let mut encode = BytesMut::new();
let mut spare = Vec::new();
loop {
if shutdown.is_cancelled() {
break;
}
if tasks.len() >= workers.count.get() {
match reap(&mut tasks, cancelled.as_mut()).await {
Turn::Delivery(joined) => spare.extend(log_worker_exit(joined)),
Turn::Ended | Turn::Shutdown => break,
}
continue;
}
let pulled = poll_fn(|cx| match stream.as_mut().poll_next(cx) {
Poll::Ready(Some(item)) => Poll::Ready(Turn::Delivery(item)),
Poll::Ready(None) => Poll::Ready(Turn::Ended),
Poll::Pending => cancelled.as_mut().poll(cx).map(|()| Turn::Shutdown),
})
.await;
match pulled {
Turn::Delivery(Ok(batch)) => {
if workers.is_sequential() {
run_batch::<_, _, C, _>(
&*handler,
batch,
&mut scratch,
&mut encode,
&name,
&state,
&delivery,
&failure,
)
.await;
} else {
let handler = Arc::clone(&handler);
let name = Arc::clone(&name);
let state = Arc::clone(&state);
let delivery = Arc::clone(&delivery);
let failure = failure.clone();
let (mut scratch, mut encode) = spare.pop().unwrap_or_default();
tasks.spawn(async move {
run_batch::<_, _, C, _>(
&*handler,
batch,
&mut scratch,
&mut encode,
&name,
&state,
&delivery,
&failure,
)
.await;
(scratch, encode)
});
}
}
Turn::Delivery(Err(err)) => {
error!(
target: "ruststream::dispatch",
error = %err,
"subscriber stream error",
);
}
Turn::Ended => {
debug!(
target: "ruststream::dispatch",
subscriber = %name,
"subscriber stream ended",
);
break;
}
Turn::Shutdown => break,
}
}
while let Some(joined) = tasks.join_next().await {
log_worker_exit(joined);
}
}))
}
#[allow(clippy::too_many_arguments)] fn spawn_batch_threads<S, H, C, St>(
mut subscriber: S,
handler: Arc<H>,
shutdown: Shutdown,
name: Arc<str>,
state: Arc<St>,
delivery: Arc<Delivery<C>>,
failure: DispatchFailure,
workers: Workers,
batch_size: NonZeroUsize,
) -> Result<JoinHandle<()>, StartThreadError>
where
S: BatchSubscriber + Send + 'static,
S::Message: Send + 'static,
S::Batch: 'static,
H: BatchHandler<S::Batch, C, St> + 'static,
C: Send + Sync + 'static,
St: Send + Sync + 'static,
{
let shared = Arc::new(pool::Shared {
handler,
name,
state,
delivery,
failure,
});
let threads = Threads::start(&shared.name, workers.count.get(), false, |own| {
let shared = Arc::clone(&shared);
let delivery = shared.delivery.on_thread(own);
let mut scratch = <H as BatchHandler<S::Batch, C, St>>::Scratch::default();
let mut encode = BytesMut::new();
async move |batch: S::Batch| {
run_batch::<_, _, C, _>(
&*shared.handler,
batch,
&mut scratch,
&mut encode,
&shared.name,
&shared.state,
&delivery,
&shared.failure,
)
.await;
}
})?;
Ok(tokio::spawn(async move {
threads
.feed(
subscriber.batches(batch_size),
&shared.name,
&shutdown,
|_| None,
)
.await;
}))
}
async fn dispatch<H, M, C, St>(
handler: &H,
slot: &mut Slot<M>,
encode: &mut BytesMut,
name: &str,
state: &St,
delivery: &Delivery<C>,
failure: &DispatchFailure,
) where
H: Handler<M, C, St>,
C: crate::BuildContext<M> + Send + Sync + 'static,
M: IncomingMessage,
St: Send + Sync,
{
#[cfg(feature = "testing")]
let watcher = delivery.hooks.coordinator().cloned();
#[cfg(feature = "testing")]
if let Some(coordinator) = &watcher {
coordinator.enqueued();
}
let (settle, continuations) = {
let msg = slot.get();
let cx = C::build(msg);
let source = FromDelivery(msg);
let mut ctx = Context::new(name, &source, state, cx, delivery)
.with_failfast(&failure.shutdown)
.with_encode_buffer(encode)
.with_decode_policy(failure.policies.decode);
#[cfg(feature = "testing")]
let result = in_harness_scope(
harness_scope(delivery),
observed!(
delivery,
AssertUnwindSafe(handler.handle(msg, &mut ctx)).catch_unwind()
),
)
.await;
#[cfg(not(feature = "testing"))]
let result = observed!(
delivery,
AssertUnwindSafe(handler.handle(msg, &mut ctx)).catch_unwind()
)
.await;
#[cfg(feature = "testing")]
let panicked = result.is_err();
let settle = match result {
Ok(s) => Some(s),
Err(payload) => {
let reason = panic_reason(payload.as_ref());
error!(
target: "ruststream::dispatch",
subscription = %name,
panic = %reason,
"handler panicked",
);
match failure.policies.panic {
FailurePolicy::FailFast => {
failure
.shutdown
.signal(name, &format!("handler panicked: {reason}"));
None
}
other => Some(
other
.settlement()
.map_or_else(super::handler::HandlerOutcome::drop, Into::into),
),
}
}
};
let continuations = match settle.as_ref() {
Some(s) if ctx.has_hooks() => Some(ctx.take_hooks_for(s.outcome())),
_ => None,
};
#[cfg(feature = "testing")]
if let Some(coordinator) = delivery.hooks.coordinator() {
coordinator.record(Record {
scope_id: delivery.scope_id,
subscription: delivery.subscription,
name: name.to_owned(),
deliveries: vec![Delivered {
raw: Bytes::copy_from_slice(msg.payload()),
value: ctx.take_carried(),
settle: settle.as_ref().map(super::handler::HandlerOutcome::outcome),
}],
panicked,
decode_failed: ctx.took_decode_failed(),
});
}
(settle, continuations)
};
if let Some(mut s) = settle {
let settling = settle_outcome(slot, s.outcome(), name, delivery, C::build as fn(&M) -> C);
settling.await;
if let Some(after) = s.take_after() {
delivery.spawn_after(after);
}
} else {
drop(slot.take());
}
if let Some(continuations) = continuations {
for fut in continuations {
delivery.spawn_after(fut);
}
}
#[cfg(feature = "testing")]
if let Some(coordinator) = &watcher {
coordinator.consumed();
}
}
#[allow(clippy::too_many_arguments)]
async fn run_batch<H, Batch, C, St>(
handler: &H,
batch: Batch,
scratch: &mut H::Scratch,
encode: &mut BytesMut,
name: &str,
state: &St,
delivery: &Delivery<C>,
failure: &DispatchFailure,
) where
H: BatchHandler<Batch, C, St>,
C: Send + Sync + 'static,
St: Send + Sync,
{
let batch = H::take(batch);
let Some(cx) = batch.context() else { return };
let empty = HeaderMap::new();
let mut ctx = Context::new(name, &empty, state, cx, delivery)
.with_failfast(&failure.shutdown)
.with_encode_buffer(encode)
.with_decode_policy(failure.policies.decode);
#[cfg(feature = "testing")]
let watcher = delivery.hooks.coordinator().cloned();
#[cfg(feature = "testing")]
let mut batch = batch;
#[cfg(feature = "testing")]
let unsettled = watcher.as_ref().map(|_| batch.unsettled());
#[cfg(feature = "testing")]
if let Some(coordinator) = &watcher {
coordinator.enqueued();
}
#[cfg(feature = "testing")]
let result = in_harness_scope(
harness_scope(delivery),
observed!(
delivery,
AssertUnwindSafe(handler.handle_batch(batch, scratch, &mut ctx)).catch_unwind()
),
)
.await;
#[cfg(not(feature = "testing"))]
let result = observed!(
delivery,
AssertUnwindSafe(handler.handle_batch(batch, scratch, &mut ctx)).catch_unwind()
)
.await;
match result {
Ok(()) => {
if ctx.has_hooks() {
for fut in ctx.take_settle_hooks() {
delivery.spawn_after(fut);
}
}
}
Err(payload) => {
let reason = panic_reason(payload.as_ref());
error!(
target: "ruststream::dispatch",
subscription = %name,
panic = %reason,
"batch handler panicked",
);
#[cfg(feature = "testing")]
if let (Some(coordinator), Some(deliveries)) = (&watcher, unsettled) {
coordinator.record(Record {
scope_id: delivery.scope_id,
subscription: delivery.subscription,
name: name.to_owned(),
deliveries,
panicked: true,
decode_failed: false,
});
}
if failure.policies.panic == FailurePolicy::FailFast {
failure
.shutdown
.signal(name, &format!("batch handler panicked: {reason}"));
}
}
}
#[cfg(feature = "testing")]
if let Some(coordinator) = &watcher {
coordinator.consumed();
}
}
#[cfg(feature = "testing")]
fn harness_scope<C>(delivery: &Delivery<C>) -> Option<HarnessScope> {
delivery
.hooks
.coordinator()
.cloned()
.map(|coordinator| HarnessScope::new(coordinator, delivery.scope_id, delivery.subscription))
}
pub(crate) async fn settle_outcome<M, C>(
slot: &mut Slot<M>,
outcome: HandlerResult,
name: &str,
delivery: &Delivery<C>,
build_cx: fn(&M) -> C,
) where
M: IncomingMessage,
C: Send + Sync + 'static,
{
let ack_result = match outcome {
HandlerResult::Ack => slot.take().ack().await,
HandlerResult::Nack { requeue: false } => slot.take().nack(false).await,
HandlerResult::Nack { requeue: true } => settle_retry(slot, name, delivery, build_cx).await,
HandlerResult::NackAfter { delay } => {
settle_nack_after(slot, name, delay, delivery, build_cx).await
}
};
if let Err(err) = ack_result {
warn!(
target: "ruststream::dispatch",
subscription = %name,
error = %err,
"ack / nack failed",
);
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Redelivery<'a> {
Subscription,
DeadLetter {
destination: &'a str,
at_cap: bool,
},
Reject,
}
fn attempt_of<M: IncomingMessage>(msg: &M) -> u64 {
msg.redelivery_count()
.unwrap_or_else(|| current_retry_count(msg.headers()) + 1)
}
fn redelivery_of<'d, M: IncomingMessage>(
msg: &M,
declaration: &'d RetryDeclaration,
) -> Redelivery<'d> {
match (declaration.max_attempts(), declaration.dead_letter()) {
(Some(cap), destination) if attempt_of(msg) >= u64::from(cap.get()) => {
destination.map_or(Redelivery::Reject, |destination| Redelivery::DeadLetter {
destination,
at_cap: true,
})
}
(None, Some(destination)) => Redelivery::DeadLetter {
destination,
at_cap: false,
},
_ => Redelivery::Subscription,
}
}
async fn settle_retry<M, C>(
slot: &mut Slot<M>,
name: &str,
delivery: &Delivery<C>,
build_cx: fn(&M) -> C,
) -> Result<(), AckError>
where
M: IncomingMessage,
C: Send + Sync + 'static,
{
if delivery.declaration.declares_nothing() {
return slot.take().nack(true).await;
}
let Some(retry) = delivery.retry.as_ref() else {
return slot.take().nack(true).await;
};
match redelivery_of(slot.get(), &delivery.declaration) {
Redelivery::Subscription if slot.get().redelivery_count().is_some() => {
slot.take().nack(true).await
}
Redelivery::Subscription => {
let destination = retry.destination.clone();
publish_copy(slot, name, destination, None, retry, delivery, build_cx).await
}
Redelivery::DeadLetter {
destination,
at_cap,
} => {
if at_cap {
warn_at_cap(name, attempt_of(slot.get()), Some(destination));
}
let destination = Arc::from(destination);
publish_copy(
slot,
name,
Some(destination),
None,
retry,
delivery,
build_cx,
)
.await
}
Redelivery::Reject => {
warn_at_cap(name, attempt_of(slot.get()), None);
slot.take().nack(false).await
}
}
}
async fn settle_nack_after<M, C>(
slot: &mut Slot<M>,
name: &str,
delay: Duration,
delivery: &Delivery<C>,
build_cx: fn(&M) -> C,
) -> Result<(), AckError>
where
M: IncomingMessage,
C: Send + Sync + 'static,
{
if slot.get().supports_nack_after() {
if let Some(retry) = delivery.retry.as_ref()
&& delivery.declaration.max_attempts().is_some()
{
match redelivery_of(slot.get(), &delivery.declaration) {
Redelivery::DeadLetter { destination, .. } => {
warn_at_cap(name, attempt_of(slot.get()), Some(destination));
let destination = Arc::from(destination);
return publish_copy(
slot,
name,
Some(destination),
None,
retry,
delivery,
build_cx,
)
.await;
}
Redelivery::Reject => {
warn_at_cap(name, attempt_of(slot.get()), None);
return slot.take().nack(false).await;
}
Redelivery::Subscription => {}
}
}
let accepted = slot.take().nack_after(delay).await;
#[cfg(feature = "testing")]
if accepted.is_ok()
&& let Some(coordinator) = delivery.hooks.coordinator()
{
coordinator.expect_redelivery(delivery.subscription, delay);
}
return accepted;
}
let Some(retry) = delivery.retry.as_ref() else {
if matches!(
redelivery_of(slot.get(), &delivery.declaration),
Redelivery::Subscription
) {
warn!(
target: "ruststream::dispatch",
subscription = %name,
"retry_after on a broker with neither native delayed redelivery nor a copy path \
of its own; requeuing immediately (the delay is dropped)",
);
}
return slot.take().nack(true).await;
};
match redelivery_of(slot.get(), &delivery.declaration) {
Redelivery::Subscription => {
let destination = retry.destination.clone();
publish_copy(
slot,
name,
destination,
Some(delay),
retry,
delivery,
build_cx,
)
.await
}
Redelivery::DeadLetter {
destination,
at_cap,
} => {
if at_cap {
warn_at_cap(name, attempt_of(slot.get()), Some(destination));
}
let destination = Arc::from(destination);
publish_copy(
slot,
name,
Some(destination),
None,
retry,
delivery,
build_cx,
)
.await
}
Redelivery::Reject => {
warn_at_cap(name, attempt_of(slot.get()), None);
slot.take().nack(false).await
}
}
}
fn warn_at_cap(name: &str, attempt: u64, destination: Option<&str>) {
if let Some(destination) = destination {
warn!(
target: "ruststream::dispatch",
subscription = %name,
attempt,
dead_letter = %destination,
"retry attempts exhausted; publishing the delivery to the dead-letter destination",
);
} else {
warn!(
target: "ruststream::dispatch",
subscription = %name,
attempt,
"retry attempts exhausted and no dead-letter destination is declared; rejecting the \
delivery",
);
}
}
async fn publish_copy<M, C>(
slot: &mut Slot<M>,
name: &str,
destination: Option<Arc<str>>,
delay: Option<Duration>,
retry: &DeferredRetry<C>,
delivery: &Delivery<C>,
build_cx: fn(&M) -> C,
) -> Result<(), AckError>
where
M: IncomingMessage,
C: Send + Sync + 'static,
{
let publisher = Arc::clone(&retry.publisher);
let subscription: Arc<str> = Arc::from(name);
let target = destination.unwrap_or_else(|| Arc::clone(&subscription));
let msg = slot.get();
let payload = Bytes::copy_from_slice(msg.payload());
let delivered = msg.headers().clone();
let mut headers = delivered.clone();
let next_count = current_retry_count(&headers) + 1;
headers.insert(Str::from_static(RETRY_COUNT_HEADER), next_count.to_string());
let context = build_cx(msg);
#[cfg(feature = "testing")]
let _hold = delivery.hooks.coordinator().map(Coordinator::hold);
match slot.take().nack(false).await {
Ok(()) | Err(AckError::Unsupported) => {}
Err(err) => return Err(err),
}
let republish = async move {
let cx = PublishContext::new(&subscription, &delivered, &context);
let copy = OutgoingMessage::new(target.as_ref(), payload.as_ref()).with_headers(headers);
if let Err(err) = publisher.publish_copy(copy, &cx).await {
warn!(
target: "ruststream::dispatch",
subscription = %subscription,
destination = %target,
error = %err,
"publishing the retry copy failed; message lost",
);
}
};
#[cfg(feature = "testing")]
let republish = in_harness_scope(harness_scope(delivery), republish);
let Some(delay) = delay.filter(|delay| !delay.is_zero()) else {
republish.await;
return Ok(());
};
#[cfg(feature = "testing")]
if let Some(coordinator) = delivery.hooks.coordinator() {
coordinator.schedule_redelivery_future(delay, republish);
return Ok(());
}
delivery.spawn_after(async move {
tokio::time::sleep(delay).await;
republish.await;
});
Ok(())
}
#[cfg(all(test, feature = "testing", feature = "memory"))]
mod harness_tests;
mod pool;
#[cfg(all(test, feature = "memory"))]
mod tests;
mod threads;