Skip to main content

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;