ruststream/runtime/dispatch.rs
1//! The per-subscriber dispatch loop: pulls messages off one subscriber and invokes its handler
2//! until shutdown is signalled or the stream ends. [`RustStream`](super::RustStream) owns the
3//! task spawning.
4
5use std::fmt;
6use std::future::{Future, poll_fn};
7use std::hash::{DefaultHasher, Hash, Hasher};
8use std::num::NonZeroUsize;
9use std::panic::AssertUnwindSafe;
10use std::pin::Pin;
11use std::sync::Arc;
12use std::task::{Context as TaskContext, Poll};
13use std::time::Duration;
14
15use bytes::{Bytes, BytesMut};
16use futures::{FutureExt, Stream};
17#[cfg(test)]
18use tokio::runtime::Handle;
19use tokio::task::{JoinError, JoinHandle, JoinSet};
20use tokio_util::task::TaskTracker;
21use tracing::{debug, error, warn};
22
23use crate::{
24 AckError, BatchSubscriber, HeaderMap, IncomingMessage, OutgoingMessage, RetryDeclaration, Str,
25 Subscriber,
26};
27
28use super::batch::{BatchHandler, TakenBatch};
29use super::context::{Context, FromDelivery};
30use super::failure::{DispatchFailure, FailurePolicy, panic_reason};
31use super::handler::{Handler, HandlerResult};
32use super::main_runtime::MainRuntime;
33#[cfg(all(test, feature = "poll-diagnostics"))]
34use super::poll_diagnostics::PollDiagnostics;
35#[cfg(feature = "poll-diagnostics")]
36use super::poll_diagnostics::PollSampler;
37use super::publish::PublishContext;
38#[cfg(all(test, feature = "memory"))]
39use super::redelivery::ErasedRetryPublisher;
40use super::redelivery::{DeferredRetry, ScopeDelivery};
41use super::shutdown::Shutdown;
42#[cfg(feature = "testing")]
43use crate::testing::coordinator::{
44 Coordinator, Delivered, HarnessScope, Record, TestHooks, in_harness_scope,
45};
46pub(crate) use threads::StartThreadError;
47use threads::Threads;
48
49/// Header carrying the framework's own retry count.
50///
51/// The runtime increments it on every copy of a delivery it publishes, and reads it back where the
52/// transport counts nothing of its own
53/// ([`IncomingMessage::redelivery_count`](crate::IncomingMessage::redelivery_count)): that is what
54/// a registration's [`max_attempts`](super::RouterWith::max_attempts) cap counts there, on the
55/// runtime's copy path and on a delay a broker crate honours by publishing a copy of its own
56/// alike. A handler can read it too, to tell a first delivery from a redelivery.
57///
58/// # Examples
59///
60/// ```
61/// # #[cfg(all(feature = "macros", feature = "json"))]
62/// # mod demo {
63/// use ruststream::prelude::*;
64/// use ruststream::runtime::RETRY_COUNT_HEADER;
65/// use serde::Deserialize;
66///
67/// #[derive(Deserialize)]
68/// struct Order {
69/// id: u64,
70/// }
71///
72/// #[subscriber("orders")]
73/// async fn reconcile(order: &Order, ctx: &mut Context<'_>) -> HandlerOutcome {
74/// // The copies the runtime publishes count themselves; the first delivery carries no header.
75/// let attempt: u64 = ctx
76/// .headers()
77/// .get_str(RETRY_COUNT_HEADER)
78/// .and_then(|v| v.parse().ok())
79/// .unwrap_or(0);
80/// if attempt >= 3 {
81/// tracing::warn!(order.id, attempt, "giving up");
82/// return HandlerOutcome::drop();
83/// }
84/// HandlerOutcome::retry()
85/// }
86/// # }
87/// # fn main() {}
88/// ```
89pub const RETRY_COUNT_HEADER: &str = "x-ruststream-retry-count";
90
91/// Parses the current [`RETRY_COUNT_HEADER`] value, defaulting to zero when absent or malformed.
92fn current_retry_count(headers: &HeaderMap) -> u64 {
93 headers
94 .get_str(RETRY_COUNT_HEADER)
95 .and_then(|v| v.parse().ok())
96 .unwrap_or(0)
97}
98
99/// Concurrency policy for one subscriber's dispatch loop, declared with the `workers(..)` or
100/// `threads(..)` macro argument (or [`Workers::sequential`] by default).
101///
102/// - `workers(n)`: `n` long-lived workers, tasks of the runtime the app runs on, process up to
103/// `n` deliveries of the subscriber concurrently; a free worker takes the next one.
104/// Back-pressure holds: the stream is not polled while `n` deliveries are in flight. Global
105/// processing order is lost by design.
106/// - `workers(n, by_key)`: the same `n` workers as sequential lanes; a delivery goes to the lane
107/// picked by hashing its [`partition_key`](crate::IncomingMessage::partition_key), so per-key
108/// ordering is preserved. Messages without a key rotate over the lanes.
109/// - `threads(n)`: `n` dedicated threads of the subscription's own, each with a current-thread
110/// runtime, for a handler that computes. The subscription's loop stays on the app's runtime and
111/// spreads deliveries over the threads round-robin; decoding, the handler and settling run on
112/// the thread a delivery was handed to, and only there. The app's runtime stays free for I/O,
113/// timers and the other subscriptions.
114/// - `threads(n, by_key)`: the same threads as sequential lanes by partition key.
115///
116/// The workers or threads and the queues that feed them are made when the subscription starts:
117/// handing a delivery to one allocates nothing.
118///
119/// The default is sequential dispatch (`workers(1)`).
120#[derive(Debug, Clone, Copy, PartialEq, Eq)]
121pub struct Workers {
122 count: usize,
123 by_key: bool,
124 placement: Placement,
125}
126
127/// Where a subscription's concurrent deliveries run.
128#[derive(Debug, Clone, Copy, PartialEq, Eq)]
129pub(crate) enum Placement {
130 /// As tasks of the runtime the app runs on.
131 Runtime,
132 /// On dedicated threads of the subscription's own.
133 Threads,
134}
135
136impl Workers {
137 /// Sequential dispatch: one delivery at a time, in stream order. The default.
138 #[must_use]
139 pub const fn sequential() -> Self {
140 Self {
141 count: 1,
142 by_key: false,
143 placement: Placement::Runtime,
144 }
145 }
146
147 /// A pool of up to `count` concurrent deliveries. One behaves like sequential dispatch.
148 #[must_use]
149 pub const fn pool(count: NonZeroUsize) -> Self {
150 Self {
151 count: count.get(),
152 by_key: false,
153 placement: Placement::Runtime,
154 }
155 }
156
157 /// `count` sequential lanes keyed by the message
158 /// [`partition_key`](crate::IncomingMessage::partition_key): per-key ordering is preserved.
159 /// One lane behaves like sequential dispatch.
160 #[must_use]
161 pub const fn keyed(count: NonZeroUsize) -> Self {
162 Self {
163 count: count.get(),
164 by_key: true,
165 placement: Placement::Runtime,
166 }
167 }
168
169 /// `count` dedicated threads for a CPU-bound handler: the subscription's deliveries are
170 /// handled on threads of its own, each running a current-thread runtime, so the handler's
171 /// computation holds up none of the app runtime's threads. The router-chain spelling of
172 /// `threads(n)`.
173 ///
174 /// # Examples
175 ///
176 /// ```
177 /// # #[cfg(all(feature = "macros", feature = "memory", feature = "json"))]
178 /// # mod demo {
179 /// use ruststream::memory::prelude::*;
180 /// use serde::Deserialize;
181 ///
182 /// #[derive(Deserialize)]
183 /// struct Image {
184 /// id: u64,
185 /// }
186 ///
187 /// #[subscriber("images.resize")]
188 /// async fn resize(image: &Image) -> HandlerOutcome {
189 /// tracing::info!(image.id, "resized");
190 /// HandlerOutcome::ack()
191 /// }
192 ///
193 /// fn images() -> Router<MemoryBroker, impl RouterDef<MemoryBroker>> {
194 /// // Resizing computes, so it runs on eight threads of its own.
195 /// Router::new()
196 /// .include(resize)
197 /// .workers(Workers::threads(nonzero!(8)))
198 /// }
199 ///
200 /// fn app() -> RustStream {
201 /// RustStream::new(AppInfo::new("images", "0.1.0")).with_broker(MemoryBroker::new(), |b| {
202 /// b.include_router(images());
203 /// })
204 /// }
205 /// # }
206 /// # fn main() {}
207 /// ```
208 #[must_use]
209 pub const fn threads(count: NonZeroUsize) -> Self {
210 Self {
211 count: count.get(),
212 by_key: false,
213 placement: Placement::Threads,
214 }
215 }
216
217 /// `count` dedicated threads as sequential lanes keyed by the message
218 /// [`partition_key`](crate::IncomingMessage::partition_key): a key always lands on the same
219 /// thread, so per-key ordering is preserved. The router-chain spelling of
220 /// `threads(n, by_key)`.
221 ///
222 /// # Examples
223 ///
224 /// ```
225 /// # #[cfg(all(feature = "macros", feature = "memory", feature = "json"))]
226 /// # mod demo {
227 /// use ruststream::memory::prelude::*;
228 /// use serde::Deserialize;
229 ///
230 /// #[derive(Deserialize)]
231 /// struct Image {
232 /// id: u64,
233 /// }
234 ///
235 /// #[subscriber("images.resize")]
236 /// async fn resize(image: &Image) -> HandlerOutcome {
237 /// tracing::info!(image.id, "resized");
238 /// HandlerOutcome::ack()
239 /// }
240 ///
241 /// fn images() -> Router<MemoryBroker, impl RouterDef<MemoryBroker>> {
242 /// // One tenant's images stay in order: a key always lands on the same thread.
243 /// Router::new()
244 /// .include(resize)
245 /// .workers(Workers::threads_keyed(nonzero!(4)))
246 /// }
247 ///
248 /// fn app() -> RustStream {
249 /// RustStream::new(AppInfo::new("images", "0.1.0")).with_broker(MemoryBroker::new(), |b| {
250 /// b.include_router(images());
251 /// })
252 /// }
253 /// # }
254 /// # fn main() {}
255 /// ```
256 #[must_use]
257 pub const fn threads_keyed(count: NonZeroUsize) -> Self {
258 Self {
259 count: count.get(),
260 by_key: true,
261 placement: Placement::Threads,
262 }
263 }
264
265 /// One worker on the app's runtime is indistinguishable from the sequential loop; one
266 /// dedicated thread is still a thread of its own.
267 pub(crate) const fn is_sequential(&self) -> bool {
268 self.count <= 1 && matches!(self.placement, Placement::Runtime)
269 }
270
271 /// The same concurrency as tasks of the runtime the loop runs on: what a `threads(n)`
272 /// subscription becomes under the test harness, which runs every subscription on the test's
273 /// own runtime.
274 #[cfg(feature = "testing")]
275 pub(crate) const fn on_runtime(self) -> Self {
276 Self {
277 placement: Placement::Runtime,
278 ..self
279 }
280 }
281}
282
283impl Default for Workers {
284 fn default() -> Self {
285 Self::sequential()
286 }
287}
288
289/// Per-subscription publish context threaded into every delivery's [`Context`]: the
290/// broker-agnostic `retry_after` fallback of this registration and the app-wide tracker for
291/// post-settle continuations.
292/// An `and_after` continuation is spawned onto `tasks` so a graceful shutdown drains it.
293pub(crate) struct Delivery<C = ()> {
294 /// The retry path of this subscription: the publisher a copy of a delivery leaves through,
295 /// with the destination the mount site or the descriptor named. `None` where the broker moves
296 /// the delivery itself, in which case a `NackAfter` on a transport with no native delayed
297 /// redelivery has nothing left to do but requeue.
298 pub(crate) retry: Option<DeferredRetry<C>>,
299 /// What the mount site declared with `max_attempts(..)` and `dead_letter(..)`. Read only
300 /// where the runtime is the one moving the delivery.
301 pub(crate) declaration: RetryDeclaration,
302 /// Per-scope task tracker for post-settle `and_after` continuations. The
303 /// dispatcher spawns each element's continuation onto it after settling, so a graceful
304 /// shutdown drains them.
305 pub(crate) tasks: TaskTracker,
306 /// The runtime the app connected its brokers on, captured when the subscription opens. Work
307 /// that must outlive the delivery (a continuation, a delayed retry copy) runs here, wherever
308 /// the handler itself runs, and a handler reaches it through
309 /// [`Context::main_runtime`](super::context::Context::main_runtime).
310 pub(crate) runtime: MainRuntime,
311 /// On a dedicated thread (`threads(n)`), the thread's own tracker: the thread waits for what
312 /// its deliveries left behind before it ends, so a shutdown does not drop it with the
313 /// thread's runtime. `None` everywhere else.
314 pub(crate) thread_tasks: Option<TaskTracker>,
315 /// The harness's recording-and-quiescence hooks for this scope. Empty (uninstalled) outside a
316 /// [`TestApp`](crate::testing::TestApp) run, so the per-delivery read is a single atomic load.
317 #[cfg(feature = "testing")]
318 pub(crate) hooks: Arc<TestHooks>,
319 /// This broker's registration index, used to scope recorded deliveries per broker.
320 #[cfg(feature = "testing")]
321 pub(crate) scope_id: usize,
322 /// The subscription of the app this delivery comes through, identifying its records.
323 #[cfg(feature = "testing")]
324 pub(crate) subscription: usize,
325 /// Times every Nth poll of this subscription's handler.
326 #[cfg(feature = "poll-diagnostics")]
327 pub(crate) poll: PollSampler,
328}
329
330/// The handler's future, observed by the poll-time diagnostics where the feature is on and handed
331/// through untouched where it is off.
332#[cfg(feature = "poll-diagnostics")]
333macro_rules! observed {
334 ($delivery:expr, $future:expr) => {
335 $delivery.poll.observe($future)
336 };
337}
338
339/// See the feature's counterpart above: without the feature, the expression itself.
340#[cfg(not(feature = "poll-diagnostics"))]
341macro_rules! observed {
342 ($delivery:expr, $future:expr) => {
343 $future
344 };
345}
346
347impl<C> Delivery<C> {
348 /// The context one subscription dispatches under: its own retry path over what the whole
349 /// scope shares.
350 #[cfg_attr(not(feature = "poll-diagnostics"), allow(unused_variables))]
351 pub(crate) fn for_subscription(
352 scope: &ScopeDelivery,
353 subscription: &str,
354 retry: Option<DeferredRetry<C>>,
355 declaration: RetryDeclaration,
356 ) -> Self {
357 Self {
358 retry,
359 declaration,
360 tasks: scope.tasks().clone(),
361 // A subscription opens inside the app's startup, on the runtime it connected on.
362 runtime: MainRuntime::current(),
363 thread_tasks: None,
364 #[cfg(feature = "testing")]
365 hooks: Arc::clone(scope.hooks()),
366 #[cfg(feature = "testing")]
367 scope_id: scope.scope_id(),
368 #[cfg(feature = "testing")]
369 subscription: scope.subscription(),
370 #[cfg(feature = "poll-diagnostics")]
371 poll: scope.poll().register(subscription),
372 }
373 }
374
375 /// A delivery context outside any scope, for tests that drive the dispatch functions
376 /// directly.
377 #[cfg(test)]
378 pub(crate) fn detached(retry: Option<DeferredRetry<C>>, tasks: TaskTracker) -> Self {
379 Self {
380 retry,
381 declaration: RetryDeclaration::new(),
382 tasks,
383 runtime: MainRuntime::new(test_runtime()),
384 thread_tasks: None,
385 #[cfg(feature = "testing")]
386 hooks: Arc::new(TestHooks::detached()),
387 #[cfg(feature = "testing")]
388 scope_id: 0,
389 #[cfg(feature = "testing")]
390 subscription: 0,
391 #[cfg(feature = "poll-diagnostics")]
392 poll: PollDiagnostics::new().register("detached"),
393 }
394 }
395
396 /// A delivery context publishing its retry copies through `publisher` at `destination`. For
397 /// tests.
398 #[cfg(all(test, feature = "memory"))]
399 pub(crate) fn deferring_to(
400 publisher: Arc<dyn ErasedRetryPublisher<C>>,
401 destination: &str,
402 tasks: TaskTracker,
403 ) -> Self {
404 Self::detached(
405 Some(DeferredRetry {
406 publisher,
407 destination: Some(Arc::from(destination)),
408 }),
409 tasks,
410 )
411 }
412
413 /// The same context with a declaration on it, as a mount site's `max_attempts(..)` and
414 /// `dead_letter(..)` leave it. For tests.
415 #[cfg(all(test, feature = "memory"))]
416 pub(crate) fn declaring(mut self, declaration: RetryDeclaration) -> Self {
417 self.declaration = declaration;
418 self
419 }
420
421 /// An empty delivery context: no deferred retry, a fresh continuation tracker. For tests.
422 #[cfg(test)]
423 pub(crate) fn empty() -> Self {
424 Self::with_tasks(TaskTracker::new())
425 }
426
427 /// An empty delivery context carrying a caller-owned continuation tracker, so a test can
428 /// observe the post-settle continuations spawned through it.
429 #[cfg(test)]
430 pub(crate) fn with_tasks(tasks: TaskTracker) -> Self {
431 Self::detached(None, tasks)
432 }
433}
434
435impl<C> Delivery<C> {
436 /// Spawns work the delivery leaves behind (a continuation, a post-settle hook, the timer of a
437 /// deferred copy) on the runtime the delivery runs on, tracked so a graceful shutdown waits
438 /// for it.
439 pub(crate) fn spawn_after<Work>(&self, work: Work)
440 where
441 Work: Future<Output = ()> + Send + 'static,
442 {
443 match &self.thread_tasks {
444 None => drop(self.tasks.spawn(work)),
445 Some(own) => drop(self.tasks.spawn(own.track_future(work))),
446 }
447 }
448
449 /// The same context for one dedicated thread, whose leftovers `own` tracks as well.
450 fn on_thread(&self, own: TaskTracker) -> Self {
451 Self {
452 retry: self.retry.as_ref().map(|retry| DeferredRetry {
453 publisher: Arc::clone(&retry.publisher),
454 destination: retry.destination.clone(),
455 }),
456 declaration: self.declaration.clone(),
457 tasks: self.tasks.clone(),
458 runtime: self.runtime.clone(),
459 thread_tasks: Some(own),
460 #[cfg(feature = "testing")]
461 hooks: Arc::clone(&self.hooks),
462 #[cfg(feature = "testing")]
463 scope_id: self.scope_id,
464 #[cfg(feature = "testing")]
465 subscription: self.subscription,
466 #[cfg(feature = "poll-diagnostics")]
467 poll: self.poll.for_thread(),
468 }
469 }
470}
471
472/// The runtime a delivery context built by a unit test answers `main_runtime` with: the test's
473/// own where it runs inside one, and a shared idle one where a synchronous test builds a context.
474#[cfg(test)]
475fn test_runtime() -> Handle {
476 static IDLE: std::sync::LazyLock<tokio::runtime::Runtime> = std::sync::LazyLock::new(|| {
477 tokio::runtime::Builder::new_current_thread()
478 .enable_all()
479 .build()
480 .expect("a current-thread runtime builds")
481 });
482 Handle::try_current().unwrap_or_else(|_| IDLE.handle().clone())
483}
484
485impl<C> fmt::Debug for Delivery<C> {
486 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
487 f.debug_struct("Delivery")
488 .field(
489 "retry_destination",
490 &self.retry.as_ref().map(|retry| &retry.destination),
491 )
492 .field("declaration", &self.declaration)
493 .field("pending_continuations", &self.tasks.len())
494 .finish_non_exhaustive()
495 }
496}
497
498/// One delivery, between the stream that yielded it and the settle that consumes it.
499///
500/// Everything on the path borrows the delivery out of the slot and the settle takes it, so a
501/// delivery is copied twice whatever its size: once out of the stream, once into
502/// [`ack(self)`](IncomingMessage::ack). The futures in between hold a pointer to the slot rather
503/// than the value, which is what keeps a pinned dispatch the same size for a 16-byte delivery and
504/// a 16-kilobyte one.
505///
506/// The inside is an `Option` because Rust cannot move a value out of a place its caller owns. A
507/// slot is never empty on the dispatch path: the sequential loop's one slot is filled by the poll
508/// that takes a delivery off the stream and emptied by the settle, and a pooled worker's is filled
509/// where it is declared; taking the delivery is the last thing that happens to it.
510pub(crate) struct Slot<M>(Option<M>);
511
512/// What a slot says when it is read after the settle has taken the delivery.
513const SETTLED: &str = "the delivery has already been settled";
514
515impl<M> Slot<M> {
516 /// The slot holding `msg`.
517 pub(crate) const fn new(msg: M) -> Self {
518 Self(Some(msg))
519 }
520
521 /// A slot with nothing in it yet: the sequential loop's own, which the stream fills in place.
522 const fn empty() -> Self {
523 Self(None)
524 }
525
526 /// Puts the next delivery in the slot.
527 fn fill(&mut self, msg: M) {
528 self.0 = Some(msg);
529 }
530
531 /// The delivery, for everything that only reads it.
532 ///
533 /// # Panics
534 ///
535 /// Panics once the settle has taken the delivery, which is the last thing the dispatch path
536 /// does with a slot.
537 pub(crate) fn get(&self) -> &M {
538 self.0.as_ref().expect(SETTLED)
539 }
540
541 /// The delivery, for the settle that consumes it.
542 ///
543 /// # Panics
544 ///
545 /// Panics if called twice; see [`Slot::get`].
546 pub(crate) fn take(&mut self) -> M {
547 self.0.take().expect(SETTLED)
548 }
549}
550
551/// Spawns a task that drives `subscriber` through `handler` until `shutdown` is triggered or the
552/// stream terminates. Each delivery is given a [`Context`] built from `name`, the message headers,
553/// shared `state`, and the `delivery` publish context.
554pub(crate) fn spawn_dispatch<S, H, C, St>(
555 mut subscriber: S,
556 handler: Arc<H>,
557 shutdown: Shutdown,
558 name: Arc<str>,
559 state: Arc<St>,
560 delivery: Arc<Delivery<C>>,
561 failure: DispatchFailure,
562) -> JoinHandle<()>
563where
564 S: Subscriber + Send + 'static,
565 H: Handler<S::Message, C, St> + 'static,
566 C: crate::BuildContext<S::Message> + Send + Sync + 'static,
567 St: Send + Sync + 'static,
568{
569 tokio::spawn(async move {
570 let mut stream = std::pin::pin!(subscriber.stream());
571 let mut cancelled = std::pin::pin!(shutdown.cancelled());
572 // One encode buffer for the whole loop: a reply that leaves through a publisher which
573 // reads the payload is written into this one every time, so the subscription allocates
574 // it once and never again.
575 let mut encode = BytesMut::new();
576 // The loop's slot, one for the whole loop: the stream's poll writes each delivery
577 // straight into it, and the settle takes it out again. A delivery returned by value
578 // through the turn and bound to a local on its way here would be copied twice more.
579 let mut slot = Slot::empty();
580 loop {
581 match turn_into(&shutdown, stream.as_mut(), cancelled.as_mut(), &mut slot).await {
582 Turn::Delivery(Ok(())) => {
583 // Pinned rather than awaited in place: a delivery's future carries the
584 // context, the handler's own state and the settle path, and awaiting the
585 // call expression makes the loop build it on the stack and copy it into
586 // its own state on every delivery.
587 let handling = std::pin::pin!(dispatch(
588 &*handler,
589 &mut slot,
590 &mut encode,
591 &name,
592 &state,
593 &delivery,
594 &failure
595 ));
596 handling.await;
597 }
598 Turn::Delivery(Err(err)) => {
599 error!(
600 target: "ruststream::dispatch",
601 error = %err,
602 "subscriber stream error",
603 );
604 }
605 Turn::Ended => {
606 debug!(
607 target: "ruststream::dispatch",
608 subscriber = %name,
609 "subscriber stream ended",
610 );
611 break;
612 }
613 Turn::Shutdown => break,
614 }
615 }
616 })
617}
618
619/// Waits for whichever comes first: the next item off `stream`, or `shutdown`. A delivery is
620/// written into the loop's slot where it lands rather than returned by value.
621///
622/// The stream is polled first and `cancelled` only where it answers `Pending`, so a turn that
623/// finds a delivery ready never touches the token's waiter list: what shutdown costs a busy
624/// subscription is the one relaxed load of [`Shutdown`]'s flag. `cancelled` is the token's own
625/// wait future, built once per subscription by the caller and pinned across the whole loop, so
626/// the turn that does park registers with the waiter list once rather than on every iteration,
627/// which is what a `select!` over `cancelled()` charges.
628///
629/// # Cancel safety
630///
631/// Both halves are cancel-safe: a delivery is written into the slot in the same poll that takes
632/// it off the stream, and a registration with the token's waiter list outlives the poll. Once
633/// `cancelled` has resolved the flag stays set - a token never un-cancels - so the check at the
634/// top answers every later turn and the resolved future is never polled again.
635///
636/// The other loops write the same wait out rather than share it: one instantiation shared with
637/// them is enough for the compiler to stop inlining it here.
638fn turn_into<'s, St, Message, Failure, Wait>(
639 shutdown: &Shutdown,
640 stream: Pin<&'s mut St>,
641 cancelled: Pin<&'s mut Wait>,
642 slot: &'s mut Slot<Message>,
643) -> TurnInto<'s, St, Message, Wait>
644where
645 St: Stream<Item = Result<Message, Failure>>,
646 Wait: Future<Output = ()>,
647{
648 TurnInto {
649 shutdown: shutdown.is_cancelled(),
650 stream,
651 cancelled,
652 slot,
653 }
654}
655
656/// The future [`turn_into`] returns: written out rather than an `async fn` so its poll carries an
657/// inline hint, which a closure inside `poll_fn` cannot.
658struct TurnInto<'s, St, Message, Wait> {
659 shutdown: bool,
660 stream: Pin<&'s mut St>,
661 cancelled: Pin<&'s mut Wait>,
662 slot: &'s mut Slot<Message>,
663}
664
665impl<St, Message, Failure, Wait> Future for TurnInto<'_, St, Message, Wait>
666where
667 St: Stream<Item = Result<Message, Failure>>,
668 Wait: Future<Output = ()>,
669{
670 type Output = Turn<Result<(), Failure>>;
671
672 #[inline]
673 fn poll(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
674 // Every field is a reference or a flag, so the future is `Unpin`.
675 let this = self.get_mut();
676 if this.shutdown {
677 return Poll::Ready(Turn::Shutdown);
678 }
679 match this.stream.as_mut().poll_next(cx) {
680 Poll::Ready(Some(Ok(msg))) => {
681 this.slot.fill(msg);
682 Poll::Ready(Turn::Delivery(Ok(())))
683 }
684 Poll::Ready(Some(Err(err))) => Poll::Ready(Turn::Delivery(Err(err))),
685 Poll::Ready(None) => Poll::Ready(Turn::Ended),
686 Poll::Pending => this.cancelled.as_mut().poll(cx).map(|()| Turn::Shutdown),
687 }
688 }
689}
690
691/// What one turn of a dispatch loop found on its subscriber.
692enum Turn<T> {
693 /// The subscriber yielded a delivery, or an error reading one.
694 Delivery(T),
695 /// The subscriber's stream ended; the loop is done.
696 Ended,
697 /// Shutdown was signalled; the loop stops without touching the stream again.
698 Shutdown,
699}
700
701/// Waits for whichever comes first: a finished worker off `tasks`, or `shutdown`.
702///
703/// [`turn_into`]'s counterpart for a full pool, in the same order and over the same pinned token
704/// future: a worker that has already finished is reaped without the token's waiter list being
705/// touched. `Ended` answers an empty set, which a full pool cannot be.
706///
707/// # Cancel safety
708///
709/// Cancel-safe, like [`turn_into`]: a `JoinSet` keeps its finished workers across polls, and the
710/// registration a parked poll made with the waiter list outlives that poll.
711async fn reap<Wait, Done>(
712 tasks: &mut JoinSet<Done>,
713 mut cancelled: Pin<&mut Wait>,
714) -> Turn<Result<Done, JoinError>>
715where
716 Wait: Future<Output = ()>,
717 Done: 'static,
718{
719 poll_fn(|cx| match tasks.poll_join_next(cx) {
720 Poll::Ready(Some(joined)) => Poll::Ready(Turn::Delivery(joined)),
721 Poll::Ready(None) => Poll::Ready(Turn::Ended),
722 Poll::Pending => cancelled.as_mut().poll(cx).map(|()| Turn::Shutdown),
723 })
724 .await
725}
726
727/// Spawns a task that drives `subscriber` through `handler` with `workers.count` long-lived
728/// workers: a pool that hands each delivery to a free worker, or with `by_key` lanes that hand it
729/// to the worker its key hashes to (see [`pool`]).
730///
731/// Sequential policies delegate to [`spawn_dispatch`]. On shutdown the stream stops being
732/// polled and the workers drain what they hold; if the app's `shutdown_timeout` aborts this task,
733/// the workers abort with it.
734// The parts are independent and each spawn site passes its own; bundling them into a struct
735// would hide that.
736#[allow(clippy::too_many_arguments)]
737pub(crate) fn spawn_dispatch_workers<S, H, C, St>(
738 subscriber: S,
739 handler: Arc<H>,
740 shutdown: Shutdown,
741 name: Arc<str>,
742 state: Arc<St>,
743 delivery: Arc<Delivery<C>>,
744 failure: DispatchFailure,
745 workers: Workers,
746) -> Result<JoinHandle<()>, StartThreadError>
747where
748 S: Subscriber + Send + 'static,
749 S::Message: Send + Sync + 'static,
750 H: Handler<S::Message, C, St> + 'static,
751 C: crate::BuildContext<S::Message> + Send + Sync + 'static,
752 St: Send + Sync + 'static,
753{
754 // Before the harness maps the placement away: the declaration is what the warning advises.
755 #[cfg(feature = "poll-diagnostics")]
756 if matches!(workers.placement, Placement::Threads) {
757 delivery.poll.on_dedicated_threads();
758 }
759 // The harness drives every subscription on the test's own runtime, so a timer the handler
760 // arms is one `advance` reaches; it does not reproduce the thread topology.
761 #[cfg(feature = "testing")]
762 let workers = if delivery.hooks.coordinator().is_some() {
763 workers.on_runtime()
764 } else {
765 workers
766 };
767 if workers.is_sequential() {
768 return Ok(spawn_dispatch(
769 subscriber, handler, shutdown, name, state, delivery, failure,
770 ));
771 }
772 if matches!(workers.placement, Placement::Threads) {
773 return spawn_dispatch_threads(
774 subscriber, handler, shutdown, name, state, delivery, failure, workers,
775 );
776 }
777 Ok(pool::spawn_dispatch_pool(
778 subscriber, handler, shutdown, name, state, delivery, failure, workers,
779 ))
780}
781
782/// Starts the dedicated threads of a `threads(n)` subscription and spawns its loop on the app's
783/// runtime (see [`threads`]).
784// See `spawn_dispatch_workers`: each part is the registration's own.
785#[allow(clippy::too_many_arguments)]
786fn spawn_dispatch_threads<S, H, C, St>(
787 mut subscriber: S,
788 handler: Arc<H>,
789 shutdown: Shutdown,
790 name: Arc<str>,
791 state: Arc<St>,
792 delivery: Arc<Delivery<C>>,
793 failure: DispatchFailure,
794 workers: Workers,
795) -> Result<JoinHandle<()>, StartThreadError>
796where
797 S: Subscriber + Send + 'static,
798 S::Message: Send + Sync + 'static,
799 H: Handler<S::Message, C, St> + 'static,
800 C: crate::BuildContext<S::Message> + Send + Sync + 'static,
801 St: Send + Sync + 'static,
802{
803 let shared = Arc::new(pool::Shared {
804 handler,
805 name,
806 state,
807 delivery,
808 failure,
809 });
810 let threads = Threads::start(&shared.name, workers.count, workers.by_key, |own| {
811 let shared = Arc::clone(&shared);
812 let delivery = shared.delivery.on_thread(own);
813 // One encode buffer per thread, for the reason the sequential loop has one.
814 let mut encode = BytesMut::new();
815 async move |msg: S::Message| {
816 let mut slot = Slot::new(msg);
817 dispatch(
818 &*shared.handler,
819 &mut slot,
820 &mut encode,
821 &shared.name,
822 &shared.state,
823 &delivery,
824 &shared.failure,
825 )
826 .await;
827 }
828 })?;
829 Ok(tokio::spawn(async move {
830 threads
831 .feed(
832 subscriber.stream(),
833 &shared.name,
834 &shutdown,
835 <S::Message as IncomingMessage>::partition_key,
836 )
837 .await;
838 }))
839}
840
841fn lane_of(key: &[u8], lanes: usize) -> usize {
842 let mut hasher = DefaultHasher::new();
843 key.hash(&mut hasher);
844 // The modulo keeps the value below `lanes`, which is a usize.
845 #[allow(clippy::cast_possible_truncation)]
846 {
847 (hasher.finish() % lanes as u64) as usize
848 }
849}
850
851/// What a finished worker handed back, or nothing and a log line where it failed.
852fn log_worker_exit<Done>(joined: Result<Done, JoinError>) -> Option<Done> {
853 match joined {
854 Ok(done) => Some(done),
855 Err(err) => {
856 error!(target: "ruststream::dispatch", error = %err, "worker task failed");
857 None
858 }
859 }
860}
861
862/// Spawns a task that drives `subscriber` through a batch `handler`, one
863/// [`BatchSubscriber::batches`] item per invocation, until `shutdown` is triggered or the stream
864/// terminates. The handler owns the batch's deliveries and settles each of them.
865///
866/// With a non-sequential `workers` policy, up to `workers.count` batches are in flight at once,
867/// each in its own task; keyed lanes do not apply at batch granularity (the macro rejects
868/// `by_key` on batch forms), so a keyed policy degrades to the plain pool.
869#[allow(clippy::too_many_arguments)] // See spawn_dispatch_workers.
870pub(crate) fn spawn_batch_dispatch<S, H, C, St>(
871 mut subscriber: S,
872 handler: Arc<H>,
873 shutdown: Shutdown,
874 name: Arc<str>,
875 state: Arc<St>,
876 delivery: Arc<Delivery<C>>,
877 failure: DispatchFailure,
878 workers: Workers,
879 batch_size: NonZeroUsize,
880) -> Result<JoinHandle<()>, StartThreadError>
881where
882 S: BatchSubscriber + Send + 'static,
883 S::Message: Send + 'static,
884 S::Batch: 'static,
885 H: BatchHandler<S::Batch, C, St> + 'static,
886 C: Send + Sync + 'static,
887 St: Send + Sync + 'static,
888{
889 // See `spawn_dispatch_workers`: the declaration is what the warning advises, and the harness
890 // runs every subscription on the test's runtime.
891 #[cfg(feature = "poll-diagnostics")]
892 if matches!(workers.placement, Placement::Threads) {
893 delivery.poll.on_dedicated_threads();
894 }
895 #[cfg(feature = "testing")]
896 let workers = if delivery.hooks.coordinator().is_some() {
897 workers.on_runtime()
898 } else {
899 workers
900 };
901 if matches!(workers.placement, Placement::Threads) {
902 return spawn_batch_threads(
903 subscriber, handler, shutdown, name, state, delivery, failure, workers, batch_size,
904 );
905 }
906 Ok(tokio::spawn(async move {
907 // The registration's own batch size, straight to the broker: whatever comes back is the
908 // batch the handler sees.
909 let mut stream = std::pin::pin!(subscriber.batches(batch_size));
910 let mut cancelled = std::pin::pin!(shutdown.cancelled());
911 let mut tasks = JoinSet::new();
912 // One decode buffer for the whole loop: the sequential path lends the same one to every
913 // batch, so the slice a handler reads is allocated once for the subscription. The encode
914 // buffer beside it is the same bargain on the way out. A pooled batch runs beside the
915 // others and cannot share these, so it takes a pair of its own from `spare`, where the
916 // pool keeps what finished workers handed back.
917 let mut scratch = <H as BatchHandler<S::Batch, C, St>>::Scratch::default();
918 let mut encode = BytesMut::new();
919 let mut spare = Vec::new();
920 loop {
921 if shutdown.is_cancelled() {
922 break;
923 }
924 if tasks.len() >= workers.count {
925 // The pool is full: reap a finished worker before polling for more.
926 match reap(&mut tasks, cancelled.as_mut()).await {
927 Turn::Delivery(joined) => spare.extend(log_worker_exit(joined)),
928 // A full pool holds a worker, so `Ended` is shutdown's twin here: either
929 // way the drain below joins whatever is still running.
930 Turn::Ended | Turn::Shutdown => break,
931 }
932 continue;
933 }
934 // The stream first, the token only where the stream has nothing: see `turn_into`, whose
935 // wait this is.
936 let pulled = poll_fn(|cx| match stream.as_mut().poll_next(cx) {
937 Poll::Ready(Some(item)) => Poll::Ready(Turn::Delivery(item)),
938 Poll::Ready(None) => Poll::Ready(Turn::Ended),
939 Poll::Pending => cancelled.as_mut().poll(cx).map(|()| Turn::Shutdown),
940 })
941 .await;
942 match pulled {
943 Turn::Delivery(Ok(batch)) => {
944 if workers.is_sequential() {
945 // Turbofish: the adapter handlers are generic over the batch context,
946 // so the spawn's own parameter names it.
947 run_batch::<_, _, C, _>(
948 &*handler,
949 batch,
950 &mut scratch,
951 &mut encode,
952 &name,
953 &state,
954 &delivery,
955 &failure,
956 )
957 .await;
958 } else {
959 let handler = Arc::clone(&handler);
960 let name = Arc::clone(&name);
961 let state = Arc::clone(&state);
962 let delivery = Arc::clone(&delivery);
963 let failure = failure.clone();
964 let (mut scratch, mut encode) = spare.pop().unwrap_or_default();
965 tasks.spawn(async move {
966 run_batch::<_, _, C, _>(
967 &*handler,
968 batch,
969 &mut scratch,
970 &mut encode,
971 &name,
972 &state,
973 &delivery,
974 &failure,
975 )
976 .await;
977 (scratch, encode)
978 });
979 }
980 }
981 Turn::Delivery(Err(err)) => {
982 error!(
983 target: "ruststream::dispatch",
984 error = %err,
985 "subscriber stream error",
986 );
987 }
988 Turn::Ended => {
989 debug!(
990 target: "ruststream::dispatch",
991 subscriber = %name,
992 "subscriber stream ended",
993 );
994 break;
995 }
996 Turn::Shutdown => break,
997 }
998 }
999 while let Some(joined) = tasks.join_next().await {
1000 log_worker_exit(joined);
1001 }
1002 }))
1003}
1004
1005/// The batch counterpart of [`spawn_dispatch_threads`]: each thread takes whole batches off its
1006/// ring, with a decode and an encode buffer of its own. Keyed lanes do not apply to batches, so a
1007/// keyed policy spreads batches as the plain one does.
1008#[allow(clippy::too_many_arguments)] // See spawn_dispatch_workers.
1009fn spawn_batch_threads<S, H, C, St>(
1010 mut subscriber: S,
1011 handler: Arc<H>,
1012 shutdown: Shutdown,
1013 name: Arc<str>,
1014 state: Arc<St>,
1015 delivery: Arc<Delivery<C>>,
1016 failure: DispatchFailure,
1017 workers: Workers,
1018 batch_size: NonZeroUsize,
1019) -> Result<JoinHandle<()>, StartThreadError>
1020where
1021 S: BatchSubscriber + Send + 'static,
1022 S::Message: Send + 'static,
1023 S::Batch: 'static,
1024 H: BatchHandler<S::Batch, C, St> + 'static,
1025 C: Send + Sync + 'static,
1026 St: Send + Sync + 'static,
1027{
1028 let shared = Arc::new(pool::Shared {
1029 handler,
1030 name,
1031 state,
1032 delivery,
1033 failure,
1034 });
1035 let threads = Threads::start(&shared.name, workers.count, false, |own| {
1036 let shared = Arc::clone(&shared);
1037 let delivery = shared.delivery.on_thread(own);
1038 let mut scratch = <H as BatchHandler<S::Batch, C, St>>::Scratch::default();
1039 let mut encode = BytesMut::new();
1040 async move |batch: S::Batch| {
1041 run_batch::<_, _, C, _>(
1042 &*shared.handler,
1043 batch,
1044 &mut scratch,
1045 &mut encode,
1046 &shared.name,
1047 &shared.state,
1048 &delivery,
1049 &shared.failure,
1050 )
1051 .await;
1052 }
1053 })?;
1054 Ok(tokio::spawn(async move {
1055 threads
1056 .feed(
1057 subscriber.batches(batch_size),
1058 &shared.name,
1059 &shutdown,
1060 |_| None,
1061 )
1062 .await;
1063 }))
1064}
1065
1066async fn dispatch<H, M, C, St>(
1067 handler: &H,
1068 slot: &mut Slot<M>,
1069 encode: &mut BytesMut,
1070 name: &str,
1071 state: &St,
1072 delivery: &Delivery<C>,
1073 failure: &DispatchFailure,
1074) where
1075 H: Handler<M, C, St>,
1076 C: crate::BuildContext<M> + Send + Sync + 'static,
1077 M: IncomingMessage,
1078 // The dispatch future is awaited inside a spawned task, so it must be `Send`: the context
1079 // borrows `&St` across the handler await, which requires `St: Sync`.
1080 St: Send + Sync,
1081{
1082 // Settling the message is what releases the harness's quiescence wait, and the post-settle
1083 // continuations are spawned after it; one in-flight token spanning the whole dispatch keeps a
1084 // `drain` from running before they exist.
1085 #[cfg(feature = "testing")]
1086 let watcher = delivery.hooks.coordinator().cloned();
1087 #[cfg(feature = "testing")]
1088 if let Some(coordinator) = &watcher {
1089 coordinator.enqueued();
1090 }
1091 // Build the broker's typed per-delivery context from the message, then attach the fail-fast
1092 // handle. The delivery is borrowed out of the slot for the whole handler phase and moves
1093 // only at the settle below.
1094 // The context lives in a block of its own and is dropped where it stands when the block ends:
1095 // a `drop(ctx)` call would move the whole context out of this future's state into the call's
1096 // argument first, a `memcpy` of the context per delivery.
1097 let (settle, continuations) = {
1098 let msg = slot.get();
1099 let cx = C::build(msg);
1100 // The delivery itself, not its header map: a handler that reads no headers never reaches
1101 // the broker's accessor, which is a parse and a buffer on most transports.
1102 let source = FromDelivery(msg);
1103 let mut ctx = Context::new(name, &source, state, cx, delivery)
1104 .with_failfast(&failure.shutdown)
1105 .with_encode_buffer(encode)
1106 .with_decode_policy(failure.policies.decode);
1107 // Catch a panicking handler so it cannot silently kill the dispatch loop (which would stop
1108 // the subscriber consuming) or leave the message unsettled. AssertUnwindSafe is required
1109 // because the future borrows `&mut ctx`; that state is discarded with the failed delivery.
1110 // Under the harness, the invocation runs in a task-local slot scope so publishes made
1111 // through injected `Out` publishers are attributed to their slot.
1112 #[cfg(feature = "testing")]
1113 let result = in_harness_scope(
1114 harness_scope(delivery),
1115 observed!(
1116 delivery,
1117 AssertUnwindSafe(handler.handle(msg, &mut ctx)).catch_unwind()
1118 ),
1119 )
1120 .await;
1121 #[cfg(not(feature = "testing"))]
1122 let result = observed!(
1123 delivery,
1124 AssertUnwindSafe(handler.handle(msg, &mut ctx)).catch_unwind()
1125 )
1126 .await;
1127 #[cfg(feature = "testing")]
1128 let panicked = result.is_err();
1129 // Resolve into a `HandlerOutcome` regardless of whether the handler panicked. `None`
1130 // means a fail-fast panic tore the service down and left the message unsettled (a broker
1131 // with redelivery hands it back after the restart).
1132 let settle = match result {
1133 Ok(s) => Some(s),
1134 Err(payload) => {
1135 let reason = panic_reason(payload.as_ref());
1136 error!(
1137 target: "ruststream::dispatch",
1138 subscription = %name,
1139 panic = %reason,
1140 "handler panicked",
1141 );
1142 match failure.policies.panic {
1143 FailurePolicy::FailFast => {
1144 failure
1145 .shutdown
1146 .signal(name, &format!("handler panicked: {reason}"));
1147 None
1148 }
1149 other => Some(
1150 other
1151 .settlement()
1152 .map_or_else(super::handler::HandlerOutcome::drop, Into::into),
1153 ),
1154 }
1155 }
1156 };
1157 // Drain the matching post-settle hooks BEFORE settling: `ctx` borrows `msg`'s headers, and
1158 // settling consumes `msg`. The drained futures own their captures. A fail-fast (no
1159 // settlement) runs no hooks. Most deliveries register none, and those pay the branch alone:
1160 // the list, its scan and the drop glue of both belong to the deliveries that did register
1161 // one.
1162 let continuations = match settle.as_ref() {
1163 Some(s) if ctx.has_hooks() => Some(ctx.take_hooks_for(s.outcome())),
1164 _ => None,
1165 };
1166 // The harness records what the handler saw and how it settled, BEFORE settling the
1167 // message: the matching decrement runs in the broker message's `Drop` (during
1168 // `settle_outcome`, or at the end of this function on the fail-fast path), so the record
1169 // is in place by the time `drive` wakes. Captured here because `settle_outcome` consumes
1170 // `msg` and dropping `ctx` clears the decode flag.
1171 #[cfg(feature = "testing")]
1172 if let Some(coordinator) = delivery.hooks.coordinator() {
1173 coordinator.record(Record {
1174 scope_id: delivery.scope_id,
1175 subscription: delivery.subscription,
1176 name: name.to_owned(),
1177 deliveries: vec![Delivered {
1178 raw: Bytes::copy_from_slice(msg.payload()),
1179 value: ctx.take_carried(),
1180 settle: settle.as_ref().map(super::handler::HandlerOutcome::outcome),
1181 }],
1182 panicked,
1183 decode_failed: ctx.took_decode_failed(),
1184 });
1185 }
1186 (settle, continuations)
1187 };
1188 if let Some(mut s) = settle {
1189 // Named for the same reason as the delivery's own future above: the settle path is
1190 // built where it is polled instead of being copied into this future's state.
1191 let settling = settle_outcome(slot, s.outcome(), name, delivery, C::build as fn(&M) -> C);
1192 settling.await;
1193 // Spawn the `and_after` continuation (if any) onto the tracked set so a graceful shutdown
1194 // drains it. At-most-once: the message is already settled, so a lost or panicking
1195 // continuation never redelivers it.
1196 if let Some(after) = s.take_after() {
1197 delivery.spawn_after(after);
1198 }
1199 } else {
1200 // A fail-fast left the delivery unsettled: it is released here, as it always was at the
1201 // end of its dispatch, rather than when the loop's slot is next filled.
1202 drop(slot.take());
1203 }
1204 // Context-registered hooks run after the message is settled: at-most-once, off the delivery
1205 // path. They ride the same app-wide tracker as an `and_after` continuation, so one drain
1206 // covers both - the harness's `drain` and the shutdown's alike.
1207 if let Some(continuations) = continuations {
1208 for fut in continuations {
1209 delivery.spawn_after(fut);
1210 }
1211 }
1212 #[cfg(feature = "testing")]
1213 if let Some(coordinator) = &watcher {
1214 coordinator.consumed();
1215 }
1216}
1217
1218/// Runs one batch through its handler under panic protection. The handler owns and settles the
1219/// batch's deliveries, so a panic there has already consumed them: the panic policy can only tear
1220/// the service down (`fail_fast`) or be logged and skipped. Per-element settlement is out of scope
1221/// (see the batch decode path for per-element decode handling). Ungated `after_settle` hooks run
1222/// once the batch has settled (per-element outcomes make a gated hook ill-defined on a batch).
1223// The loop hands this everything one batch needs, the buffers it lends included: see
1224// spawn_dispatch_workers.
1225#[allow(clippy::too_many_arguments)]
1226async fn run_batch<H, Batch, C, St>(
1227 handler: &H,
1228 batch: Batch,
1229 scratch: &mut H::Scratch,
1230 encode: &mut BytesMut,
1231 name: &str,
1232 state: &St,
1233 delivery: &Delivery<C>,
1234 failure: &DispatchFailure,
1235) where
1236 H: BatchHandler<Batch, C, St>,
1237 C: Send + Sync + 'static,
1238 St: Send + Sync,
1239{
1240 let batch = H::take(batch);
1241 // A batch spans many deliveries, so its context carries only subscription-scoped data, built
1242 // from its first delivery, or from the batch itself where the handler lends its values; the
1243 // shared app state is threaded the same way as on the single-message path. A batch with
1244 // nothing in it has nothing to settle and nothing to build a context from.
1245 let Some(cx) = batch.context() else { return };
1246 let empty = HeaderMap::new();
1247 let mut ctx = Context::new(name, &empty, state, cx, delivery)
1248 .with_failfast(&failure.shutdown)
1249 .with_encode_buffer(encode)
1250 .with_decode_policy(failure.policies.decode);
1251 // See `dispatch`: the harness scope attributes `Out` publishes to their slot, and lets the
1252 // batch settle path record the batch it applied.
1253 // A batch settles its own deliveries inside the handler (a panic settles them by dropping
1254 // them), so the last decrement lands before the batch record and the fail-fast signal below.
1255 // One extra in-flight token spans the whole dispatch, so a harness driving to quiescence
1256 // cannot return into that window.
1257 #[cfg(feature = "testing")]
1258 let watcher = delivery.hooks.coordinator().cloned();
1259 // A panicking batch settles nothing, so what it carried is captured here (the handler owns
1260 // the deliveries and a panic consumes them) to record the call the settle path never reached.
1261 // Only a harness reads that record, so a dispatch no harness watches captures nothing.
1262 #[cfg(feature = "testing")]
1263 let mut batch = batch;
1264 #[cfg(feature = "testing")]
1265 let unsettled = watcher.as_ref().map(|_| batch.unsettled());
1266 #[cfg(feature = "testing")]
1267 if let Some(coordinator) = &watcher {
1268 coordinator.enqueued();
1269 }
1270 #[cfg(feature = "testing")]
1271 let result = in_harness_scope(
1272 harness_scope(delivery),
1273 observed!(
1274 delivery,
1275 AssertUnwindSafe(handler.handle_batch(batch, scratch, &mut ctx)).catch_unwind()
1276 ),
1277 )
1278 .await;
1279 #[cfg(not(feature = "testing"))]
1280 let result = observed!(
1281 delivery,
1282 AssertUnwindSafe(handler.handle_batch(batch, scratch, &mut ctx)).catch_unwind()
1283 )
1284 .await;
1285 match result {
1286 Ok(()) => {
1287 // As on the single-message path: a batch that registered no hook pays the branch.
1288 if ctx.has_hooks() {
1289 for fut in ctx.take_settle_hooks() {
1290 delivery.spawn_after(fut);
1291 }
1292 }
1293 }
1294 Err(payload) => {
1295 let reason = panic_reason(payload.as_ref());
1296 error!(
1297 target: "ruststream::dispatch",
1298 subscription = %name,
1299 panic = %reason,
1300 "batch handler panicked",
1301 );
1302 #[cfg(feature = "testing")]
1303 if let (Some(coordinator), Some(deliveries)) = (&watcher, unsettled) {
1304 coordinator.record(Record {
1305 scope_id: delivery.scope_id,
1306 subscription: delivery.subscription,
1307 name: name.to_owned(),
1308 deliveries,
1309 panicked: true,
1310 decode_failed: false,
1311 });
1312 }
1313 if failure.policies.panic == FailurePolicy::FailFast {
1314 failure
1315 .shutdown
1316 .signal(name, &format!("batch handler panicked: {reason}"));
1317 }
1318 }
1319 }
1320 #[cfg(feature = "testing")]
1321 if let Some(coordinator) = &watcher {
1322 coordinator.consumed();
1323 }
1324}
1325
1326/// The harness scope a delivery runs under, or `None` when no [`TestApp`](crate::testing::TestApp)
1327/// is driving this app.
1328#[cfg(feature = "testing")]
1329fn harness_scope<C>(delivery: &Delivery<C>) -> Option<HarnessScope> {
1330 delivery
1331 .hooks
1332 .coordinator()
1333 .cloned()
1334 .map(|coordinator| HarnessScope::new(coordinator, delivery.scope_id, delivery.subscription))
1335}
1336
1337/// Settles one delivery by `outcome`, logging an ack / nack failure without propagating it.
1338///
1339/// The single place a settlement reaches the broker, single-message and batch paths alike: a
1340/// second one would be free to answer a [`NackAfter`](HandlerResult::NackAfter) differently, and
1341/// the retry path below is exactly the part that is easy to leave out.
1342pub(crate) async fn settle_outcome<M, C>(
1343 slot: &mut Slot<M>,
1344 outcome: HandlerResult,
1345 name: &str,
1346 delivery: &Delivery<C>,
1347 build_cx: fn(&M) -> C,
1348) where
1349 M: IncomingMessage,
1350 C: Send + Sync + 'static,
1351{
1352 let ack_result = match outcome {
1353 HandlerResult::Ack => slot.take().ack().await,
1354 HandlerResult::Nack { requeue: false } => slot.take().nack(false).await,
1355 HandlerResult::Nack { requeue: true } => settle_retry(slot, name, delivery, build_cx).await,
1356 HandlerResult::NackAfter { delay } => {
1357 settle_nack_after(slot, name, delay, delivery, build_cx).await
1358 }
1359 };
1360 if let Err(err) = ack_result {
1361 warn!(
1362 target: "ruststream::dispatch",
1363 subscription = %name,
1364 error = %err,
1365 "ack / nack failed",
1366 );
1367 }
1368}
1369
1370/// Where the next copy of one delivery goes, read off the registration's declaration and the
1371/// count the delivery carries.
1372#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1373enum Redelivery<'a> {
1374 /// Back to the subscription itself: the attempts are not spent and no destination has taken
1375 /// the copies over.
1376 Subscription,
1377 /// To the declared destination.
1378 DeadLetter {
1379 /// Where the copy is published.
1380 destination: &'a str,
1381 /// Whether it goes there because the attempts are spent, rather than because the
1382 /// registration named a destination and no cap.
1383 at_cap: bool,
1384 },
1385 /// Rejected without a copy: the attempts are spent and the registration named nowhere to send
1386 /// it. A rejection keeps the broker's own dead-letter policy in play where one is configured.
1387 Reject,
1388}
1389
1390/// How many times this message has been delivered, counting this delivery.
1391///
1392/// The broker's own count where the transport keeps one, the framework's header otherwise, never
1393/// both. A transport that counts its deliveries answers for all of them, the ones a delay or a
1394/// requeue brought round included. The header starts absent and is incremented on every copy
1395/// published for the delivery, so the first delivery counts as one either way.
1396fn attempt_of<M: IncomingMessage>(msg: &M) -> u64 {
1397 msg.redelivery_count()
1398 .unwrap_or_else(|| current_retry_count(msg.headers()) + 1)
1399}
1400
1401/// Reads the registration's declaration against this delivery.
1402fn redelivery_of<'d, M: IncomingMessage>(
1403 msg: &M,
1404 declaration: &'d RetryDeclaration,
1405) -> Redelivery<'d> {
1406 match (declaration.max_attempts(), declaration.dead_letter()) {
1407 (Some(cap), destination) if attempt_of(msg) >= u64::from(cap.get()) => {
1408 destination.map_or(Redelivery::Reject, |destination| Redelivery::DeadLetter {
1409 destination,
1410 at_cap: true,
1411 })
1412 }
1413 (None, Some(destination)) => Redelivery::DeadLetter {
1414 destination,
1415 at_cap: false,
1416 },
1417 _ => Redelivery::Subscription,
1418 }
1419}
1420
1421/// Settles an immediate retry ([`Nack { requeue: true }`](HandlerResult::Nack)).
1422///
1423/// A registration that declared neither a cap nor a destination gets the broker's own requeue,
1424/// which is what an immediate retry has always been. Under a declaration the outcome obeys it:
1425/// where the transport counts its own redeliveries the requeue stays the broker's and the count
1426/// is read from the delivery, and where it counts none the runtime publishes the copy itself so
1427/// that the framework's header carries the count forward and the cap applies to an immediate
1428/// retry as it does to a delayed one.
1429///
1430/// A registration whose copies the broker moves itself is never capped here, exactly as on the
1431/// native delayed path: the broker's requeue is the immediate retry, and the declaration reached
1432/// the subscription descriptor at startup for the broker to apply. Reading the cap here would
1433/// settle a spent delivery with a plain rejection, and on a queue that deletes a rejected
1434/// delivery (SQS) that loses the message instead of leaving it to the redrive policy.
1435///
1436/// # Errors
1437///
1438/// Returns the [`AckError`] from settling the original delivery.
1439async fn settle_retry<M, C>(
1440 slot: &mut Slot<M>,
1441 name: &str,
1442 delivery: &Delivery<C>,
1443 build_cx: fn(&M) -> C,
1444) -> Result<(), AckError>
1445where
1446 M: IncomingMessage,
1447 C: Send + Sync + 'static,
1448{
1449 if delivery.declaration.declares_nothing() {
1450 return slot.take().nack(true).await;
1451 }
1452 // The broker moves this subscription's deliveries itself: its requeue is what an immediate
1453 // retry is, and the declaration is the broker's to apply, so no cap is read here, as on the
1454 // native delayed path.
1455 let Some(retry) = delivery.retry.as_ref() else {
1456 return slot.take().nack(true).await;
1457 };
1458 match redelivery_of(slot.get(), &delivery.declaration) {
1459 // The broker counts its own redeliveries, so its requeue carries the count forward and
1460 // there is nothing for a copy to add.
1461 Redelivery::Subscription if slot.get().redelivery_count().is_some() => {
1462 slot.take().nack(true).await
1463 }
1464 Redelivery::Subscription => {
1465 let destination = retry.destination.clone();
1466 publish_copy(slot, name, destination, None, retry, delivery, build_cx).await
1467 }
1468 Redelivery::DeadLetter {
1469 destination,
1470 at_cap,
1471 } => {
1472 if at_cap {
1473 warn_at_cap(name, attempt_of(slot.get()), Some(destination));
1474 }
1475 let destination = Arc::from(destination);
1476 publish_copy(
1477 slot,
1478 name,
1479 Some(destination),
1480 None,
1481 retry,
1482 delivery,
1483 build_cx,
1484 )
1485 .await
1486 }
1487 Redelivery::Reject => {
1488 warn_at_cap(name, attempt_of(slot.get()), None);
1489 slot.take().nack(false).await
1490 }
1491 }
1492}
1493
1494/// Settles a [`NackAfter`](HandlerResult::NackAfter) outcome, choosing native delayed redelivery
1495/// or the copy the runtime publishes.
1496///
1497/// When the broker reports native support (`supports_nack_after`), this defers to
1498/// [`IncomingMessage::nack_after`] and nothing is published - unless the delivery is already at
1499/// the registration's cap. The cap is read first, from the one count the transport has (see
1500/// [`attempt_of`]): the broker's own where it keeps one, the framework's header otherwise, never
1501/// both. A delivery at the cap goes to the declared dead-letter destination or is rejected,
1502/// exactly as on the copy path.
1503/// Where the broker moves a spent delivery itself the cap is not read on this path: the
1504/// declaration is the subscription descriptor's to map onto the broker's own mechanism.
1505///
1506/// Without native support this captures the message, drops the original, and schedules a copy of it
1507/// after the delay, with the [`RETRY_COUNT_HEADER`] incremented. A zero delay publishes the copy at
1508/// once, on the dispatch path, as an immediate retry does. The copy goes to the address the
1509/// subscription's descriptor reported at startup, not to the subscription's name: the two differ
1510/// wherever a subscription is a resource of its own. It leaves through the registration's retry
1511/// slot, so the mount site's transforms and the publisher's own headers reach it as they reach any
1512/// slot publish; it carries bytes already, so no codec encodes it and the call adjusts no
1513/// per-message settings.
1514///
1515/// Once the registration's attempts are spent the copy goes to the declared dead-letter
1516/// destination instead, immediately, and with no destination declared the delivery is rejected.
1517///
1518/// A transport with no settlement at all ([`AckError::Unsupported`] from `nack`, as on MQTT at
1519/// `QoS` 0, `ZeroMQ`, or Redis pub/sub) still gets the copy: there is no original to drop and no
1520/// redelivery of the broker's own to fall back on, so the copy is the only form the retry can
1521/// take.
1522///
1523/// Where the broker moves the subscription's deliveries itself and cannot hold one back, the
1524/// delivery is requeued at once, and a warning says the delay is dropped. A delivery the
1525/// declaration sends elsewhere (its attempts are spent, or a destination takes every copy) gets
1526/// the same requeue without the warning: the broker moves it on that requeue, so no delay was due.
1527///
1528/// # Errors
1529///
1530/// Returns the [`AckError`] from settling the original when the transport does settle but this
1531/// settle failed. The delivery then stays with the broker, which redelivers it on its own timers,
1532/// and a copy on top of that would duplicate the message; the caller logs the error.
1533///
1534/// # Cancel safety
1535///
1536/// A deferred copy runs on a detached task that sleeps for `delay`. It is at-most-once over
1537/// that window: if the process exits (or the runtime is dropped) before the timer fires, the
1538/// deferred message is lost, since the original has already been dropped. Brokers that need
1539/// at-least-once delayed redelivery across a crash must provide native support.
1540async fn settle_nack_after<M, C>(
1541 slot: &mut Slot<M>,
1542 name: &str,
1543 delay: Duration,
1544 delivery: &Delivery<C>,
1545 build_cx: fn(&M) -> C,
1546) -> Result<(), AckError>
1547where
1548 M: IncomingMessage,
1549 C: Send + Sync + 'static,
1550{
1551 if slot.get().supports_nack_after() {
1552 // The cap is read before the delay reaches the broker: a native redelivery would otherwise
1553 // circle past a cap nothing in this process ever gets to apply. It is read wherever this
1554 // process publishes the subscription's copies, because a crate that honours the delay by
1555 // republishing the delivery increments the header the cap counts with where the transport
1556 // counts nothing; a broker that moves a spent delivery itself applies the declaration
1557 // itself.
1558 if let Some(retry) = delivery.retry.as_ref()
1559 && delivery.declaration.max_attempts().is_some()
1560 {
1561 match redelivery_of(slot.get(), &delivery.declaration) {
1562 Redelivery::DeadLetter { destination, .. } => {
1563 warn_at_cap(name, attempt_of(slot.get()), Some(destination));
1564 let destination = Arc::from(destination);
1565 return publish_copy(
1566 slot,
1567 name,
1568 Some(destination),
1569 None,
1570 retry,
1571 delivery,
1572 build_cx,
1573 )
1574 .await;
1575 }
1576 Redelivery::Reject => {
1577 warn_at_cap(name, attempt_of(slot.get()), None);
1578 return slot.take().nack(false).await;
1579 }
1580 Redelivery::Subscription => {}
1581 }
1582 }
1583 let accepted = slot.take().nack_after(delay).await;
1584 // The broker holds this redelivery on its own timer, out of the harness's sight, so the
1585 // harness is told when it falls due.
1586 #[cfg(feature = "testing")]
1587 if accepted.is_ok()
1588 && let Some(coordinator) = delivery.hooks.coordinator()
1589 {
1590 coordinator.expect_redelivery(delivery.subscription, delay);
1591 }
1592 return accepted;
1593 }
1594
1595 let Some(retry) = delivery.retry.as_ref() else {
1596 // A delivery the declaration sends elsewhere is the broker's to move, and this requeue is
1597 // what moves it, so it was never due a delay. Only one that comes back loses its delay.
1598 if matches!(
1599 redelivery_of(slot.get(), &delivery.declaration),
1600 Redelivery::Subscription
1601 ) {
1602 warn!(
1603 target: "ruststream::dispatch",
1604 subscription = %name,
1605 "retry_after on a broker with neither native delayed redelivery nor a copy path \
1606 of its own; requeuing immediately (the delay is dropped)",
1607 );
1608 }
1609 return slot.take().nack(true).await;
1610 };
1611
1612 match redelivery_of(slot.get(), &delivery.declaration) {
1613 Redelivery::Subscription => {
1614 let destination = retry.destination.clone();
1615 publish_copy(
1616 slot,
1617 name,
1618 destination,
1619 Some(delay),
1620 retry,
1621 delivery,
1622 build_cx,
1623 )
1624 .await
1625 }
1626 Redelivery::DeadLetter {
1627 destination,
1628 at_cap,
1629 } => {
1630 if at_cap {
1631 warn_at_cap(name, attempt_of(slot.get()), Some(destination));
1632 }
1633 // A delivery that has run out of attempts is not retried, so the delay it asked for
1634 // does not apply to the copy that carries it away.
1635 let destination = Arc::from(destination);
1636 publish_copy(
1637 slot,
1638 name,
1639 Some(destination),
1640 None,
1641 retry,
1642 delivery,
1643 build_cx,
1644 )
1645 .await
1646 }
1647 Redelivery::Reject => {
1648 warn_at_cap(name, attempt_of(slot.get()), None);
1649 slot.take().nack(false).await
1650 }
1651 }
1652}
1653
1654/// Says once per exhausted delivery where it went, because that is the point at which a message
1655/// stops coming back and an operator has to know.
1656fn warn_at_cap(name: &str, attempt: u64, destination: Option<&str>) {
1657 if let Some(destination) = destination {
1658 warn!(
1659 target: "ruststream::dispatch",
1660 subscription = %name,
1661 attempt,
1662 dead_letter = %destination,
1663 "retry attempts exhausted; publishing the delivery to the dead-letter destination",
1664 );
1665 } else {
1666 warn!(
1667 target: "ruststream::dispatch",
1668 subscription = %name,
1669 attempt,
1670 "retry attempts exhausted and no dead-letter destination is declared; rejecting the \
1671 delivery",
1672 );
1673 }
1674}
1675
1676/// Drops the original delivery and sends a copy of it to `destination`, now or after `delay`; a
1677/// zero delay is now.
1678///
1679/// `destination` is `None` only where the registration's transforms name one per delivery; the
1680/// copy then starts at the subscription's own name, which is what the transforms read as the
1681/// delivery's channel, and one of them renames it.
1682///
1683/// # Errors
1684///
1685/// Returns the [`AckError`] from dropping the original where that failed and the transport does
1686/// settle: the delivery then stays with the broker, and a copy on top of it would duplicate the
1687/// message.
1688async fn publish_copy<M, C>(
1689 slot: &mut Slot<M>,
1690 name: &str,
1691 destination: Option<Arc<str>>,
1692 delay: Option<Duration>,
1693 retry: &DeferredRetry<C>,
1694 delivery: &Delivery<C>,
1695 build_cx: fn(&M) -> C,
1696) -> Result<(), AckError>
1697where
1698 M: IncomingMessage,
1699 C: Send + Sync + 'static,
1700{
1701 let publisher = Arc::clone(&retry.publisher);
1702 let subscription: Arc<str> = Arc::from(name);
1703 let target = destination.unwrap_or_else(|| Arc::clone(&subscription));
1704
1705 // Settling consumes the message, so capture everything the copy needs first: its bytes, the
1706 // headers the handler saw, and the broker's own per-delivery context, which is what a
1707 // transform on this position reads. The bytes are lent on from here rather than handed over:
1708 // this path runs only after a delivery has failed, so a broker that keeps owned bytes copies
1709 // once there, where handing them over would cost the runtime a conversion on every retry.
1710 let msg = slot.get();
1711 let payload = Bytes::copy_from_slice(msg.payload());
1712 let delivered = msg.headers().clone();
1713 let mut headers = delivered.clone();
1714 let next_count = current_retry_count(&headers) + 1;
1715 headers.insert(Str::from_static(RETRY_COUNT_HEADER), next_count.to_string());
1716 let context = build_cx(msg);
1717
1718 // Drop the original so the broker does not also redeliver it; the copy carries the retry
1719 // forward. A transport with no settlement at all has nothing to drop and no redelivery of its
1720 // own, so there the copy is the only way the message survives: aborting on that error would
1721 // lose it. Any other settle failure leaves the delivery with the broker, which will redeliver
1722 // it on its own timers, so a copy on top would duplicate it.
1723 //
1724 // Under the harness the original's settlement closes its count before the copy's publish
1725 // opens the copy's; the hold keeps the reaction open across that gap until this function
1726 // returns, so `drive` cannot settle between the two.
1727 #[cfg(feature = "testing")]
1728 let _hold = delivery.hooks.coordinator().map(Coordinator::hold);
1729 match slot.take().nack(false).await {
1730 Ok(()) | Err(AckError::Unsupported) => {}
1731 Err(err) => return Err(err),
1732 }
1733
1734 let republish = async move {
1735 let cx = PublishContext::new(&subscription, &delivered, &context);
1736 let copy = OutgoingMessage::new(target.as_ref(), payload.as_ref()).with_headers(headers);
1737 if let Err(err) = publisher.publish_copy(copy, &cx).await {
1738 warn!(
1739 target: "ruststream::dispatch",
1740 subscription = %subscription,
1741 destination = %target,
1742 error = %err,
1743 "publishing the retry copy failed; message lost",
1744 );
1745 }
1746 };
1747
1748 // The copy leaves after the handler returned, so the harness watching the delivery is
1749 // handed on to it: the copy is one of the delivery's publishes.
1750 #[cfg(feature = "testing")]
1751 let republish = in_harness_scope(harness_scope(delivery), republish);
1752
1753 // A zero delay is an immediate copy: a timer task for it would cost a task allocation and
1754 // two scheduler turns to publish the same copy a little later.
1755 let Some(delay) = delay.filter(|delay| !delay.is_zero()) else {
1756 // An immediate copy is awaited on the dispatch path: there is no timer to wait for, and a
1757 // detached task would let the loop pull the next delivery before this one is back.
1758 republish.await;
1759 return Ok(());
1760 };
1761
1762 // Under the harness the copy is scheduled through the coordinator, the way a broker schedules
1763 // its native delayed redelivery: `TestApp::advance` then fires it and counts it before it
1764 // drives the reaction, so a test sees the copy rather than a race.
1765 #[cfg(feature = "testing")]
1766 if let Some(coordinator) = delivery.hooks.coordinator() {
1767 coordinator.schedule_redelivery_future(delay, republish);
1768 return Ok(());
1769 }
1770 // The timer is the delivery's own work, so it runs where the delivery ran; being tracked, a
1771 // graceful shutdown waits for it and the copy leaves before the broker closes.
1772 delivery.spawn_after(async move {
1773 tokio::time::sleep(delay).await;
1774 republish.await;
1775 });
1776 Ok(())
1777}
1778
1779#[cfg(all(test, feature = "testing", feature = "memory"))]
1780mod harness_tests;
1781mod pool;
1782#[cfg(all(test, feature = "memory"))]
1783mod tests;
1784mod threads;