From f6571e6ef87e8f2481be380bdc6b6c19c81a45c1 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 00:01:22 -0600 Subject: [PATCH 01/29] fix(streaming): bound gRPC stream server shutdown when a peer never speaks HTTP/2 tonic's graceful shutdown waits for every connection, and a connection that never finished its HTTP/2 handshake cannot receive GOAWAY. A health check or a half-open peer therefore held GrpcFrameTransport::shutdown, and with it Velo::shutdown, open forever. Drop the server after a short grace period so the listener closes and shutdown returns. Also move a regression doc comment back onto the test it describes. Signed-off-by: Ryan Olson --- lib/velo/src/streaming/grpc_transport.rs | 78 ++++++++++++++++++------ 1 file changed, 61 insertions(+), 17 deletions(-) diff --git a/lib/velo/src/streaming/grpc_transport.rs b/lib/velo/src/streaming/grpc_transport.rs index 85c07b31..503e37cf 100644 --- a/lib/velo/src/streaming/grpc_transport.rs +++ b/lib/velo/src/streaming/grpc_transport.rs @@ -64,6 +64,11 @@ const SESSION_ID_META: &str = "x-session-id"; /// so the heartbeat watchdog is what reports a wedged consumer, not this. const TERMINAL_ACK_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20); +/// How long `shutdown` lets open HTTP/2 connections finish GOAWAY before it +/// drops the server. A clean close takes one round trip; this bounds the case +/// where a peer never answers. +const SERVER_CLOSE_GRACE: std::time::Duration = std::time::Duration::from_secs(1); + use crate::streaming::sender::is_terminal_sentinel; // --------------------------------------------------------------------------- @@ -273,11 +278,29 @@ impl GrpcFrameTransport { tasks.spawn_until_done(async move { let server = tonic::transport::Server::builder().add_service(VeloStreamingServer::new(service)); - if let Err(e) = server - .serve_with_incoming_shutdown(TcpListenerStream::new(listener), cancel.cancelled()) - .await - { - tracing::warn!("GrpcFrameTransport server error: {}", e); + let serve = server + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), cancel.cancelled()); + // Graceful shutdown waits for every connection, and a connection + // that never finished its HTTP/2 handshake cannot be sent GOAWAY, + // so it would hold `shutdown` open forever. After the grace period + // drop the server: the listener closes, and tonic's own connection + // tasks end when their peers close. + let grace_expired = async { + cancel.cancelled().await; + tokio::time::sleep(SERVER_CLOSE_GRACE).await; + }; + tokio::select! { + result = serve => { + if let Err(e) = result { + tracing::warn!("GrpcFrameTransport server error: {}", e); + } + } + () = grace_expired => { + tracing::debug!( + "GrpcFrameTransport: connections still open after {:?}; releasing listener", + SERVER_CLOSE_GRACE + ); + } } }); @@ -551,18 +574,6 @@ mod tests { (wid, PeerInfo::new(inst, address)) } - /// Regression: the streaming RPC must not complete its response until it - /// has consumed the whole request stream. - /// - /// `connect`'s frame pump uses end-of-response as the acknowledgement that - /// its terminal sentinel reached the consumer. When the handler answered - /// with an already-empty stream, tonic sent the trailers before reading the - /// request body, so the pump had nothing to wait on: it queued its frames - /// into a 256-slot channel and tore the call down, discarding whatever had - /// not yet been flushed. On a short stream that routinely threw away *every* - /// frame including the terminal, and the server — having seen a request - /// stream that ended without a sentinel — injected `Dropped`. The consumer - /// then failed with `SenderDropped` despite a clean `finalize()`. #[tokio::test] async fn shutdown_closes_listener_binds_and_active_pumps() { let server = GrpcFrameTransport::new("127.0.0.1:0".parse().unwrap()) @@ -600,6 +611,39 @@ mod tests { .expect("listener released before shutdown returns"); } + /// A peer that opens TCP to the stream port and never sends the HTTP/2 + /// preface (a health check, a port scan, a peer that died mid-handshake) + /// must not hold shutdown open. tonic's graceful shutdown waits for every + /// connection, and hyper cannot send GOAWAY before the handshake ends, so + /// waiting for the server task alone never returns. + #[tokio::test] + async fn shutdown_does_not_wait_for_silent_connection() { + let server = GrpcFrameTransport::new("127.0.0.1:0".parse().unwrap()) + .await + .unwrap(); + let addr = server.bound_addr(); + let _silent = tokio::net::TcpStream::connect(addr).await.unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + tokio::time::timeout(Duration::from_secs(5), server.shutdown()) + .await + .expect("shutdown waited for a connection that never spoke HTTP/2"); + let _listener = tokio::net::TcpListener::bind(addr) + .await + .expect("listener released before shutdown returns"); + } + + /// Regression: the streaming RPC must not complete its response until it + /// has consumed the whole request stream. + /// + /// `connect`'s frame pump uses end-of-response as the acknowledgement that + /// its terminal sentinel reached the consumer. When the handler answered + /// with an already-empty stream, tonic sent the trailers before reading the + /// request body, so the pump had nothing to wait on: it queued its frames + /// into a 256-slot channel and tore the call down, discarding whatever had + /// not yet been flushed. On a short stream that routinely threw away *every* + /// frame including the terminal, and the server — having seen a request + /// stream that ended without a sentinel — injected `Dropped`. The consumer + /// then failed with `SenderDropped` despite a clean `finalize()`. #[tokio::test(flavor = "multi_thread")] async fn response_completes_only_after_request_stream_drains() { let server = GrpcFrameTransport::default_new().await.unwrap(); From 5a07a8b54515e087ffd979fbdb69d395ed3aa6df Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 00:03:26 -0600 Subject: [PATCH 02/29] fix(streaming): arm the unattached timeout when detach runs off-runtime detach now re-arms the anchor's unattached timeout once the Detached sentinel is queued. On the fast path that runs on the caller's thread, and a thread with no runtime (a sender moved to a binding's thread) could not spawn the timer, so an anchor nobody re-attached was never reaped. Enter the sender's runtime there. Signed-off-by: Ryan Olson --- lib/velo/src/streaming/anchor.rs | 24 ++++++++++++++++++++++++ lib/velo/src/streaming/sender.rs | 6 ++++++ 2 files changed, 30 insertions(+) diff --git a/lib/velo/src/streaming/anchor.rs b/lib/velo/src/streaming/anchor.rs index c0ddcefd..1af41071 100644 --- a/lib/velo/src/streaming/anchor.rs +++ b/lib/velo/src/streaming/anchor.rs @@ -2986,6 +2986,30 @@ mod tests { )); } + /// A sender may be moved to a thread with no runtime (a language + /// binding's thread). Detaching there must still arm the unattached + /// timeout; otherwise an anchor nobody re-attaches is never reaped. + #[tokio::test(flavor = "multi_thread")] + async fn detach_from_plain_thread_arms_unattached_timeout() { + let mgr = make_manager(); + let anchor = mgr.create_anchor::(); + anchor.set_timeout(Some(Duration::from_millis(100))); + let handle = anchor.handle(); + let (_, local_id) = handle.unpack(); + let sender = mgr.attach_stream_anchor::(handle).await.unwrap(); + std::thread::spawn(move || sender.detach().unwrap()) + .join() + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + while mgr.registry.contains_key(&local_id) { + tokio::time::sleep(Duration::from_millis(20)).await; + } + }) + .await + .expect("detached anchor was never reaped"); + drop(anchor); + } + // ----------------------------------------------------------------------- // Test 4: CancellationToken is idempotent across multiple cancel() calls // ----------------------------------------------------------------------- diff --git a/lib/velo/src/streaming/sender.rs b/lib/velo/src/streaming/sender.rs index 95b4790d..ae1ffb7b 100644 --- a/lib/velo/src/streaming/sender.rs +++ b/lib/velo/src/streaming/sender.rs @@ -411,7 +411,13 @@ impl StreamSender { let registry = self.local_registry.clone(); let (worker_id, local_id) = self.handle.unpack(); let sender_stream_id = self.sender_stream_id; + let runtime = self.runtime.clone(); send_terminal(&self.runtime, &self.tx, bytes, move || { + // The fast path runs this on the caller's thread, which may have + // no runtime; the unattached timeout must still be armed. + let _runtime = tokio::runtime::Handle::try_current() + .is_err() + .then(|| runtime.enter()); if let Some(registry) = registry && let Some(mut entry) = registry.get_mut(&local_id) && entry From df02fdb7917d8d2b13403ea875d70ce18b258e44 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 00:05:20 -0600 Subject: [PATCH 03/29] fix(streaming): let MPSC detach outlive the runtime that made the sender detach always spawned its terminal send on the runtime the sender was attached on. Once that runtime stopped, the task was dropped, detach returned ChannelClosed with room in the channel, and the consumer saw neither Detached nor Dropped. Queue the terminal with try_send first, which commits it at once, and only on a full channel spawn the wait on the caller's runtime. Drop and the cleanup guard follow the same rule. Signed-off-by: Ryan Olson --- lib/velo/src/streaming/mpsc/sender.rs | 39 ++++++++++++++++---- lib/velo/tests/streaming/mpsc_integration.rs | 32 ++++++++++++++++ 2 files changed, 63 insertions(+), 8 deletions(-) diff --git a/lib/velo/src/streaming/mpsc/sender.rs b/lib/velo/src/streaming/mpsc/sender.rs index 64b36950..528a3211 100644 --- a/lib/velo/src/streaming/mpsc/sender.rs +++ b/lib/velo/src/streaming/mpsc/sender.rs @@ -230,22 +230,39 @@ impl MpscStreamSender { /// /// Reattaching via [`crate::AnchorManager::attach_mpsc_stream_anchor`] /// allocates a fresh [`SenderId`]. Once polled, detach runs to completion - /// on the sender's runtime even if this future is dropped. Consumer - /// cancellation stops a detach that is waiting for channel space. + /// even if this future is dropped: at once when the channel has room, else + /// on the caller's runtime (the sender's own when the caller has none). + /// Consumer cancellation stops a detach that is waiting for channel space. pub async fn detach(mut self) -> Result { self.heartbeat_cancel.cancel(); let cleanup = self.cleanup_guard(); + self.sent_terminal = true; + if self.cancel_token.is_cancelled() { + return Err(SendError::ChannelClosed); + } + let sender_id = self.sender_id.0; + let bytes = cached_detached().clone(); + let bytes = match &self.channel { + SenderChannel::Local(tx) => match tx.try_send((sender_id, bytes)) { + Ok(()) => return Ok(self.handle), + Err(flume::TrySendError::Full((_, bytes))) => bytes, + Err(flume::TrySendError::Disconnected(_)) => return Err(SendError::ChannelClosed), + }, + SenderChannel::Remote(tx) => match tx.try_send(bytes) { + Ok(()) => return Ok(self.handle), + Err(flume::TrySendError::Full(bytes)) => bytes, + Err(flume::TrySendError::Disconnected(_)) => return Err(SendError::ChannelClosed), + }, + }; let channel = self.channel.clone(); let cancel = self.cancel_token.clone(); - let sender_id = self.sender_id.0; // Flume can enqueue a frame before its send future is polled Ready. // Keep one task responsible for the terminal event and cleanup, so // abandoning this caller cannot send both Detached and Dropped. // Capture only transport state: T does not need a Send bound. - let task = self.runtime.spawn(async move { + let task = self.spawn_runtime().spawn(async move { let _cleanup = cleanup; let send = async { - let bytes = cached_detached().clone(); match channel { SenderChannel::Local(tx) => tx .send_async((sender_id, bytes)) @@ -263,7 +280,6 @@ impl MpscStreamSender { result = send => result, } }); - self.sent_terminal = true; task.await.map_err(|_| SendError::ChannelClosed)??; Ok(self.handle) } @@ -286,7 +302,9 @@ impl Drop for SenderCleanup { self.sender_registry.senders.remove(&self.sender_stream_id); // The last local slot may start an unattached timeout. Drop can run // on a plain thread, so provide the runtime captured at construction. - let _entered = self.runtime.enter(); + let _entered = tokio::runtime::Handle::try_current() + .is_err() + .then(|| self.runtime.enter()); if self.local && let Some(slot) = super::anchor::remove_sender_slot( &self.mpsc_registry, @@ -301,6 +319,11 @@ impl Drop for SenderCleanup { } impl MpscStreamSender { + /// The caller's runtime first: the sender can outlive the one it was made on. + fn spawn_runtime(&self) -> tokio::runtime::Handle { + tokio::runtime::Handle::try_current().unwrap_or_else(|_| self.runtime.clone()) + } + fn cleanup_guard(&self) -> SenderCleanup { SenderCleanup { local: matches!(self.channel, SenderChannel::Local(_)), @@ -335,7 +358,7 @@ impl Drop for MpscStreamSender { // runtime worker. Drop removes the sender registry // entry, so only the anchor token can stop this wait. let tx = tx.clone(); - self.runtime.spawn(async move { + self.spawn_runtime().spawn(async move { tokio::select! { biased; _ = consumer_cancel.cancelled() => {} diff --git a/lib/velo/tests/streaming/mpsc_integration.rs b/lib/velo/tests/streaming/mpsc_integration.rs index 7c080059..f0404370 100644 --- a/lib/velo/tests/streaming/mpsc_integration.rs +++ b/lib/velo/tests/streaming/mpsc_integration.rs @@ -386,6 +386,38 @@ async fn abandoning_mpsc_detach_completes_or_stops_on_consumer_cancel() { } } +/// A sender can outlive the runtime it was attached on. Detaching it from +/// another runtime must still deliver `Detached`: work spawned on the stopped +/// runtime is dropped, so the consumer would see no terminal at all. +#[tokio::test(flavor = "multi_thread")] +async fn mpsc_detach_after_attaching_runtime_stops() { + let mgr = make_manager(); + let mut anchor = mgr.create_mpsc_anchor::(); + let handle = anchor.handle(); + let attach_mgr = Arc::clone(&mgr); + let sender = std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(attach_mgr.attach_mpsc_stream_anchor::(handle)) + // `runtime` drops here and shuts down. + }) + .join() + .unwrap() + .expect("attach"); + let sid = sender.sender_id(); + let returned = tokio::time::timeout(Duration::from_secs(2), sender.detach()) + .await + .expect("detach stalled") + .expect("detach after the attaching runtime stopped"); + assert_eq!(returned, handle); + assert!( + matches!(next_frame(&mut anchor).await, Some(Ok((id, MpscFrame::Detached))) if id == sid) + ); + anchor.cancel(); +} + /// A pending `anchor.next()` must terminate promptly when the controller /// cancels, even if sender handles are still alive. #[tokio::test(flavor = "multi_thread")] From 72c7b17b7c72446c307dc31f6e2ef1cdf61a0ede Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 00:11:39 -0600 Subject: [PATCH 04/29] perf(streaming): read the stream-task stop flag without a lock CancellationToken::is_cancelled takes a mutex, and every peer and lane of a mux checks StreamTasks::is_stopped several times per batch. Keep an AtomicBool next to the token, set under the admission lock in stop(). A microbenchmark of the check alone: 10.8 ns uncontended and 203 ns with 4 threads for the token, 0.5 ns and 1.1 ns for the flag. Signed-off-by: Ryan Olson --- lib/velo/src/streaming/tasks.rs | 33 +++++++++++++++++++++++++++++++-- 1 file changed, 31 insertions(+), 2 deletions(-) diff --git a/lib/velo/src/streaming/tasks.rs b/lib/velo/src/streaming/tasks.rs index a63fd1ff..3f49cb2e 100644 --- a/lib/velo/src/streaming/tasks.rs +++ b/lib/velo/src/streaming/tasks.rs @@ -3,6 +3,7 @@ use std::future::Future; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; @@ -11,6 +12,9 @@ use tokio_util::task::TaskTracker; #[derive(Clone)] pub(crate) struct StreamTasks { cancel: CancellationToken, + /// Mirrors `cancel` for hot-path readers: `CancellationToken::is_cancelled` + /// takes a mutex that every peer and lane of a mux would share. + stopped: Arc, tracker: TaskTracker, admission: Arc>, runtime: tokio::runtime::Handle, @@ -26,6 +30,7 @@ impl StreamTasks { pub(crate) fn new(runtime: tokio::runtime::Handle) -> Self { Self { cancel: CancellationToken::new(), + stopped: Arc::default(), tracker: TaskTracker::new(), admission: Arc::default(), runtime, @@ -61,13 +66,13 @@ impl StreamTasks { } pub(crate) fn is_stopped(&self) -> bool { - self.cancel.is_cancelled() + self.stopped.load(Ordering::Acquire) } /// Returns true only for the first stop request. pub(crate) fn stop(&self) -> bool { let _admission = self.admission.lock(); - let first = !self.is_stopped(); + let first = !self.stopped.swap(true, Ordering::AcqRel); self.cancel.cancel(); self.tracker.close(); first @@ -77,3 +82,27 @@ impl StreamTasks { self.tracker.wait().await; } } + +#[cfg(test)] +mod tests { + use super::*; + + /// `is_stopped` reads a flag, not the token, so the two must move + /// together: the first stop wins, refuses later spawns, and cancels + /// what is running. + #[tokio::test] + async fn stop_is_seen_by_flag_token_and_spawn() { + let tasks = StreamTasks::default(); + let token = tasks.cancellation_token(); + assert!(tasks.spawn(std::future::pending())); + assert!(!tasks.is_stopped()); + assert!(tasks.stop()); + assert!(!tasks.stop()); + assert!(tasks.is_stopped()); + assert!(token.is_cancelled()); + assert!(!tasks.spawn(async {})); + tokio::time::timeout(std::time::Duration::from_secs(1), tasks.wait()) + .await + .expect("a stopped tracker waits only for running tasks"); + } +} From 4eb3ccba1d9822e45395bf569832a3891501ecc8 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 00:11:41 -0600 Subject: [PATCH 05/29] refactor(streaming): give the mux one stop path in its own module MuxCore carried a cancel token that was only another name for its StreamTasks token; read the tasks instead. Stopping the tasks and closing ingress happened in four places; one stop_mux helper does it now. Batcher start and the stop path move to messenger_mux/lifecycle.rs, which brings mod.rs back under 1000 lines. The batcher keeps its own cancel handle, now documented: a stop that lands while its loop is polled exits through teardown(true), and tests use a separate token to drive that exit. Signed-off-by: Ryan Olson --- .../src/streaming/messenger_mux/lifecycle.rs | 100 ++++++++++++++++++ lib/velo/src/streaming/messenger_mux/mod.rs | 84 +-------------- .../messenger_mux/peer_batcher/mod.rs | 9 +- lib/velo/src/streaming/messenger_mux/sweep.rs | 2 +- lib/velo/src/streaming/messenger_mux/tests.rs | 2 +- 5 files changed, 114 insertions(+), 83 deletions(-) create mode 100644 lib/velo/src/streaming/messenger_mux/lifecycle.rs diff --git a/lib/velo/src/streaming/messenger_mux/lifecycle.rs b/lib/velo/src/streaming/messenger_mux/lifecycle.rs new file mode 100644 index 00000000..c5089f50 --- /dev/null +++ b/lib/velo/src/streaming/messenger_mux/lifecycle.rs @@ -0,0 +1,100 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! How the mux starts batchers and how it stops: the one stop path every +//! failure and shutdown goes through. + +use std::sync::Arc; + +use anyhow::{Result, anyhow}; + +use super::ingress::IngressRegistry; +use super::peer_batcher::{self, BatcherContext, BatcherHandle}; +use super::{MessengerMuxTransport, MuxCore, PeerLane}; +use crate::observability::MuxMetricsHandle; + +impl MessengerMuxTransport { + pub(crate) async fn shutdown(&self) { + self.core.tasks.stop(); + self.core.tasks.wait().await; + self.core.batchers.clear(); + close_ingress(&self.core.ingress, self.core.metrics.as_ref()); + self.core.drains.clear(); + } +} + +impl MuxCore { + /// The batcher for one (peer, lane), created on first use. + pub(super) fn batcher(&self, key: PeerLane) -> Result> { + if self.tasks.is_stopped() { + return Err(anyhow!("messenger mux shut down")); + } + if let Some(existing) = self.batchers.get(&key) { + return self.checked_batcher(key, existing.value()); + } + let entry = self.batchers.entry(key).or_insert_with(|| { + peer_batcher::spawn( + key, + BatcherContext { + tasks: self.tasks.clone(), + messenger: Arc::clone(&self.messenger), + config: self.config.clone(), + metrics: self.metrics.clone(), + epochs: Arc::clone(&self.epochs), + batchers: Arc::clone(&self.batchers), + ingress: Arc::clone(&self.ingress), + cancel: self.tasks.cancellation_token(), + #[cfg(test)] + hooks: self.hooks.get().cloned(), + }, + ) + }); + self.checked_batcher(key, entry.value()) + } + + /// The caller holds the registry guard, so a closed handle here cannot + /// have been removed by normal retirement between lookup and validation. + fn checked_batcher( + &self, + key: PeerLane, + handle: &Arc, + ) -> Result> { + // Spawn on a stopped runtime can drop the batcher immediately. A + // closed registered handle is a failed epoch, not normal retirement. + if self.tasks.is_stopped() || handle.is_closed() { + if stop_mux(&self.tasks, &self.ingress, self.metrics.as_ref()) { + tracing::error!(peer = %key.peer, lane = %key.lane, + "messenger mux stopped: registered batcher closed unexpectedly"); + } + return Err(anyhow!("messenger mux batcher shut down")); + } + Ok(Arc::clone(handle)) + } +} + +/// Stop the mux's tasks and close every ingress slot. Returns true only for +/// the first stop request. +pub(super) fn stop_mux( + tasks: &crate::streaming::tasks::StreamTasks, + ingress: &IngressRegistry, + metrics: Option<&MuxMetricsHandle>, +) -> bool { + let first = tasks.stop(); + close_ingress(ingress, metrics); + first +} + +pub(super) fn close_ingress(ingress: &IngressRegistry, metrics: Option<&MuxMetricsHandle>) { + let closed = ingress.shutdown(); + if let Some(metrics) = metrics { + for _ in 0..closed { + metrics.slot_closed(); + } + } +} + +impl Drop for MuxCore { + fn drop(&mut self) { + stop_mux(&self.tasks, &self.ingress, self.metrics.as_ref()); + } +} diff --git a/lib/velo/src/streaming/messenger_mux/mod.rs b/lib/velo/src/streaming/messenger_mux/mod.rs index 8cf812cc..23648477 100644 --- a/lib/velo/src/streaming/messenger_mux/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/mod.rs @@ -105,6 +105,7 @@ mod lane; mod lane_choice; #[cfg(all(test, feature = "quic"))] mod lane_tests; +mod lifecycle; pub(crate) mod peer_batcher; pub(crate) mod protocol; mod sweep; @@ -127,9 +128,8 @@ use velo_ext::{TransportKey, WorkerAddress, WorkerId}; use self::flow_control::NegotiatedLimits; use self::ingress::IngressRegistry; use self::lane_choice::LaneLoad; -use self::peer_batcher::{ - BatcherContext, BatcherHandle, BatcherMap, OpenRejected, OpenSlotRequest, -}; +use self::lifecycle::close_ingress; +use self::peer_batcher::{BatcherHandle, BatcherMap, OpenRejected, OpenSlotRequest}; use crate::messenger::{Context, Handler, Messenger}; use crate::observability::{MuxMetricsHandle, VeloMetrics}; use crate::streaming::transport::FrameTransport; @@ -224,7 +224,6 @@ struct MuxCore { /// lazily recreated must not restart its epoch, or the peer would read every /// batch of the new one as stale and discard it wholesale. epochs: Arc, - cancel: CancellationToken, tasks: crate::streaming::tasks::StreamTasks, /// Peers with credit to return, posted by draining consumers. See /// [`ingress::DrainSignal`]. @@ -266,14 +265,6 @@ struct MuxCore { } impl MessengerMuxTransport { - pub(crate) async fn shutdown(&self) { - self.core.tasks.stop(); - self.core.tasks.wait().await; - self.core.batchers.clear(); - close_ingress(&self.core.ingress, self.core.metrics.as_ref()); - self.core.drains.clear(); - } - /// Ask the sender of a claimed slot to stop, through the batcher of the /// lane its `OpenSlot` arrived on. pub(crate) fn request_stop(&self, key: PeerLane, slot: protocol::SlotId, session_id: u64) { @@ -358,7 +349,6 @@ impl MessengerMuxTransport { // Epochs start at 1 so zero is never a live epoch, which keeps a // zeroed header from reading as a legitimate one. epochs: Arc::new(AtomicU64::new(1)), - cancel: tasks.cancellation_token(), tasks, drain_tx, drain_rx, @@ -405,54 +395,6 @@ impl MessengerMuxTransport { } impl MuxCore { - /// The batcher for one (peer, lane), created on first use. - fn batcher(&self, key: PeerLane) -> Result> { - if self.tasks.is_stopped() { - return Err(anyhow!("messenger mux shut down")); - } - if let Some(existing) = self.batchers.get(&key) { - return self.checked_batcher(key, existing.value()); - } - let entry = self.batchers.entry(key).or_insert_with(|| { - peer_batcher::spawn( - key, - BatcherContext { - tasks: self.tasks.clone(), - messenger: Arc::clone(&self.messenger), - config: self.config.clone(), - metrics: self.metrics.clone(), - epochs: Arc::clone(&self.epochs), - batchers: Arc::clone(&self.batchers), - ingress: Arc::clone(&self.ingress), - cancel: self.cancel.clone(), - #[cfg(test)] - hooks: self.hooks.get().cloned(), - }, - ) - }); - self.checked_batcher(key, entry.value()) - } - - /// The caller holds the registry guard, so a closed handle here cannot - /// have been removed by normal retirement between lookup and validation. - fn checked_batcher( - &self, - key: PeerLane, - handle: &Arc, - ) -> Result> { - // Spawn on a stopped runtime can drop the batcher immediately. A - // closed registered handle is a failed epoch, not normal retirement. - if self.tasks.is_stopped() || handle.is_closed() { - if self.tasks.stop() { - tracing::error!(peer = %key.peer, lane = %key.lane, - "messenger mux stopped: registered batcher closed unexpectedly"); - } - close_ingress(&self.ingress, self.metrics.as_ref()); - return Err(anyhow!("messenger mux batcher shut down")); - } - Ok(Arc::clone(handle)) - } - /// Hand one decoded batch to the ingress lane and act on what it produced. /// /// Everything the batch produced goes back through the batcher for the @@ -586,7 +528,7 @@ impl MuxCore { /// Hand a reconcile pass's grants to the batcher of the lane the slots /// arrived on. fn return_credit(&self, key: PeerLane, replies: Vec) { - if replies.is_empty() || self.cancel.is_cancelled() { + if replies.is_empty() || self.tasks.is_stopped() { return; } let Ok(batcher) = self.batcher(key) else { @@ -691,22 +633,6 @@ impl MuxCore { } } -fn close_ingress(ingress: &IngressRegistry, metrics: Option<&MuxMetricsHandle>) { - let closed = ingress.shutdown(); - if let Some(metrics) = metrics { - for _ in 0..closed { - metrics.slot_closed(); - } - } -} - -impl Drop for MuxCore { - fn drop(&mut self) { - self.tasks.stop(); - close_ingress(&self.ingress, self.metrics.as_ref()); - } -} - /// Register the receive buffer for one `(anchor_id, session_id)` and open the /// accept window on it. /// @@ -963,7 +889,7 @@ impl MessengerMuxTransport { let core = Arc::clone(&self.core); Box::pin(async move { for _ in 0..CONNECT_ATTEMPTS { - if core.cancel.is_cancelled() { + if core.tasks.is_stopped() { return Err(anyhow!("messenger mux shut down")); } let batcher = core.batcher(key)?; diff --git a/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs b/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs index 75215960..0c3061ac 100644 --- a/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs @@ -283,6 +283,10 @@ pub(crate) struct BatcherContext { pub(crate) epochs: Arc, pub(crate) batchers: Arc, pub(crate) ingress: Arc, + /// The run loop's own exit. The mux passes `tasks.cancellation_token()`; + /// a stop that lands while the loop is being polled exits here, through + /// `teardown(true)`, rather than by the spawn wrapper dropping the loop. + /// Tests pass a separate token to drive that exit on purpose. pub(crate) cancel: CancellationToken, /// A barrier in the run loop, installed only by the tests that need to stop /// it mid-wake. See [`test_hooks`]. @@ -409,11 +413,12 @@ impl Drop for Batcher { // abort during that flush still loses credit and must stop the mux. if !self.teardown_complete { let already_cancelled = self.cancel.is_cancelled(); - if self.tasks.stop() && !already_cancelled { + if super::lifecycle::stop_mux(&self.tasks, &self.ingress, self.metrics.as_ref()) + && !already_cancelled + { tracing::error!(peer = %self.key.peer, lane = %self.key.lane, "messenger mux stopped: batcher task aborted unexpectedly"); } - super::close_ingress(&self.ingress, self.metrics.as_ref()); } // A refused spawn can drop here under the batcher-map entry guard. // Normal exits unregister in run; whole-mux shutdown clears the map. diff --git a/lib/velo/src/streaming/messenger_mux/sweep.rs b/lib/velo/src/streaming/messenger_mux/sweep.rs index 98958311..6ba117d5 100644 --- a/lib/velo/src/streaming/messenger_mux/sweep.rs +++ b/lib/velo/src/streaming/messenger_mux/sweep.rs @@ -202,7 +202,7 @@ async fn deferred_visit_due(due: Option) { /// entry is what used to leave one entry behind per tick, for good. pub(super) fn spawn_sweep(core: &Arc) { let weak = Arc::downgrade(core); - let cancel = core.cancel.clone(); + let cancel = core.tasks.cancellation_token(); let interval = core.config.credit_sweep_interval; let floor = core.config.drain_visit_floor; let drain_rx = core.drain_rx.clone(); diff --git a/lib/velo/src/streaming/messenger_mux/tests.rs b/lib/velo/src/streaming/messenger_mux/tests.rs index 5369e35c..0606b37d 100644 --- a/lib/velo/src/streaming/messenger_mux/tests.rs +++ b/lib/velo/src/streaming/messenger_mux/tests.rs @@ -212,7 +212,7 @@ fn runtime_shutdown_cancels_a_mux_with_live_batchers() { }); runtime.shutdown_timeout(Duration::from_secs(5)); assert!( - pair.producer.core.cancel.is_cancelled(), + pair.producer.core.tasks.is_stopped(), "an aborted batcher must stop retries on its closed registry entry" ); // The consumer has no batcher yet. Submitting one to its stopped runtime From 58bc11077a64756f1647a3606034b81c8fde6172 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 00:07:00 -0600 Subject: [PATCH 06/29] refactor(transports): one stop sequence for startup failure and shutdown StartupTransports::stop and VeloBackend::shutdown_now repeated the same drain, teardown and shutdown steps. Both drain calls are idempotent, so one helper serves both. Signed-off-by: Ryan Olson --- lib/velo/src/transports.rs | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/lib/velo/src/transports.rs b/lib/velo/src/transports.rs index 93697dac..216c2bca 100644 --- a/lib/velo/src/transports.rs +++ b/lib/velo/src/transports.rs @@ -145,15 +145,8 @@ struct StartupTransports { impl StartupTransports { fn stop(&mut self) { if let Some(shutdown) = self.shutdown.take() { - shutdown.begin_drain(); - for transport in self.transports.values() { - transport.begin_drain(); - } // Construction has no caller to drain work for. Use a zero budget. - shutdown.teardown_token().cancel(); - for transport in self.transports.values() { - transport.shutdown(); - } + stop_transports(&shutdown, &self.transports); } } @@ -163,6 +156,19 @@ impl StartupTransports { } } +/// Gate, tear down, and shut down every transport. Both `begin_drain` calls +/// are idempotent, so this is safe after a graceful drain already ran. +fn stop_transports(state: &ShutdownState, transports: &HashMap>) { + state.begin_drain(); + for transport in transports.values() { + transport.begin_drain(); + } + state.teardown_token().cancel(); + for transport in transports.values() { + transport.shutdown(); + } +} + impl Drop for StartupTransports { fn drop(&mut self) { self.stop(); @@ -682,13 +688,7 @@ impl VeloBackend { /// Stop transports after teardown or failed construction. pub(crate) fn shutdown_now(&self) { - if !self.shutdown_state.is_draining() { - self.begin_drain(); - } - self.shutdown_state.teardown_token().cancel(); - for transport in self.transports.values() { - transport.shutdown(); - } + stop_transports(&self.shutdown_state, &self.transports); } /// Perform a graceful 4-phase shutdown. From 449febc50a28e7c718f2145258e958a0a0ee592d Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 00:08:26 -0600 Subject: [PATCH 07/29] refactor(messenger): one helper for error replies in the server Four places sent an error response and logged a failure unless it was an admission refusal. They now share send_error_reply. fail_fast awaits the reply in the handler task, and the ordered-lane shed path spawns it on the messenger task tracker instead of an untracked task. Signed-off-by: Ryan Olson --- lib/velo/src/messenger/server/dispatcher.rs | 99 +++++++++++---------- lib/velo/src/messenger/server/mod.rs | 40 ++++----- 2 files changed, 68 insertions(+), 71 deletions(-) diff --git a/lib/velo/src/messenger/server/dispatcher.rs b/lib/velo/src/messenger/server/dispatcher.rs index 4f67ff02..4927efe2 100644 --- a/lib/velo/src/messenger/server/dispatcher.rs +++ b/lib/velo/src/messenger/server/dispatcher.rs @@ -105,19 +105,24 @@ async fn run_handler( ); fail_fast( &system, + handler.name(), message_id, response_type, "handler panicked", in_flight, - ); + ) + .await; } trace!(target: "crate::messenger::dispatcher", handler = %handler.name(), "Handler task completed"); } /// Send an error response so a waiting caller fails immediately instead of -/// hanging until its own timeout. No-op for fire-and-forget. -fn fail_fast( +/// hanging until its own timeout. No-op for fire-and-forget. Awaited inline, +/// so the caller's task tracking covers the reply; `in_flight` is held until +/// the reply is sent. +async fn fail_fast( system: &Arc, + handler: &str, message_id: ResponseId, response_type: ResponseType, reason: &'static str, @@ -126,23 +131,34 @@ fn fail_fast( if matches!(response_type, ResponseType::FireAndForget) { return; } - let backend = system.backend().clone(); - tokio::spawn(async move { - let _in_flight = in_flight; - if let Err(e) = DispatcherHub::send_error_response_static( - &backend, - message_id, - format!("Handler failed: {reason}"), - ) - .await - && !e.is::() - { - error!( - target: "crate::messenger::dispatcher", - "Failed to send handler error response: {}", e - ); - } - }); + let _in_flight = in_flight; + send_error_reply( + system.backend(), + handler, + message_id, + format!("Handler failed: {reason}"), + ) + .await; +} + +/// Send an error response to the caller and log a failure to do so. Admission +/// refusals are expected during shutdown, so they are not logged. +pub(crate) async fn send_error_reply( + backend: &VeloBackend, + handler: &str, + message_id: ResponseId, + message: String, +) { + if let Err(e) = DispatcherHub::send_error_response_static(backend, message_id, message).await + && !e.is::() + { + error!( + target: "crate::messenger::dispatcher", + handler = %handler, + message_id = %message_id, + "Failed to send error response: {e}" + ); + } } /// Dispatcher implementation that spawns handlers on a task tracker. @@ -440,13 +456,22 @@ impl ActiveMessageDispatcher for OrderedDispa {{reason=\"ordered_lane_shed\"}} for the rate" ); }); - fail_fast( - &system, - message_id, - response_type, - "ordered lane queue full", - shed.ctx.in_flight, - ); + // `dispatch` is sync, so the reply needs its own task; the + // tracker keeps it from being untracked. + let system = system.clone(); + let handler_name = self.handler.name().to_string(); + let in_flight = shed.ctx.in_flight; + system.tracker().clone().spawn(async move { + fail_fast( + &system, + &handler_name, + message_id, + response_type, + "ordered lane queue full", + in_flight, + ) + .await; + }); } } } @@ -550,15 +575,7 @@ impl DispatcherHub { let error_message = format!("Handler '{}' not found", handler_name); tokio::spawn(async move { let _in_flight = ctx.in_flight; - if let Err(e) = - Self::send_error_response_static(&backend, message_id, error_message).await - && !e.is::() - { - error!( - target: "crate::messenger::dispatcher", - "Failed to send error response for unknown handler: {}", e - ); - } + send_error_reply(&backend, &handler_name, message_id, error_message).await; }); } ResponseType::FireAndForget => { @@ -571,16 +588,6 @@ impl DispatcherHub { } } - /// Send an error response back to the sender. - /// Admission errors are already logged by the backend's error callback. - pub(crate) async fn send_error_response( - &self, - response_id: ResponseId, - error_message: String, - ) -> anyhow::Result<()> { - Self::send_error_response_static(&self.backend, response_id, error_message).await - } - /// Send an error response back to the sender (static method) async fn send_error_response_static( backend: &VeloBackend, diff --git a/lib/velo/src/messenger/server/mod.rs b/lib/velo/src/messenger/server/mod.rs index 014fa03b..1318f4e8 100644 --- a/lib/velo/src/messenger/server/mod.rs +++ b/lib/velo/src/messenger/server/mod.rs @@ -22,6 +22,7 @@ use crate::transports::{DataStreams, InboundMessage, ShutdownState, VeloBackend} use bytes::Bytes; use tokio_util::task::TaskTracker; +use dispatcher::send_error_reply; pub(crate) use dispatcher::{DispatcherHub, HandlerContext}; pub(crate) struct ActiveMessageServer { @@ -311,19 +312,14 @@ async fn create_message_handler( response_type, crate::messenger::common::messages::ResponseType::AckNack | crate::messenger::common::messages::ResponseType::Unary - ) && let Err(send_err) = hub - .send_error_response( + ) { + send_error_reply( + system.backend(), + &handler_name, message_id, format!("Failed to resolve large payload: {e}"), ) - .await - && !send_err.is::() - { - tracing::error!( - target: "crate::messenger::server", - handler = %handler_name, - "Failed to send error response: {send_err}" - ); + .await; } } } @@ -341,24 +337,18 @@ async fn create_message_handler( crate::messenger::common::messages::ResponseType::AckNack | crate::messenger::common::messages::ResponseType::Unary ) { - let hub = hub.clone(); + let backend = system.backend().clone(); + let handler_name = message.metadata.handler_name.clone(); let message_id = message.metadata.response_id; tokio::spawn(async move { let _in_flight = in_flight; - if let Err(e) = hub - .send_error_response( - message_id, - "Rendezvous resolver not configured on receiver" - .to_string(), - ) - .await - && !e.is::() - { - tracing::error!( - target: "crate::messenger::server", - "Failed to send error response for missing resolver: {e}" - ); - } + send_error_reply( + &backend, + &handler_name, + message_id, + "Rendezvous resolver not configured on receiver".to_string(), + ) + .await; }); } continue; From 2fb25beb83b6d5dfd55f121bab077e5489c63a25 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 00:09:23 -0600 Subject: [PATCH 08/29] refactor(ucx): tell connection epochs apart by the admission lock ConnHandle held a channel sender only so the pump could recognise its own epoch on cleanup. The handle is cloned on every send, so that cost a shared atomic inc and dec each time. The admission Arc is already in the handle and is unique per epoch, so compare that instead. Signed-off-by: Ryan Olson --- lib/velo/src/transports/ucx/transport.rs | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/lib/velo/src/transports/ucx/transport.rs b/lib/velo/src/transports/ucx/transport.rs index 76397a69..106161a1 100644 --- a/lib/velo/src/transports/ucx/transport.rs +++ b/lib/velo/src/transports/ucx/transport.rs @@ -130,7 +130,6 @@ impl Default for UcxConfig { /// gates admitted. An epoch is retired when the peer's endpoint fails. #[derive(Clone)] struct ConnHandle { - tx: flume::Sender, admission: Arc>, gate: AdmissionGate, _lifetime: Arc, @@ -269,8 +268,11 @@ impl UcxTransport { .entry(peer) .or_insert_with(|| { let (tx, rx) = flume::bounded::(self.config.channel_capacity); - let gate = AdmissionGate::new(tx.clone(), rt.clone()); - let epoch_tx = tx.clone(); + let gate = AdmissionGate::new(tx, rt.clone()); + // Identifies this epoch in cleanup; cheaper per send than a + // channel sender clone inside every `ConnHandle` copy. + let admission = Arc::new(Mutex::new(())); + let epoch_admission = admission.clone(); let connections = Arc::downgrade(&self.connections); let cancel = tokio_util::sync::CancellationToken::new(); let lifetime = Arc::new(cancel.clone().drop_guard()); @@ -324,7 +326,7 @@ impl UcxTransport { doorbell.ring(); } if let Some(connections) = connections.upgrade() { - if let Some((_, stale)) = connections.remove_if(&peer, |_, h| h.tx.same_channel(&epoch_tx)) { + if let Some((_, stale)) = connections.remove_if(&peer, |_, h| Arc::ptr_eq(&h.admission, &epoch_admission)) { stale.gate.fail_all(AdmissionError::ChannelClosed); } if let Some(metrics) = metrics.as_ref() { @@ -335,12 +337,11 @@ impl UcxTransport { } pending.fail_all(AdmissionError::ChannelClosed); drop(pending); - drop(epoch_tx); while let Ok(task) = rx.recv_async().await { task.fail("ucx connection closed"); } }); - ConnHandle { tx, admission: Arc::new(Mutex::new(())), gate, _lifetime: lifetime } + ConnHandle { admission, gate, _lifetime: lifetime } }) .clone(); if let Some(m) = self.shared.metrics.get() { From 2ca31ca0e32acb10c0cdd8661ce621e69906b314 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 01:11:28 -0600 Subject: [PATCH 09/29] fix(streaming): stop the mux before removing anchors at shutdown Velo::shutdown tears down the messenger transports before the anchor manager. Removing an anchor with a live mux stream closes its slot, and while the mux still ran that close was batched to the peer: TCP had no connection left, dialed, found itself cancelled, and failed the frame as a send error logged at ERROR. Stop the mux first, so the close finds no batcher and returns. The new test failed in 2 of 3 runs before the change and passed 20 of 20 after it. Signed-off-by: Ryan Olson --- lib/velo/src/lib.rs | 88 ++++++++++++++++++++++++++++++++ lib/velo/src/streaming/anchor.rs | 10 ++-- 2 files changed, 95 insertions(+), 3 deletions(-) diff --git a/lib/velo/src/lib.rs b/lib/velo/src/lib.rs index 8da496ee..85123478 100644 --- a/lib/velo/src/lib.rs +++ b/lib/velo/src/lib.rs @@ -1687,4 +1687,92 @@ mod tests { owned_transport.shutdown().await; retained.closed().await; } + + /// Shutdown must not write to peers after its transports are gone. + /// + /// `graceful_shutdown` tears down the messenger transports first. Removing + /// an anchor with a live mux stream then closes its slot, and while the mux + /// still runs that close is batched to the peer: the TCP transport has no + /// connection left, dials, finds itself cancelled, and fails the frame as + /// a send error that the default handler logs at ERROR. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn shutdown_with_live_mux_stream_sends_nothing_after_teardown() { + use crate::observability::test_helpers::MetricSnapshot; + use futures::StreamExt; + + async fn node(metrics: Option>) -> Arc { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let mut builder = Velo::builder() + .add_transport(Arc::new( + crate::transports::tcp::TcpTransportBuilder::new() + .from_listener(listener) + .unwrap() + .build() + .unwrap(), + )) + .stream_bind_addr(std::net::Ipv4Addr::LOCALHOST.into()); + if let Some(metrics) = metrics { + builder = builder.metrics(metrics); + } + builder.build().await.unwrap() + } + + let registry = prometheus::Registry::new(); + let metrics = Arc::new(VeloMetrics::register(®istry).unwrap()); + let consumer = node(Some(metrics)).await; + let producer = node(None).await; + consumer.register_peer(producer.peer_info()).unwrap(); + producer.register_peer(consumer.peer_info()).unwrap(); + for (node, peer) in [ + (&consumer, producer.instance_id()), + (&producer, consumer.instance_id()), + ] { + tokio::time::timeout( + std::time::Duration::from_secs(5), + node.wait_for_handler(peer, "_anchor_attach"), + ) + .await + .unwrap() + .unwrap(); + } + + let mut anchor = consumer.create_anchor::(); + let sender = producer + .attach_anchor::(anchor.handle()) + .await + .unwrap(); + sender.send(1).await.unwrap(); + let first = tokio::time::timeout(std::time::Duration::from_secs(5), anchor.next()) + .await + .unwrap(); + assert!(matches!( + first, + Some(Ok(crate::streaming::StreamFrame::Item(1))) + )); + + let send_errors = || { + MetricSnapshot::from_registry(®istry).counter( + "velo_transport_rejections_total", + &[("transport", "tcp"), ("reason", "send_error")], + ) + }; + let before = send_errors(); + tokio::time::timeout( + std::time::Duration::from_secs(10), + consumer.shutdown(ShutdownPolicy::WaitForever), + ) + .await + .expect("consumer shutdown hung"); + // A failed dial reports from the writer task after shutdown returns. + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + assert_eq!( + send_errors(), + before, + "shutdown sent a frame after its transport was torn down" + ); + + drop(sender); + drop(anchor); + producer.shutdown(ShutdownPolicy::WaitForever).await; + } } diff --git a/lib/velo/src/streaming/anchor.rs b/lib/velo/src/streaming/anchor.rs index 1af41071..81990b2b 100644 --- a/lib/velo/src/streaming/anchor.rs +++ b/lib/velo/src/streaming/anchor.rs @@ -1041,6 +1041,13 @@ impl AnchorManagerBuilder { impl AnchorManager { pub(crate) async fn shutdown(&self) { + // The mux stops first. `Velo::shutdown` has already torn down the + // messenger transports, and a live mux stream's anchor closes its slot + // as it is removed below; a running mux would batch that close to the + // peer over a transport with nothing left to send on. + if let Some(mux) = self.mux.get() { + mux.shutdown().await; + } // Remove entries outside shard guards: their Drop may close a mux slot. let ids: Vec<_> = self.registry.iter().map(|entry| *entry.key()).collect(); for id in ids { @@ -1065,9 +1072,6 @@ impl AnchorManager { for id in ids { self.sender_registry.cancel(id); } - if let Some(mux) = self.mux.get() { - mux.shutdown().await; - } } /// Convenience constructor with no default timeout. From 7ada25ad228d3a183f2a6397d09053b1ab96d7fc Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 01:11:53 -0600 Subject: [PATCH 10/29] docs(streaming): name the batcher's admission task as it is spawned now Signed-off-by: Ryan Olson --- lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs b/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs index 0c3061ac..7a87b0a3 100644 --- a/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs @@ -833,7 +833,7 @@ impl Batcher { /// (`docs/src/concepts/batched-streaming.md` § "Slots"), so nothing about the *sender's* admission /// order says anything about the order the receiver applies it in. /// - /// The `tokio::spawn` below always watches the admission — even an + /// The task spawned below always watches the admission — even an /// unfenced dispatch has to learn of a *failure*, which is epoch death /// whether or not the fence was ever raised. But it reports *success* to /// `singleton_resolved` only when this call actually fenced: an unfenced From a90f1f315d3e7fe50b6cdad8da95a1439c21be09 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 01:26:54 -0600 Subject: [PATCH 11/29] fix(streaming): do not enter a runtime during thread-local teardown A stream sender dropped from a thread-local destructor can outlive the thread's tokio context, and Handle::enter panics on a destroyed context; a panic there aborts the process. Enter the sender's runtime only when the thread has no runtime at all, through one helper both senders use. Otherwise the timer is not armed, as with no runtime. Signed-off-by: Ryan Olson --- lib/velo/src/streaming/mpsc/sender.rs | 4 +-- lib/velo/src/streaming/sender.rs | 4 +-- lib/velo/src/streaming/tasks.rs | 14 ++++++++ lib/velo/tests/streaming/mpsc_integration.rs | 35 ++++++++++++++++++++ 4 files changed, 51 insertions(+), 6 deletions(-) diff --git a/lib/velo/src/streaming/mpsc/sender.rs b/lib/velo/src/streaming/mpsc/sender.rs index 528a3211..fb50f3bb 100644 --- a/lib/velo/src/streaming/mpsc/sender.rs +++ b/lib/velo/src/streaming/mpsc/sender.rs @@ -302,9 +302,7 @@ impl Drop for SenderCleanup { self.sender_registry.senders.remove(&self.sender_stream_id); // The last local slot may start an unattached timeout. Drop can run // on a plain thread, so provide the runtime captured at construction. - let _entered = tokio::runtime::Handle::try_current() - .is_err() - .then(|| self.runtime.enter()); + let _entered = crate::streaming::tasks::enter_if_outside_runtime(&self.runtime); if self.local && let Some(slot) = super::anchor::remove_sender_slot( &self.mpsc_registry, diff --git a/lib/velo/src/streaming/sender.rs b/lib/velo/src/streaming/sender.rs index ae1ffb7b..209a7f1c 100644 --- a/lib/velo/src/streaming/sender.rs +++ b/lib/velo/src/streaming/sender.rs @@ -415,9 +415,7 @@ impl StreamSender { send_terminal(&self.runtime, &self.tx, bytes, move || { // The fast path runs this on the caller's thread, which may have // no runtime; the unattached timeout must still be armed. - let _runtime = tokio::runtime::Handle::try_current() - .is_err() - .then(|| runtime.enter()); + let _runtime = crate::streaming::tasks::enter_if_outside_runtime(&runtime); if let Some(registry) = registry && let Some(mut entry) = registry.get_mut(&local_id) && entry diff --git a/lib/velo/src/streaming/tasks.rs b/lib/velo/src/streaming/tasks.rs index 3f49cb2e..b835f37a 100644 --- a/lib/velo/src/streaming/tasks.rs +++ b/lib/velo/src/streaming/tasks.rs @@ -83,6 +83,20 @@ impl StreamTasks { } } +/// Enter `runtime` when the calling thread has none, so a drop or callback +/// on a plain thread can still spawn. Not when the thread's runtime context +/// is already destroyed (thread-local teardown): `Handle::enter` panics then, +/// and a panic there aborts the process. The caller's spawns then degrade +/// the way they do with no runtime at all. +pub(crate) fn enter_if_outside_runtime( + runtime: &tokio::runtime::Handle, +) -> Option> { + match tokio::runtime::Handle::try_current() { + Err(error) if error.is_missing_context() => Some(runtime.enter()), + _ => None, + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/lib/velo/tests/streaming/mpsc_integration.rs b/lib/velo/tests/streaming/mpsc_integration.rs index f0404370..0f2c6c18 100644 --- a/lib/velo/tests/streaming/mpsc_integration.rs +++ b/lib/velo/tests/streaming/mpsc_integration.rs @@ -418,6 +418,41 @@ async fn mpsc_detach_after_attaching_runtime_stops() { anchor.cancel(); } +/// A sender kept in a `thread_local!` (a language binding's thread) is +/// dropped while that thread's locals are torn down. Tokio's own context can +/// already be gone then, and `Handle::enter` panics on a destroyed context; a +/// panic in a thread-local destructor aborts the process. The drop must +/// degrade instead. +#[tokio::test(flavor = "multi_thread")] +async fn mpsc_sender_dropped_during_thread_local_teardown() { + use std::cell::RefCell; + + thread_local! { + static HELD: RefCell>> = + const { RefCell::new(None) }; + } + + let mgr = make_manager(); + let mut anchor = mgr.create_mpsc_anchor::(); + let sender = mgr + .attach_mpsc_stream_anchor::(anchor.handle()) + .await + .unwrap(); + let sid = sender.sender_id(); + std::thread::spawn(move || { + HELD.with(|held| *held.borrow_mut() = Some(sender)); + // Touch tokio's context after `HELD`, so its destructor is registered + // later and runs first. + assert!(tokio::runtime::Handle::try_current().is_err()); + }) + .join() + .expect("dropping the sender during thread-local teardown panicked"); + assert!( + matches!(next_frame(&mut anchor).await, Some(Ok((id, MpscFrame::Dropped(None)))) if id == sid) + ); + anchor.cancel(); +} + /// A pending `anchor.next()` must terminate promptly when the controller /// cancels, even if sender handles are still alive. #[tokio::test(flavor = "multi_thread")] From 2df6fb5f0df9c6f02963a65b52234a97cf4e6656 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 01:28:47 -0600 Subject: [PATCH 12/29] fix(messenger): start no reply task for a shed fire-and-forget message Moving the shed reply into a tracked task also moved the fire-and-forget check inside it, so every shed fire-and-forget message spawned a task and allocated the handler name just to return. Check before spawning: shedding is the overload path and must stay cheaper than accepting. Signed-off-by: Ryan Olson --- lib/velo/src/messenger/server/dispatcher.rs | 63 +++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/lib/velo/src/messenger/server/dispatcher.rs b/lib/velo/src/messenger/server/dispatcher.rs index 4927efe2..138b8ecc 100644 --- a/lib/velo/src/messenger/server/dispatcher.rs +++ b/lib/velo/src/messenger/server/dispatcher.rs @@ -456,6 +456,11 @@ impl ActiveMessageDispatcher for OrderedDispa {{reason=\"ordered_lane_shed\"}} for the rate" ); }); + // Fire-and-forget has no caller to tell, and the shed path + // must stay cheaper than accepting. + if matches!(response_type, ResponseType::FireAndForget) { + return; + } // `dispatch` is sync, so the reply needs its own task; the // tracker keeps it from being untracked. let system = system.clone(); @@ -705,4 +710,62 @@ mod tests { .await; messenger.closed().await; } + + struct PendingHandler; + + impl ActiveMessageHandler for PendingHandler { + fn handle(&self, _ctx: HandlerContext) -> BoxFuture<'static, ()> { + Box::pin(std::future::pending()) + } + + fn name(&self) -> &str { + "_pending" + } + } + + /// A shed fire-and-forget message has no caller to tell, so the shed path + /// must not start a reply task for it: shedding is the overload path, and + /// it must stay cheaper than accepting. The runtime is single-threaded and + /// nothing below yields, so a spawned reply task is still counted by the + /// tracker when it is read. + #[tokio::test(flavor = "current_thread")] + async fn shedding_fire_and_forget_spawns_no_reply_task() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let messenger = Messenger::builder() + .add_transport(Arc::new( + TcpTransportBuilder::new() + .from_listener(listener) + .unwrap() + .build() + .unwrap(), + )) + .build() + .await + .unwrap(); + let dispatcher = OrderedDispatcher::new( + PendingHandler, + OrderedConfig::global() + .with_max_queue_depth(Some(1)) + .with_overflow(crate::messenger::OverflowPolicy::Reject), + ); + let ctx = |n: u128| HandlerContext { + message_id: ResponseId::from_u128(n), + payload: Bytes::new(), + response_type: ResponseType::FireAndForget, + headers: None, + system: Arc::clone(&messenger), + in_flight: None, + }; + + dispatcher.dispatch(ctx(1)); + let tasks = messenger.tracker().len(); + for n in 2..10 { + dispatcher.dispatch(ctx(n)); + } + assert_eq!( + messenger.tracker().len(), + tasks, + "a shed fire-and-forget message started a task" + ); + } } From eaf265ceb0f3b2121597385a8cc49c01b8f34681 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 01:45:37 -0600 Subject: [PATCH 13/29] fix(streaming): withdraw stream feeds before the mux stops at shutdown Stopping the mux first kept shutdown from writing to torn-down transports, but stopping it retires every ingress slot by injecting Dropped. A consumer still reading through its direct feed took that as its sender's and ended with SenderDropped. Withdraw every anchor's feed first, as AnchorEntry::drop does before it closes a slot; the consumer then sees its stream end cleanly. Failed 5 of 5 runs before, passes 10 of 10 after. Signed-off-by: Ryan Olson --- lib/velo/src/lib.rs | 80 ++++++++++++++++++++++++++++++++ lib/velo/src/streaming/anchor.rs | 15 ++++-- 2 files changed, 91 insertions(+), 4 deletions(-) diff --git a/lib/velo/src/lib.rs b/lib/velo/src/lib.rs index 85123478..bf546ebf 100644 --- a/lib/velo/src/lib.rs +++ b/lib/velo/src/lib.rs @@ -1775,4 +1775,84 @@ mod tests { drop(anchor); producer.shutdown(ShutdownPolicy::WaitForever).await; } + + /// A consumer still reading when its node shuts down must see the stream + /// end, not `SenderDropped`: the sender did nothing wrong. + /// + /// Stopping the mux retires every ingress slot by injecting `Dropped` into + /// it. A consumer whose direct feed is still installed reads that record + /// as its sender's. So each anchor's feed must be withdrawn before the mux + /// stops, which is also what `AnchorEntry::drop` does before it closes a + /// slot, for the same reason. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn shutdown_ends_a_live_mux_stream_without_sender_dropped() { + use futures::StreamExt; + + async fn node() -> Arc { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + Velo::builder() + .add_transport(Arc::new( + crate::transports::tcp::TcpTransportBuilder::new() + .from_listener(listener) + .unwrap() + .build() + .unwrap(), + )) + .stream_bind_addr(std::net::Ipv4Addr::LOCALHOST.into()) + .build() + .await + .unwrap() + } + + let consumer = node().await; + let producer = node().await; + consumer.register_peer(producer.peer_info()).unwrap(); + producer.register_peer(consumer.peer_info()).unwrap(); + for (node, peer) in [ + (&consumer, producer.instance_id()), + (&producer, consumer.instance_id()), + ] { + tokio::time::timeout( + std::time::Duration::from_secs(5), + node.wait_for_handler(peer, "_anchor_attach"), + ) + .await + .unwrap() + .unwrap(); + } + + let mut anchor = consumer.create_anchor::(); + let sender = producer + .attach_anchor::(anchor.handle()) + .await + .unwrap(); + sender.send(1).await.unwrap(); + let first = tokio::time::timeout(std::time::Duration::from_secs(5), anchor.next()) + .await + .unwrap(); + assert!(matches!( + first, + Some(Ok(crate::streaming::StreamFrame::Item(1))) + )); + + let reader = tokio::spawn(async move { anchor.next().await }); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + tokio::time::timeout( + std::time::Duration::from_secs(10), + consumer.shutdown(ShutdownPolicy::WaitForever), + ) + .await + .expect("consumer shutdown hung"); + let ended = tokio::time::timeout(std::time::Duration::from_secs(5), reader) + .await + .expect("reader still waiting after shutdown") + .unwrap(); + assert!( + ended.is_none(), + "shutdown must end a live stream cleanly, got {ended:?}" + ); + + drop(sender); + producer.shutdown(ShutdownPolicy::WaitForever).await; + } } diff --git a/lib/velo/src/streaming/anchor.rs b/lib/velo/src/streaming/anchor.rs index 81990b2b..1b089439 100644 --- a/lib/velo/src/streaming/anchor.rs +++ b/lib/velo/src/streaming/anchor.rs @@ -1041,10 +1041,17 @@ impl AnchorManagerBuilder { impl AnchorManager { pub(crate) async fn shutdown(&self) { - // The mux stops first. `Velo::shutdown` has already torn down the - // messenger transports, and a live mux stream's anchor closes its slot - // as it is removed below; a running mux would batch that close to the - // peer over a transport with nothing left to send on. + // Feeds come out first: stopping the mux injects `Dropped` into every + // ingress slot, and a consumer still reading one would take it as its + // sender's and end with `SenderDropped`. Nothing here closes a slot. + for mut entry in self.registry.iter_mut() { + entry.retire_pump(); + } + // Then the mux stops, before anchors are removed. `Velo::shutdown` has + // already torn down the messenger transports, and a live mux stream's + // anchor closes its slot as it is removed below; a running mux would + // batch that close to the peer over a transport with nothing left to + // send on. if let Some(mux) = self.mux.get() { mux.shutdown().await; } From a7187b05d3d7bca805f4bd721883c2c1fa525d04 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 01:46:47 -0600 Subject: [PATCH 14/29] fix(streaming): arm the MPSC unattached timeout only on a live runtime When the last local MPSC sender leaves, its anchor re-arms the unattached timeout with a bare tokio::spawn. From a thread-local destructor, after tokio's context is gone, that panics, and a panic there aborts the process. Guard it with Handle::try_current like the SPSC timeout; with no runtime the timer does not fire. Signed-off-by: Ryan Olson --- lib/velo/src/streaming/mpsc/anchor.rs | 13 ++++++++++++- lib/velo/tests/streaming/mpsc_integration.rs | 7 ++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/lib/velo/src/streaming/mpsc/anchor.rs b/lib/velo/src/streaming/mpsc/anchor.rs index ef20ebd7..b6777b2f 100644 --- a/lib/velo/src/streaming/mpsc/anchor.rs +++ b/lib/velo/src/streaming/mpsc/anchor.rs @@ -335,8 +335,19 @@ pub(crate) fn spawn_mpsc_timeout_task( parent_cancel: &CancellationToken, ) -> CancellationToken { let tc = parent_cancel.child_token(); + // The last sender can leave from a plain thread, or from a thread-local + // destructor after tokio's own context is gone, where a bare + // `tokio::spawn` panics. The timer then does not fire, as for SPSC + // (`AnchorManager::spawn_timeout_task`). + let Ok(rt) = tokio::runtime::Handle::try_current() else { + tracing::debug!( + local_id, + "mpsc anchor: no runtime to arm the unattached timeout on; it will not fire" + ); + return tc; + }; let tc_clone = tc.clone(); - tokio::spawn(async move { + rt.spawn(async move { tokio::select! { _ = tc_clone.cancelled() => {} _ = tokio::time::sleep(timeout) => { diff --git a/lib/velo/tests/streaming/mpsc_integration.rs b/lib/velo/tests/streaming/mpsc_integration.rs index 0f2c6c18..b916aba0 100644 --- a/lib/velo/tests/streaming/mpsc_integration.rs +++ b/lib/velo/tests/streaming/mpsc_integration.rs @@ -432,8 +432,13 @@ async fn mpsc_sender_dropped_during_thread_local_teardown() { const { RefCell::new(None) }; } + // With an unattached timeout, the last sender leaving re-arms it, which + // spawns: that spawn must degrade the same way. let mgr = make_manager(); - let mut anchor = mgr.create_mpsc_anchor::(); + let mut anchor = mgr.create_mpsc_anchor_with_config::(MpscAnchorConfig { + unattached_timeout: Some(Duration::from_secs(60)), + ..Default::default() + }); let sender = mgr .attach_mpsc_stream_anchor::(anchor.handle()) .await From d41522f3f6ecae869e0e80d3374e9493f6932a59 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 02:02:13 -0600 Subject: [PATCH 15/29] fix(streaming): end live MPSC mux streams cleanly at shutdown Withdrawing SPSC feeds before the mux stops did not cover MPSC, whose slot records reach the consumer through a pump: the pump forwarded the Dropped that slot retirement injects, as the sender's own. Cancelling the pumps first is not enough on its own, since a stopped pump releases its slot, and a running mux would send that close over the torn-down transport. So shutdown stops the mux's tasks first (nothing more is written to a peer), then takes streams off their slots (SPSC feeds withdrawn, MPSC pumps cancelled), then retires the slots. Signed-off-by: Ryan Olson --- lib/velo/src/lib.rs | 95 +++++++++++++++++++ lib/velo/src/streaming/anchor.rs | 28 ++++-- .../src/streaming/messenger_mux/lifecycle.rs | 8 ++ 3 files changed, 122 insertions(+), 9 deletions(-) diff --git a/lib/velo/src/lib.rs b/lib/velo/src/lib.rs index bf546ebf..ec026636 100644 --- a/lib/velo/src/lib.rs +++ b/lib/velo/src/lib.rs @@ -1855,4 +1855,99 @@ mod tests { drop(sender); producer.shutdown(ShutdownPolicy::WaitForever).await; } + + /// The MPSC form of the test above. An MPSC slot's records reach the + /// consumer through a pump, so stopping the mux injects `Dropped` that a + /// live pump forwards as the sender's own `Dropped`. The pumps must stop + /// before the mux does. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn shutdown_ends_a_live_mpsc_mux_stream_without_dropped() { + use futures::StreamExt; + + use crate::observability::test_helpers::MetricSnapshot; + + async fn node(metrics: Option>) -> Arc { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let mut builder = Velo::builder() + .add_transport(Arc::new( + crate::transports::tcp::TcpTransportBuilder::new() + .from_listener(listener) + .unwrap() + .build() + .unwrap(), + )) + .stream_bind_addr(std::net::Ipv4Addr::LOCALHOST.into()); + if let Some(metrics) = metrics { + builder = builder.metrics(metrics); + } + builder.build().await.unwrap() + } + + let registry = prometheus::Registry::new(); + let metrics = Arc::new(VeloMetrics::register(®istry).unwrap()); + let consumer = node(Some(metrics)).await; + let producer = node(None).await; + consumer.register_peer(producer.peer_info()).unwrap(); + producer.register_peer(consumer.peer_info()).unwrap(); + for (node, peer) in [ + (&consumer, producer.instance_id()), + (&producer, consumer.instance_id()), + ] { + tokio::time::timeout( + std::time::Duration::from_secs(5), + node.wait_for_handler(peer, "_mpsc_anchor_attach"), + ) + .await + .unwrap() + .unwrap(); + } + + let mut anchor = consumer.create_mpsc_anchor::(); + let sender = producer + .attach_mpsc_anchor::(anchor.handle()) + .await + .unwrap(); + sender.send(1).await.unwrap(); + let first = tokio::time::timeout(std::time::Duration::from_secs(5), anchor.next()) + .await + .unwrap(); + assert!(matches!( + first, + Some(Ok((_, crate::streaming::mpsc::MpscFrame::Item(1)))) + )); + + let send_errors = || { + MetricSnapshot::from_registry(®istry).counter( + "velo_transport_rejections_total", + &[("transport", "tcp"), ("reason", "send_error")], + ) + }; + let before = send_errors(); + let reader = tokio::spawn(async move { anchor.next().await }); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + tokio::time::timeout( + std::time::Duration::from_secs(10), + consumer.shutdown(ShutdownPolicy::WaitForever), + ) + .await + .expect("consumer shutdown hung"); + let ended = tokio::time::timeout(std::time::Duration::from_secs(5), reader) + .await + .expect("reader still waiting after shutdown") + .unwrap(); + assert!( + ended.is_none(), + "shutdown must end a live MPSC stream cleanly, got {ended:?}" + ); + // A stopped pump releases its slot; that must not reach the wire. + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + assert_eq!( + send_errors(), + before, + "shutdown sent a frame after its transport was torn down" + ); + + drop(sender); + producer.shutdown(ShutdownPolicy::WaitForever).await; + } } diff --git a/lib/velo/src/streaming/anchor.rs b/lib/velo/src/streaming/anchor.rs index 1b089439..63f71cf0 100644 --- a/lib/velo/src/streaming/anchor.rs +++ b/lib/velo/src/streaming/anchor.rs @@ -1041,18 +1041,28 @@ impl AnchorManagerBuilder { impl AnchorManager { pub(crate) async fn shutdown(&self) { - // Feeds come out first: stopping the mux injects `Dropped` into every - // ingress slot, and a consumer still reading one would take it as its - // sender's and end with `SenderDropped`. Nothing here closes a slot. + // `Velo::shutdown` has already torn down the messenger transports, so + // the mux stops sending first. Every slot close below (an anchor's + // drop, a stopped MPSC pump's release) then finds no batcher, rather + // than writing to a transport with nothing left to send on. + let mux = self.mux.get(); + if let Some(mux) = mux { + mux.stop_sending(); + } + // Then streams come off their slots: retiring the slots injects + // `Dropped` into each, and a consumer still reading one, directly or + // through an MPSC pump, would take it as its sender's. for mut entry in self.registry.iter_mut() { entry.retire_pump(); } - // Then the mux stops, before anchors are removed. `Velo::shutdown` has - // already torn down the messenger transports, and a live mux stream's - // anchor closes its slot as it is removed below; a running mux would - // batch that close to the peer over a transport with nothing left to - // send on. - if let Some(mux) = self.mux.get() { + for entry in self.mpsc_registry.iter() { + for slot in entry.senders.values() { + if let Some(pump) = &slot.pump_token { + pump.cancel(); + } + } + } + if let Some(mux) = mux { mux.shutdown().await; } // Remove entries outside shard guards: their Drop may close a mux slot. diff --git a/lib/velo/src/streaming/messenger_mux/lifecycle.rs b/lib/velo/src/streaming/messenger_mux/lifecycle.rs index c5089f50..dbeecaa6 100644 --- a/lib/velo/src/streaming/messenger_mux/lifecycle.rs +++ b/lib/velo/src/streaming/messenger_mux/lifecycle.rs @@ -14,6 +14,14 @@ use super::{MessengerMuxTransport, MuxCore, PeerLane}; use crate::observability::MuxMetricsHandle; impl MessengerMuxTransport { + /// Stop the mux's tasks, so nothing more is written to any peer, but + /// leave its slots open: [`Self::shutdown`] retires them. Shutdown calls + /// this first so streams can be detached from their slots before + /// retirement injects `Dropped` into them. + pub(crate) fn stop_sending(&self) { + self.core.tasks.stop(); + } + pub(crate) async fn shutdown(&self) { self.core.tasks.stop(); self.core.tasks.wait().await; From 5c3f1bb6aac1c8a562982ae3bce64a5d05a052c8 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 02:02:26 -0600 Subject: [PATCH 16/29] docs(velo): describe what Velo::shutdown leaves in memory plainly Signed-off-by: Ryan Olson --- README.md | 2 +- lib/velo/src/lib.rs | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 787b229d..20a8a710 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ NOTE: Velo is experimental. Its design, development, and tests are still in prog Typed streams support graceful stop, cancellation, and ordered finalization, including prebound tickets. See [Stream lifecycle](docs/src/concepts/streaming.md). -Use `Velo::shutdown(policy)` to close resources and join owned tasks while Tokio keeps running. This release retains the legacy ownership graph in memory. Accepted requests remain counted until their replies are submitted, including error replies. See [Shutdown and drain](docs/src/concepts/shutdown.md). +Use `Velo::shutdown(policy)` to close resources and join owned tasks while Tokio keeps running. The instance and its handles stay valid afterwards; their memory is released when the last handle is dropped. Accepted requests remain counted until their replies are submitted, including error replies. See [Shutdown and drain](docs/src/concepts/shutdown.md). Batched streams give each stream a window of 32 records in flight by default. A larger window lets a saturated consumer node queue a new stream's first record behind every other stream's backlog. See [Change the credit window](docs/src/guides/tune-batched-streaming.md#change-the-credit-window). diff --git a/lib/velo/src/lib.rs b/lib/velo/src/lib.rs index ec026636..4b0f64b1 100644 --- a/lib/velo/src/lib.rs +++ b/lib/velo/src/lib.rs @@ -801,7 +801,8 @@ impl Velo { /// remain the caller's responsibility. Stream watchdogs and heartbeats are /// cancelled; application handlers that exceed `policy` can still be running. /// Use this when an instance is removed while its Tokio runtime stays alive. - /// This closes resources but retains the legacy ownership graph in memory. + /// It closes resources, but the instance and its handles stay valid; their + /// memory is released when the last handle is dropped. pub async fn shutdown(&self, policy: ShutdownPolicy) { self.graceful_shutdown(policy).await; self.anchor_manager.shutdown().await; From 31fc99a2a7d97f7e46526875dade0b7f707ad9f4 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 02:12:20 -0600 Subject: [PATCH 17/29] docs(velo): say what shutdown does to a stream with a local sender Shutdown cancels a sender on the same instance: its cancellation token fires and later sends fail. The reader ends when the application drops or finalizes that sender; shutdown does not end the reader itself, because that needs a check on every read. A test pins this for SPSC and MPSC streams. Signed-off-by: Ryan Olson --- README.md | 2 +- docs/src/concepts/shutdown.md | 2 + lib/velo/src/lib.rs | 75 ++++++++++++++++++++++++++++++++++- 3 files changed, 77 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 20a8a710..11296502 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ NOTE: Velo is experimental. Its design, development, and tests are still in prog Typed streams support graceful stop, cancellation, and ordered finalization, including prebound tickets. See [Stream lifecycle](docs/src/concepts/streaming.md). -Use `Velo::shutdown(policy)` to close resources and join owned tasks while Tokio keeps running. The instance and its handles stay valid afterwards; their memory is released when the last handle is dropped. Accepted requests remain counted until their replies are submitted, including error replies. See [Shutdown and drain](docs/src/concepts/shutdown.md). +Use `Velo::shutdown(policy)` to close resources and join owned tasks while Tokio keeps running. The instance and its handles stay valid after shutdown, and their memory is released when the last handle is dropped. Accepted requests remain counted until their replies are submitted, including error replies. See [Shutdown and drain](docs/src/concepts/shutdown.md). Batched streams give each stream a window of 32 records in flight by default. A larger window lets a saturated consumer node queue a new stream's first record behind every other stream's backlog. See [Change the credit window](docs/src/guides/tune-batched-streaming.md#change-the-credit-window). diff --git a/docs/src/concepts/shutdown.md b/docs/src/concepts/shutdown.md index a99db6f0..dc9de4bc 100644 --- a/docs/src/concepts/shutdown.md +++ b/docs/src/concepts/shutdown.md @@ -44,6 +44,8 @@ To let open streams finish, call `begin_drain`, wait until your streams end, and Use `Velo::shutdown(policy)` when an application removes a Velo instance but keeps its Tokio runtime. It first runs `graceful_shutdown`, then cancels live anchors and senders, stops the builder-owned per-stream listener, and joins the messenger receive loops and streaming tasks. Pending remote event waits fail at teardown, and their subscription tasks stop. Local event completion remains available. Custom frame transports remain the caller's responsibility. +If the sender of a stream is on the same instance, shutdown cancels that sender. The `cancellation_token` of the sender fires, and later sends fail. The reader of the stream ends when the application drops or finalizes the sender. Shutdown does not end the reader itself. That needs a check on every read. + Call shutdown explicitly. Dropping a `Velo` value or its last `Arc` does not perform a graceful drain. This release keeps the existing strong references between the runtime and its services. A retained Messenger keeps streaming services available after Velo is dropped. Explicit shutdown closes resources and joins owned tasks, but does not release this ownership graph from memory. diff --git a/lib/velo/src/lib.rs b/lib/velo/src/lib.rs index 4b0f64b1..627261cb 100644 --- a/lib/velo/src/lib.rs +++ b/lib/velo/src/lib.rs @@ -800,8 +800,11 @@ impl Velo { /// listener and pump tasks owned by the builder. Custom frame transports /// remain the caller's responsibility. Stream watchdogs and heartbeats are /// cancelled; application handlers that exceed `policy` can still be running. + /// If the sender of a stream is on this instance, shutdown cancels that + /// sender: its `cancellation_token` fires, and later sends fail. The reader + /// ends when the application drops or finalizes the sender. /// Use this when an instance is removed while its Tokio runtime stays alive. - /// It closes resources, but the instance and its handles stay valid; their + /// It closes resources. The instance and its handles stay valid, and their /// memory is released when the last handle is dropped. pub async fn shutdown(&self, policy: ShutdownPolicy) { self.graceful_shutdown(policy).await; @@ -1951,4 +1954,74 @@ mod tests { drop(sender); producer.shutdown(ShutdownPolicy::WaitForever).await; } + + /// What `Velo::shutdown` does to a stream whose sender is on this node + /// and still held by the application: the sender is cancelled (its + /// `cancellation_token` fires and later sends fail), and its reader ends + /// when the application drops it. Shutdown does not end the reader itself: + /// a consumer poll that watched for it would cost every read a check. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn shutdown_cancels_a_held_local_sender_and_its_reader_ends_on_drop() { + use futures::StreamExt; + + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let velo = Velo::builder() + .add_transport(Arc::new( + crate::transports::tcp::TcpTransportBuilder::new() + .from_listener(listener) + .unwrap() + .build() + .unwrap(), + )) + .build() + .await + .unwrap(); + + let mut anchor = velo.create_anchor::(); + let sender = velo.attach_anchor::(anchor.handle()).await.unwrap(); + let mut mpsc_anchor = velo.create_mpsc_anchor::(); + let mpsc_sender = velo + .attach_mpsc_anchor::(mpsc_anchor.handle()) + .await + .unwrap(); + let reader = tokio::spawn(async move { anchor.next().await }); + let mpsc_reader = tokio::spawn(async move { mpsc_anchor.next().await }); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + + tokio::time::timeout( + std::time::Duration::from_secs(10), + velo.shutdown(ShutdownPolicy::WaitForever), + ) + .await + .expect("shutdown hung"); + assert!(sender.cancellation_token().is_cancelled()); + assert!(mpsc_sender.cancellation_token().is_cancelled()); + assert!(sender.send(1).await.is_err()); + assert!(mpsc_sender.send(1).await.is_err()); + + drop(sender); + drop(mpsc_sender); + let ended = tokio::time::timeout(std::time::Duration::from_secs(2), reader) + .await + .expect("SPSC reader still waiting after its sender was dropped") + .unwrap(); + assert!( + matches!( + ended, + Some(Err(crate::streaming::StreamError::SenderDropped)) + ), + "got {ended:?}" + ); + let ended = tokio::time::timeout(std::time::Duration::from_secs(2), mpsc_reader) + .await + .expect("MPSC reader still waiting after its sender was dropped") + .unwrap(); + assert!( + matches!( + ended, + Some(Ok((_, crate::streaming::mpsc::MpscFrame::Dropped(_)))) + ), + "got {ended:?}" + ); + } } From 9415e19042083bbeb4d2994e1b7dd54db5d5736d Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 02:25:53 -0600 Subject: [PATCH 18/29] fix(streaming): retire mux slots only once shutdown has detached streams Shutdown stops the mux's tasks, takes every stream off its slot, then retires the slots. But a batcher parked at an await outside its cancel arm is dropped by its task wrapper when the stop lands, and its Drop retired every slot at once, so a reader still on its slot could take the injected Dropped as its sender's. Only the first stop of a failed mux retires slots now; shutdown and MuxCore::drop retire them explicitly. deliver_batch re-closes a late claim once the slots are retired, keyed on a flag shutdown sets before retiring, not on the tasks having stopped. Signed-off-by: Ryan Olson --- .../src/streaming/messenger_mux/lifecycle.rs | 30 ++++++++++-- lib/velo/src/streaming/messenger_mux/mod.rs | 17 +++++-- lib/velo/src/streaming/messenger_mux/tests.rs | 47 +++++++++++++++++++ 3 files changed, 85 insertions(+), 9 deletions(-) diff --git a/lib/velo/src/streaming/messenger_mux/lifecycle.rs b/lib/velo/src/streaming/messenger_mux/lifecycle.rs index dbeecaa6..33cfc229 100644 --- a/lib/velo/src/streaming/messenger_mux/lifecycle.rs +++ b/lib/velo/src/streaming/messenger_mux/lifecycle.rs @@ -5,6 +5,7 @@ //! failure and shutdown goes through. use std::sync::Arc; +use std::sync::atomic::Ordering; use anyhow::{Result, anyhow}; @@ -26,7 +27,7 @@ impl MessengerMuxTransport { self.core.tasks.stop(); self.core.tasks.wait().await; self.core.batchers.clear(); - close_ingress(&self.core.ingress, self.core.metrics.as_ref()); + self.core.retire_slots(); self.core.drains.clear(); } } @@ -80,15 +81,31 @@ impl MuxCore { } } -/// Stop the mux's tasks and close every ingress slot. Returns true only for -/// the first stop request. +impl MuxCore { + /// Close every ingress slot, injecting `Dropped` into each. The flag goes + /// first; `deliver_batch` reads it after a claim (see there). + fn retire_slots(&self) { + self.slots_retired.store(true, Ordering::SeqCst); + close_ingress(&self.ingress, self.metrics.as_ref()); + } +} + +/// Stop a mux that failed: stop its tasks and close every ingress slot. +/// Returns true only for the first stop request. +/// +/// A stop that was not the first leaves the slots alone. Shutdown stops the +/// tasks first and retires the slots only after it has taken every stream off +/// them; a batcher its stop drops mid-wake ends here, and closing the slots +/// now would hand a live reader `Dropped` as its sender's. pub(super) fn stop_mux( tasks: &crate::streaming::tasks::StreamTasks, ingress: &IngressRegistry, metrics: Option<&MuxMetricsHandle>, ) -> bool { let first = tasks.stop(); - close_ingress(ingress, metrics); + if first { + close_ingress(ingress, metrics); + } first } @@ -103,6 +120,9 @@ pub(super) fn close_ingress(ingress: &IngressRegistry, metrics: Option<&MuxMetri impl Drop for MuxCore { fn drop(&mut self) { - stop_mux(&self.tasks, &self.ingress, self.metrics.as_ref()); + // Unconditional: a shutdown abandoned after `stop_sending` left the + // slots open. + self.tasks.stop(); + self.retire_slots(); } } diff --git a/lib/velo/src/streaming/messenger_mux/mod.rs b/lib/velo/src/streaming/messenger_mux/mod.rs index 23648477..c7ed813c 100644 --- a/lib/velo/src/streaming/messenger_mux/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/mod.rs @@ -116,7 +116,7 @@ mod tests; use std::collections::VecDeque; use std::sync::Arc; -use std::sync::atomic::AtomicU64; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::time::Duration; use anyhow::{Result, anyhow}; @@ -225,6 +225,9 @@ struct MuxCore { /// batch of the new one as stale and discard it wholesale. epochs: Arc, tasks: crate::streaming::tasks::StreamTasks, + /// Set once the slots are retired, which is later than the tasks stop: + /// shutdown takes streams off their slots in between. + slots_retired: AtomicBool, /// Peers with credit to return, posted by draining consumers. See /// [`ingress::DrainSignal`]. drain_tx: flume::Sender, @@ -350,6 +353,7 @@ impl MessengerMuxTransport { // zeroed header from reading as a legitimate one. epochs: Arc::new(AtomicU64::new(1)), tasks, + slots_retired: AtomicBool::new(false), drain_tx, drain_rx, drains: DashMap::new(), @@ -422,12 +426,17 @@ impl MuxCore { } } - // A batch already in flight can claim a bind after shutdown visited - // its peer. Retire that late claim before returning from this handler. - if self.tasks.is_stopped() { + // A batch already in flight can claim a bind after shutdown retired + // the slots. Retire that late claim before returning from this + // handler. Shutdown sets the flag before it retires, so a claim this + // check misses is one that retirement still covers. + if self.slots_retired.load(Ordering::SeqCst) { close_ingress(&self.ingress, self.metrics.as_ref()); return; } + if self.tasks.is_stopped() { + return; + } if outcome.replies.is_empty() && outcome.grants.is_empty() diff --git a/lib/velo/src/streaming/messenger_mux/tests.rs b/lib/velo/src/streaming/messenger_mux/tests.rs index 0606b37d..f197fdbf 100644 --- a/lib/velo/src/streaming/messenger_mux/tests.rs +++ b/lib/velo/src/streaming/messenger_mux/tests.rs @@ -293,6 +293,53 @@ fn aborting_a_retiring_batchers_final_flush_stops_the_mux() { assert!(receiver.is_disconnected()); } +/// Once shutdown has stopped the mux's tasks, the slots stay open until +/// `shutdown` retires them, even when the stop drops a batcher mid-wake. +/// +/// `AnchorManager::shutdown` stops the mux, then takes every stream off its +/// slot, then retires the slots: retirement injects `Dropped`, and a reader +/// still on its slot would take that as its sender's. A batcher parked at an +/// await outside its own cancel arm is dropped by its task wrapper when the +/// stop lands, and its `Drop` used to retire every slot on the spot, inside +/// that window. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_batcher_dropped_by_stop_sending_leaves_slots_to_shutdown() { + use super::test_support::{StallingTransport, stalling_address}; + + let (transport, _wire) = StallingTransport::new(tokio::runtime::Handle::current()); + let messenger = Messenger::builder() + .add_transport(transport) + .build() + .await + .unwrap(); + let peer = velo_ext::InstanceId::new_v4(); + messenger + .register_peer(velo_ext::PeerInfo::new(peer, stalling_address())) + .unwrap(); + let mux = MessengerMuxTransport::new(messenger, MuxConfig::default(), None).unwrap(); + let hooks = Arc::new(TestHooks::default()); + assert!(mux.core.hooks.set(hooks.clone()).is_ok()); + let receiver = mux.bind(1, 1).await.unwrap(); + let batcher = mux.core.batcher(lane0(peer.worker_id())).unwrap(); + + hooks.pause(); + assert!(batcher.reply(&[peer_batcher::ReplyRecord::CloseSlot { + slot: protocol::SlotId::from_raw(7), + reason: protocol::CloseReason::UnknownSlot, + }])); + hooks.wait_until_parked().await; + mux.stop_sending(); + eventually(|| batcher.is_closed()).await; + assert!( + !receiver.is_disconnected(), + "a batcher dropped by the stop retired the slots before shutdown did" + ); + + hooks.release(); + mux.shutdown().await; + assert!(receiver.is_disconnected()); +} + #[test] fn a_secondary_runtime_does_not_own_the_mux_batchers() { let owner = tokio::runtime::Builder::new_multi_thread() From a03a120802e952df9d0bb58a008643648b97b516 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 02:30:14 -0600 Subject: [PATCH 19/29] fix(streaming): finish every shutdown state change before its first await AnchorManager::shutdown withdrew feeds before it awaited the mux's tasks but removed anchors only after. A caller that drops the future during that await (a timeout around Velo::shutdown) left anchors with no feed in the registry, and their readers waiting. Remove anchors, cancel MPSC entries and local senders before the await; the mux is already stopped, so none of it sends. Document the phases and the two rules they keep, and update the mux lifecycle module doc to match. Signed-off-by: Ryan Olson --- lib/velo/src/lib.rs | 36 +++++++++++++++++++ lib/velo/src/streaming/anchor.rs | 31 ++++++++++------ .../src/streaming/messenger_mux/lifecycle.rs | 8 +++-- 3 files changed, 63 insertions(+), 12 deletions(-) diff --git a/lib/velo/src/lib.rs b/lib/velo/src/lib.rs index 627261cb..dd5d0a39 100644 --- a/lib/velo/src/lib.rs +++ b/lib/velo/src/lib.rs @@ -2024,4 +2024,40 @@ mod tests { "got {ended:?}" ); } + + /// Every state change in shutdown happens before its first await. + /// + /// A caller can bound `Velo::shutdown` with a timeout and drop it while it + /// waits for the mux's tasks. Streams must already be off their slots and + /// their anchors removed by then; a dropped future that had withdrawn the + /// feeds but not removed the anchors would leave their readers waiting on + /// nothing. The runtime is single-threaded, so the mux's own tasks have + /// not run and its join is still pending at the first poll. + #[tokio::test(flavor = "current_thread")] + async fn abandoning_anchor_shutdown_still_removes_every_anchor() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let velo = Velo::builder() + .add_transport(Arc::new( + crate::transports::tcp::TcpTransportBuilder::new() + .from_listener(listener) + .unwrap() + .build() + .unwrap(), + )) + .build() + .await + .unwrap(); + let _anchor = velo.create_anchor::(); + let _mpsc_anchor = velo.create_mpsc_anchor::(); + { + let shutdown = velo.anchor_manager.shutdown(); + futures::pin_mut!(shutdown); + assert!( + futures::poll!(shutdown.as_mut()).is_pending(), + "the test needs shutdown to wait at its first poll" + ); + } + assert!(velo.anchor_manager.registry.is_empty()); + assert!(velo.anchor_manager.mpsc_registry.is_empty()); + } } diff --git a/lib/velo/src/streaming/anchor.rs b/lib/velo/src/streaming/anchor.rs index 63f71cf0..91682646 100644 --- a/lib/velo/src/streaming/anchor.rs +++ b/lib/velo/src/streaming/anchor.rs @@ -1040,18 +1040,29 @@ impl AnchorManagerBuilder { } impl AnchorManager { + /// End every stream this manager holds. Called by `Velo::shutdown` after + /// the messenger transports are torn down. + /// + /// Two rules hold throughout: nothing is sent to a peer, because the + /// transports are gone, and no mux slot is retired while a stream still + /// reads from it, because retirement injects `Dropped` that the reader + /// would take as its sender's. Hence the phases: + /// + /// 1. Stop the mux's tasks. A slot close after this finds no batcher and + /// sends nothing. The slots stay open. + /// 2. Take streams off their slots: SPSC feeds withdrawn, MPSC pumps + /// cancelled. + /// 3. Remove anchors and MPSC entries, and cancel local senders. + /// 4. Join the mux's tasks and retire its slots. + /// + /// Steps 1-3 do not await, so a caller that drops this future (a timeout + /// around `Velo::shutdown`) still leaves no stream waiting on a slot or an + /// anchor; `MuxCore::drop` retires the slots then. pub(crate) async fn shutdown(&self) { - // `Velo::shutdown` has already torn down the messenger transports, so - // the mux stops sending first. Every slot close below (an anchor's - // drop, a stopped MPSC pump's release) then finds no batcher, rather - // than writing to a transport with nothing left to send on. let mux = self.mux.get(); if let Some(mux) = mux { mux.stop_sending(); } - // Then streams come off their slots: retiring the slots injects - // `Dropped` into each, and a consumer still reading one, directly or - // through an MPSC pump, would take it as its sender's. for mut entry in self.registry.iter_mut() { entry.retire_pump(); } @@ -1062,9 +1073,6 @@ impl AnchorManager { } } } - if let Some(mux) = mux { - mux.shutdown().await; - } // Remove entries outside shard guards: their Drop may close a mux slot. let ids: Vec<_> = self.registry.iter().map(|entry| *entry.key()).collect(); for id in ids { @@ -1089,6 +1097,9 @@ impl AnchorManager { for id in ids { self.sender_registry.cancel(id); } + if let Some(mux) = mux { + mux.shutdown().await; + } } /// Convenience constructor with no default timeout. diff --git a/lib/velo/src/streaming/messenger_mux/lifecycle.rs b/lib/velo/src/streaming/messenger_mux/lifecycle.rs index 33cfc229..19920c72 100644 --- a/lib/velo/src/streaming/messenger_mux/lifecycle.rs +++ b/lib/velo/src/streaming/messenger_mux/lifecycle.rs @@ -1,8 +1,12 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! How the mux starts batchers and how it stops: the one stop path every -//! failure and shutdown goes through. +//! How the mux starts batchers and how it stops. +//! +//! A mux that fails stops in one step: `stop_mux` stops its tasks and retires +//! its slots. Shutdown uses two steps, so that streams can come off their +//! slots in between: `stop_sending` stops the tasks, and `shutdown` retires +//! the slots. `MuxCore::drop` retires them if a shutdown was abandoned. use std::sync::Arc; use std::sync::atomic::Ordering; From 4591c3b464bdd015d1606292f42793a6ca893f98 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Sun, 4 Oct 2026 02:41:58 -0600 Subject: [PATCH 20/29] docs(velo): do not promise that shutdown memory is freed on drop The messenger, its server and its dispatcher hold each other, so dropping the last handle does not free the instance. The book already says shutdown does not release the ownership graph; the README and the Velo::shutdown rustdoc now agree with it. Signed-off-by: Ryan Olson --- README.md | 2 +- lib/velo/src/lib.rs | 4 ++-- lib/velo/src/streaming/anchor.rs | 4 ++-- lib/velo/src/streaming/messenger_mux/lifecycle.rs | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 11296502..16c4a331 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ NOTE: Velo is experimental. Its design, development, and tests are still in prog Typed streams support graceful stop, cancellation, and ordered finalization, including prebound tickets. See [Stream lifecycle](docs/src/concepts/streaming.md). -Use `Velo::shutdown(policy)` to close resources and join owned tasks while Tokio keeps running. The instance and its handles stay valid after shutdown, and their memory is released when the last handle is dropped. Accepted requests remain counted until their replies are submitted, including error replies. See [Shutdown and drain](docs/src/concepts/shutdown.md). +Use `Velo::shutdown(policy)` to close resources and join owned tasks while Tokio keeps running. The instance and its handles stay valid after shutdown, and shutdown does not release them from memory. Accepted requests remain counted until their replies are submitted, including error replies. See [Shutdown and drain](docs/src/concepts/shutdown.md). Batched streams give each stream a window of 32 records in flight by default. A larger window lets a saturated consumer node queue a new stream's first record behind every other stream's backlog. See [Change the credit window](docs/src/guides/tune-batched-streaming.md#change-the-credit-window). diff --git a/lib/velo/src/lib.rs b/lib/velo/src/lib.rs index dd5d0a39..2cf4332f 100644 --- a/lib/velo/src/lib.rs +++ b/lib/velo/src/lib.rs @@ -804,8 +804,8 @@ impl Velo { /// sender: its `cancellation_token` fires, and later sends fail. The reader /// ends when the application drops or finalizes the sender. /// Use this when an instance is removed while its Tokio runtime stays alive. - /// It closes resources. The instance and its handles stay valid, and their - /// memory is released when the last handle is dropped. + /// It closes resources. The instance and its handles stay valid, and + /// shutdown does not release them from memory. pub async fn shutdown(&self, policy: ShutdownPolicy) { self.graceful_shutdown(policy).await; self.anchor_manager.shutdown().await; diff --git a/lib/velo/src/streaming/anchor.rs b/lib/velo/src/streaming/anchor.rs index 91682646..8125a98c 100644 --- a/lib/velo/src/streaming/anchor.rs +++ b/lib/velo/src/streaming/anchor.rs @@ -1055,9 +1055,9 @@ impl AnchorManager { /// 3. Remove anchors and MPSC entries, and cancel local senders. /// 4. Join the mux's tasks and retire its slots. /// - /// Steps 1-3 do not await, so a caller that drops this future (a timeout + /// Steps 1-3 do not await. A caller that drops this future (a timeout /// around `Velo::shutdown`) still leaves no stream waiting on a slot or an - /// anchor; `MuxCore::drop` retires the slots then. + /// anchor. The slots then stay open, with no reader, until the mux drops. pub(crate) async fn shutdown(&self) { let mux = self.mux.get(); if let Some(mux) = mux { diff --git a/lib/velo/src/streaming/messenger_mux/lifecycle.rs b/lib/velo/src/streaming/messenger_mux/lifecycle.rs index 19920c72..d42412c2 100644 --- a/lib/velo/src/streaming/messenger_mux/lifecycle.rs +++ b/lib/velo/src/streaming/messenger_mux/lifecycle.rs @@ -6,7 +6,7 @@ //! A mux that fails stops in one step: `stop_mux` stops its tasks and retires //! its slots. Shutdown uses two steps, so that streams can come off their //! slots in between: `stop_sending` stops the tasks, and `shutdown` retires -//! the slots. `MuxCore::drop` retires them if a shutdown was abandoned. +//! the slots. If a shutdown was abandoned, `MuxCore::drop` retires them. use std::sync::Arc; use std::sync::atomic::Ordering; From 6dbb88f372b24a5f6bb84b300e5e001c8c719adc Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Mon, 5 Oct 2026 08:15:00 -0600 Subject: [PATCH 21/29] fix(streaming): clear mux binds before retiring slots IngressRegistry::shutdown walked the slot tables and only then cleared the binds. A batch handler claiming a bind between the two put a live slot in a table already walked, and that slot was never retired; its reader waited forever. The last change made this reachable on the failed-mux path too, where deliver_batch no longer re-closed the claim. Clear binds first. A claim takes its bind under the peer table's lock, in a table it inserted before the claim, so it either finds no bind or leaves a slot the walk still retires. The slots_retired flag and the late-claim re-close in deliver_batch are no longer needed. Signed-off-by: Ryan Olson --- .../streaming/messenger_mux/ingress/mod.rs | 19 ++++++++++++- .../streaming/messenger_mux/ingress/tests.rs | 28 +++++++++++++++++++ .../src/streaming/messenger_mux/lifecycle.rs | 16 ++--------- lib/velo/src/streaming/messenger_mux/mod.rs | 18 +++--------- 4 files changed, 53 insertions(+), 28 deletions(-) diff --git a/lib/velo/src/streaming/messenger_mux/ingress/mod.rs b/lib/velo/src/streaming/messenger_mux/ingress/mod.rs index 954642bb..db620869 100644 --- a/lib/velo/src/streaming/messenger_mux/ingress/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/ingress/mod.rs @@ -98,6 +98,9 @@ impl Drop for BindEntry { } } +#[cfg(test)] +pub(crate) type ShutdownHook = Box; + /// Registry of binds and per-peer slot tables. #[derive(Default)] pub(crate) struct IngressRegistry { @@ -110,6 +113,10 @@ pub(crate) struct IngressRegistry { /// Calls into `close_consumer_gone`, each of which takes a peer's lock. #[cfg(test)] consumer_gone_calls: std::sync::atomic::AtomicUsize, + /// Runs inside [`Self::shutdown`], between its two steps, so a test can + /// land a claim there. + #[cfg(test)] + pub(crate) shutdown_hook: std::sync::OnceLock, /// Per-table "a credit-return visit is already queued" flags, read and set /// by draining consumers without taking the table's mutex. See /// [`DrainSignal`]. @@ -445,13 +452,23 @@ impl IngressRegistry { /// /// Used when the transport itself goes away, so a consumer never waits out /// its heartbeat watchdog for a sender that has already been dismantled. + /// + /// Binds go first. A batch handler can be claiming one concurrently, and + /// it claims under its peer table's lock, in a table it inserted before + /// the claim. So a claim either finds no bind, or put its slot in a table + /// the walk below has yet to lock. The other order let a claim land in a + /// table already walked, and that slot was never retired. pub(crate) fn shutdown(&self) -> usize { + self.binds.clear(); + #[cfg(test)] + if let Some(hook) = self.shutdown_hook.get() { + hook(self); + } let mut closed = 0; for entry in self.peers.iter() { let mut state = lock(entry.value()); closed += retire_epoch(&mut state, None); } - self.binds.clear(); closed } } diff --git a/lib/velo/src/streaming/messenger_mux/ingress/tests.rs b/lib/velo/src/streaming/messenger_mux/ingress/tests.rs index 75b6dbce..24eabafa 100644 --- a/lib/velo/src/streaming/messenger_mux/ingress/tests.rs +++ b/lib/velo/src/streaming/messenger_mux/ingress/tests.rs @@ -2056,3 +2056,31 @@ fn the_lanes_of_one_peer_share_one_byte_budget() { "a retired epoch's holds must go back to the shared budget" ); } + +/// A slot claimed while the registry shuts down is retired with the rest. +/// +/// A batch handler can be inside `handle_batch` while the mux fails or shuts +/// down. Shutdown walked the slot tables and only then cleared the binds, so +/// an `OpenSlot` landing between the two claimed a bind into a table already +/// walked: a live slot nobody would retire, and a reader that waited on it +/// forever. The hook lands the claim at exactly that point. +#[test] +fn a_claim_during_shutdown_is_retired_or_refused() { + let (registry, consumer, config) = bound(); + let claim_config = config.clone(); + assert!( + registry + .shutdown_hook + .set(Box::new(move |registry| { + open(registry, &claim_config, SlotId::new(0, 0).unwrap(), 1); + })) + .is_ok() + ); + registry.shutdown(); + assert_eq!( + registry.live_slots(peer()), + 0, + "a claimed slot outlived shutdown" + ); + assert!(consumer.rx.is_disconnected()); +} diff --git a/lib/velo/src/streaming/messenger_mux/lifecycle.rs b/lib/velo/src/streaming/messenger_mux/lifecycle.rs index d42412c2..c5715d18 100644 --- a/lib/velo/src/streaming/messenger_mux/lifecycle.rs +++ b/lib/velo/src/streaming/messenger_mux/lifecycle.rs @@ -9,7 +9,6 @@ //! the slots. If a shutdown was abandoned, `MuxCore::drop` retires them. use std::sync::Arc; -use std::sync::atomic::Ordering; use anyhow::{Result, anyhow}; @@ -31,7 +30,7 @@ impl MessengerMuxTransport { self.core.tasks.stop(); self.core.tasks.wait().await; self.core.batchers.clear(); - self.core.retire_slots(); + close_ingress(&self.core.ingress, self.core.metrics.as_ref()); self.core.drains.clear(); } } @@ -85,15 +84,6 @@ impl MuxCore { } } -impl MuxCore { - /// Close every ingress slot, injecting `Dropped` into each. The flag goes - /// first; `deliver_batch` reads it after a claim (see there). - fn retire_slots(&self) { - self.slots_retired.store(true, Ordering::SeqCst); - close_ingress(&self.ingress, self.metrics.as_ref()); - } -} - /// Stop a mux that failed: stop its tasks and close every ingress slot. /// Returns true only for the first stop request. /// @@ -113,7 +103,7 @@ pub(super) fn stop_mux( first } -pub(super) fn close_ingress(ingress: &IngressRegistry, metrics: Option<&MuxMetricsHandle>) { +fn close_ingress(ingress: &IngressRegistry, metrics: Option<&MuxMetricsHandle>) { let closed = ingress.shutdown(); if let Some(metrics) = metrics { for _ in 0..closed { @@ -127,6 +117,6 @@ impl Drop for MuxCore { // Unconditional: a shutdown abandoned after `stop_sending` left the // slots open. self.tasks.stop(); - self.retire_slots(); + close_ingress(&self.ingress, self.metrics.as_ref()); } } diff --git a/lib/velo/src/streaming/messenger_mux/mod.rs b/lib/velo/src/streaming/messenger_mux/mod.rs index c7ed813c..43a0025d 100644 --- a/lib/velo/src/streaming/messenger_mux/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/mod.rs @@ -116,7 +116,7 @@ mod tests; use std::collections::VecDeque; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::atomic::AtomicU64; use std::time::Duration; use anyhow::{Result, anyhow}; @@ -128,7 +128,6 @@ use velo_ext::{TransportKey, WorkerAddress, WorkerId}; use self::flow_control::NegotiatedLimits; use self::ingress::IngressRegistry; use self::lane_choice::LaneLoad; -use self::lifecycle::close_ingress; use self::peer_batcher::{BatcherHandle, BatcherMap, OpenRejected, OpenSlotRequest}; use crate::messenger::{Context, Handler, Messenger}; use crate::observability::{MuxMetricsHandle, VeloMetrics}; @@ -225,9 +224,6 @@ struct MuxCore { /// batch of the new one as stale and discard it wholesale. epochs: Arc, tasks: crate::streaming::tasks::StreamTasks, - /// Set once the slots are retired, which is later than the tasks stop: - /// shutdown takes streams off their slots in between. - slots_retired: AtomicBool, /// Peers with credit to return, posted by draining consumers. See /// [`ingress::DrainSignal`]. drain_tx: flume::Sender, @@ -353,7 +349,6 @@ impl MessengerMuxTransport { // zeroed header from reading as a legitimate one. epochs: Arc::new(AtomicU64::new(1)), tasks, - slots_retired: AtomicBool::new(false), drain_tx, drain_rx, drains: DashMap::new(), @@ -426,14 +421,9 @@ impl MuxCore { } } - // A batch already in flight can claim a bind after shutdown retired - // the slots. Retire that late claim before returning from this - // handler. Shutdown sets the flag before it retires, so a claim this - // check misses is one that retirement still covers. - if self.slots_retired.load(Ordering::SeqCst) { - close_ingress(&self.ingress, self.metrics.as_ref()); - return; - } + // A batch in flight when the mux stopped has nobody to reply to. A + // bind it claimed is retired with the rest: `IngressRegistry::shutdown` + // clears binds before it walks the slot tables. if self.tasks.is_stopped() { return; } From 05db2599ff2f0d01fca8e6a4fcd4a2487f527068 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Mon, 5 Oct 2026 08:15:01 -0600 Subject: [PATCH 22/29] docs(streaming): state the window before AnchorManager::shutdown stops the mux The mux runs until AnchorManager::shutdown stops it, which is after graceful_shutdown has torn the transports down; its batchers can hand batches to a transport that is going away in that window. Say so rather than claim nothing is sent. Signed-off-by: Ryan Olson --- lib/velo/src/streaming/anchor.rs | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/lib/velo/src/streaming/anchor.rs b/lib/velo/src/streaming/anchor.rs index 8125a98c..2fce520e 100644 --- a/lib/velo/src/streaming/anchor.rs +++ b/lib/velo/src/streaming/anchor.rs @@ -1043,10 +1043,10 @@ impl AnchorManager { /// End every stream this manager holds. Called by `Velo::shutdown` after /// the messenger transports are torn down. /// - /// Two rules hold throughout: nothing is sent to a peer, because the - /// transports are gone, and no mux slot is retired while a stream still - /// reads from it, because retirement injects `Dropped` that the reader - /// would take as its sender's. Hence the phases: + /// Two rules hold from step 1 on. Nothing is sent to a peer, because the + /// transports are gone. No mux slot is retired while a stream still reads + /// from it, because retirement injects `Dropped` that the reader would + /// take as its sender's. Hence the steps: /// /// 1. Stop the mux's tasks. A slot close after this finds no batcher and /// sends nothing. The slots stay open. @@ -1058,6 +1058,13 @@ impl AnchorManager { /// Steps 1-3 do not await. A caller that drops this future (a timeout /// around `Velo::shutdown`) still leaves no stream waiting on a slot or an /// anchor. The slots then stay open, with no reader, until the mux drops. + /// + /// Known gap: the mux runs until step 1, which comes after transport + /// teardown. While `graceful_shutdown` tears the transports down, a + /// batcher can still hand a batch (credit, a local sender's records) to a + /// transport that is going away. Stopping the mux before teardown needs + /// `Velo::shutdown` to be split around this function, which is not done + /// here. pub(crate) async fn shutdown(&self) { let mux = self.mux.get(); if let Some(mux) = mux { From cf0c64a9a3a0337129997be2dc205373e293c70e Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Mon, 5 Oct 2026 08:28:58 -0600 Subject: [PATCH 23/29] docs(streaming): name the mux stop paths in the comments that list them Signed-off-by: Ryan Olson --- lib/velo/src/streaming/messenger_mux/ingress/mod.rs | 5 +++-- lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs | 7 ++++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/lib/velo/src/streaming/messenger_mux/ingress/mod.rs b/lib/velo/src/streaming/messenger_mux/ingress/mod.rs index db620869..a833f3b3 100644 --- a/lib/velo/src/streaming/messenger_mux/ingress/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/ingress/mod.rs @@ -450,8 +450,9 @@ impl IngressRegistry { /// Tear down every slot of every peer, injecting `Dropped` into each. /// - /// Used when the transport itself goes away, so a consumer never waits out - /// its heartbeat watchdog for a sender that has already been dismantled. + /// Used when the mux stops (shutdown, or a failed mux), so a consumer + /// never waits out its heartbeat watchdog for a sender that has already + /// been dismantled. /// /// Binds go first. A batch handler can be claiming one concurrently, and /// it claims under its peer table's lock, in a table it inserted before diff --git a/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs b/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs index 7a87b0a3..e7b0b852 100644 --- a/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs @@ -1137,9 +1137,10 @@ impl Batcher { // that terminates only if a closed batcher is never the registered // one. The retire path holds it because the sweep removes the entry // before posting `retire`; this order is what holds it on the other - // exit. Nothing writes after cancel today — it comes only from - // `MuxCore::drop` — which is why the invariant is kept structural - // rather than argued from the callers. + // exit. Cancel comes from `stop_sending`, `stop_mux` and + // `MuxCore::drop`. A writer that runs after it (a slot close during + // anchor removal at shutdown) is refused by `MuxCore::batcher`, and + // this order keeps the invariant even if that guard were missed. if unregister { let handle = Arc::clone(&self.handle); self.batchers From cc282eb939cc84121247418539e7d04231c6c5bb Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Mon, 5 Oct 2026 08:30:36 -0600 Subject: [PATCH 24/29] test(velo): share the node and peer wiring of the shutdown tests Three shutdown tests each built the same TCP node and the same two-peer setup inline. Move that into two helpers in the test module. The tests keep their assertions, timeouts and comments. Signed-off-by: Ryan Olson --- lib/velo/src/lib.rs | 141 ++++++++++++++------------------------------ 1 file changed, 44 insertions(+), 97 deletions(-) diff --git a/lib/velo/src/lib.rs b/lib/velo/src/lib.rs index 2cf4332f..5a1cca9d 100644 --- a/lib/velo/src/lib.rs +++ b/lib/velo/src/lib.rs @@ -1692,39 +1692,31 @@ mod tests { retained.closed().await; } - /// Shutdown must not write to peers after its transports are gone. - /// - /// `graceful_shutdown` tears down the messenger transports first. Removing - /// an anchor with a live mux stream then closes its slot, and while the mux - /// still runs that close is batched to the peer: the TCP transport has no - /// connection left, dials, finds itself cancelled, and fails the frame as - /// a send error that the default handler logs at ERROR. - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] - async fn shutdown_with_live_mux_stream_sends_nothing_after_teardown() { - use crate::observability::test_helpers::MetricSnapshot; - use futures::StreamExt; - - async fn node(metrics: Option>) -> Arc { - let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); - let mut builder = Velo::builder() - .add_transport(Arc::new( - crate::transports::tcp::TcpTransportBuilder::new() - .from_listener(listener) - .unwrap() - .build() - .unwrap(), - )) - .stream_bind_addr(std::net::Ipv4Addr::LOCALHOST.into()); - if let Some(metrics) = metrics { - builder = builder.metrics(metrics); - } - builder.build().await.unwrap() + /// A node with one TCP transport on a loopback port and a loopback stream listener. + async fn tcp_stream_node(metrics: Option>) -> Arc { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let mut builder = Velo::builder() + .add_transport(Arc::new( + crate::transports::tcp::TcpTransportBuilder::new() + .from_listener(listener) + .unwrap() + .build() + .unwrap(), + )) + .stream_bind_addr(std::net::Ipv4Addr::LOCALHOST.into()); + if let Some(metrics) = metrics { + builder = builder.metrics(metrics); } + builder.build().await.unwrap() + } - let registry = prometheus::Registry::new(); - let metrics = Arc::new(VeloMetrics::register(®istry).unwrap()); - let consumer = node(Some(metrics)).await; - let producer = node(None).await; + /// A (consumer, producer) pair that know each other and both serve `attach_handler`. + async fn connected_pair( + consumer_metrics: Option>, + attach_handler: &str, + ) -> (Arc, Arc) { + let consumer = tcp_stream_node(consumer_metrics).await; + let producer = tcp_stream_node(None).await; consumer.register_peer(producer.peer_info()).unwrap(); producer.register_peer(consumer.peer_info()).unwrap(); for (node, peer) in [ @@ -1733,12 +1725,30 @@ mod tests { ] { tokio::time::timeout( std::time::Duration::from_secs(5), - node.wait_for_handler(peer, "_anchor_attach"), + node.wait_for_handler(peer, attach_handler), ) .await .unwrap() .unwrap(); } + (consumer, producer) + } + + /// Shutdown must not write to peers after its transports are gone. + /// + /// `graceful_shutdown` tears down the messenger transports first. Removing + /// an anchor with a live mux stream then closes its slot, and while the mux + /// still runs that close is batched to the peer: the TCP transport has no + /// connection left, dials, finds itself cancelled, and fails the frame as + /// a send error that the default handler logs at ERROR. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn shutdown_with_live_mux_stream_sends_nothing_after_teardown() { + use crate::observability::test_helpers::MetricSnapshot; + use futures::StreamExt; + + let registry = prometheus::Registry::new(); + let metrics = Arc::new(VeloMetrics::register(®istry).unwrap()); + let (consumer, producer) = connected_pair(Some(metrics), "_anchor_attach").await; let mut anchor = consumer.create_anchor::(); let sender = producer @@ -1792,38 +1802,7 @@ mod tests { async fn shutdown_ends_a_live_mux_stream_without_sender_dropped() { use futures::StreamExt; - async fn node() -> Arc { - let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); - Velo::builder() - .add_transport(Arc::new( - crate::transports::tcp::TcpTransportBuilder::new() - .from_listener(listener) - .unwrap() - .build() - .unwrap(), - )) - .stream_bind_addr(std::net::Ipv4Addr::LOCALHOST.into()) - .build() - .await - .unwrap() - } - - let consumer = node().await; - let producer = node().await; - consumer.register_peer(producer.peer_info()).unwrap(); - producer.register_peer(consumer.peer_info()).unwrap(); - for (node, peer) in [ - (&consumer, producer.instance_id()), - (&producer, consumer.instance_id()), - ] { - tokio::time::timeout( - std::time::Duration::from_secs(5), - node.wait_for_handler(peer, "_anchor_attach"), - ) - .await - .unwrap() - .unwrap(); - } + let (consumer, producer) = connected_pair(None, "_anchor_attach").await; let mut anchor = consumer.create_anchor::(); let sender = producer @@ -1870,41 +1849,9 @@ mod tests { use crate::observability::test_helpers::MetricSnapshot; - async fn node(metrics: Option>) -> Arc { - let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); - let mut builder = Velo::builder() - .add_transport(Arc::new( - crate::transports::tcp::TcpTransportBuilder::new() - .from_listener(listener) - .unwrap() - .build() - .unwrap(), - )) - .stream_bind_addr(std::net::Ipv4Addr::LOCALHOST.into()); - if let Some(metrics) = metrics { - builder = builder.metrics(metrics); - } - builder.build().await.unwrap() - } - let registry = prometheus::Registry::new(); let metrics = Arc::new(VeloMetrics::register(®istry).unwrap()); - let consumer = node(Some(metrics)).await; - let producer = node(None).await; - consumer.register_peer(producer.peer_info()).unwrap(); - producer.register_peer(consumer.peer_info()).unwrap(); - for (node, peer) in [ - (&consumer, producer.instance_id()), - (&producer, consumer.instance_id()), - ] { - tokio::time::timeout( - std::time::Duration::from_secs(5), - node.wait_for_handler(peer, "_mpsc_anchor_attach"), - ) - .await - .unwrap() - .unwrap(); - } + let (consumer, producer) = connected_pair(Some(metrics), "_mpsc_anchor_attach").await; let mut anchor = consumer.create_mpsc_anchor::(); let sender = producer From c4abc477871d4ce3b3463c473307a91ff7c0b271 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Mon, 5 Oct 2026 08:49:18 -0600 Subject: [PATCH 25/29] fix(streaming): fail a pre-bind that races mux shutdown instead of panicking prebind_anchor binds a mux slot and then takes the drain signal the bind parked. Mux shutdown clears parked drain signals, so a public, sync prebind_anchor call on another thread could find it gone and panic on an expect. Release the bind and return None, as a pre-bind that finds the mux stopped does. Signed-off-by: Ryan Olson --- lib/velo/src/streaming/anchor.rs | 9 ++++--- .../src/streaming/messenger_mux/lane_tests.rs | 25 +++++++++++++++++++ 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/lib/velo/src/streaming/anchor.rs b/lib/velo/src/streaming/anchor.rs index 2fce520e..c423e454 100644 --- a/lib/velo/src/streaming/anchor.rs +++ b/lib/velo/src/streaming/anchor.rs @@ -1400,9 +1400,12 @@ impl AnchorManager { let lane = mux.choose_lane(None, lane_key); let lane_index = lane.lane(); let receiver = mux.bind_on_lane(local_id, routing_session_id, lane).ok()?; - let drain = mux - .take_drain_signal(local_id, routing_session_id) - .expect("prebind parks a drain signal for the pair it just registered"); + // Missing only when mux shutdown cleared it after the bind; the + // pre-bind then fails like one that found the mux stopped. + let Some(drain) = mux.take_drain_signal(local_id, routing_session_id) else { + mux.release_bind(local_id, routing_session_id); + return None; + }; // Negotiated against the local mux alone. There is no peer here to // intersect with — that is the whole point — so the terms are this diff --git a/lib/velo/src/streaming/messenger_mux/lane_tests.rs b/lib/velo/src/streaming/messenger_mux/lane_tests.rs index a238a8f9..c62ae6d9 100644 --- a/lib/velo/src/streaming/messenger_mux/lane_tests.rs +++ b/lib/velo/src/streaming/messenger_mux/lane_tests.rs @@ -1143,3 +1143,28 @@ async fn a_lane_that_refuses_admission_fails_only_its_own_streams() { .await .expect("the streams on the other lanes finish in order"); } + +/// A pre-bind that races shutdown fails; it does not panic. +/// +/// `prebind_anchor` is a public sync call an application can make on any +/// thread. Mux shutdown clears the parked drain signals, so one can vanish +/// between the pre-bind's bind and its take. The bind hook plays that +/// shutdown. +#[tokio::test(flavor = "multi_thread")] +async fn a_prebind_racing_shutdown_returns_none() { + let node = Node::new(1).await; + let mux = node.mux(); + let core = Arc::downgrade(&mux.core); + assert!( + mux.core + .bind_hook + .set(Box::new(move |_| { + if let Some(core) = core.upgrade() { + core.drains.clear(); + } + })) + .is_ok() + ); + let anchor = node.velo.create_anchor::(); + assert!(node.velo.prebind_anchor(anchor.handle()).is_none()); +} From 0ce261ad23da13d9624f0e49eee0eb3e21fc3c68 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Mon, 5 Oct 2026 08:50:28 -0600 Subject: [PATCH 26/29] fix(streaming): let a cancelled SPSC reader pump stop while its channel is full A consumer that holds its anchor without reading fills the anchor channel, and the reader pump then waited on the send without watching its cancel token. Velo::shutdown cancels the pump and removes the anchor, but the application still holds the receiver, so the pump kept its transport receiver and channel alive. Race the send against the cancel on the full-channel path, as the MPSC pump already does; the try_send fast path is unchanged. Also say in the Velo::shutdown doc that reader pumps are cancelled, not joined. Signed-off-by: Ryan Olson --- lib/velo/src/lib.rs | 5 ++-- lib/velo/src/streaming/control/pump.rs | 13 +++++++-- lib/velo/src/streaming/control/tests.rs | 35 +++++++++++++++++++++++++ 3 files changed, 49 insertions(+), 4 deletions(-) diff --git a/lib/velo/src/lib.rs b/lib/velo/src/lib.rs index 5a1cca9d..1cf3ce9f 100644 --- a/lib/velo/src/lib.rs +++ b/lib/velo/src/lib.rs @@ -796,8 +796,9 @@ impl Velo { /// Drain messenger work and stop this instance's streaming services. /// - /// This also cancels live streams and joins the receive loops and streaming - /// listener and pump tasks owned by the builder. Custom frame transports + /// This also cancels live streams, and joins the receive loops and the + /// streaming transport tasks owned by the builder. Stream reader pumps are + /// cancelled, not joined. Custom frame transports /// remain the caller's responsibility. Stream watchdogs and heartbeats are /// cancelled; application handlers that exceed `policy` can still be running. /// If the sender of a stream is on this instance, shutdown cancels that diff --git a/lib/velo/src/streaming/control/pump.rs b/lib/velo/src/streaming/control/pump.rs index d9ba2dd2..f0802d13 100644 --- a/lib/velo/src/streaming/control/pump.rs +++ b/lib/velo/src/streaming/control/pump.rs @@ -240,8 +240,17 @@ pub(crate) async fn reader_pump( if let Some(m) = metrics.as_ref() { m.record_reader_pump_backpressure(); } - if frame_tx.send_async(b).await.is_err() { - break; // consumer dropped + // Raced against the cancel, as the MPSC pump + // does: a consumer that holds its anchor + // without reading keeps this send waiting, + // and shutdown must still end the pump. + let sent = tokio::select! { + biased; + _ = &mut cancelled => false, + sent = frame_tx.send_async(b) => sent.is_ok(), + }; + if !sent { + break; // cancelled, or consumer dropped } } Err(flume::TrySendError::Disconnected(_)) => break, diff --git a/lib/velo/src/streaming/control/tests.rs b/lib/velo/src/streaming/control/tests.rs index ae44c4d6..f374477f 100644 --- a/lib/velo/src/streaming/control/tests.rs +++ b/lib/velo/src/streaming/control/tests.rs @@ -2200,3 +2200,38 @@ fn attach_messages_from_before_lanes_read_as_lane_zero() { ); } } + +/// Cancelling a reader pump stops it even while it waits for room in a full +/// anchor channel. +/// +/// A consumer that holds its anchor without reading fills the channel, and +/// the pump then waits on the send. `Velo::shutdown` cancels the pump's token +/// and removes the anchor, but the anchor's receiver lives on with the +/// application, so the send would never fail. A pump that did not race the +/// cancel would hold its transport receiver and channel until the +/// application read or dropped the anchor. +#[tokio::test] +async fn reader_pump_cancel_ends_a_send_waiting_on_a_full_channel() { + let manager = make_test_manager(); + let (transport_tx, transport_rx) = flume::bounded::>(4); + let (frame_tx, _frame_rx) = flume::bounded::>(1); + let cancel = tokio_util::sync::CancellationToken::new(); + let pump = tokio::spawn(reader_pump( + transport_rx, + frame_tx, + cancel.clone(), + manager.anchor_context(), + 999, + std::time::Duration::from_secs(60), + )); + let item = rmp_serde::to_vec(&crate::streaming::frame::StreamFrame::Item(1u32)).unwrap(); + transport_tx.send_async(item.clone()).await.unwrap(); + transport_tx.send_async(item).await.unwrap(); + // The first item fills the channel; the pump now waits to forward the second. + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + cancel.cancel(); + tokio::time::timeout(std::time::Duration::from_secs(2), pump) + .await + .expect("a cancelled pump kept waiting on a full channel") + .unwrap(); +} From 63242f80ef72f3c75ef61f287f05903127d3b7e2 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Mon, 5 Oct 2026 08:52:34 -0600 Subject: [PATCH 27/29] fix(streaming): stop calling a batcher ended by runtime shutdown an abort Nothing aborts a batcher on purpose: one that ends without its teardown has panicked, which tokio reports, or its runtime shut down before Velo::shutdown ran. The ERROR claiming an unexpected abort misled the second case. Log a warning that names both causes. Signed-off-by: Ryan Olson --- lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs b/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs index e7b0b852..9df98f16 100644 --- a/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs @@ -416,8 +416,12 @@ impl Drop for Batcher { if super::lifecycle::stop_mux(&self.tasks, &self.ingress, self.metrics.as_ref()) && !already_cancelled { - tracing::error!(peer = %self.key.peer, lane = %self.key.lane, - "messenger mux stopped: batcher task aborted unexpectedly"); + // Nothing aborts a batcher on purpose. This is a panic, which + // tokio reports itself, or a runtime dropped before + // `Velo::shutdown` ran, which is not an error of its own. + tracing::warn!(peer = %self.key.peer, lane = %self.key.lane, + "messenger mux stopped: a batcher task ended without its teardown \ + (a panic, or the runtime shut down before Velo::shutdown)"); } } // A refused spawn can drop here under the batcher-map entry guard. From 65cfc8c0f371d58ba0d6fc94e4af070f44fc3c8c Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Mon, 5 Oct 2026 08:52:34 -0600 Subject: [PATCH 28/29] test(streaming): pin that a refused stream task drops outside the lock spawn_until_done returns a refused future with the admission lock released only because a parameter drops after the body's locals. A batcher's Drop stops the service and takes that lock again, so a change that dropped the future under the lock would deadlock. Say so at the return, and add a test whose refused future stops the service from its Drop; dropping it under the lock fails the test. Signed-off-by: Ryan Olson --- lib/velo/src/streaming/tasks.rs | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/lib/velo/src/streaming/tasks.rs b/lib/velo/src/streaming/tasks.rs index b835f37a..fb56347d 100644 --- a/lib/velo/src/streaming/tasks.rs +++ b/lib/velo/src/streaming/tasks.rs @@ -51,6 +51,8 @@ impl StreamTasks { let tracked = { let _admission = self.admission.lock(); if self.is_stopped() { + // `future` is a parameter, so it drops after `_admission`. + // Its Drop may stop this service and take the lock again. return false; } self.tracker.track_future(future) @@ -119,4 +121,33 @@ mod tests { .await .expect("a stopped tracker waits only for running tasks"); } + + /// A refused future is dropped after the admission lock is released. Its + /// `Drop` can stop the service (a batcher's does), and that takes the + /// same, non-reentrant lock. + #[tokio::test] + async fn a_refused_future_may_stop_the_service_from_its_drop() { + struct StopOnDrop(StreamTasks); + impl Drop for StopOnDrop { + fn drop(&mut self) { + self.0.stop(); + } + } + + let tasks = StreamTasks::default(); + tasks.stop(); + let (done_tx, done_rx) = std::sync::mpsc::channel(); + let refused = tasks.clone(); + std::thread::spawn(move || { + let guard = StopOnDrop(refused.clone()); + let spawned = refused.spawn_until_done(async move { + let _guard = guard; + }); + done_tx.send(spawned).unwrap(); + }); + let spawned = done_rx + .recv_timeout(std::time::Duration::from_secs(5)) + .expect("dropping a refused future deadlocked on the admission lock"); + assert!(!spawned); + } } From 36d4e8c804a6c3bf251fa05ca8ad339779f294b2 Mon Sep 17 00:00:00 2001 From: Ryan Olson Date: Mon, 5 Oct 2026 09:04:58 -0600 Subject: [PATCH 29/29] fix(streaming): count a pre-bind that finds the mux stopped as failed The two exits taken when the mux is stopped returned None without the debug line and the error count every other failed pre-bind records, so an operator could not tell a shut-down mux from zero-RTT being off. Record and log them like the rest. Signed-off-by: Ryan Olson --- lib/velo/src/streaming/anchor.rs | 28 +++++++++++--- .../src/streaming/messenger_mux/lane_tests.rs | 37 ++++++++++++++++--- 2 files changed, 55 insertions(+), 10 deletions(-) diff --git a/lib/velo/src/streaming/anchor.rs b/lib/velo/src/streaming/anchor.rs index c423e454..08347bfc 100644 --- a/lib/velo/src/streaming/anchor.rs +++ b/lib/velo/src/streaming/anchor.rs @@ -1399,11 +1399,29 @@ impl AnchorManager { // own pre-binds. Chosen before the bind so the bind is counted on it. let lane = mux.choose_lane(None, lane_key); let lane_index = lane.lane(); - let receiver = mux.bind_on_lane(local_id, routing_session_id, lane).ok()?; - // Missing only when mux shutdown cleared it after the bind; the - // pre-bind then fails like one that found the mux stopped. - let Some(drain) = mux.take_drain_signal(local_id, routing_session_id) else { - mux.release_bind(local_id, routing_session_id); + // The bind fails once the mux is stopped. Its drain signal is missing + // only when mux shutdown cleared it after the bind, which is the same + // case, a moment later. + let bound = mux + .bind_on_lane(local_id, routing_session_id, lane) + .ok() + .and_then( + |receiver| match mux.take_drain_signal(local_id, routing_session_id) { + Some(drain) => Some((receiver, drain)), + None => { + mux.release_bind(local_id, routing_session_id); + None + } + }, + ); + let Some((receiver, drain)) = bound else { + tracing::debug!(%handle, "prebind_anchor: the mux is shut down; no ticket minted"); + self.record_streaming_operation( + StreamingOp::Prebind, + HandlerOutcome::Error, + "unknown", + started, + ); return None; }; diff --git a/lib/velo/src/streaming/messenger_mux/lane_tests.rs b/lib/velo/src/streaming/messenger_mux/lane_tests.rs index c62ae6d9..189e23a1 100644 --- a/lib/velo/src/streaming/messenger_mux/lane_tests.rs +++ b/lib/velo/src/streaming/messenger_mux/lane_tests.rs @@ -1144,7 +1144,8 @@ async fn a_lane_that_refuses_admission_fails_only_its_own_streams() { .expect("the streams on the other lanes finish in order"); } -/// A pre-bind that races shutdown fails; it does not panic. +/// A pre-bind that races shutdown fails, and is counted as a failed pre-bind; +/// it does not panic. /// /// `prebind_anchor` is a public sync call an application can make on any /// thread. Mux shutdown clears the parked drain signals, so one can vanish @@ -1152,8 +1153,27 @@ async fn a_lane_that_refuses_admission_fails_only_its_own_streams() { /// shutdown. #[tokio::test(flavor = "multi_thread")] async fn a_prebind_racing_shutdown_returns_none() { - let node = Node::new(1).await; - let mux = node.mux(); + use crate::observability::test_helpers::MetricSnapshot; + + let registry = prometheus::Registry::new(); + let metrics = Arc::new(crate::observability::VeloMetrics::register(®istry).unwrap()); + let velo = Velo::builder() + .add_transport(Arc::new( + QuicTransportBuilder::new() + .bind_addr("127.0.0.1:0".parse().expect("loopback")) + .build() + .expect("quic transport"), + ) as Arc) + .stream_bind_addr(std::net::Ipv4Addr::LOCALHOST.into()) + .metrics(metrics) + .build() + .await + .expect("build velo"); + let mux = velo + .anchor_manager() + .mux_handle() + .and_then(|mux| mux.upgrade()) + .expect("velo installs the mux by default"); let core = Arc::downgrade(&mux.core); assert!( mux.core @@ -1165,6 +1185,13 @@ async fn a_prebind_racing_shutdown_returns_none() { })) .is_ok() ); - let anchor = node.velo.create_anchor::(); - assert!(node.velo.prebind_anchor(anchor.handle()).is_none()); + let anchor = velo.create_anchor::(); + assert!(velo.prebind_anchor(anchor.handle()).is_none()); + assert_eq!( + MetricSnapshot::from_registry(®istry).counter( + "velo_streaming_anchor_operations_total", + &[("operation", "prebind"), ("outcome", "error")], + ), + 1.0 + ); }