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