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