diff --git a/Cargo.lock b/Cargo.lock index f5f8bc5b..5a0c2509 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3289,7 +3289,7 @@ checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" [[package]] name = "velo" -version = "0.13.0" +version = "0.14.0" dependencies = [ "anyhow", "async-nats", diff --git a/Cargo.toml b/Cargo.toml index 1916cda4..ba0dd924 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,7 @@ exclude = ["examples"] resolver = "3" [workspace.package] -version = "0.13.0" +version = "0.14.0" edition = "2024" description = "Velo Active Messenging Library" authors = ["NVIDIA Inc. "] diff --git a/README.md b/README.md index 2cfa5bf3..27b342f4 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,8 @@ NOTE: Velo is experimental. Its design, development, and tests are still in prog **Documentation: [the Velo book](https://ai-dynamo.github.io/velo/)**. The source is in [`docs/src/`](docs/src/SUMMARY.md). +Typed streams support graceful stop, cancellation, and ordered finalization, including prebound tickets. See [Stream lifecycle](docs/src/concepts/streaming.md). + ## Crates | Crate | For | Contents | diff --git a/docs/src/concepts/batched-streaming.md b/docs/src/concepts/batched-streaming.md index 305a8649..45ee630b 100644 --- a/docs/src/concepts/batched-streaming.md +++ b/docs/src/concepts/batched-streaming.md @@ -1,6 +1,6 @@ # Batched streaming -The messenger mux carries every stream to one peer over the Messenger connection that already exists to that peer. It packs the records for that peer into `_stream_batch` active messages. Its transport key is `messenger-mux-v1`. The mux is opt-in and is negotiated per attach. Senders do not change: `StreamSender::send` stages a record, and the layer below it decides when to write. +The messenger mux carries every stream to one peer over the Messenger connection that already exists to that peer. It packs the records for that peer into `_stream_batch` active messages. Its transport key is `messenger-mux-v2`. The mux is opt-in and is negotiated per attach. Senders do not change: `StreamSender::send` stages a record, and the layer below it decides when to write. This chapter describes how the mux works. The [Tune batched streaming](../guides/tune-batched-streaming.md) guide tells you how to configure it. The [Batched streaming design](../development/batched-streaming-design.md) chapter records why it works this way and which alternatives were rejected. @@ -85,12 +85,12 @@ _stream_batch payload: [16 B batch header][record_count x record] batch header: - [u8 mux_version = 1][u8 flags][u16 record_count][u64 peer_epoch][u32 batch_seq] + [u8 mux_version = 2][u8 flags][u16 record_count][u64 peer_epoch][u32 batch_seq] record: [u8 record_type][u32 slot][u32 frame_seq][u32 len][len bytes body] -record_type: 0 = Data, 1 = OpenSlot, 2 = CloseSlot, 3 = CreditUpdate, 4 = SlotHeartbeat +record_type: 0 = Data, 1 = OpenSlot, 2 = CloseSlot, 3 = CreditUpdate, 4 = SlotHeartbeat, 5 = LifecycleSlot ``` Every multi-byte field is big-endian, in the header and in each record. Senders write `flags` as zero, and receivers ignore unknown bits. @@ -105,6 +105,7 @@ Record bodies: - **`OpenSlot`** carries `[u64 anchor_id][u64 session_id]`. This is the 16-byte attach handshake, moved into a record. - **`CloseSlot`** carries `[u8 reason]`: `0` terminal sent, `1` peer gone, `2` unknown slot, `3` protocol error. - **`CreditUpdate`** carries `[u32 delta]` from receiver to sender. +- **`LifecycleSlot`** carries the 64-bit session identity from `OpenSlot` and a one-byte action (0 = stop, 1 = cancel). Stop leaves the slot open for remaining output. Cancel takes precedence for the same session. Both the queued signal and the live producer check session identity, so a delayed signal cannot target a reused slot even if its compact generation has wrapped. Early stop is retained by the pre-bind until `OpenSlot` claims it. - **`SlotHeartbeat`** has no body. The decoder accepts it, but no sender emits it. See [Heartbeats](#heartbeats). `CloseSlot` travels in both directions and has no direction bit. The reason carries the direction. `TerminalSent` and `PeerGone` travel from slot owner to receiver. `UnknownSlot` and `ProtocolError` travel from receiver to slot owner. Both sides can hold a slot at the same dense index, and the reason tells them apart. @@ -356,19 +357,19 @@ At serving depth, `Manual` is both higher and repeatable. A batcher that writes ## Negotiation and compatibility -The attach selects the mux. No wire magic exists. `AnchorAttachRequest` and `MpscAnchorAttachRequest` carry `supported_transport_keys` (with `#[serde(default)]`, so an older sender deserializes as advertising nothing). The attach handler intersects the sender's keys with its installed transports. It picks `messenger-mux-v1` only when both sides name it. Otherwise it answers with its default key. +The attach selects the mux. No wire magic exists. `AnchorAttachRequest` and `MpscAnchorAttachRequest` carry `supported_transport_keys` (with `#[serde(default)]`, so an older sender deserializes as advertising nothing). The attach handler intersects the sender's keys with its installed transports. It picks `messenger-mux-v2` only when both sides name it. Otherwise it answers with its default key. The sender reads the answer: -- A key other than `messenger-mux-v1` is the per-stream path, and the credit fields do not apply. Every older receiver answers this way. -- `messenger-mux-v1` with a window opens a slot that already holds that window. -- `messenger-mux-v1` with no window (`initial_credit` zero) is refused. No shipped receiver answers this, and a node with a mux cannot advertise a zero window. A fallback to another transport reaches nothing that listens, and the stream hangs until the watchdog fires. +- A key other than `messenger-mux-v2` is the per-stream path, and the credit fields do not apply. Every older receiver answers this way. +- `messenger-mux-v2` with a window opens a slot that already holds that window. +- `messenger-mux-v2` with no window (`initial_credit` zero) is refused. No shipped receiver answers this, and a node with a mux cannot advertise a zero window. A fallback to another transport reaches nothing that listens, and the stream hangs until the watchdog fires. -A node with the mux enabled registers both `messenger-mux-v1` and its configured per-stream transport, so it still serves older peers. `resolve_transport` fails on an unknown key in a non-empty registry, so a receiver that answers the mux on its own breaks every older sender. SPSC and MPSC anchors both negotiate the mux. +A node with the mux enabled registers both `messenger-mux-v2` and its configured per-stream transport, so it still serves older peers. `resolve_transport` fails on an unknown key in a non-empty registry, so a receiver that answers the mux on its own breaks every older sender. SPSC and MPSC anchors both negotiate the mux. `StreamSender::negotiated_transport()` returns the key that the attach settled on. It returns `None` for a same-worker attach, which uses no transport. Compare it with the public constant `MESSENGER_MUX_KEY`. -`MuxConfig::enabled = false` is the rollback. The node stops advertising `messenger-mux-v1`, and the next attach negotiates the per-stream path with no code or wire change. See [Zero-RTT stream setup](#zero-rtt-stream-setup) for the order when tickets are in use. +`MuxConfig::enabled = false` is the rollback. The node stops advertising `messenger-mux-v2`, and the next attach negotiates the per-stream path with no code or wire change. See [Zero-RTT stream setup](#zero-rtt-stream-setup) for the order when tickets are in use. ## Observability @@ -381,3 +382,7 @@ The mux series all start with `velo_streaming_mux_`. The most important ones: - `velo_messenger_ordered_lane_wait_seconds{handler="_stream_batch"}` shows the ingress lane wait, one sample per batch. The [Metrics reference](../appendix/metrics.md) lists every series. [Stream saturation](../operations/saturation.md) explains how to read them under load. + +### Peer loss + +A batcher with live producer slots checks its selected message transport every five seconds. A failed check closes the producer slots and triggers their cancellation and stop tokens, including idle producers. The check is per peer, not per token or stream. It does not change the selected transport or permit a fallback. diff --git a/docs/src/concepts/streaming.md b/docs/src/concepts/streaming.md index db14dce1..39a76e9c 100644 --- a/docs/src/concepts/streaming.md +++ b/docs/src/concepts/streaming.md @@ -38,7 +38,9 @@ while let Some(frame) = anchor.next().await { | `SenderError(String)` | The producer could not serialize an item | | `TransportError(String)` | The network failed during delivery | -The consumer can cancel the stream with `anchor.cancel()` or with a cloned `StreamController`. The producer sees the cancel through `sender.cancellation_token()`. +The consumer can request graceful stop with `anchor.controller().request_stop()`. The producer observes `sender.stop_token()` and can send its final output, then call `finalize()`. Stop is idempotent and does not discard buffered output. + +The consumer can cancel the stream with `anchor.cancel()` or a cloned `StreamController`. Cancel triggers both producer tokens and ends delivery. Dropping the consumer also cancels. Stop and cancel made before attach or ticket open are retained; the producer observes them when it opens, even if it sends no data. Transport failure remains a stream error. ## Many producers (MPSC) @@ -89,4 +91,4 @@ sender.finalize()?; Call `prebind_anchor` from a runtime context, because it starts tasks. A ticket that the worker never opens does not stay forever. After a 60 s accept window, the bind is reclaimed and the anchor sees `SenderDropped`. `velo_streaming_unclaimed_bind_reaped_total` counts these tickets. -A sender opened from a ticket never sent `_anchor_attach`, so it has no cancel handle, and `sender.cancellation_token()` never fires. Use the result of `send` and `finalize` instead. A consumer that has gone shows as a send error. +Ticket senders support `stop_token()` and `cancellation_token()` through messenger mux version 2. The existing slot-open exchange carries the stream session identity; no extra stream registration is needed. A peer epoch and slot generation reject stale lifecycle signals. Applications that require this lifecycle must require `messenger-mux-v2`; legacy attach fallback does not establish that capability. diff --git a/docs/src/development/batched-streaming-design.md b/docs/src/development/batched-streaming-design.md index 070a2b2b..f080789f 100644 --- a/docs/src/development/batched-streaming-design.md +++ b/docs/src/development/batched-streaming-design.md @@ -126,7 +126,7 @@ The price of `Manual` is also known. A per-pass flush caps a batch at that pass' ### Negotiate per attach and register both transports -The receiver answers `messenger-mux-v1` only when the sender named it. `resolve_transport` fails on an unknown key, so a receiver that chooses the mux on its own breaks every older sender. A node with the mux registers both keys. The wire version lives in the key, so incompatible versions never pair. A zero window with the mux key is refused, because a fallback connects to a transport where nothing listens. +The receiver answers `messenger-mux-v2` only when the sender named it. `resolve_transport` fails on an unknown key, so a receiver that chooses the mux on its own breaks every older sender. A node with the mux registers both keys. The wire version lives in the key, so incompatible versions never pair. A zero window with the mux key is refused, because a fallback connects to a transport where nothing listens. ### Credit return: the pump posts, the reconcile decides @@ -194,7 +194,7 @@ The mux is a user of the Messenger send path, not something a transport implemen The additions this work made to `velo-ext` were shaped to last. `Transport::max_message_size(target) -> Option` has a default, and `None` means unknown, which costs the caller a conservative budget. `SendOutcome::Pending(SendAdmission)` changed a published enum, so it shipped once, as a coordinated `velo-ext` and `velo` release. The egress recorders on `TransportObservability` have no-op defaults. -The cost of deferral is bounded. Out-of-tree `FrameTransport`s get no multiplexing, and nothing else breaks. A receiver offers `messenger-mux-v1` only when the mux is installed and enabled, so a deployment without the mux degrades correctly. +The cost of deferral is bounded. Out-of-tree `FrameTransport`s get no multiplexing, and nothing else breaks. A receiver offers `messenger-mux-v2` only when the mux is installed and enabled, so a deployment without the mux degrades correctly. Methods that the mux needs outside the trait (`advertised_limits`, `connect_negotiated`, `take_drain_signal`, `prebind`, `release_bind`, `close_claimed_slot`) are inherent on the concrete `MessengerMuxTransport`. `AnchorManager` already holds the concrete type, so none of them needs a trait change. diff --git a/docs/src/guides/tune-batched-streaming.md b/docs/src/guides/tune-batched-streaming.md index e45d0e21..69e92035 100644 --- a/docs/src/guides/tune-batched-streaming.md +++ b/docs/src/guides/tune-batched-streaming.md @@ -24,7 +24,7 @@ let velo = Velo::builder() .await?; ``` -Enable the mux on both nodes of a pair. An attach uses the mux only when both sides advertise `messenger-mux-v1`. Every other pair uses the per-stream path. +Enable the mux on both nodes of a pair. An attach uses the mux only when both sides advertise `messenger-mux-v2`. Every other pair uses the per-stream path. Call `messenger_mux` once per `Velo` instance. A second call returns an error. @@ -86,7 +86,7 @@ All settings are fields of `MuxConfig`. Always build it with `..Default::default | Field | Default | What it controls | |---|---|---| -| `enabled` | `false` | Installs the mux and advertises `messenger-mux-v1`. Setting it back to `false` is the rollback. | +| `enabled` | `false` | Installs the mux and advertises `messenger-mux-v2`. Setting it back to `false` is the rollback. | | `max_batch_bytes` | 60 KiB | The configured cap on one batch. The eager budget and the 64 KiB coalescing threshold also clamp it. | | `initial_credit` | 256 | Data credit C per slot. Each slot buffer holds C+1 records. Zero is refused at build time. | | `slot_byte_budget` | 1 MiB | Bytes one slot can hold in flight, and the cap on its withheld queue. Zero means the default. | @@ -165,7 +165,7 @@ let sender = match envelope.ticket { }; ``` -A zero-RTT sender has no cancel handle, so its `cancellation_token` never fires. When the consumer drops the anchor, the producer's next `send` returns an error. An idle producer also receives a close from the consumer. +A ticket sender observes `stop_token()` and `cancellation_token()` through the mux's slot lifecycle. Stop leaves the stream open for final output. Cancel or consumer drop wakes an idle producer and ends delivery; it does not need a later send. The ticket stays valid for the 60-second accept window. After the window, the consumer reaps the bind and sees `SenderDropped`. diff --git a/docs/src/operations/benchmarking.md b/docs/src/operations/benchmarking.md index 6188224c..86a44a02 100644 --- a/docs/src/operations/benchmarking.md +++ b/docs/src/operations/benchmarking.md @@ -115,7 +115,7 @@ Each rule below exists because a measurement without it was wrong. 5. **Compare reps at a matched backlog draw.** On the external rig, one or more mocker processes fall behind during the opening burst and keep that backlog. The number of processes that hold the backlog (the "holders") sets throughput, ITL and end-to-end latency for every arm. Compare two arms only at the same holder count. 6. **Do not fix the arm order.** A matrix that always ran the same arm second put that arm in a degraded band in every rep. Its request-level numbers measured run position, not the setting. 7. **Change one setting per arm.** A pair that differed in two settings attributed a result to the wrong one. -8. **Assert the transport per request.** Fail any request whose negotiated key is not `messenger-mux-v1`. A silent fallback measures the per-stream path under a mux label. +8. **Assert the transport per request.** Fail any request whose negotiated key is not `messenger-mux-v2`. A silent fallback measures the per-stream path under a mux label. 9. **Record the build.** Write the velo commit and its dirty state into each rep's metadata. If the installed build and the checkout differ, fail the rep. 10. **Report latency beside throughput.** A change that adds latency can raise throughput by letting the consumer catch up. Watch TTFT in particular, because any windowed flush policy can make it worse. 11. **Prove a regression test before you trust it.** Revert the fix and make sure that the test fails. A test that passes with the fix reverted proves nothing. diff --git a/docs/src/operations/response-plane-performance.md b/docs/src/operations/response-plane-performance.md index f0c4c19f..2a3e81d6 100644 --- a/docs/src/operations/response-plane-performance.md +++ b/docs/src/operations/response-plane-performance.md @@ -164,13 +164,15 @@ UCX did not beat TCP, and it did not hold its tail. The failure chain: 5. The reader-pump watchdog (3 × 5 s) killed 1,722 streams. Every log line showed `anchor_frame_tx_len=0` and `transport_rx_len=0`. `closed_slot` drops were 1,722 × 255 exactly, so no record had arrived before each kill. 6. The worker was healthy the whole time. It completed all 27,702 of its requests. -The defect is that the UCX send path has no backpressure edge beyond its shared ring: +The measured revision had no UCX backpressure edge beyond its shared ring: - `ucp_am_send_nbx` never refuses. An exhausted endpoint returns a request pointer, not an error. - `inflight_ops` counts posted operations but does not gate admission. - The per-peer admission gate queues without a bound. - All peers share one 1,024-entry ring into the progress thread. TCP gives each connection a 256-entry channel and a blocking write. +Version 0.14 bounds queued and outstanding sends per peer. A send permit is released on UCX completion or failure, and peer retirement cancels pending admission. The progress thread can still process control and endpoint cleanup when a data peer is full. This change requires a new saturation measurement; it does not change the historical results below. + Because admission never blocks, the on-admission flush policy never parks, and batches collapse. UCX carried 6.97 to 8.65 records per batch against 18 to 29 over TCP with the same settings. That is 2.6 to 4.1 times the message rate for the same records. These alternatives were checked and refuted: @@ -181,7 +183,7 @@ These alternatives were checked and refuted: - UCX scheduling is not unfair. The failed peer carried eight times the load of the others. - Spinning progress threads do not starve the CPU. There were 9 progress threads on 288 cores. -No tuning setting reaches this mechanism. The fix has two parts. The UCX transport needs a backpressure edge (gate admission on in-flight operations or on a per-peer ring share). Liveness needs a heartbeat path that data cannot block, or a watchdog that can tell a starved stream from a dead one. The UCX inbound path now records frames like every other transport. The other two fixes are not built. [RDMA performance](rdma-performance.md) covers the UCX transport outside the response plane. +No tuning setting reaches this mechanism. The fix has two parts. The UCX transport needs a backpressure edge (gate admission on in-flight operations or on a per-peer ring share). Liveness needs a heartbeat path that data cannot block, or a watchdog that can tell a starved stream from a dead one. The UCX inbound path records frames like every other transport. Version 0.14 adds the bounded admission edge. Saturation measurements must still check whether data backlog delays stream heartbeats beyond the watchdog. [RDMA performance](rdma-performance.md) covers the UCX transport outside the response plane. ## Instrumentation cost diff --git a/docs/src/operations/saturation.md b/docs/src/operations/saturation.md index 55bbc389..dcaeebdf 100644 --- a/docs/src/operations/saturation.md +++ b/docs/src/operations/saturation.md @@ -2,7 +2,7 @@ If streaming traffic shows sender errors, missed deadlines, or a `Dropped` frame with no producer crash, use this runbook. The usual cause is saturation. The producer generates frames faster than the consumer drains them. Velo reports saturation through Prometheus counters and one log line. This page tells you which signals to read and what to do. -The signals differ between the per-stream path and the messenger mux. Read [The per-stream cascade](#the-per-stream-cascade) for streams with one TCP or gRPC connection each. Read [Saturation under the mux](#saturation-under-the-mux) for streams that negotiated `messenger-mux-v1`. +The signals differ between the per-stream path and the messenger mux. Read [The per-stream cascade](#the-per-stream-cascade) for streams with one TCP or gRPC connection each. Read [Saturation under the mux](#saturation-under-the-mux) for streams that negotiated `messenger-mux-v2`. ## The per-stream cascade diff --git a/examples/Cargo.lock b/examples/Cargo.lock index 07829466..9b0f113f 100644 --- a/examples/Cargo.lock +++ b/examples/Cargo.lock @@ -2728,7 +2728,7 @@ checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" [[package]] name = "velo" -version = "0.13.0" +version = "0.14.0" dependencies = [ "anyhow", "async-nats", diff --git a/examples/examples/soak/main.rs b/examples/examples/soak/main.rs index adb15422..f2c2c80a 100644 --- a/examples/examples/soak/main.rs +++ b/examples/examples/soak/main.rs @@ -37,7 +37,7 @@ use crate::harness::tier::Tier; /// /// The messenger-backed transport is not an option here. The AM-backed /// `VeloFrameTransport` was deleted for known correctness problems under -/// multi-stream concurrency, and its replacement — `messenger-mux-v1`, see +/// multi-stream concurrency, and its replacement — `messenger-mux-v2`, see /// `docs/src/concepts/batched-streaming.md` — is not yet exposed as a soak /// harness option. #[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)] diff --git a/lib/velo/src/lib.rs b/lib/velo/src/lib.rs index df5ffe69..6f3f1373 100644 --- a/lib/velo/src/lib.rs +++ b/lib/velo/src/lib.rs @@ -289,7 +289,7 @@ impl VeloBuilder { } /// Install the batched, multiplexed streaming transport - /// (`messenger-mux-v1`), described in `docs/src/concepts/batched-streaming.md`. + /// (`messenger-mux-v2`), described in `docs/src/concepts/batched-streaming.md`. /// /// **Opt-in, and the mux is not the default transport.** /// [`MuxConfig::enabled`](crate::streaming::MuxConfig::enabled) defaults to @@ -302,7 +302,7 @@ impl VeloBuilder { /// advertised. So a canary is one node with the flag on, talking the mux to /// other canaries and the legacy path to everything else, and **rollback is /// the same flag**: set it back to `false` and the node stops advertising - /// `messenger-mux-v1`, so the next attach negotiates the legacy path. No + /// `messenger-mux-v2`, so the next attach negotiates the legacy path. No /// code change, no wire change, and no coordination with peers, because a /// key that is never advertised is never selected. /// @@ -409,7 +409,7 @@ impl VeloBuilder { // Step 5: Build the mux, if it was switched on. It joins the registry // *beside* the legacy transport rather than replacing it: negotiation - // answers `messenger-mux-v1` only to peers that advertised it, and + // answers `messenger-mux-v2` only to peers that advertised it, and // every other peer is still answered — and must still be served — on // the legacy key. let mux = match self.mux_config.filter(|config| config.enabled) { @@ -924,6 +924,7 @@ impl Velo { } /// Wait for a specific handler to become available on a remote instance. + /// Uses cached availability; refreshes only when the handler is not known. pub async fn wait_for_handler( &self, instance_id: InstanceId, @@ -994,9 +995,8 @@ impl Velo { /// `_anchor_attach` round trip, because `ticket` already carries what one /// would have returned. /// - /// No attach means no `StreamCancelHandle`, so the returned sender's - /// `cancellation_token` never fires; a dropped consumer surfaces as a - /// send error instead. + /// The mux carries stop and cancel through the ticket's session and slot + /// identity. Both producer tokens work even before the first item is sent. pub async fn open_anchor_stream( &self, handle: StreamAnchorHandle, diff --git a/lib/velo/src/messenger/client/mod.rs b/lib/velo/src/messenger/client/mod.rs index 21f34901..4cf446b6 100644 --- a/lib/velo/src/messenger/client/mod.rs +++ b/lib/velo/src/messenger/client/mod.rs @@ -247,6 +247,11 @@ impl ActiveMessageClient { self.peer_registry.has_handler_info(instance_id) } + /// Check the cached handler list for this exact peer instance. + pub(crate) fn has_cached_handler(&self, instance_id: InstanceId, handler: &str) -> bool { + self.peer_registry.handler_exists(instance_id, handler) + } + /// Check if we can send a message directly (fast path) pub(crate) fn can_send_directly(&self, target: InstanceId, handler: &str) -> bool { // 1. Peer must be registered diff --git a/lib/velo/src/messenger/messenger.rs b/lib/velo/src/messenger/messenger.rs index ca9b16d4..ffdb800a 100644 --- a/lib/velo/src/messenger/messenger.rs +++ b/lib/velo/src/messenger/messenger.rs @@ -479,6 +479,7 @@ impl Messenger { } /// Wait for a specific handler to become available on a remote instance. + /// Uses cached availability; refreshes only when the handler is not known. pub async fn wait_for_handler( &self, instance_id: InstanceId, @@ -487,11 +488,14 @@ impl Messenger { const MAX_ATTEMPTS: u32 = 10; const DELAY: std::time::Duration = std::time::Duration::from_millis(100); + if self.client.has_cached_handler(instance_id, handler_name) { + return Ok(()); + } + for _ in 0..MAX_ATTEMPTS { self.refresh_handlers(instance_id).await?; - let handlers = self.available_handlers(instance_id).await?; - if handlers.contains(&handler_name.to_string()) { + if self.client.has_cached_handler(instance_id, handler_name) { return Ok(()); } @@ -838,11 +842,6 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_system_and_event_handlers_are_available_immediately_after_startup() { - test_transport_registry() - .lock() - .expect("transport registry poisoned") - .clear(); - let (transport_a, transport_b) = make_transport_pair(); let a = Messenger::builder() .add_transport(transport_a) @@ -874,6 +873,68 @@ mod tests { ); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_wait_for_handler_reuses_cache_and_refreshes_missing_handlers() { + use crate::observability::test_helpers::MetricSnapshot; + + tokio::time::timeout(Duration::from_secs(5), async { + let registry = prometheus::Registry::new(); + let metrics = Arc::new(VeloMetrics::register(®istry).unwrap()); + let handshakes = || { + MetricSnapshot::from_registry(®istry).counter( + "velo_messenger_client_resolution_total", + &[("path", "handshake"), ("outcome", "attempt")], + ) + }; + let (transport_a, transport_b) = make_transport_pair(); + let a = Messenger::builder() + .add_transport(transport_a) + .metrics(metrics) + .build() + .await + .unwrap(); + let b = Messenger::builder() + .add_transport(transport_b.clone()) + .build() + .await + .unwrap(); + + a.register_peer(b.peer_info()).unwrap(); + a.wait_for_handler(b.instance_id(), "_hello").await.unwrap(); + assert_eq!(handshakes(), 1.0, "unknown peers need the initial hello"); + assert!(b.client.is_peer_registered(a.instance_id())); + + // Dynamo registers the response peer again for each request. + for _ in 0..3 { + a.register_peer(b.peer_info()).unwrap(); + a.wait_for_handler(b.instance_id(), "_hello").await.unwrap(); + } + assert_eq!(handshakes(), 1.0, "cached handlers need no round trip"); + + b.register_streaming_handler(Handler::am_handler("_late", |_ctx| Ok(())).build()) + .unwrap(); + a.wait_for_handler(b.instance_id(), "_late").await.unwrap(); + assert_eq!(handshakes(), 2.0, "a missing cached handler must refresh"); + + // Reuse the address with a new instance, as after a peer restart. + b.graceful_shutdown(crate::ShutdownPolicy::Timeout(Duration::from_secs(1))) + .await; + let restarted = Messenger::builder() + .add_transport(transport_b) + .build() + .await + .unwrap(); + assert_ne!(b.instance_id(), restarted.instance_id()); + a.register_peer(restarted.peer_info()).unwrap(); + a.wait_for_handler(restarted.instance_id(), "_hello") + .await + .unwrap(); + assert_eq!(handshakes(), 3.0, "new instances need a new handshake"); + }) + .await + .expect("handler discovery must complete"); + } + /// Exercise the `.await_capacity()` chain on each public builder /// (AmSend, AmSync, Unary, TypedUnary). Each delegation is a one-line /// pass-through to `MessageBuilder::await_capacity`; this test pins that diff --git a/lib/velo/src/observability.rs b/lib/velo/src/observability.rs index 472654a8..c557ad8b 100644 --- a/lib/velo/src/observability.rs +++ b/lib/velo/src/observability.rs @@ -865,7 +865,7 @@ impl BatcherWake { } } -/// Collectors for the `messenger-mux-v1` streaming transport. +/// Collectors for the `messenger-mux-v2` streaming transport. /// /// Bound once per mux transport via [`VeloMetrics::bind_mux`] so the hot paths /// (the per-peer batcher and the `_stream_batch` ingress lane) hold concrete diff --git a/lib/velo/src/streaming.rs b/lib/velo/src/streaming.rs index dd1246ab..48d5ff41 100644 --- a/lib/velo/src/streaming.rs +++ b/lib/velo/src/streaming.rs @@ -30,7 +30,7 @@ pub mod frame; #[cfg(feature = "grpc")] pub mod grpc_transport; pub mod handle; -/// Batched, multiplexed streaming over the Messenger (`messenger-mux-v1`). +/// Batched, multiplexed streaming over the Messenger (`messenger-mux-v2`). /// /// The transport itself is internal — it is opt-in and selected when a /// stream's terms are decided, at attach or at pre-bind for zero-RTT setup, so diff --git a/lib/velo/src/streaming/anchor.rs b/lib/velo/src/streaming/anchor.rs index 6e8aea8d..13212cde 100644 --- a/lib/velo/src/streaming/anchor.rs +++ b/lib/velo/src/streaming/anchor.rs @@ -194,6 +194,7 @@ pub(crate) struct AnchorEntry { /// ordinary attach path, which is every anchor whose application sends no /// ticket. pub prebind: Option, + pub stop_requested: bool, } // --------------------------------------------------------------------------- @@ -223,19 +224,9 @@ pub(crate) struct AnchorEntry { /// learns a `StreamCancelHandle` either, and this is the only prompt path /// left. /// -/// Three races follow from `is_claimed` being an unlocked `OnceLock` read -/// rather than something the registry's shard lock also covers, and none of -/// them leak: an `OpenSlot` claiming between [`PreBind::drop`]'s read and its -/// `release_bind` leaves a live slot with no close posted, discovered on the -/// producer's next record or heartbeat; one claiming right after -/// [`AnchorManager::adopt_prebind`] reads `is_claimed` for its verdict hands -/// the attaching sender terms already in use, whose `OpenSlot` is then -/// refused `UnknownSlot`; and one claiming between the co-located branch of -/// [`AnchorManager::attach_stream_anchor`]'s guard read and its release below -/// is evicted correctly rather than left as a second writer — that branch -/// cancels this bind's pump before the release runs, so the claim has -/// already lost its only route to the anchor's channel by the time -/// [`PreBind::drop`] posts the close. +/// Claim and lifecycle changes are serialized by `DrainSignal`. An early stop +/// or cancel is therefore applied when OpenSlot claims the bind. Adoption still +/// checks the claim before it transfers ownership; a competing open is refused. pub(crate) struct PreBind { /// The anchor this slot was bound for. Half of the bind's key; the other /// half is `ticket.routing_session_id`. @@ -317,8 +308,10 @@ impl Drop for PreBind { let Some(mux) = self.mux.upgrade() else { return; }; - match self.drain.claimed() { - Some((peer, slot)) => mux.close_claimed_slot(peer, slot), + match self.drain.cancel() { + Some((peer, slot)) => { + mux.cancel_claimed_session(peer, slot, self.ticket.routing_session_id) + } None => mux.release_bind(self.anchor_id, self.ticket.routing_session_id), } } @@ -336,7 +329,17 @@ struct SenderIdentity { sender_stream_id: u64, cancel_token: CancellationToken, poison_tx: flume::Sender<()>, - poison_rx: flume::Receiver<()>, + stop_token: CancellationToken, + registry: Arc, + armed: bool, +} + +impl Drop for SenderIdentity { + fn drop(&mut self) { + if self.armed { + self.registry.senders.remove(&self.sender_stream_id); + } + } } /// What an incoming `_anchor_attach` may do with a pre-bound slot. @@ -358,6 +361,7 @@ pub(crate) enum PrebindAdoption { /// Wrapped in `Arc` so `StreamController` can outlive `StreamAnchor` being /// moved into StreamExt combinators. struct StreamControllerInner { + worker_id: velo_ext::WorkerId, local_id: u64, registry: Arc>, /// Sibling MPSC registry — held so the shared gauge update includes MPSC @@ -385,6 +389,32 @@ pub struct StreamController { } impl StreamController { + /// Request graceful stop without closing the response stream. + /// An early request is retained until a producer attaches or opens its ticket. + pub fn request_stop(&self) { + let Some(mut entry) = self.inner.registry.get_mut(&self.inner.local_id) else { + return; + }; + if entry.stop_requested { + return; + } + entry.stop_requested = true; + if let Some(prebind) = &entry.prebind { + if let Some((peer, slot)) = prebind.drain.request_stop() + && let Some(mux) = prebind.mux.upgrade() + { + mux.request_stop(peer, slot, prebind.ticket.routing_session_id); + } + } else if let Some(handle) = entry.stream_cancel_handle { + crate::streaming::control::request_sender_stop( + handle, + self.inner.worker_id, + &self.inner.sender_registry, + self.inner.messenger.as_ref(), + ); + } + } + /// Cancel the stream: remove the anchor from the registry and send a /// `_stream_cancel` AM to the sender's worker (fire-and-forget). /// @@ -432,7 +462,10 @@ impl StreamController { // the entry was already removed (e.g. finalize/detach ran first). if let Some(handle) = stream_cancel_handle { let (sender_worker_id, sender_stream_id) = handle.unpack(); - if let Some((_, entry)) = self.inner.sender_registry.senders.remove(&sender_stream_id) { + if sender_worker_id == self.inner.worker_id + && let Some((_, entry)) = + self.inner.sender_registry.senders.remove(&sender_stream_id) + { drop(entry.rx_closer.lock().unwrap().take()); entry.cancel_token.cancel(); } @@ -541,6 +574,7 @@ impl StreamAnchor { metrics, } = ctx; let inner = Arc::new(StreamControllerInner { + worker_id: handle.unpack().0, local_id, registry: registry.clone(), mpsc_registry: mpsc_registry.clone(), @@ -895,7 +929,7 @@ pub struct AnchorManager { #[builder(setter(skip), default = "std::sync::OnceLock::new()")] pub(crate) messenger_lock: std::sync::OnceLock>, - /// The `messenger-mux-v1` transport, when one is installed. + /// The `messenger-mux-v2` transport, when one is installed. /// /// Held as its concrete type rather than only as a registry entry because /// negotiation needs things the `FrameTransport` trait does not carry: the @@ -994,6 +1028,7 @@ impl AnchorManager { heartbeat_interval, stream_cancel_handle: None, // populated on attach prebind: None, // populated by `prebind_anchor` + stop_requested: false, }; self.registry.insert(local_id, entry); @@ -1100,7 +1135,7 @@ impl AnchorManager { /// /// Separate from the transport registry, which the mux also joins: the /// registry answers "can I `connect()` on this key", while this answers - /// "may I offer, and drive, `messenger-mux-v1`". Both are needed and they + /// "may I offer, and drive, `messenger-mux-v2`". Both are needed and they /// are set together by the builder. pub(crate) fn install_mux( &self, @@ -1249,6 +1284,9 @@ impl AnchorManager { // slot to a sender that used it instead of the // ticket -- see `PumpContext::prebound`. let prebound = Arc::new(AtomicBool::new(true)); + if entry.stop_requested { + drain.request_stop(); + } entry.prebind = Some(PreBind { anchor_id: local_id, ticket: ticket.clone(), @@ -1427,6 +1465,14 @@ impl AnchorManager { let prebind = entry.prebind.take().expect("the verdict says it is there"); entry.attachment = true; entry.stream_cancel_handle = Some(req.stream_cancel_handle); + if entry.stop_requested { + crate::streaming::control::request_sender_stop( + req.stream_cancel_handle, + self.worker_id, + &self.sender_registry, + self.messenger_lock.get(), + ); + } (PrebindAdoption::Adopted(prebind.adopt()), None) } } @@ -1488,13 +1534,15 @@ impl AnchorManager { /// `CreditUpdate` used to cost. async fn connect_streaming( &self, - key: &velo_ext::TransportKey, + ticket: &crate::streaming::control::StreamOpenTicket, peer: velo_ext::WorkerId, anchor_id: u64, - session_id: u64, - initial_credit: u32, - slot_byte_budget: u32, + lifecycle: Option<(CancellationToken, CancellationToken)>, ) -> Result>, AttachError> { + let key = &ticket.streaming_transport_key; + let session_id = ticket.routing_session_id; + let initial_credit = ticket.initial_credit; + let slot_byte_budget = ticket.slot_byte_budget; match crate::streaming::negotiation::choose(key, initial_credit, slot_byte_budget) { Ok(crate::streaming::negotiation::Connect::Mux(limits)) => { let mux = self.mux.get().ok_or_else(|| { @@ -1504,7 +1552,7 @@ impl AnchorManager { )) })?; Ok(mux - .connect_negotiated(peer, anchor_id, session_id, limits) + .connect_controlled(peer, anchor_id, session_id, limits, lifecycle) .await?) } Ok(crate::streaming::negotiation::Connect::Legacy) => { @@ -1749,6 +1797,12 @@ impl AnchorManager { &self.sender_registry, )))?; + messenger.register_streaming_handler( + crate::streaming::control::create_stream_stop_handler(Arc::clone( + &self.sender_registry, + )), + )?; + // MPSC handlers — share the same SenderRegistry so `_stream_cancel` // covers both SPSC and MPSC senders uniformly. messenger.register_streaming_handler( @@ -1905,17 +1959,14 @@ impl AnchorManager { /// already released or has expired still returns `Ok` here and only /// surfaces later, as a send failure on the returned `StreamSender`. /// - /// No `_anchor_attach` means no `StreamCancelHandle` either, so the - /// returned sender's - /// [`cancellation_token`](crate::streaming::sender::StreamSender::cancellation_token) - /// never fires; a dropped consumer surfaces as a send error instead. + /// Slot lifecycle signals reach the producer even before its first send. /// /// # Errors /// - [`AttachError::TransportError`] if `handle` names this node's own /// worker id (a ticket is for a *different* worker to open; a /// same-worker sender should use /// [`attach_stream_anchor`](Self::attach_stream_anchor) instead), if - /// the ticket names `messenger-mux-v1` but this node has no mux + /// the ticket names `messenger-mux-v2` but this node has no mux /// installed, if the ticket names a transport key this node cannot /// resolve, or if staging the local slot fails after retrying a /// batcher that keeps retiring underneath it. @@ -1967,11 +2018,24 @@ impl AnchorManager { /// Allocate the sender-side identity one stream is opened under. fn new_sender_identity(&self) -> SenderIdentity { let (poison_tx, poison_rx) = flume::bounded::<()>(1); + let sender_stream_id = self.next_sender_stream_id.fetch_add(1, Ordering::Relaxed) + 1; + let cancel_token = CancellationToken::new(); + let stop_token = cancel_token.child_token(); + self.sender_registry.senders.insert( + sender_stream_id, + crate::streaming::control::SenderEntry { + cancel_token: cancel_token.clone(), + stop_token: stop_token.clone(), + rx_closer: std::sync::Mutex::new(Some(poison_rx)), + }, + ); SenderIdentity { - sender_stream_id: self.next_sender_stream_id.fetch_add(1, Ordering::Relaxed) + 1, - cancel_token: CancellationToken::new(), + sender_stream_id, + cancel_token, + stop_token, poison_tx, - poison_rx, + registry: self.sender_registry.clone(), + armed: true, } } @@ -1986,7 +2050,7 @@ impl AnchorManager { &self, handle: StreamAnchorHandle, ticket: &crate::streaming::control::StreamOpenTicket, - identity: SenderIdentity, + mut identity: SenderIdentity, ) -> Result, AttachError> { let started = Instant::now(); let (handle_worker_id, local_id) = handle.unpack(); @@ -1995,12 +2059,10 @@ impl AnchorManager { // bound streaming transport, then connect by WorkerId. let frame_tx = match self .connect_streaming( - &ticket.streaming_transport_key, + ticket, handle_worker_id, local_id, - ticket.routing_session_id, - ticket.initial_credit, - ticket.slot_byte_budget, + Some((identity.cancel_token.clone(), identity.stop_token.clone())), ) .await { @@ -2016,21 +2078,10 @@ impl AnchorManager { } }; - let SenderIdentity { - sender_stream_id, - cancel_token, - poison_tx, - poison_rx, - } = identity; - - // Register SenderEntry for _stream_cancel routing - self.sender_registry.senders.insert( - sender_stream_id, - crate::streaming::control::SenderEntry { - cancel_token: cancel_token.clone(), - rx_closer: std::sync::Mutex::new(Some(poison_rx)), - }, - ); + let sender_stream_id = identity.sender_stream_id; + let cancel_token = identity.cancel_token.clone(); + let poison_tx = identity.poison_tx.clone(); + identity.armed = false; // Build StreamSender: frame_tx from the transport (not a local registry // frame_tx). No local AnchorEntry is created for the remote anchor. @@ -2189,9 +2240,13 @@ impl AnchorManager { let (poison_tx, poison_rx) = flume::bounded::<()>(1); let sender_entry = crate::streaming::control::SenderEntry { + stop_token: cancel_token.child_token(), cancel_token: cancel_token.clone(), rx_closer: std::sync::Mutex::new(Some(poison_rx)), }; + if entry.stop_requested { + sender_entry.stop_token.cancel(); + } self.sender_registry .senders .insert(sender_stream_id, sender_entry); @@ -2388,6 +2443,7 @@ impl AnchorManager { // Register SenderEntry outside the shard lock. let sender_entry = crate::streaming::control::SenderEntry { + stop_token: cancel_token.child_token(), cancel_token: cancel_token.clone(), rx_closer: std::sync::Mutex::new(Some(poison_rx)), }; @@ -2482,16 +2538,21 @@ impl AnchorManager { }; let frame_tx = self .connect_streaming( - &streaming_transport_key, + &crate::streaming::control::StreamOpenTicket { + streaming_transport_key: streaming_transport_key.clone(), + heartbeat_interval_ms, + routing_session_id: connect_session_id, + initial_credit, + slot_byte_budget, + }, handle_worker_id, local_id, - connect_session_id, - initial_credit, - slot_byte_budget, + None, ) .await?; let sender_entry = crate::streaming::control::SenderEntry { + stop_token: cancel_token.child_token(), cancel_token: cancel_token.clone(), rx_closer: std::sync::Mutex::new(Some(poison_rx)), }; diff --git a/lib/velo/src/streaming/control.rs b/lib/velo/src/streaming/control.rs index 42b6f6e5..179b86c2 100644 --- a/lib/velo/src/streaming/control.rs +++ b/lib/velo/src/streaming/control.rs @@ -128,6 +128,8 @@ pub struct StreamCancelRequest { pub struct SenderEntry { /// Fires when `_stream_cancel` is received — user-facing via `cancellation_token()`. pub cancel_token: tokio_util::sync::CancellationToken, + /// Graceful stop leaves the response channel open. + pub stop_token: tokio_util::sync::CancellationToken, /// Drop this to signal cancellation to `StreamSender::send()` via /// `poison_tx.is_disconnected()`. Wrapped in `Mutex>` so the @@ -180,6 +182,60 @@ pub fn create_stream_cancel_handler( .build() } +/// Send a graceful stop through the identity established by attach. +/// +/// A sender on `local_worker` is stopped through the local registry, as +/// `StreamController::cancel` does: a same-worker attach registers it there, +/// and an active message to the local worker is not guaranteed to resolve. +pub(crate) fn request_sender_stop( + handle: StreamCancelHandle, + local_worker: velo_ext::WorkerId, + registry: &SenderRegistry, + messenger: Option<&Arc>, +) { + let (worker, sender_stream_id) = handle.unpack(); + if worker == local_worker { + if let Some(entry) = registry.senders.get(&sender_stream_id) { + entry.stop_token.cancel(); + } + return; + } + // Stream IDs are local to a worker. Do not cancel an unrelated local sender. + if let Some(messenger) = messenger { + // On the messenger's runtime, not the caller's: `request_stop` is + // synchronous and can run on a thread with no runtime, where the stop + // would be dropped after the anchor had already recorded it. + let rt = messenger.runtime().clone(); + let messenger = Arc::clone(messenger); + rt.spawn(async move { + let payload = serde_json::to_vec(&StreamCancelRequest { sender_stream_id }) + .expect("stream identity"); + let _ = messenger + .am_send_streaming("_stream_stop") + .expect("stream stop handler") + .raw_payload(bytes::Bytes::from(payload)) + .worker(worker) + .send() + .await; + }); + } else if let Some(entry) = registry.senders.get(&sender_stream_id) { + entry.stop_token.cancel(); + } +} + +pub(crate) fn create_stream_stop_handler( + registry: Arc, +) -> crate::messenger::Handler { + crate::messenger::Handler::am_handler("_stream_stop", move |ctx: crate::messenger::Context| { + let req = serde_json::from_slice::(&ctx.payload)?; + if let Some(entry) = registry.senders.get(&req.sender_stream_id) { + entry.stop_token.cancel(); + } + Ok(()) + }) + .build() +} + // --------------------------------------------------------------------------- // Request / Response types // --------------------------------------------------------------------------- @@ -202,7 +258,7 @@ pub struct AnchorAttachRequest { /// asked to `connect()` on. /// /// The receiver intersects this with its own installed set and prefers - /// `messenger-mux-v1` when it appears in both. `#[serde(default)]` means a + /// `messenger-mux-v2` when it appears in both. `#[serde(default)]` means a /// sender that predates negotiation deserializes as one advertising /// nothing, which is exactly right: an empty list can never intersect, so /// absent a pre-bound slot on the anchor, such a sender is always answered @@ -430,7 +486,7 @@ pub fn create_anchor_attach_handler(manager: Arc) -> crate::messe .fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1; // Which transport this attach rides is decided here, from what - // the sender advertised: `messenger-mux-v1` when both sides + // the sender advertised: `messenger-mux-v2` when both sides // named it, and otherwise exactly the local default this // handler answered with before negotiation existed. let selection = manager.select_streaming_transport(&req.supported_transport_keys); @@ -512,6 +568,15 @@ pub fn create_anchor_attach_handler(manager: Arc) -> crate::messe // Mark as attached and store cancel handle for upstream cancel routing entry.attachment = true; entry.stream_cancel_handle = Some(req.stream_cancel_handle); + if entry.stop_requested { + request_sender_stop( + req.stream_cancel_handle, + // The anchor's worker is this worker. + req.handle.unpack().0, + &manager.sender_registry, + manager.messenger_lock.get(), + ); + } // Drop shard lock before spawning drop(occ); diff --git a/lib/velo/src/streaming/control/tests.rs b/lib/velo/src/streaming/control/tests.rs index 78781775..6b122ace 100644 --- a/lib/velo/src/streaming/control/tests.rs +++ b/lib/velo/src/streaming/control/tests.rs @@ -180,6 +180,7 @@ async fn reader_pump_unclaimed_bind_reap_increments_counter() { heartbeat_interval: std::time::Duration::from_secs(5), stream_cancel_handle: None, prebind: None, + stop_requested: false, }, ); @@ -463,7 +464,7 @@ fn an_attach_request_round_trips_its_advertised_keys() { session_id: 3, stream_cancel_handle: StreamCancelHandle::pack(velo_ext::WorkerId::from_u64(4), 5), supported_transport_keys: vec![ - velo_ext::TransportKey::new("messenger-mux-v1"), + velo_ext::TransportKey::new("messenger-mux-v2"), velo_ext::TransportKey::new("tcp-stream"), ], }; @@ -476,7 +477,7 @@ fn an_attach_request_round_trips_its_advertised_keys() { .iter() .map(velo_ext::TransportKey::as_str) .collect::>(), - ["messenger-mux-v1", "tcp-stream"], + ["messenger-mux-v2", "tcp-stream"], ); } @@ -823,6 +824,7 @@ fn make_pump_test_infra() -> ( heartbeat_interval: Duration::from_secs(5), stream_cancel_handle: None, prebind: None, + stop_requested: false, }, ); @@ -912,6 +914,7 @@ fn make_prebind_pump_test_infra( heartbeat_interval: heartbeat_deadline, stream_cancel_handle: None, prebind: None, + stop_requested: false, }, ); @@ -1494,6 +1497,7 @@ async fn reader_pump_does_not_count_a_blocked_forward_as_heartbeat_silence() { heartbeat_interval: deadline, stream_cancel_handle: None, prebind: None, + stop_requested: false, }, ); let ctx = crate::streaming::anchor::AnchorContext { @@ -1703,6 +1707,7 @@ async fn test_child_token_reattach_pump_survives() { heartbeat_interval: Duration::from_secs(5), stream_cancel_handle: None, prebind: None, + stop_requested: false, }, ); @@ -1883,7 +1888,7 @@ fn attach_response_golden_encoding_unchanged() { "handle": {"hi": 1, "lo": 2}, "session_id": 3, "stream_cancel_handle": {"hi": 4, "lo": 5}, - "supported_transport_keys": ["messenger-mux-v1"] + "supported_transport_keys": ["messenger-mux-v2"] }"#; let decoded: AnchorAttachRequest = serde_json::from_str(ticketless_request).expect("a ticketless request must deserialize"); @@ -1894,14 +1899,14 @@ fn attach_response_golden_encoding_unchanged() { .iter() .map(velo_ext::TransportKey::as_str) .collect::>(), - ["messenger-mux-v1"] + ["messenger-mux-v2"] ); // The response keeps its five fields and gains none. Compared as a value // rather than as bytes because the field *set* is the invariant; rmp-serde // writes named fields, so an added one would show up here as an extra key. let response = AnchorAttachResponse::Ok { - streaming_transport_key: velo_ext::TransportKey::new("messenger-mux-v1"), + streaming_transport_key: velo_ext::TransportKey::new("messenger-mux-v2"), heartbeat_interval_ms: 1234, routing_session_id: 42, initial_credit: 64, @@ -1933,7 +1938,7 @@ fn attach_response_golden_encoding_unchanged() { // has no legacy sender to default for (see `StreamOpenTicket`'s own doc), // so `heartbeat_interval_ms` is set explicitly rather than left absent. let ticket: StreamOpenTicket = serde_json::from_str( - r#"{"streaming_transport_key":"messenger-mux-v1","heartbeat_interval_ms":1500,"routing_session_id":7,"initial_credit":8,"slot_byte_budget":0}"#, + r#"{"streaming_transport_key":"messenger-mux-v2","heartbeat_interval_ms":1500,"routing_session_id":7,"initial_credit":8,"slot_byte_budget":0}"#, ) .expect("a fully-populated ticket must deserialize"); assert_eq!(ticket.routing_session_id, 7); @@ -1955,25 +1960,25 @@ fn attach_response_golden_encoding_unchanged() { /// -- all wrong answers reached without error. #[test] fn a_ticket_missing_a_minted_field_fails_rather_than_silently_defaulting() { - let missing_routing_session_id = r#"{"streaming_transport_key":"messenger-mux-v1","initial_credit":8,"slot_byte_budget":4096}"#; + let missing_routing_session_id = r#"{"streaming_transport_key":"messenger-mux-v2","initial_credit":8,"slot_byte_budget":4096}"#; assert!( serde_json::from_str::(missing_routing_session_id).is_err(), "a ticket missing routing_session_id must not silently decode as session 0" ); - let missing_initial_credit = r#"{"streaming_transport_key":"messenger-mux-v1","routing_session_id":7,"slot_byte_budget":4096}"#; + let missing_initial_credit = r#"{"streaming_transport_key":"messenger-mux-v2","routing_session_id":7,"slot_byte_budget":4096}"#; assert!( serde_json::from_str::(missing_initial_credit).is_err(), "a ticket missing initial_credit must not silently decode as 'not offering the mux'" ); - let missing_slot_byte_budget = r#"{"streaming_transport_key":"messenger-mux-v1","routing_session_id":7,"initial_credit":8}"#; + let missing_slot_byte_budget = r#"{"streaming_transport_key":"messenger-mux-v2","routing_session_id":7,"initial_credit":8}"#; assert!( serde_json::from_str::(missing_slot_byte_budget).is_err(), "a ticket missing slot_byte_budget must not silently decode as 'use the default'" ); - let missing_heartbeat = r#"{"streaming_transport_key":"messenger-mux-v1","routing_session_id":7,"initial_credit":8,"slot_byte_budget":4096}"#; + let missing_heartbeat = r#"{"streaming_transport_key":"messenger-mux-v2","routing_session_id":7,"initial_credit":8,"slot_byte_budget":4096}"#; assert!( serde_json::from_str::(missing_heartbeat).is_err(), "a ticket missing heartbeat_interval_ms must not silently decode at 5000ms, which can \ diff --git a/lib/velo/src/streaming/messenger_mux/config.rs b/lib/velo/src/streaming/messenger_mux/config.rs index c4486628..a24adb59 100644 --- a/lib/velo/src/streaming/messenger_mux/config.rs +++ b/lib/velo/src/streaming/messenger_mux/config.rs @@ -135,7 +135,7 @@ pub struct MuxConfig { /// /// **Defaults to `false`, and stays that way** — the mux is opt-in, not the /// default transport. This flag is also the rollback: set it back to - /// `false` and the node stops registering `messenger-mux-v1` and stops + /// `false` and the node stops registering `messenger-mux-v2` and stops /// advertising it on attach, so the next attach negotiates the legacy path /// with no code change and no wire change. That is what makes a canary /// safe, and why activation is config-only. diff --git a/lib/velo/src/streaming/messenger_mux/flow_control/mod.rs b/lib/velo/src/streaming/messenger_mux/flow_control/mod.rs index bbfb9a25..630bac2b 100644 --- a/lib/velo/src/streaming/messenger_mux/flow_control/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/flow_control/mod.rs @@ -181,9 +181,10 @@ impl CreditClass { /// `reader_pump`'s `DETECTION_MULTIPLIER` stops firing on a saturated slot. pub(crate) const fn of(record_type: RecordType, is_terminal: bool) -> Self { match record_type { - RecordType::OpenSlot | RecordType::CloseSlot | RecordType::CreditUpdate => { - Self::Control - } + RecordType::OpenSlot + | RecordType::CloseSlot + | RecordType::CreditUpdate + | RecordType::LifecycleSlot => Self::Control, RecordType::Data if is_terminal => Self::Terminal, RecordType::Data | RecordType::SlotHeartbeat => Self::Data, } diff --git a/lib/velo/src/streaming/messenger_mux/ingress/drain.rs b/lib/velo/src/streaming/messenger_mux/ingress/drain.rs index 7f6ddd9a..5cb6a0fa 100644 --- a/lib/velo/src/streaming/messenger_mux/ingress/drain.rs +++ b/lib/velo/src/streaming/messenger_mux/ingress/drain.rs @@ -44,6 +44,8 @@ pub(crate) struct DrainSignal { /// Whose bind this turned out to be, and where its drains are posted. All /// of it arrives together when an `OpenSlot` claims the bind. claim: OnceLock, + // Serializes claim with early stop/cancel; neither side can miss the other. + lifecycle: std::sync::Mutex, /// Records this slot's pump has taken out of the buffer since the last /// [`IngressSlot::reconcile`](super::slot::IngressSlot::reconcile) swapped /// it to zero. @@ -68,6 +70,7 @@ impl DrainSignal { pub(crate) fn new(wake: flume::Sender) -> Self { Self { claim: OnceLock::new(), + lifecycle: std::sync::Mutex::new(0), drained: AtomicU32::new(0), listed: AtomicBool::new(false), wake, @@ -83,13 +86,29 @@ impl DrainSignal { slot: SlotId, pending: Arc, lane: flume::Sender, - ) { + ) -> u8 { + let lifecycle = self.lifecycle.lock().unwrap(); let _ = self.claim.set(SlotClaim { peer, slot, pending, lane, }); + *lifecycle + } + + pub(crate) fn request_stop(&self) -> Option<(WorkerId, SlotId)> { + let mut state = self.lifecycle.lock().unwrap(); + if *state != 0 { + return None; + } + *state = 1; + self.claimed() + } + + pub(crate) fn cancel(&self) -> Option<(WorkerId, SlotId)> { + *self.lifecycle.lock().unwrap() = 2; + self.claimed() } /// The peer and slot that claimed this bind, once one has. diff --git a/lib/velo/src/streaming/messenger_mux/ingress/mod.rs b/lib/velo/src/streaming/messenger_mux/ingress/mod.rs index 035028f2..ed25111c 100644 --- a/lib/velo/src/streaming/messenger_mux/ingress/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/ingress/mod.rs @@ -179,6 +179,7 @@ pub(crate) struct BatchOutcome { pub(crate) grants: Vec<(SlotId, u32)>, /// `CloseSlot`s addressed to slots *we* own, likewise. pub(crate) peer_closes: Vec<(SlotId, CloseReason)>, + pub(crate) peer_stops: Vec<(SlotId, u64, bool)>, /// Slots this batch created. pub(crate) opened: usize, /// Slots this batch retired. @@ -211,6 +212,7 @@ impl IngressRegistry { peer: WorkerId, id: SlotId, metrics: Option<&MuxMetricsHandle>, + session_id: Option, ) -> Option { let entry = self.peers.get(&peer)?; let mut state = lock(entry.value()); @@ -227,6 +229,10 @@ impl IngressRegistry { { return None; } + let live = state.slots[id.index() as usize].as_ref()?; + if session_id.is_some_and(|session| live.session_id != session) { + return None; + } let mut outcome = BatchOutcome::default(); finish_close( &mut state, @@ -235,9 +241,16 @@ impl IngressRegistry { metrics, &mut outcome, ); - (outcome.closed > 0).then_some(ReplyRecord::CloseSlot { - slot: id, - reason: CloseReason::UnknownSlot, + (outcome.closed > 0).then_some(match session_id { + Some(session_id) => ReplyRecord::LifecycleSlot { + slot: id, + session_id, + cancel: true, + }, + None => ReplyRecord::CloseSlot { + slot: id, + reason: CloseReason::UnknownSlot, + }, }) } @@ -586,6 +599,9 @@ fn apply_record( RecordBody::CloseSlot { reason } => { close_slot(state, ctx, record.slot, record.frame_seq, reason, outcome); } + RecordBody::LifecycleSlot { session_id, cancel } => { + outcome.peer_stops.push((record.slot, session_id, cancel)) + } RecordBody::Data(body) => deliver(state, ctx, record, body.to_vec(), outcome), RecordBody::SlotHeartbeat => deliver(state, ctx, record, heartbeat_frame(), outcome), } @@ -680,14 +696,14 @@ fn open_slot( // The bind now has an owner, so its drain signal can start counting and // posting. Before this point it is inert: nothing has been delivered on // this slot, so nothing can have drained. - bind.drain.claimed_by( + let lifecycle = bind.drain.claimed_by( ctx.peer, id, ctx.registry.pending_wake(ctx.peer), state.drained_tx.clone(), ); - let slot = IngressSlot::new( + let mut slot = IngressSlot::new( id, bind.frame_tx, Arc::clone(&bind.drain), @@ -695,8 +711,23 @@ fn open_slot( ctx.config.slot_byte_budget, record.frame_seq.saturating_add(1), ); + slot.session_id = session_id; state.slots[index] = Some(slot); outcome.opened += 1; + if lifecycle == 2 { + finish_close(state, id, CloseReason::UnknownSlot, ctx.metrics, outcome); + outcome.replies.push(ReplyRecord::LifecycleSlot { + slot: id, + session_id, + cancel: true, + }); + } else if lifecycle == 1 { + outcome.replies.push(ReplyRecord::LifecycleSlot { + slot: id, + session_id, + cancel: false, + }); + } } fn close_slot( diff --git a/lib/velo/src/streaming/messenger_mux/ingress/slot.rs b/lib/velo/src/streaming/messenger_mux/ingress/slot.rs index d21f4ea6..6978077f 100644 --- a/lib/velo/src/streaming/messenger_mux/ingress/slot.rs +++ b/lib/velo/src/streaming/messenger_mux/ingress/slot.rs @@ -51,6 +51,7 @@ pub(super) enum Applied { /// One receive-side slot. pub(super) struct IngressSlot { + pub(super) session_id: u64, /// Index and generation this slot answers to. pub(super) id: SlotId, /// The mux-owned `C + 1`-deep buffer handed to the anchor by `bind`. @@ -118,6 +119,7 @@ impl IngressSlot { first_seq: u32, ) -> Self { Self { + session_id: 0, id, frame_tx, drain, diff --git a/lib/velo/src/streaming/messenger_mux/mod.rs b/lib/velo/src/streaming/messenger_mux/mod.rs index 4cbfbeaf..a6cbed32 100644 --- a/lib/velo/src/streaming/messenger_mux/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/mod.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Batched, multiplexed streaming over the Messenger — the `messenger-mux-v1` +//! Batched, multiplexed streaming over the Messenger — the `messenger-mux-v2` //! transport described by `docs/src/concepts/batched-streaming.md`. //! //! Today one stream owns one connection: X concurrent streams to one peer means @@ -142,7 +142,7 @@ pub use self::config::{AutoFlush, FlushPolicy, MuxConfig}; /// [`StreamSender::negotiated_transport`](crate::streaming::StreamSender::negotiated_transport) /// is compared against — a caller that had to spell the string itself would be /// re-deriving the one value negotiation is keyed on. -pub const MESSENGER_MUX_KEY: &str = "messenger-mux-v1"; +pub const MESSENGER_MUX_KEY: &str = "messenger-mux-v2"; /// The active-message handler every batch travels through. pub(crate) const STREAM_BATCH_HANDLER: &str = "_stream_batch"; @@ -180,7 +180,7 @@ const CONNECT_ATTEMPTS: usize = 3; /// scale with slot count. It needs to absorb a burst across distinct peers. const DRAIN_WAKE_CAPACITY: usize = 1024; -/// The `messenger-mux-v1` [`FrameTransport`]. +/// The `messenger-mux-v2` [`FrameTransport`]. /// /// Holds no listener and no connections. `connect` allocates a slot on the /// peer's batcher; `bind` registers a buffer the peer's `OpenSlot` will claim. @@ -232,6 +232,17 @@ struct MuxCore { } impl MessengerMuxTransport { + pub(crate) fn request_stop(&self, peer: WorkerId, slot: protocol::SlotId, session_id: u64) { + self.core.return_credit( + peer, + vec![peer_batcher::ReplyRecord::LifecycleSlot { + slot, + session_id, + cancel: false, + }], + ); + } + /// Take the [`ingress::DrainSignal`] `bind` parked for this pair. /// /// Called once by the attach path, between `bind` returning and the pump @@ -380,12 +391,18 @@ impl MuxCore { } } - if outcome.replies.is_empty() && outcome.grants.is_empty() && outcome.peer_closes.is_empty() + if outcome.replies.is_empty() + && outcome.grants.is_empty() + && outcome.peer_closes.is_empty() + && outcome.peer_stops.is_empty() { return; } let batcher = self.batcher(peer); + for (slot, session_id, cancel) in outcome.peer_stops { + batcher.peer_stopped(slot, session_id, cancel); + } for (slot, delta) in outcome.grants { batcher.grant(slot, delta); } @@ -485,7 +502,7 @@ impl MuxCore { /// sends none. The reply is what that idle producer needs, since the /// fault that carries the same news to it otherwise rides on the next /// record it sends. - fn close_claimed_slot(&self, peer: WorkerId, slot: protocol::SlotId) { + fn close_claimed_slot(&self, peer: WorkerId, slot: protocol::SlotId, session_id: Option) { // Checked before touching the ingress table, not after: this runs // from a `Drop` that can land on a thread with no runtime under it // (`StreamController::cancel` guards its own spawn the same way and @@ -504,9 +521,9 @@ impl MuxCore { ); return; } - let Some(reply) = self - .ingress - .close_consumer_gone(peer, slot, self.metrics.as_ref()) + let Some(reply) = + self.ingress + .close_consumer_gone(peer, slot, self.metrics.as_ref(), session_id) else { return; }; @@ -701,8 +718,18 @@ impl MessengerMuxTransport { /// Idempotent, and silent where there is nothing to close: a stream that /// ended on its own terminal retired the slot then, so the ordinary end of /// a stream costs no extra record on the wire. + pub(crate) fn cancel_claimed_session( + &self, + peer: WorkerId, + slot: protocol::SlotId, + session_id: u64, + ) { + self.core.close_claimed_slot(peer, slot, Some(session_id)); + } + + #[cfg(test)] pub(crate) fn close_claimed_slot(&self, peer: WorkerId, slot: protocol::SlotId) { - self.core.close_claimed_slot(peer, slot); + self.core.close_claimed_slot(peer, slot, None); } /// Binds registered and neither claimed nor released. @@ -768,6 +795,17 @@ impl MessengerMuxTransport { anchor_id: u64, session_id: u64, limits: NegotiatedLimits, + ) -> BoxFuture<'_, Result>>> { + self.connect_controlled(peer, anchor_id, session_id, limits, None) + } + + pub(crate) fn connect_controlled( + &self, + peer: WorkerId, + anchor_id: u64, + session_id: u64, + limits: NegotiatedLimits, + lifecycle: Option<(CancellationToken, CancellationToken)>, ) -> BoxFuture<'_, Result>>> { let core = Arc::clone(&self.core); Box::pin(async move { @@ -787,6 +825,7 @@ impl MessengerMuxTransport { anchor_id, session_id, inlet: inlet_rx, + lifecycle: lifecycle.clone(), credit: limits.open_credit(), slot_byte_budget: limits.slot_byte_budget(), ack: ack_tx, diff --git a/lib/velo/src/streaming/messenger_mux/peer_batcher/control.rs b/lib/velo/src/streaming/messenger_mux/peer_batcher/control.rs index 0be4c5c2..b123e53f 100644 --- a/lib/velo/src/streaming/messenger_mux/peer_batcher/control.rs +++ b/lib/velo/src/streaming/messenger_mux/peer_batcher/control.rs @@ -65,6 +65,7 @@ use crate::observability::MuxMetricsHandle; /// Coalesced control for one slot **this** batcher owns. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub(super) struct OwnedControl { + pub(super) lifecycle: Option<(u64, bool)>, /// Credit granted since the batcher last looked. pub(super) credit: u32, /// The receiver asked us to abandon the slot. @@ -227,6 +228,10 @@ struct ControlState { /// matching means the key names a stale generation — one this index has /// since moved past via a reopen, or one it never had at all. live_generations: Vec>, + live_sessions: Vec, + // One coalesced lifecycle signal per locally owned consumer session. + // A session is never coalesced with a later reuse of the same slot id. + lifecycle_replies: HashMap<(u32, u64), bool>, /// Entries refused: an index `mine` never allocated, or a `rejects` key /// past its cap. An ordinary stale-generation race is not counted here — /// see `entry_mine`. @@ -242,6 +247,7 @@ impl ControlState { && self.peers.is_empty() && self.resolutions.is_empty() && self.rejects.is_empty() + && self.lifecycle_replies.is_empty() } /// The sweep evicted this batcher from the registry. @@ -263,7 +269,11 @@ impl ControlState { /// own allocation history rather than by anything a peer can grow. #[cfg(test)] fn len(&self) -> usize { - self.mine.len() + self.peers.len() + self.resolutions.len() + self.rejects.len() + self.mine.len() + + self.peers.len() + + self.resolutions.len() + + self.rejects.len() + + self.lifecycle_replies.len() } /// Take everything pending, leaving the state empty. @@ -291,6 +301,7 @@ impl ControlState { flush: std::mem::take(&mut self.flush), mine, peers, + lifecycle_replies: std::mem::take(&mut self.lifecycle_replies), } } @@ -377,6 +388,7 @@ pub(super) struct DrainedControl { pub(super) flush: bool, pub(super) mine: HashMap, pub(super) peers: HashMap, + pub(super) lifecycle_replies: HashMap<(u32, u64), bool>, } /// The state plus the wakeup that tells the batcher to look at it. @@ -470,13 +482,25 @@ impl ControlInbox { /// open just published, never against the one before it. Called before /// the peer can possibly have learned `slot`'s id, so by the time a /// legitimate grant for it can arrive, this has always already run. + #[cfg(test)] pub(super) fn note_allocated(&self, slot: SlotId) { + self.note_allocated_session(slot, 0); + } + + pub(super) fn note_allocated_session(&self, slot: SlotId, session_id: u64) { let mut state = self.lock(); let index = slot.index() as usize; if index >= state.live_generations.len() { state.live_generations.resize(index + 1, None); + state.live_sessions.resize(index + 1, 0); } state.live_generations[index] = Some(slot.generation()); + state.live_sessions[index] = session_id; + // Drop a pending lifecycle signal from a prior session even when the + // compact generation wrapped back to the same value. + if let Some(pending) = state.mine.get_mut(&slot.raw()) { + pending.lifecycle = None; + } } /// An inbound `CreditUpdate` for a slot we own. @@ -488,6 +512,18 @@ impl ControlInbox { }); } + pub(super) fn peer_stopped(&self, slot: SlotId, session_id: u64, cancel: bool) { + self.mutate(|state| { + if state.live_sessions.get(slot.index() as usize) != Some(&session_id) { + return; + } + if let Some(entry) = state.entry_mine(slot) { + let prior_cancel = entry.lifecycle.is_some_and(|(_, cancelled)| cancelled); + entry.lifecycle = Some((session_id, cancel || prior_cancel)); + } + }); + } + /// The receiver asked us to abandon a slot we own. pub(super) fn peer_closed(&self, slot: SlotId, reason: CloseReason) { self.mutate(|state| { @@ -521,6 +557,16 @@ impl ControlInbox { self.mutate(|state| { for record in records { match *record { + ReplyRecord::LifecycleSlot { + slot, + session_id, + cancel, + } => { + *state + .lifecycle_replies + .entry((slot.raw(), session_id)) + .or_default() |= cancel; + } ReplyRecord::CreditUpdate { slot, delta } => { let entry = state.entry_peer(slot); entry.credit = entry.credit.saturating_add(delta); @@ -610,6 +656,46 @@ mod tests { inbox } + #[test] + fn lifecycle_checks_session_after_generation_reuse() { + let inbox = ControlInbox::default(); + let id = slot(0, 0); + inbox.note_allocated_session(id, 11); + inbox.peer_stopped(id, 11, true); + // The compact slot generation can eventually wrap. The full session + // still prevents both queued and newly arriving stale cancellation. + inbox.note_allocated_session(id, 12); + inbox.peer_stopped(id, 11, true); + inbox.peer_stopped(id, 12, false); + assert_eq!( + inbox.take().unwrap().mine[&id.raw()].lifecycle, + Some((12, false)) + ); + inbox.peer_stopped(id, 12, true); + inbox.peer_stopped(id, 12, false); + assert_eq!( + inbox.take().unwrap().mine[&id.raw()].lifecycle, + Some((12, true)) + ); + } + + #[test] + fn lifecycle_replies_keep_sessions_separate_and_cancel_wins() { + let inbox = ControlInbox::default(); + let id = slot(0, 0); + for (session_id, cancel) in [(11, false), (12, false), (11, true), (11, false)] { + assert!(inbox.reply(&[ReplyRecord::LifecycleSlot { + slot: id, + session_id, + cancel + }])); + } + let replies = inbox.take().unwrap().lifecycle_replies; + assert_eq!(replies.len(), 2); + assert!(replies[&(id.raw(), 11)]); + assert!(!replies[&(id.raw(), 12)]); + } + #[test] fn credit_accumulates_into_one_entry() { let inbox = inbox_with(8); 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 849f3ccc..5434ae83 100644 --- a/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/peer_batcher/mod.rs @@ -64,6 +64,7 @@ mod writer; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering}; +use std::time::Duration; use dashmap::DashMap; use futures::future::FutureExt; @@ -105,6 +106,7 @@ pub(crate) struct OpenSlotRequest { pub(crate) anchor_id: u64, pub(crate) session_id: u64, pub(crate) inlet: flume::Receiver>, + pub(crate) lifecycle: Option<(CancellationToken, CancellationToken)>, /// The ledger the slot opens with — the window the receiver advertised on /// its attach response, already granted. /// @@ -140,6 +142,11 @@ const OPEN_QUEUE_DEPTH: usize = 64; /// enforced. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum ReplyRecord { + LifecycleSlot { + slot: SlotId, + session_id: u64, + cancel: bool, + }, /// Additional data credit for the peer's slot. CreditUpdate { slot: SlotId, delta: u32 }, /// Tell the peer to abandon a slot the ingress holds and is closing. @@ -173,6 +180,10 @@ pub(crate) struct BatcherHandle { } impl BatcherHandle { + pub(crate) fn peer_stopped(&self, slot: SlotId, session_id: u64, cancel: bool) { + self.control.peer_stopped(slot, session_id, cancel); + } + /// Queue an attach, waiting if this batcher already has `OPEN_QUEUE_DEPTH` /// of them outstanding. pub(crate) async fn open_slot( @@ -334,6 +345,7 @@ enum Work { /// the decision it leads to is [`FlushGate::should_flush`] reading the /// deadline as state. Linger, + Probe, } /// Whether [`Batcher::fire_singleton`] may skip the fence on a synchronously @@ -383,6 +395,11 @@ struct Batcher { impl Batcher { async fn run(mut self, opens: flume::Receiver) { let cancel = self.cancel.clone(); + let mut probe = tokio::time::interval_at( + tokio::time::Instant::now() + Duration::from_secs(5), + Duration::from_secs(5), + ); + probe.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); let control = Arc::clone(&self.control); loop { let deadline = self.gate.deadline(); @@ -398,6 +415,7 @@ impl Batcher { // Drained by the pass below between the wake and the take. None => continue, }, + _ = probe.tick(), if self.slots.live() > 0 => Work::Probe, Some((index, item)) = self.streams.next() => Work::Slot(index, item), () = linger_until(deadline) => Work::Linger, }; @@ -408,7 +426,7 @@ impl Batcher { Work::Slot(_, SlotItem::InletClosed) => BatcherWake::InletClosed, Work::Open(_) => BatcherWake::Open, Work::Control(_) => BatcherWake::Control, - Work::Linger => BatcherWake::Linger, + Work::Linger | Work::Probe => BatcherWake::Linger, }); } self.dispatch(work).await; @@ -494,6 +512,11 @@ impl Batcher { Work::Open(request) => self.on_open_slot(request).await, Work::Control(drained) => self.on_control(drained).await, Work::Linger => {} + Work::Probe => { + if !self.writer.peer_is_alive().await { + self.epoch_death(); + } + } } } @@ -507,6 +530,13 @@ impl Batcher { if drained.flush { self.gate.kick(); } + for ((raw, session_id), cancel) in drained.lifecycle_replies { + let slot = SlotId::from_raw(raw); + self.push_reply(RecordType::LifecycleSlot, |encoder| { + encoder.push_lifecycle(slot, session_id, cancel) + }) + .await; + } for (raw, entry) in drained.peers { self.on_reply(SlotId::from_raw(raw), entry).await; } @@ -552,6 +582,18 @@ impl Batcher { self.on_peer_closed(slot, reason); return; } + if let Some((session_id, cancel)) = entry.lifecycle + && let Some(live) = self.slots.get_mut_checked(slot) + && live.session_id == session_id + { + if cancel { + self.close_local(slot.index()); + return; + } + if let Some((_, stop)) = &live.lifecycle { + stop.cancel(); + } + } let mut touched = false; if let Some(live) = self.slots.get_mut_checked(slot) { if entry.credit > 0 { @@ -577,6 +619,7 @@ impl Batcher { anchor_id, session_id, inlet, + lifecycle, credit, slot_byte_budget, ack, @@ -596,7 +639,10 @@ impl Batcher { // information the control inbox has no other way to see, published // here because `entry_mine`'s bound needs it and nothing shorter than // this call site can hand it over. - self.control.note_allocated(id); + let live = self.slots.get_mut_checked(id).expect("allocated slot"); + live.lifecycle = lifecycle; + live.session_id = session_id; + self.control.note_allocated_session(id, session_id); self.streams.push(stream); self.publish_live_slots(); if let Some(metrics) = &self.metrics { diff --git a/lib/velo/src/streaming/messenger_mux/peer_batcher/records.rs b/lib/velo/src/streaming/messenger_mux/peer_batcher/records.rs index 830d8b78..82f2a37c 100644 --- a/lib/velo/src/streaming/messenger_mux/peer_batcher/records.rs +++ b/lib/velo/src/streaming/messenger_mux/peer_batcher/records.rs @@ -275,6 +275,9 @@ impl Batcher { } if terminal { + if let Some(slot) = self.slots.get_mut(index) { + slot.lifecycle = None; + } self.close_local(index); } } @@ -321,6 +324,9 @@ impl Batcher { // `fire_singleton`, has to stay a choice made here rather than inside // `fire_singleton` itself. if terminal { + if let Some(slot) = self.slots.get_mut(index) { + slot.lifecycle = None; + } self.close_local(index); } } diff --git a/lib/velo/src/streaming/messenger_mux/peer_batcher/slot_stream.rs b/lib/velo/src/streaming/messenger_mux/peer_batcher/slot_stream.rs index f1090569..75e72b1d 100644 --- a/lib/velo/src/streaming/messenger_mux/peer_batcher/slot_stream.rs +++ b/lib/velo/src/streaming/messenger_mux/peer_batcher/slot_stream.rs @@ -231,6 +231,11 @@ pub(super) struct WithheldOverflow { /// One live egress slot. pub(super) struct EgressSlot { + pub(super) session_id: u64, + pub(super) lifecycle: Option<( + tokio_util::sync::CancellationToken, + tokio_util::sync::CancellationToken, + )>, /// Index and generation as they travel on the wire. pub(super) id: SlotId, /// What this side may still send. Opens at the window the peer advertised @@ -321,6 +326,9 @@ impl EgressSlot { /// under — would then be rejected as stale: the deferred close would never /// be written and the slot would sit in `live_slots` for good. pub(super) fn disconnect(&self) { + if let Some((cancel, _)) = &self.lifecycle { + cancel.cancel(); + } self.gate.close(); } @@ -407,6 +415,8 @@ impl EgressSlots { }; self.entries[index as usize] = Some(EgressSlot { + session_id: 0, + lifecycle: None, id, credit, next_seq: 0, @@ -439,7 +449,7 @@ impl EgressSlots { let Some(slot) = self.entries.get_mut(index as usize).and_then(Option::take) else { return false; }; - slot.gate.close(); + slot.disconnect(); self.generations[index as usize] = slot.id.generation().wrapping_add(1); self.free.push(index); self.live -= 1; diff --git a/lib/velo/src/streaming/messenger_mux/peer_batcher/tests/support.rs b/lib/velo/src/streaming/messenger_mux/peer_batcher/tests/support.rs index b1f5e79f..3a2be634 100644 --- a/lib/velo/src/streaming/messenger_mux/peer_batcher/tests/support.rs +++ b/lib/velo/src/streaming/messenger_mux/peer_batcher/tests/support.rs @@ -274,6 +274,7 @@ impl Harness { anchor_id, session_id, inlet: inlet_rx, + lifecycle: None, // The default is deliberately starved. On the attach path a // slot opens holding the window its peer advertised, but most // arms below are about what the batcher does once a slot has @@ -499,6 +500,7 @@ impl StalledHarness { anchor_id, session_id, inlet: inlet_rx, + lifecycle: None, credit: SlotCredit::new(credit), slot_byte_budget: self.config.slot_byte_budget, ack: ack_tx, diff --git a/lib/velo/src/streaming/messenger_mux/peer_batcher/writer.rs b/lib/velo/src/streaming/messenger_mux/peer_batcher/writer.rs index 0a51f513..d5beccca 100644 --- a/lib/velo/src/streaming/messenger_mux/peer_batcher/writer.rs +++ b/lib/velo/src/streaming/messenger_mux/peer_batcher/writer.rs @@ -113,6 +113,19 @@ impl BatchWriter { } } + pub(super) async fn peer_is_alive(&mut self) -> bool { + let Some(peer) = self.peer_instance() else { + return false; + }; + matches!( + self.messenger + .backend() + .check_peer_health(peer, std::time::Duration::from_secs(5)) + .await, + Ok(()) | Err(crate::transports::HealthCheckError::NeverConnected) + ) + } + /// The epoch every batch this writer opens is stamped with. pub(super) const fn epoch(&self) -> u64 { self.epoch diff --git a/lib/velo/src/streaming/messenger_mux/protocol/mod.rs b/lib/velo/src/streaming/messenger_mux/protocol/mod.rs index 7e6b5957..90850c20 100644 --- a/lib/velo/src/streaming/messenger_mux/protocol/mod.rs +++ b/lib/velo/src/streaming/messenger_mux/protocol/mod.rs @@ -34,9 +34,9 @@ use std::ops::Range; /// Mux wire version carried in every batch header. /// /// Bumped only for a change that an older peer cannot skip past. Negotiation -/// (`messenger-mux-v1`) already keeps unequal versions from meeting; the +/// (`messenger-mux-v2`) already keeps unequal versions from meeting; the /// field is the belt to that pair of braces. -pub(crate) const MUX_VERSION: u8 = 1; +pub(crate) const MUX_VERSION: u8 = 2; /// Encoded size of a batch header. pub(crate) const BATCH_HEADER_LEN: usize = 16; @@ -254,16 +254,20 @@ pub(crate) enum RecordType { CreditUpdate = 3, /// Per-slot liveness beat. No body. SlotHeartbeat = 4, + /// Lifecycle request, receiver to producer: `[u64 session_id][u8 action]`, + /// where action 0 is a graceful stop and 1 is a cancel. + LifecycleSlot = 5, } /// The label value `velo_streaming_mux_records_sent_total` files each /// [`RecordType`] under, indexed by [`RecordType::count_index`]. -pub(crate) const RECORD_TYPE_LABELS: [&str; 5] = [ +pub(crate) const RECORD_TYPE_LABELS: [&str; 6] = [ "data", "open_slot", "close_slot", "credit_update", "slot_heartbeat", + "lifecycle", ]; /// Number of [`RecordType`] variants — the width of a per-type count array. @@ -294,6 +298,7 @@ impl RecordType { Self::CloseSlot => 2, Self::CreditUpdate => 3, Self::SlotHeartbeat => 4, + Self::LifecycleSlot => 5, } } @@ -305,6 +310,7 @@ impl RecordType { 2 => Some(Self::CloseSlot), 3 => Some(Self::CreditUpdate), 4 => Some(Self::SlotHeartbeat), + 5 => Some(Self::LifecycleSlot), _ => None, } } @@ -324,7 +330,10 @@ impl RecordType { /// the watchdog kill that `docs/src/operations/saturation.md` documents. #[cfg(test)] pub(crate) const fn is_control(self) -> bool { - matches!(self, Self::OpenSlot | Self::CloseSlot | Self::CreditUpdate) + matches!( + self, + Self::OpenSlot | Self::CloseSlot | Self::CreditUpdate | Self::LifecycleSlot + ) } } @@ -377,6 +386,8 @@ pub(crate) enum RecordBody<'a> { CreditUpdate { delta: u32 }, /// Liveness only. SlotHeartbeat, + /// Graceful stop keeps the slot open for remaining output. + LifecycleSlot { session_id: u64, cancel: bool }, } impl RecordBody<'_> { @@ -389,6 +400,7 @@ impl RecordBody<'_> { Self::CloseSlot { .. } => RecordType::CloseSlot, Self::CreditUpdate { .. } => RecordType::CreditUpdate, Self::SlotHeartbeat => RecordType::SlotHeartbeat, + Self::LifecycleSlot { .. } => RecordType::LifecycleSlot, } } } @@ -508,6 +520,18 @@ pub(crate) struct BatchEncoder { } impl BatchEncoder { + pub(crate) fn push_lifecycle( + &mut self, + slot: SlotId, + session_id: u64, + cancel: bool, + ) -> Result<(), EncodeError> { + self.push(RecordType::LifecycleSlot, slot, 0, 9, |body| { + body.extend_from_slice(&session_id.to_be_bytes()); + body.extend_from_slice(&[u8::from(cancel)]); + }) + } + /// Opens a batch in a fresh buffer. pub(crate) fn new(peer_epoch: u64, batch_seq: u32) -> Self { Self::with_buffer(BytesMut::new(), peer_epoch, batch_seq) @@ -810,6 +834,15 @@ fn decode_body( } Ok(RecordBody::CreditUpdate { delta }) } + RecordType::LifecycleSlot => { + if body.len() != 9 || body[8] > 1 { + return Err(mismatch(9)); + } + Ok(RecordBody::LifecycleSlot { + session_id: read_u64(body, 0).expect("length checked"), + cancel: body[8] == 1, + }) + } RecordType::SlotHeartbeat => { if !body.is_empty() { return Err(mismatch(0)); diff --git a/lib/velo/src/streaming/messenger_mux/protocol/tests.rs b/lib/velo/src/streaming/messenger_mux/protocol/tests.rs index 1e393245..182cc415 100644 --- a/lib/velo/src/streaming/messenger_mux/protocol/tests.rs +++ b/lib/velo/src/streaming/messenger_mux/protocol/tests.rs @@ -609,20 +609,20 @@ fn an_unknown_record_type_is_rejected_with_its_offset() { encoder.push_heartbeat(slot(), 1).expect("push"); }); let second = BATCH_HEADER_LEN + RECORD_HEADER_LEN; - batch[second] = 5; + batch[second] = 6; assert_eq!( first_error(&batch), Some(DecodeError::UnknownRecordType { offset: second, - value: 5, + value: 6, }) ); - for value in 0..=4u8 { + for value in 0..=5u8 { assert!(RecordType::from_u8(value).is_some(), "type {value}"); } - for value in 5..=u8::MAX { + for value in 6..=u8::MAX { assert_eq!(RecordType::from_u8(value), None, "type {value}"); } } diff --git a/lib/velo/src/streaming/messenger_mux/tests.rs b/lib/velo/src/streaming/messenger_mux/tests.rs index e1deb1a6..7cff62a4 100644 --- a/lib/velo/src/streaming/messenger_mux/tests.rs +++ b/lib/velo/src/streaming/messenger_mux/tests.rs @@ -205,7 +205,7 @@ async fn eventually(mut predicate: impl FnMut() -> bool) { /// falls back to the legacy path, with no error anywhere to say why. #[test] fn the_negotiated_key_is_the_string_that_shipped() { - assert_eq!(MESSENGER_MUX_KEY, "messenger-mux-v1"); + assert_eq!(MESSENGER_MUX_KEY, "messenger-mux-v2"); } #[tokio::test(flavor = "multi_thread")] diff --git a/lib/velo/src/streaming/mpsc/control.rs b/lib/velo/src/streaming/mpsc/control.rs index 3e151631..47564b96 100644 --- a/lib/velo/src/streaming/mpsc/control.rs +++ b/lib/velo/src/streaming/mpsc/control.rs @@ -995,7 +995,7 @@ mod tests { handle: StreamAnchorHandle::pack_mpsc(WorkerId::from_u64(1), 2), session_id: 3, stream_cancel_handle: StreamCancelHandle::pack(WorkerId::from_u64(4), 5), - supported_transport_keys: vec![TransportKey::new("messenger-mux-v1")], + supported_transport_keys: vec![TransportKey::new("messenger-mux-v2")], }; let decoded: MpscAnchorAttachRequest = rmp_serde::from_slice(&rmp_serde::to_vec(&req).expect("encode")).expect("decode"); @@ -1005,11 +1005,11 @@ mod tests { .iter() .map(TransportKey::as_str) .collect::>(), - ["messenger-mux-v1"], + ["messenger-mux-v2"], ); let resp = MpscAnchorAttachResponse::Ok { - streaming_transport_key: TransportKey::new("messenger-mux-v1"), + streaming_transport_key: TransportKey::new("messenger-mux-v2"), heartbeat_interval_ms: 5000, sender_id: 7, routing_session_id: 8, diff --git a/lib/velo/src/streaming/negotiation.rs b/lib/velo/src/streaming/negotiation.rs index 88ea0c3e..aaa9e741 100644 --- a/lib/velo/src/streaming/negotiation.rs +++ b/lib/velo/src/streaming/negotiation.rs @@ -14,7 +14,7 @@ //! 3. The sender decides what the answer obliges it to do ([`choose`]). //! //! Every asymmetry here exists to keep a mixed deployment working. The receiver -//! prefers `messenger-mux-v1` **only** when the sender named it, because +//! prefers `messenger-mux-v2` **only** when the sender named it, because //! `resolve_transport` hard-errors on a key it does not know and a receiver that //! answered the mux unilaterally would break every older sender. A node with the //! mux enabled therefore registers both it and the configured legacy transport, @@ -131,7 +131,7 @@ pub(crate) enum Connect { /// /// The mux key with no window is the case that cannot be honoured. It is /// unreachable from any shipped peer: no version before this one answers -/// `messenger-mux-v1`, and this one refuses to build a mux at zero credit. So it +/// `messenger-mux-v2`, and this one refuses to build a mux at zero credit. So it /// means a peer that bound a mux receiver and then told us to ignore it, and /// there is no safe reading of that — connecting over any other transport would /// reach nothing the peer is listening on and hang until the anchor's watchdog diff --git a/lib/velo/src/streaming/sender.rs b/lib/velo/src/streaming/sender.rs index cdf57f98..466d81a2 100644 --- a/lib/velo/src/streaming/sender.rs +++ b/lib/velo/src/streaming/sender.rs @@ -138,6 +138,7 @@ pub struct StreamSender { registry: Arc>, /// User-facing cancellation signal: fires when _stream_cancel is received. cancel_token: CancellationToken, + stop_token: CancellationToken, /// Key in the sender-side registry for cleanup and for the _stream_cancel handler. sender_stream_id: u64, /// Sender-side registry shared with the _stream_cancel handler. @@ -196,6 +197,11 @@ impl StreamSender { sender_registry, poison_tx, } = cancel; + let stop_token = sender_registry + .senders + .get(&sender_stream_id) + .map(|entry| entry.stop_token.clone()) + .unwrap_or_else(|| cancel_token.child_token()); let heartbeat_cancel = CancellationToken::new(); // Spawn heartbeat background task @@ -224,6 +230,7 @@ impl StreamSender { tx, handle, heartbeat_cancel, + stop_token, sent_terminal: false, registry, cancel_token, @@ -271,20 +278,17 @@ impl StreamSender { /// # } /// ``` /// - /// This fires only for a stream opened by attach. A sender built by - /// [`AnchorManager::open_anchor_stream`](crate::streaming::AnchorManager::open_anchor_stream) - /// (or [`Velo::open_anchor_stream`](crate::Velo::open_anchor_stream)) never - /// sent an attach, so the consumer never learned a `StreamCancelHandle` to - /// reach it by, and this token never cancels for it -- a dropped consumer - /// surfaces there as `SendError::ChannelClosed` on the next `send` - /// instead, posted promptly once the pre-bind's reclamation runs. The - /// `tokio::select!` above is still correct for that case: the token side - /// just never fires, so the loop exits through the `send` error the way - /// the corrected example above does. + /// Also fires for ticket opens, including an idle producer whose slot closes. pub fn cancellation_token(&self) -> CancellationToken { self.cancel_token.clone() } + /// Fires when the consumer requests graceful stop or cancels the stream. + /// The producer can send buffered output and finalize after graceful stop. + pub fn stop_token(&self) -> CancellationToken { + self.stop_token.clone() + } + /// Send a typed item through the channel. /// /// Serializes `StreamFrame::Item(item)` via `rmp_serde` and pushes the @@ -296,7 +300,7 @@ impl StreamSender { /// - [`SendError::ChannelClosed`] if the receiver has been dropped. pub async fn send(&self, item: T) -> Result<(), SendError> { // Check if the poison channel has been disconnected (rx_closer dropped by _stream_cancel handler). - if self.poison_tx.is_disconnected() { + if self.cancel_token.is_cancelled() || self.poison_tx.is_disconnected() { return Err(SendError::ChannelClosed); } let bytes = rmp_serde::to_vec(&StreamFrame::Item(item)) @@ -311,10 +315,11 @@ impl StreamSender { if let Some(m) = self.metrics.as_ref() { m.record_producer_send_backpressure(); } - self.tx - .send_async(b) - .await - .map_err(|_| SendError::ChannelClosed) + tokio::select! { + biased; + _ = self.cancel_token.cancelled() => Err(SendError::ChannelClosed), + result = self.tx.send_async(b) => result.map_err(|_| SendError::ChannelClosed), + } } Err(flume::TrySendError::Disconnected(_)) => Err(SendError::ChannelClosed), } @@ -548,6 +553,7 @@ mod tests { // Insert the SenderEntry into the registry (simulating what attach_stream_anchor does) let entry = crate::streaming::control::SenderEntry { + stop_token: cancel_token.child_token(), cancel_token: cancel_token.clone(), rx_closer: std::sync::Mutex::new(Some(poison_rx)), }; diff --git a/lib/velo/src/transports.rs b/lib/velo/src/transports.rs index bb17b4c0..47440935 100644 --- a/lib/velo/src/transports.rs +++ b/lib/velo/src/transports.rs @@ -250,6 +250,20 @@ impl VeloBackend { .map(|entry| entry.value().key()) } + /// Probe the selected transport without changing its priority or fallback. + pub(crate) async fn check_peer_health( + &self, + target: InstanceId, + timeout: std::time::Duration, + ) -> Result<(), transport::HealthCheckError> { + let transport = self + .primary_transport + .get(&target) + .map(|entry| entry.value().clone()) + .ok_or(transport::HealthCheckError::PeerNotRegistered)?; + transport.check_health(target, timeout).await + } + /// Largest `header + payload` the peer's primary transport will carry to /// `target` in one send, or `None` when that cannot be established. /// diff --git a/lib/velo/src/transports/ucx/tests.rs b/lib/velo/src/transports/ucx/tests.rs index 4f56d5b6..6dc53297 100644 --- a/lib/velo/src/transports/ucx/tests.rs +++ b/lib/velo/src/transports/ucx/tests.rs @@ -401,6 +401,89 @@ async fn shutdown_fails_queued_sends() { b.transport.shutdown(); } +// A held command stands in for a UCX send that has not completed. This makes +// pressure and cleanup deterministic without timing a network transfer. +#[tokio::test] +async fn send_budget_bounds_admission_and_releases_on_peer_retirement() { + use super::super::address::{AM_ID_BASE, BLOB_VERSION, UcxEndpoint}; + let transport = UcxTransport::new( + "ucx".into(), + UcxConfig { + channel_capacity: 1, + ..Default::default() + }, + ); + transport + .runtime + .set(tokio::runtime::Handle::current()) + .unwrap(); + let peer = InstanceId::new_v4(); + transport.shared.peers.insert( + peer, + UcxEndpoint { + v: BLOB_VERSION, + am_id_base: AM_ID_BASE, + eager_max: 1024, + incarnation: 1, + worker_addr: vec![], + }, + ); + let rx = transport.ring_rx.lock().unwrap().take().unwrap(); + let errors = CountingErrors::new(); + let send = || { + transport.send_message( + peer, + Bytes::new(), + Bytes::from_static(b"data"), + MessageType::Message, + errors.clone(), + ) + }; + assert!(send().is_admitted()); + let held = tokio::time::timeout(T, rx.recv_async()) + .await + .unwrap() + .unwrap(); + // One command holds the permit, one may wait for it, and one fits in the + // channel. The following ticket must wait before admission. + let ticket = loop { + if let SendOutcome::Pending(ticket) = send() { + break ticket; + } + }; + assert_eq!( + transport.connections.get(&peer).unwrap().gate.queued_len(), + 1 + ); + send(); + assert_eq!(errors.count(), 1, "excess queue entry must fail explicitly"); + ticket.cancel(); + assert_eq!( + transport.connections.get(&peer).unwrap().gate.queued_len(), + 0 + ); + let next = send(); + transport.shared.failed_peers.insert(peer, ()); + transport.reap_failed_connection(peer); + if let SendOutcome::Pending(ticket) = next { + assert!(tokio::time::timeout(T, ticket).await.unwrap().is_err()); + } + // Completion releases the outstanding permit and its retained buffers. + drop(held); + tokio::time::timeout(T, async { + while errors.count() < 2 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert!(transport.connections.is_empty()); + assert!( + rx.is_empty(), + "retired sends must not reach the progress thread" + ); +} + // --------------------------------------------------------------------------- // RMA // --------------------------------------------------------------------------- diff --git a/lib/velo/src/transports/ucx/transport.rs b/lib/velo/src/transports/ucx/transport.rs index 242b6f73..2bfe30bc 100644 --- a/lib/velo/src/transports/ucx/transport.rs +++ b/lib/velo/src/transports/ucx/transport.rs @@ -56,7 +56,8 @@ pub struct UcxConfig { /// makes loaded submission a plain ring push. Default 20 µs. pub spin_us: u64, /// Capacity of the command ring between senders and the progress thread. - /// The per-peer admission gates queue (and preserve order) beyond it. + /// Also bounds each peer's queued and outstanding sends. Admission waits + /// for UCX completion when a peer exhausts its send permits. pub channel_capacity: usize, /// Override for `UCX_TLS` (e.g. `"rc_verbs,ud_verbs,self"` or `"tcp"`), /// applied only when the environment does not already set it. Note an @@ -120,7 +121,9 @@ impl Default for UcxConfig { /// gates admitted. An epoch is retired when the peer's endpoint fails. #[derive(Clone)] struct ConnHandle { - gate: AdmissionGate, + admission: Arc>, + gate: AdmissionGate, + _lifetime: Arc, } /// UCX messenger transport. See the module docs for the model. @@ -239,8 +242,65 @@ impl UcxTransport { let handle = self .connections .entry(peer) - .or_insert_with(|| ConnHandle { - gate: AdmissionGate::new(self.ring_tx.clone(), rt.clone()), + .or_insert_with(|| { + let (tx, rx) = flume::bounded::(self.config.channel_capacity); + let gate = AdmissionGate::new(tx, rt.clone()); + let cancel = tokio_util::sync::CancellationToken::new(); + let lifetime = Arc::new(cancel.clone().drop_guard()); + let permits = Arc::new(tokio::sync::Semaphore::new(self.config.channel_capacity)); + let ring = self.ring_tx.clone(); + let doorbell = self.shared.doorbell.clone(); + let pending = gate.clone(); + rt.spawn(async move { + loop { + let mut task = tokio::select! { + biased; + _ = cancel.cancelled() => break, + task = rx.recv_async() => match task { Ok(task) => task, Err(_) => break }, + }; + let permit = tokio::select! { + biased; + _ = cancel.cancelled() => { + task.fail("ucx connection closed"); + break; + } + permit = permits.clone().acquire_owned() => permit.expect("send permits remain open"), + }; + task.permit = Some(permit); + match ring.try_send(Cmd::Send(task)) { + Ok(()) => {} + Err(flume::TrySendError::Disconnected(command)) => { + command.refuse_for_shutdown(); + break; + } + Err(flume::TrySendError::Full(Cmd::Send(task))) => { + let header = task.header.clone(); + let payload = task.payload.clone(); + let on_error = task.on_error.clone(); + tokio::select! { + biased; + _ = cancel.cancelled() => { + on_error.on_error(header, payload, "ucx connection closed".into()); + break; + } + result = ring.send_async(Cmd::Send(task)) => { + if let Err(error) = result { + error.0.refuse_for_shutdown(); + break; + } + } + } + } + Err(flume::TrySendError::Full(_)) => unreachable!("only sends enter the peer queue"), + } + doorbell.ring(); + } + pending.fail_all(AdmissionError::ChannelClosed); + while let Ok(task) = rx.try_recv() { + task.fail("ucx connection closed"); + } + }); + ConnHandle { admission: Arc::new(Mutex::new(())), gate, _lifetime: lifetime } }) .clone(); if let Some(m) = self.shared.metrics.get() { @@ -250,28 +310,22 @@ impl UcxTransport { } fn admit(&self, handle: &ConnHandle, task: SendTask) -> SendOutcome { - match handle.gate.send(Cmd::Send(task)) { - SendOutcome::Admitted => { - // The frame is on the ring: wake the progress thread now. - self.shared.doorbell.ring(); - SendOutcome::Admitted - } - SendOutcome::Pending(admission) => { - if let Some(m) = self.shared.metrics.get() { - m.record_send_backpressure(); - } - // The ring push happens later, from the gate's driver task — - // ring the doorbell when it actually lands, so a queued frame - // does not wait out the park timeout. The park's bounded - // timeout remains the backstop. - let doorbell = Arc::clone(&self.shared.doorbell); - SendOutcome::Pending(admission.on_resolved(move |result| { - if result.is_ok() { - doorbell.ring(); - } - })) - } + // Serialize the capacity check with admission. The gate's driver can + // only reduce the queue while this lock is held. + let admission = handle.admission.lock().unwrap_or_else(|e| e.into_inner()); + if handle.gate.queued_len() >= self.config.channel_capacity { + drop(admission); + task.fail("ucx peer admission queue full"); + return SendOutcome::Admitted; + } + let outcome = handle.gate.send(task); + drop(admission); + if matches!(&outcome, SendOutcome::Pending(_)) + && let Some(m) = self.shared.metrics.get() + { + m.record_send_backpressure(); } + outcome } /// Retire a peer's epoch after its endpoint failed: queued frames belong @@ -353,6 +407,7 @@ impl Transport for UcxTransport { let task = SendTask { peer: instance_id, msg_type: message_type, + permit: None, header, payload, on_error, diff --git a/lib/velo/src/transports/ucx/worker.rs b/lib/velo/src/transports/ucx/worker.rs index 7b6effa8..c4c91205 100644 --- a/lib/velo/src/transports/ucx/worker.rs +++ b/lib/velo/src/transports/ucx/worker.rs @@ -241,6 +241,7 @@ use super::transport::UcxConfig; /// One message staged for the progress thread. pub(crate) struct SendTask { pub peer: InstanceId, + pub permit: Option, pub msg_type: MessageType, pub header: Bytes, pub payload: Bytes, @@ -513,6 +514,7 @@ enum OpKind { /// A velo frame: failure reports through `on_error` with the original /// buffers, per the `Transport` contract. Frame { + _permit: Option, header: Bytes, payload: Bytes, on_error: Arc, @@ -534,6 +536,7 @@ impl OpState { { match state.kind { OpKind::Frame { + _permit, header, payload, on_error, @@ -1488,6 +1491,7 @@ impl WorkerState { let reply = matches!(task.msg_type, MessageType::Message); let op = Arc::new(OpState { kind: OpKind::Frame { + _permit: task.permit, header: task.header.clone(), payload: task.payload.clone(), on_error: task.on_error, diff --git a/lib/velo/tests/streaming/endurance_remote.rs b/lib/velo/tests/streaming/endurance_remote.rs index 17f7c9b7..d5e24dfc 100644 --- a/lib/velo/tests/streaming/endurance_remote.rs +++ b/lib/velo/tests/streaming/endurance_remote.rs @@ -10,7 +10,7 @@ //! **raw TCP** (`TcpFrameTransport`). //! //! The messenger-backed variants that used to sit beside them went with the -//! deleted `VeloFrameTransport`. Their replacement rides `messenger-mux-v1`, +//! deleted `VeloFrameTransport`. Their replacement rides `messenger-mux-v2`, //! which is crate-internal until attach-time negotiation can select it, so the //! mux endurance scenarios land with that stage rather than here. //! diff --git a/lib/velo/tests/streaming/mux_negotiation/mod.rs b/lib/velo/tests/streaming/mux_negotiation/mod.rs index 0cc7b051..945e5887 100644 --- a/lib/velo/tests/streaming/mux_negotiation/mod.rs +++ b/lib/velo/tests/streaming/mux_negotiation/mod.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Attach-time negotiation of `messenger-mux-v1`, over two real `Velo` nodes. +//! Attach-time negotiation of `messenger-mux-v2`, over two real `Velo` nodes. //! //! The unit tests beside `streaming/negotiation.rs` pin the decision table; this //! file checks that two nodes actually reach those verdicts across a wire, and diff --git a/lib/velo/tests/streaming/mux_negotiation/zero_rtt.rs b/lib/velo/tests/streaming/mux_negotiation/zero_rtt.rs index c4058a81..a099865b 100644 --- a/lib/velo/tests/streaming/mux_negotiation/zero_rtt.rs +++ b/lib/velo/tests/streaming/mux_negotiation/zero_rtt.rs @@ -19,6 +19,146 @@ use super::*; use velo::streaming::AttachError; use velo::streaming::control::{AnchorAttachRequest, StreamOpenTicket}; +/// Stop leaves output usable; cancel wakes an idle producer on both open paths. +#[tokio::test(flavor = "multi_thread")] +async fn native_lifecycle_survives_early_stop_and_escalation() { + let (consumer, producer) = pair(Some(mux_config()), Some(mux_config())).await; + for ticket_open in [false, true] { + for early in [false, true] { + let mut anchor = consumer.velo.create_anchor::(); + let controller = anchor.controller(); + if early { + controller.request_stop(); + } + let sender = if ticket_open { + let ticket = consumer.velo.prebind_anchor(anchor.handle()).unwrap(); + producer + .velo + .open_anchor_stream::(anchor.handle(), ship(ticket)) + .await + .unwrap() + } else { + producer + .velo + .attach_anchor::(anchor.handle()) + .await + .unwrap() + }; + controller.request_stop(); + controller.request_stop(); + tokio::time::timeout(Duration::from_secs(5), sender.stop_token().cancelled()) + .await + .unwrap(); + assert!(!sender.cancellation_token().is_cancelled()); + sender.send(42).await.unwrap(); + assert!(matches!( + tokio::time::timeout(Duration::from_secs(5), anchor.next()) + .await + .unwrap(), + Some(Ok(StreamFrame::Item(42))) + )); + controller.cancel(); + tokio::time::timeout( + Duration::from_secs(5), + sender.cancellation_token().cancelled(), + ) + .await + .unwrap(); + assert!(sender.stop_token().is_cancelled()); + assert!(sender.send(43).await.is_err()); + } + } +} + +/// A stop requested from a thread without a Tokio runtime still reaches a +/// remote sender. +/// +/// `request_stop` is synchronous, so a caller can be on any thread: a Python +/// binding's thread, or a drop on a plain thread. The stop to a remote sender +/// goes out as an active message, which needs a runtime to send it. Taken from +/// the calling thread, there was none, and the stop was dropped. The anchor had +/// already recorded it, so a second call returned early, and the sender never +/// stopped. The messenger's own runtime sends it. +#[tokio::test(flavor = "multi_thread")] +async fn a_stop_requested_off_runtime_reaches_a_remote_sender() { + for mux in [Some(mux_config()), None] { + let (consumer, producer) = pair(mux.clone(), mux).await; + let anchor = consumer.velo.create_anchor::(); + let controller = anchor.controller(); + let sender = producer + .velo + .attach_anchor::(transfer(anchor.handle())) + .await + .expect("remote attach"); + + std::thread::spawn(move || controller.request_stop()) + .join() + .expect("request_stop on a plain thread"); + + tokio::time::timeout(Duration::from_secs(5), sender.stop_token().cancelled()) + .await + .expect("the remote sender never saw the stop"); + assert!(!sender.cancellation_token().is_cancelled()); + drop(anchor); + } +} + +/// A stop reaches a sender on the anchor's own worker. +/// +/// A same-worker attach puts the sender in the local sender registry, the way +/// `cancel` finds it. A stop sent as an active message has to resolve its own +/// worker through the messenger instead, and that send's result is ignored. +#[tokio::test(flavor = "multi_thread")] +async fn a_stop_reaches_a_sender_on_the_same_worker() { + for early in [false, true] { + let node = node(Some(mux_config())).await; + let anchor = node.velo.create_anchor::(); + let controller = anchor.controller(); + if early { + controller.request_stop(); + } + let sender = node + .velo + .attach_anchor::(anchor.handle()) + .await + .expect("same-worker attach"); + if !early { + controller.request_stop(); + } + tokio::time::timeout(Duration::from_secs(5), sender.stop_token().cancelled()) + .await + .unwrap_or_else(|_| panic!("early={early}: the same-worker sender never saw the stop")); + assert!(!sender.cancellation_token().is_cancelled()); + drop(anchor); + } +} + +/// An unused or claimed ticket must cancel without a subsequent producer send. +#[tokio::test(flavor = "multi_thread")] +async fn ticket_cancel_wakes_idle_producer() { + let (consumer, producer) = pair(Some(mux_config()), Some(mux_config())).await; + for before_open in [false, true] { + let anchor = consumer.velo.create_anchor::(); + let ticket = consumer.velo.prebind_anchor(anchor.handle()).unwrap(); + if before_open { + anchor.controller().cancel(); + } + let sender = producer + .velo + .open_anchor_stream::(anchor.handle(), ship(ticket)) + .await + .unwrap(); + drop(anchor); + tokio::time::timeout( + Duration::from_secs(5), + sender.cancellation_token().cancelled(), + ) + .await + .unwrap(); + assert!(sender.stop_token().is_cancelled()); + } +} + /// A ticket as the worker receives it: through the encoding, never by /// reference. An application carries one in the request envelope it already /// sends, so anything that failed to serialise would fail there and not here.