Skip to main content

camel_processor/
wire_tap.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::atomic::{AtomicBool, Ordering};
4use std::sync::{Arc, Mutex};
5use std::task::{Context, Poll};
6use std::time::Duration;
7
8use async_trait::async_trait;
9use tokio::sync::{Semaphore, TryAcquireError};
10use tokio_util::sync::CancellationToken;
11use tokio_util::task::TaskTracker;
12use tower::{Service, ServiceExt};
13
14use camel_api::{CamelError, Exchange, StepLifecycle, StepShutdownReason};
15
16/// Configuration for [`WireTapService`].
17///
18/// Default concurrency bound is 20 (Camel-faithful flat-semaphore).
19/// `shutdown_grace` defaults to 5 seconds.
20#[derive(Clone)]
21pub struct WireTapConfig {
22    /// Maximum number of concurrent tap tasks. `None` means unlimited.
23    pub max_concurrent: Option<usize>,
24    /// Grace period for in-flight tap tasks to complete on shutdown.
25    /// A value of zero means "skip drain, cancel immediately".
26    pub shutdown_grace: std::time::Duration,
27}
28
29impl Default for WireTapConfig {
30    fn default() -> Self {
31        Self {
32            max_concurrent: Some(20),
33            shutdown_grace: std::time::Duration::from_secs(5),
34        }
35    }
36}
37
38impl WireTapConfig {
39    /// Validate the config, panicking on invalid states.
40    ///
41    /// `shutdown_grace` of zero is valid (means "skip drain, cancel immediately").
42    pub fn validate(&self) {
43        if self.max_concurrent == Some(0) {
44            panic!("max_concurrent must be > 0 when set");
45        }
46    }
47
48    /// Create a config with a bounded concurrency limit.
49    pub fn bounded(max_concurrent: usize) -> Self {
50        assert!(max_concurrent > 0, "max_concurrent must be > 0");
51        Self {
52            max_concurrent: Some(max_concurrent),
53            shutdown_grace: std::time::Duration::from_secs(5),
54        }
55    }
56}
57
58/// Mutable admission-gate state guarded by [`WireTapShared::inner`].
59///
60/// The `Mutex` over this struct serializes the "check `open` → register task"
61/// critical section so a `shutdown()` racing with a `call()` cannot orphan a
62/// task: either `call()` registers under the lock (and `shutdown` drains it via
63/// `tracker.wait`), or `shutdown` closes admission first (and `call()` returns
64/// early without registering). There is no `await` point while the lock is held.
65#[derive(Debug)]
66struct WireTapSharedInner {
67    open: bool,
68    tracker: TaskTracker,
69    cancel: CancellationToken,
70    semaphore: Option<Arc<Semaphore>>,
71    shutdown_grace: Duration,
72}
73
74/// Shared admission gate, liveness tracker, and cancellation token for a
75/// [`WireTapService`] and its clones.
76///
77/// All clones of a `WireTapService` share the SAME `Arc<WireTapShared>`: per-
78/// request clones drop an `Arc` ref but do NOT close admission or cancel taps.
79/// Only the last-ref drop (canonical-service teardown) fires [`Drop`], which
80/// cancels every in-flight tap. This is defense-in-depth alongside the runtime-
81/// driven `StepLifecycle::shutdown` path (ADR-0022 mandates shutdown-before-drop,
82/// but Drop guarantees cleanup if the runtime fails to call shutdown).
83#[derive(Debug)]
84struct WireTapShared {
85    inner: Mutex<WireTapSharedInner>,
86}
87
88impl Drop for WireTapShared {
89    fn drop(&mut self) {
90        // Defense-in-depth for the cancel-and-drain contract: when the last
91        // `Arc<WireTapShared>` ref drops (canonical-service teardown), cancel
92        // every in-flight tap so spawned tasks unwind promptly via the
93        // `cancel.cancelled()` select branch in `run_tap`. The runtime calls
94        // `StepLifecycle::shutdown` before drop per ADR-0022, but Drop
95        // guarantees cleanup if it does not.
96        //
97        // Cancel BEFORE the inner fields drop: once `cancel.cancel()` fires the
98        // cancellation state is latched into the token (and all clones held by
99        // in-flight tasks), so the subsequent `CancellationToken::drop` and
100        // `TaskTracker::drop` (which detaches rather than aborts) do not race
101        // the cancellation signal.
102        self.inner
103            .lock()
104            .expect("WireTapShared mutex poisoned") // allow-unwrap
105            .cancel
106            .cancel();
107    }
108}
109
110pub struct WireTapService {
111    tap_endpoint: camel_api::BoxProcessor,
112    shared: Arc<WireTapShared>,
113}
114
115// The shared admission gate, liveness tracker, and cancellation token live in
116// `Arc<WireTapShared>`: each clone gets a new ref to the SAME shared state.
117// This is required because the route pipeline clones the `BoxProcessor` per
118// request (`BoxCloneService` contract) and drops the clone once `call()`'s
119// immediate-return future resolves. Per-clone state would close admission on
120// every request drop. Sharing keeps the gate open until canonical teardown.
121impl Clone for WireTapService {
122    fn clone(&self) -> Self {
123        Self {
124            tap_endpoint: self.tap_endpoint.clone(),
125            shared: Arc::clone(&self.shared),
126        }
127    }
128}
129
130impl WireTapService {
131    /// Create a new `WireTapService` with default (bounded-20) concurrency.
132    pub fn new(tap_endpoint: camel_api::BoxProcessor) -> Self {
133        Self::with_config(tap_endpoint, WireTapConfig::default())
134    }
135
136    /// Create a new `WireTapService` from a [`WireTapConfig`].
137    pub fn with_config(tap_endpoint: camel_api::BoxProcessor, config: WireTapConfig) -> Self {
138        config.validate();
139        let semaphore = config
140            .max_concurrent
141            .map(|limit| Arc::new(Semaphore::new(limit)));
142        let shared = Arc::new(WireTapShared {
143            inner: Mutex::new(WireTapSharedInner {
144                open: true,
145                tracker: TaskTracker::new(),
146                cancel: CancellationToken::new(),
147                semaphore,
148                shutdown_grace: config.shutdown_grace,
149            }),
150        });
151        Self {
152            tap_endpoint,
153            shared,
154        }
155    }
156
157    /// Test-only accessor for the count of currently-tracked (detached) tap
158    /// tasks. This counts ONLY detached tasks registered with the
159    /// [`TaskTracker`]; the inline CallerRuns path is NOT counted. Therefore
160    /// the bound invariant (`bound + 1` total concurrent execution, where the
161    /// +1 is the inline tap) is not observable through this accessor.
162    #[cfg(test)]
163    pub(crate) fn in_flight_count(&self) -> usize {
164        self.shared
165            .inner
166            .lock()
167            .expect("WireTapShared mutex poisoned") // allow-unwrap
168            .tracker
169            .len()
170    }
171}
172
173/// A lifecycle handle for a [`WireTapService`] implementing graceful-drain-
174/// then-abort teardown via [`StepLifecycle::shutdown`].
175///
176/// Obtain via [`WireTapService::lifecycle`]. The handle shares the same
177/// [`Arc<WireTapShared>`] as the service, so `shutdown` observes the live
178/// admission gate and task tracker.
179#[derive(Debug)]
180pub struct WireTapLifecycle {
181    shared: Arc<WireTapShared>,
182    shutdown_called: AtomicBool,
183}
184
185#[async_trait]
186impl StepLifecycle for WireTapLifecycle {
187    fn name(&self) -> &'static str {
188        "wiretap"
189    }
190
191    async fn shutdown(&self, _reason: StepShutdownReason) -> Result<(), CamelError> {
192        // Idempotency gate.
193        if self.shutdown_called.swap(true, Ordering::SeqCst) {
194            return Ok(());
195        }
196
197        // (a) Close admission, close the tracker, clone handles out of the lock.
198        let (tracker, cancel, grace) = {
199            let mut guard = self
200                .shared
201                .inner
202                .lock()
203                .expect("WireTapShared mutex poisoned"); // allow-unwrap
204            guard.open = false;
205            guard.tracker.close();
206            (
207                guard.tracker.clone(),
208                guard.cancel.clone(),
209                guard.shutdown_grace,
210            )
211            // MutexGuard dropped here — NO .await held across the guard.
212        };
213
214        // (b) If zero grace, skip drain — go straight to cancel.
215        // (c) DRAIN FIRST: await in-flight taps that complete naturally
216        //     within the grace period.
217        if !grace.is_zero() {
218            let _ = tokio::time::timeout(grace, tracker.wait()).await;
219        }
220
221        // (d) CANCEL: abort any stragglers that did not drain within grace.
222        cancel.cancel();
223
224        // (e) Await cancel-completions (tasks that abort on token fire).
225        let _ = tracker.wait().await;
226
227        Ok(())
228    }
229}
230
231impl WireTapService {
232    /// Obtain a shared lifecycle handle for this service's admission gate and
233    /// task tracker. Callers can invoke [`StepLifecycle::shutdown`] on the
234    /// returned handle for graceful-drain-then-abort teardown.
235    pub fn lifecycle(&self) -> Arc<dyn StepLifecycle> {
236        Arc::new(WireTapLifecycle {
237            shared: Arc::clone(&self.shared),
238            shutdown_called: AtomicBool::new(false),
239        })
240    }
241}
242
243impl Service<Exchange> for WireTapService {
244    type Response = Exchange;
245    type Error = CamelError;
246    type Future = Pin<Box<dyn Future<Output = Result<Exchange, CamelError>> + Send>>;
247
248    /// Always ready (ADR-0019): the main route never blocks on tap readiness.
249    /// Tap endpoint readiness is driven inside [`run_tap`] on the fire-and-
250    /// forgetget path; a tap readiness error is logged and suppressed, never
251    /// propagated to the main exchange.
252    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
253        Poll::Ready(Ok(()))
254    }
255
256    fn call(&mut self, exchange: Exchange) -> Self::Future {
257        let tap_endpoint = self.tap_endpoint.clone();
258        let tap_exchange = exchange.clone();
259
260        // Admission critical section: hold the lock across "check open →
261        // admit-or-inline decision → register tracked task" so a racing
262        // `shutdown` cannot close the tracker between the open-check and the
263        // task registration. There is NO await point while the lock is held;
264        // `try_acquire_owned`, `tracker.spawn`, and the open-check are all sync.
265        let inner = self
266            .shared
267            .inner
268            .lock()
269            .expect("WireTapShared mutex poisoned"); // allow-unwrap
270        if !inner.open {
271            tracing::warn!("WireTap admission closed, dropping tap");
272            drop(inner);
273            return Box::pin(async move { Ok(exchange) });
274        }
275
276        match &inner.semaphore {
277            Some(sem) => match Arc::clone(sem).try_acquire_owned() {
278                Ok(permit) => {
279                    // Admit: register a detached tracked task holding the permit.
280                    // The OwnedSemaphorePermit is MOVED into the task body and
281                    // lives for the task's lifetime, releasing on completion.
282                    let cancel = inner.cancel.clone();
283                    inner.tracker.spawn(async move {
284                        let _permit = permit;
285                        run_tap(tap_endpoint, tap_exchange, cancel).await;
286                    });
287                    drop(inner);
288                    Box::pin(async move { Ok(exchange) })
289                }
290                Err(TryAcquireError::NoPermits) => {
291                    // Saturated: run the tap INLINE on the caller's future. No
292                    // permit is acquired or held, so total concurrent execution
293                    // transiently reaches `bound + 1` (this inline tap alongside
294                    // the `bound` detached permit-holders). The caller is
295                    // back-pressured until the inline tap finishes (CallerRuns).
296                    let cancel = inner.cancel.clone();
297                    drop(inner);
298                    Box::pin(async move {
299                        run_tap(tap_endpoint, tap_exchange, cancel).await;
300                        Ok(exchange)
301                    })
302                }
303                Err(TryAcquireError::Closed) => {
304                    tracing::warn!("WireTap semaphore closed, dropping tap");
305                    drop(inner);
306                    Box::pin(async move { Ok(exchange) })
307                }
308            },
309            None => {
310                // Unbounded: register a detached tracked task with no permit.
311                let cancel = inner.cancel.clone();
312                inner.tracker.spawn(async move {
313                    run_tap(tap_endpoint, tap_exchange, cancel).await;
314                });
315                drop(inner);
316                Box::pin(async move { Ok(exchange) })
317            }
318        }
319    }
320}
321
322/// Single private helper shared by the detached path and the inline CallerRuns
323/// path. Drives the tap endpoint to readiness then calls it, racing against the
324/// shared `cancel` token so shutdown/abort unwinds promptly. Tap readiness and
325/// processing errors are logged at `warn!` (category handler-owned per
326/// ADR-0012) and suppressed — the main exchange proceeds unchanged.
327async fn run_tap(
328    mut tap_endpoint: camel_api::BoxProcessor,
329    tap_exchange: Exchange,
330    cancel: CancellationToken,
331) {
332    // Readiness phase: cancel races against `tap_endpoint.ready()`.
333    {
334        let ready_fut = tap_endpoint.ready();
335        tokio::pin!(ready_fut);
336        let ready_result = tokio::select! {
337            biased;
338            _ = cancel.cancelled() => { return; }
339            r = &mut ready_fut => r,
340        };
341        if let Err(e) = ready_result {
342            // log-policy: handler-owned
343            tracing::warn!("WireTap endpoint poll_ready failed: {}", e);
344            return;
345        }
346    }
347    // Call phase: tap_endpoint is now Ready; cancel races against the call.
348    {
349        let call_fut = tap_endpoint.call(tap_exchange);
350        tokio::pin!(call_fut);
351        let call_result = tokio::select! {
352            biased;
353            _ = cancel.cancelled() => { return; }
354            r = &mut call_fut => r,
355        };
356        if let Err(e) = call_result {
357            // log-policy: handler-owned
358            tracing::warn!("WireTap processing error: {}", e);
359        }
360    }
361}
362
363/// A Tower layer that produces `WireTapService` instances.
364pub struct WireTapLayer {
365    tap_endpoint: camel_api::BoxProcessor,
366    config: WireTapConfig,
367}
368
369impl WireTapLayer {
370    /// Create a new WireTapLayer with the given tap endpoint processor (default bounded-20 concurrency).
371    pub fn new(tap_endpoint: camel_api::BoxProcessor) -> Self {
372        Self {
373            tap_endpoint,
374            config: WireTapConfig::default(),
375        }
376    }
377
378    /// Create a new WireTapLayer with bounded concurrency.
379    pub fn bounded(tap_endpoint: camel_api::BoxProcessor, max_concurrent: usize) -> Self {
380        Self {
381            tap_endpoint,
382            config: WireTapConfig::bounded(max_concurrent),
383        }
384    }
385}
386
387impl<S> tower::Layer<S> for WireTapLayer {
388    type Service = WireTapService;
389
390    fn layer(&self, _inner: S) -> Self::Service {
391        WireTapService::with_config(self.tap_endpoint.clone(), self.config.clone())
392    }
393}
394
395#[cfg(test)]
396mod tests {
397    use super::*;
398    use camel_api::{BoxProcessor, BoxProcessorExt, Message};
399    use std::sync::Arc;
400    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
401    use tower::ServiceExt;
402
403    // --- Existing tests retained / adapted to the new shared-state model ---
404
405    #[tokio::test]
406    async fn test_wire_tap_returns_original_immediately() {
407        let tap_processor = BoxProcessor::from_fn(|ex| Box::pin(async move { Ok(ex) }));
408
409        let mut wire_tap = WireTapService::new(tap_processor);
410        let exchange = Exchange::new(Message::new("test message"));
411
412        let result = wire_tap
413            .ready()
414            .await
415            .unwrap()
416            .call(exchange)
417            .await
418            .unwrap();
419
420        assert_eq!(result.input.body.as_text(), Some("test message"));
421    }
422
423    #[tokio::test]
424    async fn test_wire_tap_endpoint_receives_clone() {
425        let received_count = Arc::new(AtomicUsize::new(0));
426        let count_clone = received_count.clone();
427
428        let tap_processor = BoxProcessor::from_fn(move |ex| {
429            let count = count_clone.clone();
430            Box::pin(async move {
431                count.fetch_add(1, Ordering::SeqCst);
432                Ok(ex)
433            })
434        });
435
436        let mut wire_tap = WireTapService::new(tap_processor);
437        let exchange = Exchange::new(Message::new("test"));
438
439        let _result = wire_tap
440            .ready()
441            .await
442            .unwrap()
443            .call(exchange)
444            .await
445            .unwrap();
446
447        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
448
449        assert_eq!(received_count.load(Ordering::SeqCst), 1);
450    }
451
452    #[tokio::test]
453    async fn test_wire_tap_isolates_errors() {
454        let tap_processor = BoxProcessor::from_fn(|_ex| {
455            Box::pin(async move { Err(CamelError::ProcessorError("tap error".into())) })
456        });
457
458        let mut wire_tap = WireTapService::new(tap_processor);
459        let exchange = Exchange::new(Message::new("test"));
460
461        let result = wire_tap.ready().await.unwrap().call(exchange).await;
462
463        assert!(result.is_ok());
464        assert_eq!(result.unwrap().input.body.as_text(), Some("test"));
465    }
466
467    #[tokio::test]
468    async fn test_wire_tap_layer() {
469        use tower::Layer;
470
471        let tap_processor = BoxProcessor::from_fn(|ex| Box::pin(async move { Ok(ex) }));
472
473        let layer = super::WireTapLayer::new(tap_processor);
474        let inner = camel_api::IdentityProcessor;
475        let mut svc = layer.layer(inner);
476
477        let exchange = Exchange::new(Message::new("test"));
478        let result = svc.ready().await.unwrap().call(exchange).await.unwrap();
479
480        assert_eq!(result.input.body.as_text(), Some("test"));
481    }
482
483    #[tokio::test]
484    async fn test_wiretap_bounded_concurrency() {
485        // Under the new CallerRuns admission model, when the bound is saturated
486        // the next call runs its tap INLINE on the caller's future (without
487        // acquiring a permit). The transient peak concurrent execution is
488        // therefore `bound + 1` (the inline tap alongside the `bound` detached
489        // permit-holders), matching the spec invariant. The old `<= bound`
490        // assertion reflected the leaky spawn-then-acquire model.
491        let concurrent = Arc::new(AtomicUsize::new(0));
492        let max_concurrent = Arc::new(AtomicUsize::new(0));
493
494        let c = Arc::clone(&concurrent);
495        let mc = Arc::clone(&max_concurrent);
496        let tap_processor = BoxProcessor::from_fn(move |ex| {
497            let c = Arc::clone(&c);
498            let mc = Arc::clone(&mc);
499            Box::pin(async move {
500                let current = c.fetch_add(1, Ordering::SeqCst) + 1;
501                mc.fetch_max(current, Ordering::SeqCst);
502                tokio::time::sleep(std::time::Duration::from_millis(50)).await;
503                c.fetch_sub(1, Ordering::SeqCst);
504                Ok(ex)
505            })
506        });
507
508        let config = super::WireTapConfig::bounded(2);
509        let mut svc = super::WireTapService::with_config(tap_processor, config);
510
511        for _ in 0..3 {
512            let ex = Exchange::new(Message::new("test"));
513            let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
514        }
515
516        tokio::time::sleep(std::time::Duration::from_millis(300)).await;
517
518        let observed_max = max_concurrent.load(Ordering::SeqCst);
519        // CallerRuns allows `bound + 1` (the inline tap).
520        assert!(
521            observed_max <= 3,
522            "max concurrency was {observed_max}, expected <= bound+1 (=3) under CallerRuns"
523        );
524    }
525
526    #[tokio::test]
527    async fn test_wire_tap_survives_per_request_clone_drop() {
528        // Regression for the clone-abort bug (rc-vq91): the route pipeline
529        // clones the BoxProcessor per request and drops the clone once call()'s
530        // immediate-return future resolves. With per-clone state, that drop
531        // would close admission. Sharing `Arc<WireTapShared>` keeps the gate
532        // open across clone drops; only the last-ref drop fires cancellation.
533        let completed = Arc::new(AtomicUsize::new(0));
534        let completed_clone = completed.clone();
535
536        let tap_processor = BoxProcessor::from_fn(move |ex| {
537            let c = completed_clone.clone();
538            Box::pin(async move {
539                tokio::time::sleep(std::time::Duration::from_millis(150)).await;
540                c.fetch_add(1, Ordering::SeqCst);
541                Ok(ex)
542            })
543        });
544
545        let canonical = WireTapService::new(tap_processor);
546
547        for _ in 0..3 {
548            let mut clone = canonical.clone();
549            let _ = clone
550                .ready()
551                .await
552                .unwrap()
553                .call(Exchange::new(Message::new("req")))
554                .await
555                .unwrap();
556        }
557
558        let _ = tokio::time::timeout(std::time::Duration::from_secs(2), async {
559            while completed.load(Ordering::SeqCst) < 3 {
560                tokio::time::sleep(std::time::Duration::from_millis(20)).await;
561            }
562        })
563        .await;
564        assert_eq!(
565            completed.load(Ordering::SeqCst),
566            3,
567            "all tap tasks must complete despite per-request clone drops"
568        );
569    }
570
571    #[test]
572    fn test_wiretap_config_default_is_bounded_20() {
573        let cfg = WireTapConfig::default();
574        assert_eq!(cfg.max_concurrent, Some(20));
575        assert_eq!(cfg.shutdown_grace, std::time::Duration::from_secs(5));
576    }
577
578    #[test]
579    fn test_wiretap_config_bounded_zero_panics() {
580        let result = std::panic::catch_unwind(|| WireTapConfig::bounded(0));
581        assert!(result.is_err());
582        if let Err(payload) = result {
583            let msg = payload
584                .downcast_ref::<&str>()
585                .expect("panic payload should be &str");
586            assert!(
587                msg.contains("max_concurrent"),
588                "panic message should contain 'max_concurrent', got: {msg}"
589            );
590        }
591    }
592
593    #[test]
594    fn test_wiretap_config_validate_rejects_zero_bound() {
595        let cfg = WireTapConfig {
596            max_concurrent: Some(0),
597            shutdown_grace: std::time::Duration::from_secs(5),
598        };
599        let result = std::panic::catch_unwind(|| cfg.validate());
600        assert!(result.is_err());
601        let payload = result.unwrap_err();
602        let msg = payload
603            .downcast_ref::<&str>()
604            .expect("panic payload should be &str");
605        assert!(
606            msg.contains("max_concurrent"),
607            "panic message should contain 'max_concurrent', got: {msg}"
608        );
609    }
610
611    #[tokio::test]
612    async fn test_wire_tap_drop_aborts_spawned_tasks() {
613        // Under the new shared-state model, dropping the canonical service
614        // drops the last `Arc<WireTapShared>` ref, firing `WireTapShared::drop`
615        // which cancels the token. The spawned tap's `run_tap` selects on
616        // `cancel.cancelled()` and returns promptly, so the 10s sleep is
617        // aborted and `task_completed` stays false.
618        let task_started = Arc::new(AtomicBool::new(false));
619        let task_completed = Arc::new(AtomicBool::new(false));
620        let started_clone = task_started.clone();
621        let completed_clone = task_completed.clone();
622
623        let tap_processor = BoxProcessor::from_fn(move |_ex| {
624            let started = started_clone.clone();
625            let completed = completed_clone.clone();
626            Box::pin(async move {
627                started.store(true, Ordering::SeqCst);
628                tokio::time::sleep(std::time::Duration::from_secs(10)).await;
629                completed.store(true, Ordering::SeqCst);
630                Ok(Exchange::default())
631            })
632        });
633
634        let mut service = WireTapService::new(tap_processor);
635        let _ = service
636            .ready()
637            .await
638            .unwrap()
639            .call(Exchange::default())
640            .await
641            .unwrap();
642
643        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
644        assert!(
645            task_started.load(Ordering::SeqCst),
646            "tap task should be running"
647        );
648        assert!(
649            !task_completed.load(Ordering::SeqCst),
650            "task should not have completed yet"
651        );
652
653        drop(service);
654
655        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
656
657        assert!(
658            !task_completed.load(Ordering::SeqCst),
659            "task should have been aborted, not completed"
660        );
661    }
662
663    // --- New tests for the bounded-admission + TaskTracker + cancellation model ---
664
665    #[tokio::test]
666    async fn test_wiretap_bounded_detached_count_never_exceeds_bound() {
667        // Detached tracked count must stay `<= bound`. The inline CallerRuns tap
668        // is NOT tracked so it cannot be observed here (transient total
669        // execution may briefly reach `bound + 1` — that invariant is exercised
670        // by `test_wiretap_bounded_concurrency` above).
671        let tap_processor = BoxProcessor::from_fn(|_ex| {
672            Box::pin(async move {
673                tokio::time::sleep(std::time::Duration::from_millis(50)).await;
674                Ok(Exchange::default())
675            })
676        });
677
678        let canonical = WireTapService::with_config(tap_processor, WireTapConfig::bounded(2));
679        let max_seen = Arc::new(AtomicUsize::new(0));
680        let stop = Arc::new(AtomicBool::new(false));
681
682        // Background sampler: continuously polls in_flight_count() via a shared
683        // clone and tracks the peak observed detached task count.
684        let sampler_svc = canonical.clone();
685        let sampler_max = Arc::clone(&max_seen);
686        let sampler_stop = Arc::clone(&stop);
687        let sampler = tokio::spawn(async move {
688            while !sampler_stop.load(Ordering::SeqCst) {
689                let n = sampler_svc.in_flight_count();
690                sampler_max.fetch_max(n, Ordering::SeqCst);
691                tokio::task::yield_now().await;
692            }
693        });
694
695        // Fire 5 call() futures from spawned callers under contention.
696        let mut callers = Vec::new();
697        for _ in 0..5 {
698            let mut caller_svc = canonical.clone();
699            callers.push(tokio::spawn(async move {
700                let _ = caller_svc
701                    .ready()
702                    .await
703                    .unwrap()
704                    .call(Exchange::new(Message::new("x")))
705                    .await;
706            }));
707        }
708        for h in callers {
709            let _ = h.await;
710        }
711
712        stop.store(true, Ordering::SeqCst);
713        let _ = sampler.await;
714
715        let observed = max_seen.load(Ordering::SeqCst);
716        assert!(
717            observed <= 2,
718            "detached tracked task count peaked at {observed}, expected <= bound (=2)"
719        );
720    }
721
722    #[tokio::test]
723    async fn test_wiretap_caller_backpressured_when_saturated() {
724        // CallerRuns: when bound is saturated, the next call's tap runs INLINE
725        // on the caller's future. The caller is back-pressured until the inline
726        // tap finishes. The leaky spawn-then-acquire version would resolve the
727        // call immediately regardless of the tap's progress.
728        use tokio::sync::Notify;
729
730        let notify = Arc::new(Notify::new());
731        let tap_notify = Arc::clone(&notify);
732        let tap_processor = BoxProcessor::from_fn(move |_ex| {
733            let n = Arc::clone(&tap_notify);
734            Box::pin(async move {
735                n.notified().await;
736                Ok(Exchange::default())
737            })
738        });
739
740        let mut svc = WireTapService::with_config(tap_processor, WireTapConfig::bounded(1));
741
742        // First call: acquires the sole permit, spawns detached tap awaiting Notify.
743        let _ = svc
744            .ready()
745            .await
746            .unwrap()
747            .call(Exchange::default())
748            .await
749            .unwrap();
750        // Yield to let the spawned tap actually register its `notified()` waiter.
751        tokio::time::sleep(std::time::Duration::from_millis(20)).await;
752
753        // Second call: try_acquire fails (NoPermits). CallerRuns path runs the
754        // tap inline on this future, which awaits Notify.
755        let mut svc2 = svc.clone();
756        let mut fut2 = Box::pin(
757            svc2.ready()
758                .await
759                .unwrap()
760                .call(Exchange::new(Message::new("inline"))),
761        );
762
763        // Race fut2 against a 50ms sleep: fut2 should still be Pending (it is
764        // running the tap inline, awaiting Notify).
765        let pending_after_50ms = tokio::select! {
766            r = &mut fut2 => {
767                panic!(
768                    "fut2 should be Pending after 50ms under CallerRuns back-pressure; resolved early: {:?}",
769                    r.is_ok()
770                );
771            }
772            _ = tokio::time::sleep(std::time::Duration::from_millis(50)) => true,
773        };
774        assert!(
775            pending_after_50ms,
776            "fut2 should be Pending (inline tap awaiting Notify) after 50ms under CallerRuns back-pressure"
777        );
778
779        // Release all waiters: notify_waiters wakes both the detached tap 1 and
780        // the inline tap (fut2). fut2 resolves Ok.
781        notify.notify_waiters();
782        let result = fut2.await;
783        assert!(
784            result.is_ok(),
785            "fut2 should resolve Ok after notify_waiters"
786        );
787    }
788
789    #[tokio::test]
790    async fn test_wiretap_unbounded_none_path_detaches_without_permit() {
791        // Explicit unbounded (max_concurrent = None) path: tasks detach into the
792        // tracker without acquiring a permit. The tracker count rises above 0
793        // then drains to 0 as tasks complete.
794        let tap_processor = BoxProcessor::from_fn(|_ex| {
795            Box::pin(async move {
796                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
797                Ok(Exchange::default())
798            })
799        });
800
801        let mut svc = WireTapService::with_config(
802            tap_processor,
803            WireTapConfig {
804                max_concurrent: None,
805                shutdown_grace: std::time::Duration::from_secs(5),
806            },
807        );
808
809        let sampler_svc = svc.clone();
810        let peak = Arc::new(AtomicUsize::new(0));
811        let peak_clone = Arc::clone(&peak);
812        let done = Arc::new(AtomicBool::new(false));
813        let done_clone = Arc::clone(&done);
814        let sampler = tokio::spawn(async move {
815            while !done_clone.load(Ordering::SeqCst) {
816                let n = sampler_svc.in_flight_count();
817                peak_clone.fetch_max(n, Ordering::SeqCst);
818                tokio::task::yield_now().await;
819            }
820        });
821
822        for _ in 0..50 {
823            let ex = Exchange::new(Message::new("x"));
824            let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
825        }
826
827        // Poll until the tracker drains to 0 within 2s.
828        let drained = tokio::time::timeout(std::time::Duration::from_secs(2), async {
829            loop {
830                if svc.in_flight_count() == 0 {
831                    return;
832                }
833                tokio::time::sleep(std::time::Duration::from_millis(5)).await;
834            }
835        })
836        .await
837        .is_ok();
838        done.store(true, Ordering::SeqCst);
839        let _ = sampler.await;
840
841        assert!(drained, "unbounded path tasks should drain to 0 within 2s");
842        assert!(
843            peak.load(Ordering::SeqCst) > 0,
844            "unbounded path should have observed tracked tasks (peak > 0)"
845        );
846    }
847
848    #[tokio::test]
849    async fn test_wiretap_no_unbounded_task_growth_across_bursts() {
850        // Regression for the leaky spawn-then-acquire model: completed tasks
851        // MUST decrement the tracker's len() so subsequent bursts do not
852        // accumulate.
853        let tap_processor =
854            BoxProcessor::from_fn(|_ex| Box::pin(async move { Ok(Exchange::default()) }));
855
856        let svc = WireTapService::with_config(tap_processor, WireTapConfig::default());
857
858        let drain_to_zero = |svc: &WireTapService| {
859            let s = svc.clone();
860            async move {
861                tokio::time::timeout(std::time::Duration::from_secs(2), async {
862                    loop {
863                        if s.in_flight_count() == 0 {
864                            return;
865                        }
866                        tokio::time::sleep(std::time::Duration::from_millis(2)).await;
867                    }
868                })
869                .await
870                .is_ok()
871            }
872        };
873
874        // Burst 1.
875        let mut callers = Vec::new();
876        for _ in 0..1000 {
877            let mut s = svc.clone();
878            callers.push(tokio::spawn(async move {
879                let _ = s.ready().await.unwrap().call(Exchange::default()).await;
880            }));
881        }
882        for h in callers {
883            let _ = h.await;
884        }
885        assert!(
886            drain_to_zero(&svc).await,
887            "burst 1 must drain to in_flight_count == 0 within 2s"
888        );
889
890        // Burst 2.
891        let mut callers = Vec::new();
892        for _ in 0..1000 {
893            let mut s = svc.clone();
894            callers.push(tokio::spawn(async move {
895                let _ = s.ready().await.unwrap().call(Exchange::default()).await;
896            }));
897        }
898        for h in callers {
899            let _ = h.await;
900        }
901        assert!(
902            drain_to_zero(&svc).await,
903            "burst 2 must drain to in_flight_count == 0 within 2s (no accumulation across bursts)"
904        );
905    }
906
907    // --- Tracing capture helper for warn-log assertions ---
908
909    /// `MakeWriter` that appends formatted events to a shared `Vec<u8>` sink.
910    /// Used by the warn-suppression tests to assert that a `warn!` record was
911    /// emitted. The sink collects the ANSI-stripped fmt layer output.
912    #[derive(Clone)]
913    struct CapturingWriter {
914        sink: Arc<Mutex<Vec<u8>>>,
915    }
916
917    impl std::io::Write for CapturingWriter {
918        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
919            self.sink.lock().unwrap().extend_from_slice(buf); // allow-unwrap: test-only
920            Ok(buf.len())
921        }
922        fn flush(&mut self) -> std::io::Result<()> {
923            Ok(())
924        }
925    }
926
927    impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for CapturingWriter {
928        type Writer = CapturingWriter;
929        fn make_writer(&'a self) -> Self::Writer {
930            self.clone()
931        }
932    }
933
934    fn capture_sink() -> (Arc<Mutex<Vec<u8>>>, impl tracing::Subscriber) {
935        let sink: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
936        let writer = CapturingWriter {
937            sink: Arc::clone(&sink),
938        };
939        let subscriber = tracing_subscriber::fmt()
940            .with_writer(writer)
941            .with_ansi(false)
942            .finish();
943        (sink, subscriber)
944    }
945
946    /// Custom Service whose `poll_ready` always returns `Err`. Used to exercise
947    /// the tap-readiness-error suppression path.
948    #[derive(Clone)]
949    struct ReadyFailingSvc {
950        err_msg: &'static str,
951    }
952
953    impl Service<Exchange> for ReadyFailingSvc {
954        type Response = Exchange;
955        type Error = CamelError;
956        type Future = Pin<Box<dyn Future<Output = Result<Exchange, CamelError>> + Send>>;
957        fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
958            Poll::Ready(Err(CamelError::ProcessorError(self.err_msg.into())))
959        }
960        fn call(&mut self, ex: Exchange) -> Self::Future {
961            Box::pin(async move { Ok(ex) })
962        }
963    }
964
965    #[tokio::test]
966    async fn test_wiretap_tap_readiness_error_suppressed_with_log() {
967        let tap: camel_api::BoxProcessor = tower::util::BoxCloneService::new(ReadyFailingSvc {
968            err_msg: "ready-boom",
969        });
970        let mut svc = WireTapService::new(tap);
971
972        let (sink, subscriber) = capture_sink();
973        let exchange = Exchange::new(Message::new("main"));
974
975        // set_default propagates to tasks spawned via tokio within this scope.
976        let _guard = tracing::subscriber::set_default(subscriber);
977        // Parallel tests race tracing's per-callsite interest cache against
978        // this thread-local subscriber; force a rebuild so the callsites
979        // below re-evaluate against it (bd rc-u9hs).
980        tracing::callsite::rebuild_interest_cache();
981        let result = svc.ready().await.unwrap().call(exchange).await;
982
983        assert!(result.is_ok(), "tap readiness error must be suppressed");
984        assert_eq!(result.unwrap().input.body.as_text(), Some("main"));
985
986        // Give the spawned tap task time to run ready() and log warn.
987        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
988        drop(_guard);
989
990        let captured = String::from_utf8(sink.lock().unwrap().clone()).unwrap(); // allow-unwrap: test-only
991        assert!(
992            captured.contains("ready-boom"),
993            "a warn! record mentioning the readiness error should have been emitted; got: {captured}"
994        );
995    }
996
997    #[tokio::test]
998    async fn test_wiretap_tap_processing_error_suppressed_with_log() {
999        let tap_processor = BoxProcessor::from_fn(|_ex| {
1000            Box::pin(async move { Err(CamelError::ProcessorError("call-boom".into())) })
1001        });
1002        let mut svc = WireTapService::new(tap_processor);
1003
1004        let (sink, subscriber) = capture_sink();
1005        let exchange = Exchange::new(Message::new("main"));
1006
1007        let _guard = tracing::subscriber::set_default(subscriber);
1008        // Parallel tests race tracing's per-callsite interest cache against
1009        // this thread-local subscriber; force a rebuild so the callsites
1010        // below re-evaluate against it (bd rc-u9hs).
1011        tracing::callsite::rebuild_interest_cache();
1012        let result = svc.ready().await.unwrap().call(exchange).await;
1013
1014        assert!(result.is_ok(), "tap processing error must be suppressed");
1015        assert_eq!(result.unwrap().input.body.as_text(), Some("main"));
1016
1017        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
1018        drop(_guard);
1019
1020        let captured = String::from_utf8(sink.lock().unwrap().clone()).unwrap(); // allow-unwrap: test-only
1021        assert!(
1022            captured.contains("call-boom"),
1023            "a warn! record mentioning the processing error should have been emitted; got: {captured}"
1024        );
1025    }
1026
1027    #[tokio::test]
1028    async fn test_wiretap_poll_ready_always_ready() {
1029        // poll_ready returns Ready(Ok(())) unconditionally (ADR-0019), even
1030        // when the tap endpoint's own readiness would fail.
1031        let tap: camel_api::BoxProcessor = tower::util::BoxCloneService::new(ReadyFailingSvc {
1032            err_msg: "would-fail",
1033        });
1034        let mut svc = WireTapService::new(tap);
1035
1036        let waker = futures::task::noop_waker();
1037        let mut cx = Context::from_waker(&waker);
1038        let poll = svc.poll_ready(&mut cx);
1039        assert!(
1040            matches!(poll, Poll::Ready(Ok(()))),
1041            "poll_ready must be Ready(Ok(())) unconditionally (ADR-0019), got opposite"
1042        );
1043    }
1044
1045    // --- WireTapLifecycle + StepLifecycle shutdown tests (Task 4) ---
1046
1047    #[tokio::test]
1048    async fn test_wiretap_shutdown_drains_fast_aborts_slow() {
1049        let fast_done = Arc::new(AtomicBool::new(false));
1050        let slow_done = Arc::new(AtomicBool::new(false));
1051        let call_idx = Arc::new(AtomicUsize::new(0));
1052
1053        let fd = fast_done.clone();
1054        let sd = slow_done.clone();
1055        let ci = call_idx.clone();
1056        let tap_processor = BoxProcessor::from_fn(move |ex| {
1057            let fd = fd.clone();
1058            let sd = sd.clone();
1059            let ci = ci.clone();
1060            Box::pin(async move {
1061                let n = ci.fetch_add(1, Ordering::SeqCst);
1062                if n == 0 {
1063                    tokio::time::sleep(Duration::from_millis(10)).await;
1064                    fd.store(true, Ordering::SeqCst);
1065                } else {
1066                    tokio::time::sleep(std::time::Duration::from_secs(10)).await;
1067                    sd.store(true, Ordering::SeqCst);
1068                }
1069                Ok(ex)
1070            })
1071        });
1072
1073        let config = WireTapConfig {
1074            max_concurrent: Some(20),
1075            shutdown_grace: Duration::from_millis(200),
1076        };
1077        let mut svc = WireTapService::with_config(tap_processor, config);
1078
1079        let _ = svc
1080            .ready()
1081            .await
1082            .unwrap()
1083            .call(Exchange::new(Message::new("fast")))
1084            .await
1085            .unwrap();
1086        let _ = svc
1087            .ready()
1088            .await
1089            .unwrap()
1090            .call(Exchange::new(Message::new("slow")))
1091            .await
1092            .unwrap();
1093
1094        tokio::time::sleep(Duration::from_millis(20)).await;
1095
1096        let lifecycle = svc.lifecycle();
1097        let start = tokio::time::Instant::now();
1098        lifecycle
1099            .shutdown(StepShutdownReason::RouteStop)
1100            .await
1101            .unwrap();
1102        let elapsed = start.elapsed();
1103
1104        assert!(
1105            fast_done.load(Ordering::SeqCst),
1106            "fast tap should drain before grace expires"
1107        );
1108        assert!(
1109            !slow_done.load(Ordering::SeqCst),
1110            "slow tap should be aborted after grace, not complete"
1111        );
1112        assert!(
1113            elapsed < Duration::from_millis(500),
1114            "shutdown took {:?}, expected < 500ms",
1115            elapsed
1116        );
1117    }
1118
1119    #[tokio::test]
1120    async fn test_wiretap_shutdown_idempotent() {
1121        let slow_done = Arc::new(AtomicBool::new(false));
1122        let sd = slow_done.clone();
1123        let tap_processor = BoxProcessor::from_fn(move |ex| {
1124            let sd = sd.clone();
1125            Box::pin(async move {
1126                tokio::time::sleep(std::time::Duration::from_secs(10)).await;
1127                sd.store(true, Ordering::SeqCst);
1128                Ok(ex)
1129            })
1130        });
1131
1132        let config = WireTapConfig {
1133            max_concurrent: Some(20),
1134            shutdown_grace: Duration::from_millis(50),
1135        };
1136        let mut svc = WireTapService::with_config(tap_processor, config);
1137
1138        let _ = svc
1139            .ready()
1140            .await
1141            .unwrap()
1142            .call(Exchange::new(Message::new("slow")))
1143            .await
1144            .unwrap();
1145
1146        tokio::time::sleep(Duration::from_millis(20)).await;
1147
1148        let lifecycle = svc.lifecycle();
1149        lifecycle
1150            .shutdown(StepShutdownReason::RouteStop)
1151            .await
1152            .unwrap();
1153
1154        let start = tokio::time::Instant::now();
1155        let result = lifecycle.shutdown(StepShutdownReason::HotSwap).await;
1156        let elapsed = start.elapsed();
1157
1158        assert!(result.is_ok(), "second shutdown must return Ok");
1159        assert!(
1160            elapsed < Duration::from_millis(100),
1161            "second shutdown must return promptly, took {:?}",
1162            elapsed
1163        );
1164        assert!(
1165            !slow_done.load(Ordering::SeqCst),
1166            "slow tap must be aborted, not completed"
1167        );
1168    }
1169
1170    #[tokio::test]
1171    async fn test_wiretap_calls_after_close_rejected() {
1172        let tap_invoked = Arc::new(AtomicBool::new(false));
1173        let ti = tap_invoked.clone();
1174        let tap_processor = BoxProcessor::from_fn(move |ex| {
1175            let ti = ti.clone();
1176            Box::pin(async move {
1177                ti.store(true, Ordering::SeqCst);
1178                Ok(ex)
1179            })
1180        });
1181
1182        let mut svc = WireTapService::new(tap_processor);
1183        let lifecycle = svc.lifecycle();
1184        lifecycle
1185            .shutdown(StepShutdownReason::RouteStop)
1186            .await
1187            .unwrap();
1188
1189        let result = svc
1190            .ready()
1191            .await
1192            .unwrap()
1193            .call(Exchange::new(Message::new("post-close")))
1194            .await;
1195
1196        assert!(
1197            result.is_ok(),
1198            "call after close must return Ok(original exchange)"
1199        );
1200        assert!(
1201            !tap_invoked.load(Ordering::SeqCst),
1202            "tap must not be invoked after admission closed"
1203        );
1204    }
1205
1206    #[tokio::test]
1207    async fn test_wiretap_cancellation_while_pending_readiness() {
1208        // Service whose poll_ready returns Pending indefinitely, so the
1209        // spawned task blocks in run_tap's ready() phase. Shutdown cancels
1210        // the token, the biased select! picks it up, and the task exits
1211        // cleanly without reaching call().
1212        #[derive(Clone)]
1213        struct ForeverPendingSvc {
1214            called: Arc<AtomicBool>,
1215        }
1216
1217        impl Service<Exchange> for ForeverPendingSvc {
1218            type Response = Exchange;
1219            type Error = CamelError;
1220            type Future = Pin<Box<dyn Future<Output = Result<Exchange, CamelError>> + Send>>;
1221
1222            fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
1223                Poll::Pending
1224            }
1225
1226            fn call(&mut self, ex: Exchange) -> Self::Future {
1227                self.called.store(true, Ordering::SeqCst);
1228                Box::pin(async move { Ok(ex) })
1229            }
1230        }
1231
1232        let called = Arc::new(AtomicBool::new(false));
1233        let tap: camel_api::BoxProcessor = tower::util::BoxCloneService::new(ForeverPendingSvc {
1234            called: called.clone(),
1235        });
1236        let mut svc = WireTapService::new(tap);
1237
1238        let _ = svc
1239            .ready()
1240            .await
1241            .unwrap()
1242            .call(Exchange::new(Message::new("hanging")))
1243            .await
1244            .unwrap();
1245
1246        tokio::time::sleep(Duration::from_millis(20)).await;
1247
1248        let lifecycle = svc.lifecycle();
1249        let result = lifecycle.shutdown(StepShutdownReason::RouteStop).await;
1250
1251        assert!(
1252            result.is_ok(),
1253            "shutdown must succeed even with pending readiness: {:?}",
1254            result
1255        );
1256        assert!(
1257            !called.load(Ordering::SeqCst),
1258            "tap call() must never be reached — cancelled during readiness phase"
1259        );
1260    }
1261
1262    #[tokio::test]
1263    async fn test_wiretap_zero_grace_immediate_cancel() {
1264        let slow_done = Arc::new(AtomicBool::new(false));
1265        let sd = slow_done.clone();
1266        let tap_processor = BoxProcessor::from_fn(move |ex| {
1267            let sd = sd.clone();
1268            Box::pin(async move {
1269                tokio::time::sleep(std::time::Duration::from_secs(10)).await;
1270                sd.store(true, Ordering::SeqCst);
1271                Ok(ex)
1272            })
1273        });
1274
1275        let config = WireTapConfig {
1276            max_concurrent: Some(20),
1277            shutdown_grace: Duration::ZERO,
1278        };
1279        let mut svc = WireTapService::with_config(tap_processor, config);
1280
1281        let _ = svc
1282            .ready()
1283            .await
1284            .unwrap()
1285            .call(Exchange::new(Message::new("slow")))
1286            .await
1287            .unwrap();
1288
1289        tokio::time::sleep(Duration::from_millis(20)).await;
1290
1291        let lifecycle = svc.lifecycle();
1292        let start = tokio::time::Instant::now();
1293        lifecycle
1294            .shutdown(StepShutdownReason::RouteStop)
1295            .await
1296            .unwrap();
1297        let elapsed = start.elapsed();
1298
1299        assert!(
1300            !slow_done.load(Ordering::SeqCst),
1301            "slow tap must be aborted immediately (zero grace)"
1302        );
1303        assert!(
1304            elapsed < Duration::from_millis(200),
1305            "zero-grace shutdown must return quickly, took {:?}",
1306            elapsed
1307        );
1308    }
1309
1310    #[tokio::test(flavor = "multi_thread")]
1311    async fn test_wiretap_admission_shutdown_no_orphan_task() {
1312        // Stress test: fire concurrent call()s and shutdown() across many
1313        // randomized iterations. Verify in_flight_count() == 0 after shutdown
1314        // completes — no orphan task escaped the tracker.
1315        const ITERATIONS: usize = 200;
1316
1317        for _ in 0..ITERATIONS {
1318            let tap_processor = BoxProcessor::from_fn(|_ex| {
1319                Box::pin(async move {
1320                    tokio::time::sleep(Duration::from_millis(1)).await;
1321                    Ok(Exchange::default())
1322                })
1323            });
1324
1325            let svc = WireTapService::new(tap_processor);
1326            let lifecycle = svc.lifecycle();
1327
1328            // Fire several concurrent callers.
1329            let mut handles = Vec::new();
1330            for _ in 0..4 {
1331                let mut c = svc.clone();
1332                handles.push(tokio::spawn(async move {
1333                    let _ = c.ready().await.unwrap().call(Exchange::default()).await;
1334                }));
1335            }
1336
1337            // Yield to let spawns register in the tracker.
1338            tokio::task::yield_now().await;
1339            tokio::time::sleep(Duration::from_millis(1)).await;
1340
1341            // Shutdown concurrently with callers still in-flight.
1342            lifecycle
1343                .shutdown(StepShutdownReason::RouteStop)
1344                .await
1345                .unwrap();
1346
1347            for h in handles {
1348                let _ = h.await;
1349            }
1350
1351            // Poll until tracker drains (already waited in shutdown, but
1352            // defensive check).
1353            let drained = tokio::time::timeout(Duration::from_secs(2), async {
1354                loop {
1355                    if svc.in_flight_count() == 0 {
1356                        return;
1357                    }
1358                    tokio::time::sleep(Duration::from_millis(5)).await;
1359                }
1360            })
1361            .await
1362            .is_ok();
1363
1364            assert!(
1365                drained,
1366                "iteration: in_flight_count must drain to 0 after shutdown"
1367            );
1368        }
1369    }
1370}