use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::Duration;
use async_trait::async_trait;
use tokio::sync::{Semaphore, TryAcquireError};
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use tower::{Service, ServiceExt};
use camel_api::{CamelError, Exchange, StepLifecycle, StepShutdownReason};
#[derive(Clone)]
pub struct WireTapConfig {
pub max_concurrent: Option<usize>,
pub shutdown_grace: std::time::Duration,
}
impl Default for WireTapConfig {
fn default() -> Self {
Self {
max_concurrent: Some(20),
shutdown_grace: std::time::Duration::from_secs(5),
}
}
}
impl WireTapConfig {
pub fn validate(&self) {
if self.max_concurrent == Some(0) {
panic!("max_concurrent must be > 0 when set");
}
}
pub fn bounded(max_concurrent: usize) -> Self {
assert!(max_concurrent > 0, "max_concurrent must be > 0");
Self {
max_concurrent: Some(max_concurrent),
shutdown_grace: std::time::Duration::from_secs(5),
}
}
}
#[derive(Debug)]
struct WireTapSharedInner {
open: bool,
tracker: TaskTracker,
cancel: CancellationToken,
semaphore: Option<Arc<Semaphore>>,
shutdown_grace: Duration,
}
#[derive(Debug)]
struct WireTapShared {
inner: Mutex<WireTapSharedInner>,
}
impl Drop for WireTapShared {
fn drop(&mut self) {
self.inner
.lock()
.expect("WireTapShared mutex poisoned") .cancel
.cancel();
}
}
pub struct WireTapService {
tap_endpoint: camel_api::BoxProcessor,
shared: Arc<WireTapShared>,
}
impl Clone for WireTapService {
fn clone(&self) -> Self {
Self {
tap_endpoint: self.tap_endpoint.clone(),
shared: Arc::clone(&self.shared),
}
}
}
impl WireTapService {
pub fn new(tap_endpoint: camel_api::BoxProcessor) -> Self {
Self::with_config(tap_endpoint, WireTapConfig::default())
}
pub fn with_config(tap_endpoint: camel_api::BoxProcessor, config: WireTapConfig) -> Self {
config.validate();
let semaphore = config
.max_concurrent
.map(|limit| Arc::new(Semaphore::new(limit)));
let shared = Arc::new(WireTapShared {
inner: Mutex::new(WireTapSharedInner {
open: true,
tracker: TaskTracker::new(),
cancel: CancellationToken::new(),
semaphore,
shutdown_grace: config.shutdown_grace,
}),
});
Self {
tap_endpoint,
shared,
}
}
#[cfg(test)]
pub(crate) fn in_flight_count(&self) -> usize {
self.shared
.inner
.lock()
.expect("WireTapShared mutex poisoned") .tracker
.len()
}
}
#[derive(Debug)]
pub struct WireTapLifecycle {
shared: Arc<WireTapShared>,
shutdown_called: AtomicBool,
}
#[async_trait]
impl StepLifecycle for WireTapLifecycle {
fn name(&self) -> &'static str {
"wiretap"
}
async fn shutdown(&self, _reason: StepShutdownReason) -> Result<(), CamelError> {
if self.shutdown_called.swap(true, Ordering::SeqCst) {
return Ok(());
}
let (tracker, cancel, grace) = {
let mut guard = self
.shared
.inner
.lock()
.expect("WireTapShared mutex poisoned"); guard.open = false;
guard.tracker.close();
(
guard.tracker.clone(),
guard.cancel.clone(),
guard.shutdown_grace,
)
};
if !grace.is_zero() {
let _ = tokio::time::timeout(grace, tracker.wait()).await;
}
cancel.cancel();
let _ = tracker.wait().await;
Ok(())
}
}
impl WireTapService {
pub fn lifecycle(&self) -> Arc<dyn StepLifecycle> {
Arc::new(WireTapLifecycle {
shared: Arc::clone(&self.shared),
shutdown_called: AtomicBool::new(false),
})
}
}
impl Service<Exchange> for WireTapService {
type Response = Exchange;
type Error = CamelError;
type Future = Pin<Box<dyn Future<Output = Result<Exchange, CamelError>> + Send>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, exchange: Exchange) -> Self::Future {
let tap_endpoint = self.tap_endpoint.clone();
let tap_exchange = exchange.clone();
let inner = self
.shared
.inner
.lock()
.expect("WireTapShared mutex poisoned"); if !inner.open {
tracing::warn!("WireTap admission closed, dropping tap");
drop(inner);
return Box::pin(async move { Ok(exchange) });
}
match &inner.semaphore {
Some(sem) => match Arc::clone(sem).try_acquire_owned() {
Ok(permit) => {
let cancel = inner.cancel.clone();
inner.tracker.spawn(async move {
let _permit = permit;
run_tap(tap_endpoint, tap_exchange, cancel).await;
});
drop(inner);
Box::pin(async move { Ok(exchange) })
}
Err(TryAcquireError::NoPermits) => {
let cancel = inner.cancel.clone();
drop(inner);
Box::pin(async move {
run_tap(tap_endpoint, tap_exchange, cancel).await;
Ok(exchange)
})
}
Err(TryAcquireError::Closed) => {
tracing::warn!("WireTap semaphore closed, dropping tap");
drop(inner);
Box::pin(async move { Ok(exchange) })
}
},
None => {
let cancel = inner.cancel.clone();
inner.tracker.spawn(async move {
run_tap(tap_endpoint, tap_exchange, cancel).await;
});
drop(inner);
Box::pin(async move { Ok(exchange) })
}
}
}
}
async fn run_tap(
mut tap_endpoint: camel_api::BoxProcessor,
tap_exchange: Exchange,
cancel: CancellationToken,
) {
{
let ready_fut = tap_endpoint.ready();
tokio::pin!(ready_fut);
let ready_result = tokio::select! {
biased;
_ = cancel.cancelled() => { return; }
r = &mut ready_fut => r,
};
if let Err(e) = ready_result {
tracing::warn!("WireTap endpoint poll_ready failed: {}", e);
return;
}
}
{
let call_fut = tap_endpoint.call(tap_exchange);
tokio::pin!(call_fut);
let call_result = tokio::select! {
biased;
_ = cancel.cancelled() => { return; }
r = &mut call_fut => r,
};
if let Err(e) = call_result {
tracing::warn!("WireTap processing error: {}", e);
}
}
}
pub struct WireTapLayer {
tap_endpoint: camel_api::BoxProcessor,
config: WireTapConfig,
}
impl WireTapLayer {
pub fn new(tap_endpoint: camel_api::BoxProcessor) -> Self {
Self {
tap_endpoint,
config: WireTapConfig::default(),
}
}
pub fn bounded(tap_endpoint: camel_api::BoxProcessor, max_concurrent: usize) -> Self {
Self {
tap_endpoint,
config: WireTapConfig::bounded(max_concurrent),
}
}
}
impl<S> tower::Layer<S> for WireTapLayer {
type Service = WireTapService;
fn layer(&self, _inner: S) -> Self::Service {
WireTapService::with_config(self.tap_endpoint.clone(), self.config.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
use camel_api::{BoxProcessor, BoxProcessorExt, Message};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use tower::ServiceExt;
#[tokio::test]
async fn test_wire_tap_returns_original_immediately() {
let tap_processor = BoxProcessor::from_fn(|ex| Box::pin(async move { Ok(ex) }));
let mut wire_tap = WireTapService::new(tap_processor);
let exchange = Exchange::new(Message::new("test message"));
let result = wire_tap
.ready()
.await
.unwrap()
.call(exchange)
.await
.unwrap();
assert_eq!(result.input.body.as_text(), Some("test message"));
}
#[tokio::test]
async fn test_wire_tap_endpoint_receives_clone() {
let received_count = Arc::new(AtomicUsize::new(0));
let count_clone = received_count.clone();
let tap_processor = BoxProcessor::from_fn(move |ex| {
let count = count_clone.clone();
Box::pin(async move {
count.fetch_add(1, Ordering::SeqCst);
Ok(ex)
})
});
let mut wire_tap = WireTapService::new(tap_processor);
let exchange = Exchange::new(Message::new("test"));
let _result = wire_tap
.ready()
.await
.unwrap()
.call(exchange)
.await
.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
assert_eq!(received_count.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_wire_tap_isolates_errors() {
let tap_processor = BoxProcessor::from_fn(|_ex| {
Box::pin(async move { Err(CamelError::ProcessorError("tap error".into())) })
});
let mut wire_tap = WireTapService::new(tap_processor);
let exchange = Exchange::new(Message::new("test"));
let result = wire_tap.ready().await.unwrap().call(exchange).await;
assert!(result.is_ok());
assert_eq!(result.unwrap().input.body.as_text(), Some("test"));
}
#[tokio::test]
async fn test_wire_tap_layer() {
use tower::Layer;
let tap_processor = BoxProcessor::from_fn(|ex| Box::pin(async move { Ok(ex) }));
let layer = super::WireTapLayer::new(tap_processor);
let inner = camel_api::IdentityProcessor;
let mut svc = layer.layer(inner);
let exchange = Exchange::new(Message::new("test"));
let result = svc.ready().await.unwrap().call(exchange).await.unwrap();
assert_eq!(result.input.body.as_text(), Some("test"));
}
#[tokio::test]
async fn test_wiretap_bounded_concurrency() {
let concurrent = Arc::new(AtomicUsize::new(0));
let max_concurrent = Arc::new(AtomicUsize::new(0));
let c = Arc::clone(&concurrent);
let mc = Arc::clone(&max_concurrent);
let tap_processor = BoxProcessor::from_fn(move |ex| {
let c = Arc::clone(&c);
let mc = Arc::clone(&mc);
Box::pin(async move {
let current = c.fetch_add(1, Ordering::SeqCst) + 1;
mc.fetch_max(current, Ordering::SeqCst);
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
c.fetch_sub(1, Ordering::SeqCst);
Ok(ex)
})
});
let config = super::WireTapConfig::bounded(2);
let mut svc = super::WireTapService::with_config(tap_processor, config);
for _ in 0..3 {
let ex = Exchange::new(Message::new("test"));
let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
}
tokio::time::sleep(std::time::Duration::from_millis(300)).await;
let observed_max = max_concurrent.load(Ordering::SeqCst);
assert!(
observed_max <= 3,
"max concurrency was {observed_max}, expected <= bound+1 (=3) under CallerRuns"
);
}
#[tokio::test]
async fn test_wire_tap_survives_per_request_clone_drop() {
let completed = Arc::new(AtomicUsize::new(0));
let completed_clone = completed.clone();
let tap_processor = BoxProcessor::from_fn(move |ex| {
let c = completed_clone.clone();
Box::pin(async move {
tokio::time::sleep(std::time::Duration::from_millis(150)).await;
c.fetch_add(1, Ordering::SeqCst);
Ok(ex)
})
});
let canonical = WireTapService::new(tap_processor);
for _ in 0..3 {
let mut clone = canonical.clone();
let _ = clone
.ready()
.await
.unwrap()
.call(Exchange::new(Message::new("req")))
.await
.unwrap();
}
let _ = tokio::time::timeout(std::time::Duration::from_secs(2), async {
while completed.load(Ordering::SeqCst) < 3 {
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
}
})
.await;
assert_eq!(
completed.load(Ordering::SeqCst),
3,
"all tap tasks must complete despite per-request clone drops"
);
}
#[test]
fn test_wiretap_config_default_is_bounded_20() {
let cfg = WireTapConfig::default();
assert_eq!(cfg.max_concurrent, Some(20));
assert_eq!(cfg.shutdown_grace, std::time::Duration::from_secs(5));
}
#[test]
fn test_wiretap_config_bounded_zero_panics() {
let result = std::panic::catch_unwind(|| WireTapConfig::bounded(0));
assert!(result.is_err());
if let Err(payload) = result {
let msg = payload
.downcast_ref::<&str>()
.expect("panic payload should be &str");
assert!(
msg.contains("max_concurrent"),
"panic message should contain 'max_concurrent', got: {msg}"
);
}
}
#[test]
fn test_wiretap_config_validate_rejects_zero_bound() {
let cfg = WireTapConfig {
max_concurrent: Some(0),
shutdown_grace: std::time::Duration::from_secs(5),
};
let result = std::panic::catch_unwind(|| cfg.validate());
assert!(result.is_err());
let payload = result.unwrap_err();
let msg = payload
.downcast_ref::<&str>()
.expect("panic payload should be &str");
assert!(
msg.contains("max_concurrent"),
"panic message should contain 'max_concurrent', got: {msg}"
);
}
#[tokio::test]
async fn test_wire_tap_drop_aborts_spawned_tasks() {
let task_started = Arc::new(AtomicBool::new(false));
let task_completed = Arc::new(AtomicBool::new(false));
let started_clone = task_started.clone();
let completed_clone = task_completed.clone();
let tap_processor = BoxProcessor::from_fn(move |_ex| {
let started = started_clone.clone();
let completed = completed_clone.clone();
Box::pin(async move {
started.store(true, Ordering::SeqCst);
tokio::time::sleep(std::time::Duration::from_secs(10)).await;
completed.store(true, Ordering::SeqCst);
Ok(Exchange::default())
})
});
let mut service = WireTapService::new(tap_processor);
let _ = service
.ready()
.await
.unwrap()
.call(Exchange::default())
.await
.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
assert!(
task_started.load(Ordering::SeqCst),
"tap task should be running"
);
assert!(
!task_completed.load(Ordering::SeqCst),
"task should not have completed yet"
);
drop(service);
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
assert!(
!task_completed.load(Ordering::SeqCst),
"task should have been aborted, not completed"
);
}
#[tokio::test]
async fn test_wiretap_bounded_detached_count_never_exceeds_bound() {
let tap_processor = BoxProcessor::from_fn(|_ex| {
Box::pin(async move {
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
Ok(Exchange::default())
})
});
let canonical = WireTapService::with_config(tap_processor, WireTapConfig::bounded(2));
let max_seen = Arc::new(AtomicUsize::new(0));
let stop = Arc::new(AtomicBool::new(false));
let sampler_svc = canonical.clone();
let sampler_max = Arc::clone(&max_seen);
let sampler_stop = Arc::clone(&stop);
let sampler = tokio::spawn(async move {
while !sampler_stop.load(Ordering::SeqCst) {
let n = sampler_svc.in_flight_count();
sampler_max.fetch_max(n, Ordering::SeqCst);
tokio::task::yield_now().await;
}
});
let mut callers = Vec::new();
for _ in 0..5 {
let mut caller_svc = canonical.clone();
callers.push(tokio::spawn(async move {
let _ = caller_svc
.ready()
.await
.unwrap()
.call(Exchange::new(Message::new("x")))
.await;
}));
}
for h in callers {
let _ = h.await;
}
stop.store(true, Ordering::SeqCst);
let _ = sampler.await;
let observed = max_seen.load(Ordering::SeqCst);
assert!(
observed <= 2,
"detached tracked task count peaked at {observed}, expected <= bound (=2)"
);
}
#[tokio::test]
async fn test_wiretap_caller_backpressured_when_saturated() {
use tokio::sync::Notify;
let notify = Arc::new(Notify::new());
let tap_notify = Arc::clone(¬ify);
let tap_processor = BoxProcessor::from_fn(move |_ex| {
let n = Arc::clone(&tap_notify);
Box::pin(async move {
n.notified().await;
Ok(Exchange::default())
})
});
let mut svc = WireTapService::with_config(tap_processor, WireTapConfig::bounded(1));
let _ = svc
.ready()
.await
.unwrap()
.call(Exchange::default())
.await
.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
let mut svc2 = svc.clone();
let mut fut2 = Box::pin(
svc2.ready()
.await
.unwrap()
.call(Exchange::new(Message::new("inline"))),
);
let pending_after_50ms = tokio::select! {
r = &mut fut2 => {
panic!(
"fut2 should be Pending after 50ms under CallerRuns back-pressure; resolved early: {:?}",
r.is_ok()
);
}
_ = tokio::time::sleep(std::time::Duration::from_millis(50)) => true,
};
assert!(
pending_after_50ms,
"fut2 should be Pending (inline tap awaiting Notify) after 50ms under CallerRuns back-pressure"
);
notify.notify_waiters();
let result = fut2.await;
assert!(
result.is_ok(),
"fut2 should resolve Ok after notify_waiters"
);
}
#[tokio::test]
async fn test_wiretap_unbounded_none_path_detaches_without_permit() {
let tap_processor = BoxProcessor::from_fn(|_ex| {
Box::pin(async move {
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
Ok(Exchange::default())
})
});
let mut svc = WireTapService::with_config(
tap_processor,
WireTapConfig {
max_concurrent: None,
shutdown_grace: std::time::Duration::from_secs(5),
},
);
let sampler_svc = svc.clone();
let peak = Arc::new(AtomicUsize::new(0));
let peak_clone = Arc::clone(&peak);
let done = Arc::new(AtomicBool::new(false));
let done_clone = Arc::clone(&done);
let sampler = tokio::spawn(async move {
while !done_clone.load(Ordering::SeqCst) {
let n = sampler_svc.in_flight_count();
peak_clone.fetch_max(n, Ordering::SeqCst);
tokio::task::yield_now().await;
}
});
for _ in 0..50 {
let ex = Exchange::new(Message::new("x"));
let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
}
let drained = tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
if svc.in_flight_count() == 0 {
return;
}
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
}
})
.await
.is_ok();
done.store(true, Ordering::SeqCst);
let _ = sampler.await;
assert!(drained, "unbounded path tasks should drain to 0 within 2s");
assert!(
peak.load(Ordering::SeqCst) > 0,
"unbounded path should have observed tracked tasks (peak > 0)"
);
}
#[tokio::test]
async fn test_wiretap_no_unbounded_task_growth_across_bursts() {
let tap_processor =
BoxProcessor::from_fn(|_ex| Box::pin(async move { Ok(Exchange::default()) }));
let svc = WireTapService::with_config(tap_processor, WireTapConfig::default());
let drain_to_zero = |svc: &WireTapService| {
let s = svc.clone();
async move {
tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
if s.in_flight_count() == 0 {
return;
}
tokio::time::sleep(std::time::Duration::from_millis(2)).await;
}
})
.await
.is_ok()
}
};
let mut callers = Vec::new();
for _ in 0..1000 {
let mut s = svc.clone();
callers.push(tokio::spawn(async move {
let _ = s.ready().await.unwrap().call(Exchange::default()).await;
}));
}
for h in callers {
let _ = h.await;
}
assert!(
drain_to_zero(&svc).await,
"burst 1 must drain to in_flight_count == 0 within 2s"
);
let mut callers = Vec::new();
for _ in 0..1000 {
let mut s = svc.clone();
callers.push(tokio::spawn(async move {
let _ = s.ready().await.unwrap().call(Exchange::default()).await;
}));
}
for h in callers {
let _ = h.await;
}
assert!(
drain_to_zero(&svc).await,
"burst 2 must drain to in_flight_count == 0 within 2s (no accumulation across bursts)"
);
}
#[derive(Clone)]
struct CapturingWriter {
sink: Arc<Mutex<Vec<u8>>>,
}
impl std::io::Write for CapturingWriter {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.sink.lock().unwrap().extend_from_slice(buf); Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for CapturingWriter {
type Writer = CapturingWriter;
fn make_writer(&'a self) -> Self::Writer {
self.clone()
}
}
fn capture_sink() -> (Arc<Mutex<Vec<u8>>>, impl tracing::Subscriber) {
let sink: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
let writer = CapturingWriter {
sink: Arc::clone(&sink),
};
let subscriber = tracing_subscriber::fmt()
.with_writer(writer)
.with_ansi(false)
.finish();
(sink, subscriber)
}
#[derive(Clone)]
struct ReadyFailingSvc {
err_msg: &'static str,
}
impl Service<Exchange> for ReadyFailingSvc {
type Response = Exchange;
type Error = CamelError;
type Future = Pin<Box<dyn Future<Output = Result<Exchange, CamelError>> + Send>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Err(CamelError::ProcessorError(self.err_msg.into())))
}
fn call(&mut self, ex: Exchange) -> Self::Future {
Box::pin(async move { Ok(ex) })
}
}
#[tokio::test]
async fn test_wiretap_tap_readiness_error_suppressed_with_log() {
let tap: camel_api::BoxProcessor = tower::util::BoxCloneService::new(ReadyFailingSvc {
err_msg: "ready-boom",
});
let mut svc = WireTapService::new(tap);
let (sink, subscriber) = capture_sink();
let exchange = Exchange::new(Message::new("main"));
let _guard = tracing::subscriber::set_default(subscriber);
let result = svc.ready().await.unwrap().call(exchange).await;
assert!(result.is_ok(), "tap readiness error must be suppressed");
assert_eq!(result.unwrap().input.body.as_text(), Some("main"));
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
drop(_guard);
let captured = String::from_utf8(sink.lock().unwrap().clone()).unwrap(); assert!(
captured.contains("ready-boom"),
"a warn! record mentioning the readiness error should have been emitted; got: {captured}"
);
}
#[tokio::test]
async fn test_wiretap_tap_processing_error_suppressed_with_log() {
let tap_processor = BoxProcessor::from_fn(|_ex| {
Box::pin(async move { Err(CamelError::ProcessorError("call-boom".into())) })
});
let mut svc = WireTapService::new(tap_processor);
let (sink, subscriber) = capture_sink();
let exchange = Exchange::new(Message::new("main"));
let _guard = tracing::subscriber::set_default(subscriber);
let result = svc.ready().await.unwrap().call(exchange).await;
assert!(result.is_ok(), "tap processing error must be suppressed");
assert_eq!(result.unwrap().input.body.as_text(), Some("main"));
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
drop(_guard);
let captured = String::from_utf8(sink.lock().unwrap().clone()).unwrap(); assert!(
captured.contains("call-boom"),
"a warn! record mentioning the processing error should have been emitted; got: {captured}"
);
}
#[tokio::test]
async fn test_wiretap_poll_ready_always_ready() {
let tap: camel_api::BoxProcessor = tower::util::BoxCloneService::new(ReadyFailingSvc {
err_msg: "would-fail",
});
let mut svc = WireTapService::new(tap);
let waker = futures::task::noop_waker();
let mut cx = Context::from_waker(&waker);
let poll = svc.poll_ready(&mut cx);
assert!(
matches!(poll, Poll::Ready(Ok(()))),
"poll_ready must be Ready(Ok(())) unconditionally (ADR-0019), got opposite"
);
}
#[tokio::test]
async fn test_wiretap_shutdown_drains_fast_aborts_slow() {
let fast_done = Arc::new(AtomicBool::new(false));
let slow_done = Arc::new(AtomicBool::new(false));
let call_idx = Arc::new(AtomicUsize::new(0));
let fd = fast_done.clone();
let sd = slow_done.clone();
let ci = call_idx.clone();
let tap_processor = BoxProcessor::from_fn(move |ex| {
let fd = fd.clone();
let sd = sd.clone();
let ci = ci.clone();
Box::pin(async move {
let n = ci.fetch_add(1, Ordering::SeqCst);
if n == 0 {
tokio::time::sleep(Duration::from_millis(10)).await;
fd.store(true, Ordering::SeqCst);
} else {
tokio::time::sleep(std::time::Duration::from_secs(10)).await;
sd.store(true, Ordering::SeqCst);
}
Ok(ex)
})
});
let config = WireTapConfig {
max_concurrent: Some(20),
shutdown_grace: Duration::from_millis(200),
};
let mut svc = WireTapService::with_config(tap_processor, config);
let _ = svc
.ready()
.await
.unwrap()
.call(Exchange::new(Message::new("fast")))
.await
.unwrap();
let _ = svc
.ready()
.await
.unwrap()
.call(Exchange::new(Message::new("slow")))
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(20)).await;
let lifecycle = svc.lifecycle();
let start = tokio::time::Instant::now();
lifecycle
.shutdown(StepShutdownReason::RouteStop)
.await
.unwrap();
let elapsed = start.elapsed();
assert!(
fast_done.load(Ordering::SeqCst),
"fast tap should drain before grace expires"
);
assert!(
!slow_done.load(Ordering::SeqCst),
"slow tap should be aborted after grace, not complete"
);
assert!(
elapsed < Duration::from_millis(500),
"shutdown took {:?}, expected < 500ms",
elapsed
);
}
#[tokio::test]
async fn test_wiretap_shutdown_idempotent() {
let slow_done = Arc::new(AtomicBool::new(false));
let sd = slow_done.clone();
let tap_processor = BoxProcessor::from_fn(move |ex| {
let sd = sd.clone();
Box::pin(async move {
tokio::time::sleep(std::time::Duration::from_secs(10)).await;
sd.store(true, Ordering::SeqCst);
Ok(ex)
})
});
let config = WireTapConfig {
max_concurrent: Some(20),
shutdown_grace: Duration::from_millis(50),
};
let mut svc = WireTapService::with_config(tap_processor, config);
let _ = svc
.ready()
.await
.unwrap()
.call(Exchange::new(Message::new("slow")))
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(20)).await;
let lifecycle = svc.lifecycle();
lifecycle
.shutdown(StepShutdownReason::RouteStop)
.await
.unwrap();
let start = tokio::time::Instant::now();
let result = lifecycle.shutdown(StepShutdownReason::HotSwap).await;
let elapsed = start.elapsed();
assert!(result.is_ok(), "second shutdown must return Ok");
assert!(
elapsed < Duration::from_millis(100),
"second shutdown must return promptly, took {:?}",
elapsed
);
assert!(
!slow_done.load(Ordering::SeqCst),
"slow tap must be aborted, not completed"
);
}
#[tokio::test]
async fn test_wiretap_calls_after_close_rejected() {
let tap_invoked = Arc::new(AtomicBool::new(false));
let ti = tap_invoked.clone();
let tap_processor = BoxProcessor::from_fn(move |ex| {
let ti = ti.clone();
Box::pin(async move {
ti.store(true, Ordering::SeqCst);
Ok(ex)
})
});
let mut svc = WireTapService::new(tap_processor);
let lifecycle = svc.lifecycle();
lifecycle
.shutdown(StepShutdownReason::RouteStop)
.await
.unwrap();
let result = svc
.ready()
.await
.unwrap()
.call(Exchange::new(Message::new("post-close")))
.await;
assert!(
result.is_ok(),
"call after close must return Ok(original exchange)"
);
assert!(
!tap_invoked.load(Ordering::SeqCst),
"tap must not be invoked after admission closed"
);
}
#[tokio::test]
async fn test_wiretap_cancellation_while_pending_readiness() {
#[derive(Clone)]
struct ForeverPendingSvc {
called: Arc<AtomicBool>,
}
impl Service<Exchange> for ForeverPendingSvc {
type Response = Exchange;
type Error = CamelError;
type Future = Pin<Box<dyn Future<Output = Result<Exchange, CamelError>> + Send>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Pending
}
fn call(&mut self, ex: Exchange) -> Self::Future {
self.called.store(true, Ordering::SeqCst);
Box::pin(async move { Ok(ex) })
}
}
let called = Arc::new(AtomicBool::new(false));
let tap: camel_api::BoxProcessor = tower::util::BoxCloneService::new(ForeverPendingSvc {
called: called.clone(),
});
let mut svc = WireTapService::new(tap);
let _ = svc
.ready()
.await
.unwrap()
.call(Exchange::new(Message::new("hanging")))
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(20)).await;
let lifecycle = svc.lifecycle();
let result = lifecycle.shutdown(StepShutdownReason::RouteStop).await;
assert!(
result.is_ok(),
"shutdown must succeed even with pending readiness: {:?}",
result
);
assert!(
!called.load(Ordering::SeqCst),
"tap call() must never be reached — cancelled during readiness phase"
);
}
#[tokio::test]
async fn test_wiretap_zero_grace_immediate_cancel() {
let slow_done = Arc::new(AtomicBool::new(false));
let sd = slow_done.clone();
let tap_processor = BoxProcessor::from_fn(move |ex| {
let sd = sd.clone();
Box::pin(async move {
tokio::time::sleep(std::time::Duration::from_secs(10)).await;
sd.store(true, Ordering::SeqCst);
Ok(ex)
})
});
let config = WireTapConfig {
max_concurrent: Some(20),
shutdown_grace: Duration::ZERO,
};
let mut svc = WireTapService::with_config(tap_processor, config);
let _ = svc
.ready()
.await
.unwrap()
.call(Exchange::new(Message::new("slow")))
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(20)).await;
let lifecycle = svc.lifecycle();
let start = tokio::time::Instant::now();
lifecycle
.shutdown(StepShutdownReason::RouteStop)
.await
.unwrap();
let elapsed = start.elapsed();
assert!(
!slow_done.load(Ordering::SeqCst),
"slow tap must be aborted immediately (zero grace)"
);
assert!(
elapsed < Duration::from_millis(200),
"zero-grace shutdown must return quickly, took {:?}",
elapsed
);
}
#[tokio::test(flavor = "multi_thread")]
async fn test_wiretap_admission_shutdown_no_orphan_task() {
const ITERATIONS: usize = 200;
for _ in 0..ITERATIONS {
let tap_processor = BoxProcessor::from_fn(|_ex| {
Box::pin(async move {
tokio::time::sleep(Duration::from_millis(1)).await;
Ok(Exchange::default())
})
});
let svc = WireTapService::new(tap_processor);
let lifecycle = svc.lifecycle();
let mut handles = Vec::new();
for _ in 0..4 {
let mut c = svc.clone();
handles.push(tokio::spawn(async move {
let _ = c.ready().await.unwrap().call(Exchange::default()).await;
}));
}
tokio::task::yield_now().await;
tokio::time::sleep(Duration::from_millis(1)).await;
lifecycle
.shutdown(StepShutdownReason::RouteStop)
.await
.unwrap();
for h in handles {
let _ = h.await;
}
let drained = tokio::time::timeout(Duration::from_secs(2), async {
loop {
if svc.in_flight_count() == 0 {
return;
}
tokio::time::sleep(Duration::from_millis(5)).await;
}
})
.await
.is_ok();
assert!(
drained,
"iteration: in_flight_count must drain to 0 after shutdown"
);
}
}
}