From 64056e453d5129ef4b2972fc31929aee46edc36c Mon Sep 17 00:00:00 2001 From: danilo-najkov-db Date: Fri, 25 Sep 2026 15:58:04 +0000 Subject: [PATCH 1/2] [Rust] Extract shared mux core Signed-off-by: danilo-najkov-db --- rust/NEXT_CHANGELOG.md | 3 + .../src/builder/stream_builder/multiplexed.rs | 188 +++++++-- rust/sdk/src/multiplexed_stream.rs | 367 +----------------- rust/sdk/src/multiplexed_stream/core.rs | 365 +++++++++++++++++ rust/sdk/src/multiplexed_stream/lane.rs | 94 +++++ 5 files changed, 634 insertions(+), 383 deletions(-) create mode 100644 rust/sdk/src/multiplexed_stream/core.rs create mode 100644 rust/sdk/src/multiplexed_stream/lane.rs diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index f0dbdba7..780c44a0 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -41,6 +41,9 @@ resume-watermark reconciliation after a lost acknowledgment, and validation for setup responses, acknowledgment bounds, and offset overflow. +- Extracted a private transport-generic mux core and lane contract, preserving + existing gRPC multiplexed-stream behavior and construction guarantees. + ### Breaking Changes ### Deprecations diff --git a/rust/sdk/src/builder/stream_builder/multiplexed.rs b/rust/sdk/src/builder/stream_builder/multiplexed.rs index 3841a417..c36ac506 100644 --- a/rust/sdk/src/builder/stream_builder/multiplexed.rs +++ b/rust/sdk/src/builder/stream_builder/multiplexed.rs @@ -176,48 +176,162 @@ impl<'a> MultiplexedStreamBuilder<'a> { }, ); - // `join_all` polls every open concurrently and preserves input order, - // so completion timing cannot change the assigned stream indices. - let results = join_all(opens).await; - let mut indexed_streams = Vec::with_capacity(self.stream_count); - let mut first_error = None; - for (stream_index, result) in results.into_iter().enumerate() { - match result { - Ok(stream) => indexed_streams.push((stream_index, stream)), - Err(err) => { - error!(stream_index, error = %err, "Failed to create multiplexed sub-stream"); - first_error.get_or_insert(err); - } + let streams = collect_opened(opens, |mut stream: ZerobusStream| async move { + stream.close().await + }) + .await?; + debug_assert_eq!(streams.len(), self.stream_count); + // This is one user-initiated logical stream creation. Counting each + // internal lane would make the churn warning flag intentional muxes. + crate::client_warnings::record_stream_creation(&table_properties.table_name); + Ok(MultiplexedStream::from_streams(streams)) + } +} + +/// Preserves preassigned indices, waits for every open, and closes all successes +/// before returning the lowest-index construction error. Cancelling this future +/// drops pending opens and successfully constructed lanes together. +async fn collect_opened( + opens: impl IntoIterator, + close: C, +) -> ZerobusResult> +where + O: std::future::Future>, + C: Fn(S) -> F, + F: std::future::Future>, +{ + // `join_all` polls every open concurrently and preserves input order, + // so completion timing cannot change the assigned stream indices. + let results = join_all(opens).await; + let mut indexed_streams = Vec::new(); + let mut first_error = None; + for (stream_index, result) in results.into_iter().enumerate() { + match result { + Ok(stream) => indexed_streams.push((stream_index, stream)), + Err(err) => { + error!(stream_index, error = %err, "Failed to create multiplexed sub-stream"); + first_error.get_or_insert(err); } } + } - if let Some(first_error) = first_error { - let cleanup_results = join_all(indexed_streams.into_iter().map( - |(successful_stream_index, mut stream)| async move { - (successful_stream_index, stream.close().await) - }, - )) - .await; - for (successful_stream_index, result) in cleanup_results { - if let Err(err) = result { - warn!( - stream_index = successful_stream_index, - error = %err, - "Failed to clean up multiplexed sub-stream after construction failure" - ); - } + if let Some(first_error) = first_error { + let cleanup_results = join_all(indexed_streams.into_iter().map( + |(successful_stream_index, stream)| { + let close = &close; + async move { (successful_stream_index, close(stream).await) } + }, + )) + .await; + for (successful_stream_index, result) in cleanup_results { + if let Err(err) = result { + warn!( + stream_index = successful_stream_index, + error = %err, + "Failed to clean up multiplexed sub-stream after construction failure" + ); } - return Err(first_error); } + return Err(first_error); + } - let streams: Vec<_> = indexed_streams - .into_iter() - .map(|(_stream_index, stream)| stream) - .collect(); - debug_assert_eq!(streams.len(), self.stream_count); - // This is one user-initiated logical stream creation. Counting each - // internal lane would make the churn warning flag intentional muxes. - crate::client_warnings::record_stream_creation(&table_properties.table_name); - Ok(MultiplexedStream::from_streams(streams)) + let streams: Vec<_> = indexed_streams + .into_iter() + .map(|(_stream_index, stream)| stream) + .collect(); + Ok(streams) +} + +#[cfg(test)] +mod construction_tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + use tokio::sync::Notify; + + #[derive(Debug)] + struct Opened(usize, Arc); + impl Drop for Opened { + fn drop(&mut self) { + self.1.fetch_add(1, Ordering::SeqCst); + } + } + + #[tokio::test(start_paused = true)] + async fn completion_order_does_not_change_lane_indices() { + let dropped = Arc::new(AtomicUsize::new(0)); + let opens = (0..3).map(|i| { + let dropped = dropped.clone(); + async move { + tokio::time::sleep(Duration::from_millis((3 - i) as u64 * 100)).await; + Ok(Opened(i, dropped)) + } + }); + let lanes = collect_opened(opens, |_| async { Ok(()) }).await.unwrap(); + assert_eq!(lanes.iter().map(|s| s.0).collect::>(), [0, 1, 2]); + assert_eq!(dropped.load(Ordering::SeqCst), 0); + drop(lanes); + assert_eq!(dropped.load(Ordering::SeqCst), 3); + } + + #[tokio::test(start_paused = true)] + async fn lowest_index_failure_wins_after_all_opens_and_cleanup_settle() { + let dropped = Arc::new(AtomicUsize::new(0)); + let cleaned = Arc::new(AtomicUsize::new(0)); + let opens = (0..4).map(|i| { + let dropped = dropped.clone(); + async move { + tokio::time::sleep(Duration::from_millis((4 - i) as u64 * 100)).await; + if i % 2 == 0 { + Ok(Opened(i, dropped)) + } else { + Err(ZerobusError::InvalidArgument(format!("lane-{i}"))) + } + } + }); + let result = collect_opened(opens, |lane| { + let cleaned = cleaned.clone(); + async move { + tokio::time::sleep(Duration::from_secs(1)).await; + cleaned.fetch_add(1, Ordering::SeqCst); + drop(lane); + Err(ZerobusError::InvalidArgument("cleanup error".into())) + } + }) + .await; + assert!(matches!(result, Err(ZerobusError::InvalidArgument(msg)) if msg == "lane-1")); + assert_eq!(cleaned.load(Ordering::SeqCst), 2); + assert_eq!(dropped.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn cancellation_drops_completed_and_in_progress_opens() { + let dropped = Arc::new(AtomicUsize::new(0)); + let reached = Arc::new(Notify::new()); + let task = { + let dropped = dropped.clone(); + let reached = reached.clone(); + tokio::spawn(async move { + let opens = (0..3).map(|i| { + let dropped = dropped.clone(); + let reached = reached.clone(); + async move { + let lane = Opened(i, dropped); + if i == 1 { + reached.notify_one(); + std::future::pending::<()>().await; + } + if i == 2 { + tokio::time::sleep(Duration::from_secs(3600)).await; + } + Ok(lane) + } + }); + collect_opened(opens, |_| async { Ok(()) }).await + }) + }; + reached.notified().await; + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + assert_eq!(dropped.load(Ordering::SeqCst), 3); } } diff --git a/rust/sdk/src/multiplexed_stream.rs b/rust/sdk/src/multiplexed_stream.rs index aed4191e..3c079d65 100644 --- a/rust/sdk/src/multiplexed_stream.rs +++ b/rust/sdk/src/multiplexed_stream.rs @@ -18,20 +18,17 @@ //! [`get_unacked_records`](MultiplexedStream::get_unacked_records) or //! [`get_unacked_batches`](MultiplexedStream::get_unacked_batches). -use futures::future::join_all; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::sync::{Arc, OnceLock}; -use std::time::Duration; -use tokio_util::sync::CancellationToken; -use tracing::{error, info, warn}; +use std::sync::Arc; + +mod core; +mod lane; +use core::MuxCore; use crate::{ AckCallback, DynamicRecord, EncodedBatch, EncodedRecord, MessageDescriptor, OffsetId, - PreparedInput, ZerobusError, ZerobusResult, ZerobusStream, + PreparedInput, ZerobusResult, ZerobusStream, }; -const CAPACITY_WAIT_TIMEOUT: Duration = Duration::from_secs(30); - /// Number of bits reserved for the stream index. /// 6 bits supports up to 64 sub-streams. const STREAM_BITS: u32 = 6; @@ -146,14 +143,7 @@ pub(crate) fn multiplexed_ack_callback( /// /// This API is in Beta. pub struct MultiplexedStream { - streams: Vec, - round_robin_counter: AtomicUsize, - is_closed: AtomicBool, - closed_token: CancellationToken, - failure: OnceLock, - /// Completed close-time flush result; retries resume teardown without - /// flushing a transport whose supervisor may already have stopped. - close_flush_result: Option>, + core: MuxCore, } impl MultiplexedStream { @@ -170,199 +160,20 @@ impl MultiplexedStream { } pub(crate) fn from_streams(streams: Vec) -> Self { - assert!( - !streams.is_empty(), - "MultiplexedStream requires at least one sub-stream" - ); - assert!( - streams.len() <= (1 << STREAM_BITS), - "MultiplexedStream supports at most {} sub-streams", - 1 << STREAM_BITS - ); Self { - streams, - round_robin_counter: AtomicUsize::new(0), - is_closed: AtomicBool::new(false), - closed_token: CancellationToken::new(), - failure: OnceLock::new(), - close_flush_result: None, + core: MuxCore::from_streams(streams), } } /// Returns the schema descriptor configured with [`crate::StreamBuilder::dynamic_proto`]. /// Returns an error if this is not a dynamic-protobuf stream. pub fn message_descriptor(&self) -> ZerobusResult { - self.streams[0].message_descriptor() + self.core.first().message_descriptor() } /// Creates an empty record using this mux's dynamic-protobuf schema. pub fn new_record(&self) -> ZerobusResult { - self.streams[0].new_record() - } - - #[allow(clippy::result_large_err)] - fn check_closed(&self) -> ZerobusResult<()> { - if let Some(error) = self.failure.get() { - return Err(error.clone()); - } - if self.is_closed_fast() { - return Err(self.closed_error()); - } - Ok(()) - } - - fn closed_error(&self) -> ZerobusError { - self.failure.get().cloned().unwrap_or_else(|| { - ZerobusError::InvalidStateError("MultiplexedStream is closed".to_string()) - }) - } - - fn is_closed_fast(&self) -> bool { - self.is_closed.load(Ordering::Relaxed) - } - - fn first_closed_stream(&self) -> Option { - self.streams.iter().position(ZerobusStream::is_closed) - } - - async fn lane_terminal_error(&self, idx: usize, fallback: ZerobusError) -> ZerobusError { - self.streams[idx].terminal_error().await.unwrap_or(fallback) - } - - fn shutdown_on_failure(&self, trigger_index: usize, cause: &ZerobusError) { - if self.is_closed_fast() || self.failure.set(cause.clone()).is_err() { - return; - } - self.is_closed.store(true, Ordering::Relaxed); - self.closed_token.cancel(); - - error!( - trigger_stream_index = trigger_index, - cause = %cause, - num_streams = self.streams.len(), - "MultiplexedStream poisoned due to sub-stream failure" - ); - } - - // TODO: if the picked sub-stream is at capacity, try the next one before - // falling back to waiting. - fn pick_substream(&self) -> usize { - self.round_robin_counter.fetch_add(1, Ordering::Relaxed) % self.streams.len() - } - - async fn reserve_capacity( - &self, - stream: &ZerobusStream, - idx: usize, - ) -> ZerobusResult { - let started_at = tokio::time::Instant::now(); - let timeout_ms = CAPACITY_WAIT_TIMEOUT.as_millis(); - let table_name = stream.table_properties.table_name.as_str(); - let max_inflight_requests = stream.options.max_inflight_requests; - - self.check_closed()?; - - let wait_for_reservation = async { - let reservation = stream.reserve_capacity(); - tokio::pin!(reservation); - - if let Ok(result) = tokio::time::timeout(Duration::from_secs(1), &mut reservation).await - { - return result; - } - let waited_ms = started_at.elapsed().as_millis(); - warn!( - stream_index = idx, - table_name, - waited_ms, - timeout_ms, - max_inflight_requests, - "Backpressure: sub-stream at capacity, waiting for drain" - ); - reservation.await - }; - - let result = tokio::select! { - result = tokio::time::timeout(CAPACITY_WAIT_TIMEOUT, wait_for_reservation) => result, - _ = self.closed_token.cancelled() => { - let waited_ms = started_at.elapsed().as_millis(); - if let Some(failure) = self.failure.get() { - return Err(failure.clone()); - } - warn!( - stream_index = idx, - table_name, - waited_ms, - max_inflight_requests, - "Multiplexed capacity wait cancelled by shutdown" - ); - return Err(ZerobusError::InvalidStateError( - format!( - "MultiplexedStream closed after {waited_ms} ms while waiting for capacity on sub-stream {idx} for table {table_name} (max_inflight_requests: {max_inflight_requests})" - ), - )); - } - }; - - match result { - Ok(Ok(reservation)) => Ok(reservation), - Ok(Err(e)) => Err(self.handle_lane_error(idx, e).await), - Err(_) => { - self.check_closed()?; - if stream.is_closed() { - return Err(self - .handle_lane_error( - idx, - ZerobusError::ConnectionTimeout(format!( - "Timed out waiting for capacity on multiplexed sub-stream {idx}" - )), - ) - .await); - } - let waited_ms = started_at.elapsed().as_millis(); - warn!( - stream_index = idx, - table_name, - waited_ms, - timeout_ms, - max_inflight_requests, - "Timed out waiting for multiplexed sub-stream capacity" - ); - Err(ZerobusError::ConnectionTimeout(format!( - "Timed out after {waited_ms} ms waiting for capacity on multiplexed sub-stream {idx} for table {table_name} (configured timeout: {timeout_ms} ms, max_inflight_requests: {max_inflight_requests})" - ))) - } - } - } - - async fn enqueue_reserved( - &self, - stream: &ZerobusStream, - idx: usize, - encoded_batch: EncodedBatch, - ) -> ZerobusResult { - let reservation = self.reserve_capacity(stream, idx).await?; - let enqueue_result = stream - .enqueue_reserved_admitted(encoded_batch, reservation, || self.check_closed()) - .await; - - match enqueue_result { - Ok(off) => Ok(MessageId::new(idx, off)), - Err(e) => Err(self.handle_lane_error(idx, e).await), - } - } - - // Payload errors and wait timeouts leave the lane alive; only terminal - // lane errors poison the mux. - async fn handle_lane_error(&self, idx: usize, e: ZerobusError) -> ZerobusError { - if self.streams[idx].is_closed() { - let cause = self.lane_terminal_error(idx, e).await; - self.shutdown_on_failure(idx, &cause); - cause - } else { - warn!(stream_index = idx, error = %e, "Sub-stream operation errored but lane remains alive"); - e - } + self.core.first().new_record() } /// Ingests a single record into the next sub-stream (round-robin). @@ -375,11 +186,9 @@ impl MultiplexedStream { &self, payload: impl Into, ) -> ZerobusResult { - self.check_closed()?; - let idx = self.pick_substream(); - let stream = &self.streams[idx]; - let encoded_batch = stream.prepare_record(payload)?; - self.enqueue_reserved(stream, idx, encoded_batch).await + self.core + .ingest(|stream| stream.prepare_record(payload)) + .await } /// Ingests a batch of records into a single sub-stream (round-robin). @@ -393,15 +202,13 @@ impl MultiplexedStream { I: IntoIterator, T: Into, { - self.check_closed()?; + self.core.check_closed()?; let records: Vec = payload.into_iter().collect(); if records.is_empty() { return Ok(None); } - let idx = self.pick_substream(); - let stream = &self.streams[idx]; - let encoded_batch = stream.prepare_records(records)?; - self.enqueue_reserved(stream, idx, encoded_batch) + self.core + .ingest(|stream| stream.prepare_records(records)) .await .map(Some) } @@ -415,53 +222,7 @@ impl MultiplexedStream { /// otherwise the first flush error is returned. Additional errors are /// logged. pub async fn flush(&self) -> ZerobusResult<()> { - if self.is_closed_fast() && self.failure.get().is_none() { - return Err(self.closed_error()); - } - if self.failure.get().is_none() { - if let Some(idx) = self.first_closed_stream() { - let error = self - .lane_terminal_error( - idx, - ZerobusError::InvalidStateError(format!( - "MultiplexedStream sub-stream {idx} is closed" - )), - ) - .await; - self.shutdown_on_failure(idx, &error); - } - } - - let results = join_all(self.streams.iter().map(ZerobusStream::flush)).await; - let mut first_error: Option = None; - let mut first_terminal: Option<(usize, ZerobusError)> = None; - for (i, result) in results.into_iter().enumerate() { - if let Err(e) = result { - if self.streams[i].is_closed() && first_terminal.is_none() { - let terminal_error = self.lane_terminal_error(i, e.clone()).await; - first_terminal = Some((i, terminal_error)); - } - if first_error.is_none() { - first_error = Some(e); - } else { - warn!( - stream_index = i, - error = %e, - "Additional sub-stream flush error" - ); - } - } - } - if let Some((i, error)) = first_terminal { - self.shutdown_on_failure(i, &error); - } else if let Some(error) = &first_error { - warn!(error = %error, "flush errored but sub-streams still alive"); - } - self.failure - .get() - .cloned() - .or(first_error) - .map_or(Ok(()), Err) + self.core.flush().await } /// Waits for server acknowledgment of the record or batch behind a @@ -471,20 +232,7 @@ impl MultiplexedStream { /// Only the lane that owns this message can complete or fail the wait; a /// failure on another lane does not make an acknowledged record retryable. pub async fn wait_for_message_id(&self, message_id: MessageId) -> ZerobusResult<()> { - let idx = message_id.stream_index(); - if idx >= self.streams.len() { - return Err(ZerobusError::InvalidArgument(format!( - "Invalid stream index {} in message id", - idx - ))); - } - match self.streams[idx] - .wait_for_offset(message_id.sub_offset()) - .await - { - Ok(()) => Ok(()), - Err(e) => Err(self.handle_lane_error(idx, e).await), - } + self.core.wait_for_message_id(message_id).await } /// Flushes and closes all sub-streams, releasing their resources. @@ -496,67 +244,13 @@ impl MultiplexedStream { /// Retrying a cancelled close resumes teardown once its flush attempt has /// completed, retaining both the flush result and terminal lane failures. pub async fn close(&mut self) -> ZerobusResult<()> { - info!("Closing MultiplexedStream"); - - self.is_closed.store(true, Ordering::Relaxed); - self.closed_token.cancel(); - - if self.close_flush_result.is_none() { - if self.failure.get().is_none() { - if let Some(idx) = self.first_closed_stream() { - if let Some(error) = self.streams[idx].terminal_error().await { - let _ = self.failure.set(error); - } - } - } - let mut first_error = None; - // All lanes stay live through this flush attempt. Cache its result - // before any supervisor is stopped or callback drain can be cancelled. - let results = join_all(self.streams.iter().map(ZerobusStream::flush)).await; - for (i, result) in results.into_iter().enumerate() { - if let Err(error) = result { - if self.failure.get().is_none() && self.streams[i].is_closed() { - if let Some(cause) = self.streams[i].terminal_error().await { - let _ = self.failure.set(cause); - } - } - if first_error.is_none() { - first_error = Some(error); - } else { - warn!(stream_index = i, error = %error, "Additional sub-stream flush error during close"); - } - } - } - self.close_flush_result = Some(first_error.map_or(Ok(()), Err)); - } - - let failure = &self.failure; - join_all(self.streams.iter_mut().enumerate().map(|(i, stream)| async move { - if let Some(error) = stream.close_after_flush().await { - // Publish immediately, not after join_all: another lane's - // callbacks may still be draining when this close is cancelled. - if let Err(error) = failure.set(error) { - warn!(stream_index = i, error = %error, "Additional terminal lane error during close"); - } - } - stream.shutdown_callbacks().await; - })) - .await; - - if let Some(error) = self.failure.get() { - Err(error.clone()) - } else { - self.close_flush_result - .as_ref() - .expect("close flush completed") - .clone() - } + self.core.close().await } /// Returns whether the mux is closed — either via [`close`](Self::close) /// or because a mux operation observed a sub-stream failure. pub fn is_closed(&self) -> bool { - self.is_closed_fast() + self.core.is_closed() } /// Returns records that were ingested but not acknowledged. @@ -576,26 +270,7 @@ impl MultiplexedStream { /// so results are always complete. Any error from close is swallowed — if records can /// still be recovered, they will be returned. pub async fn get_unacked_batches(&mut self) -> ZerobusResult> { - let _ = self.close().await; - let mut all_batches = Vec::new(); - for stream in &self.streams { - all_batches.extend(stream.get_unacked_batches().await?); - } - Ok(all_batches) - } -} - -impl Drop for MultiplexedStream { - fn drop(&mut self) { - self.is_closed.store(true, Ordering::Relaxed); - self.closed_token.cancel(); - // Fire cancellation on every sub-stream in parallel so their - // background tasks can start unwinding concurrently. The Vec drop - // below then runs each `ZerobusStream::Drop`, which aborts any - // JoinHandles that haven't already exited. - for stream in &self.streams { - stream.signal_shutdown(); - } + self.core.get_unacked_batches().await } } diff --git a/rust/sdk/src/multiplexed_stream/core.rs b/rust/sdk/src/multiplexed_stream/core.rs new file mode 100644 index 00000000..6a3f1e81 --- /dev/null +++ b/rust/sdk/src/multiplexed_stream/core.rs @@ -0,0 +1,365 @@ +//! Shared mux routing and lifecycle. Transport-specific invariants stay in `MuxLane`. +use super::lane::MuxLane; +use super::{MessageId, STREAM_BITS}; +use crate::{ZerobusError, ZerobusResult}; +use futures::future::join_all; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::OnceLock; +use std::time::Duration; +use tokio_util::sync::CancellationToken; +use tracing::{error, info, warn}; + +const CAPACITY_WAIT_TIMEOUT: Duration = Duration::from_secs(30); + +pub(super) struct MuxCore { + streams: Vec, + round_robin_counter: AtomicUsize, + is_closed: AtomicBool, + closed_token: CancellationToken, + failure: OnceLock, + // Cache completed pre-close work so cancelled teardown can resume. + close_flush_result: Option>, +} + +impl MuxCore { + pub(super) fn from_streams(streams: Vec) -> Self { + assert!( + !streams.is_empty(), + "MultiplexedStream requires at least one sub-stream" + ); + assert!( + streams.len() <= (1 << STREAM_BITS), + "MultiplexedStream supports at most {} sub-streams", + 1 << STREAM_BITS + ); + Self { + streams, + round_robin_counter: AtomicUsize::new(0), + is_closed: AtomicBool::new(false), + closed_token: CancellationToken::new(), + failure: OnceLock::new(), + close_flush_result: None, + } + } + + #[allow(clippy::result_large_err)] + pub(super) fn check_closed(&self) -> ZerobusResult<()> { + if let Some(error) = self.failure.get() { + return Err(error.clone()); + } + if self.is_closed_fast() { + return Err(self.closed_error()); + } + Ok(()) + } + + fn closed_error(&self) -> ZerobusError { + self.failure.get().cloned().unwrap_or_else(|| { + ZerobusError::InvalidStateError("MultiplexedStream is closed".to_string()) + }) + } + + fn is_closed_fast(&self) -> bool { + self.is_closed.load(Ordering::Relaxed) + } + + fn first_closed_stream(&self) -> Option { + self.streams.iter().position(MuxLane::is_closed) + } + + async fn lane_terminal_error(&self, idx: usize, fallback: ZerobusError) -> ZerobusError { + self.streams[idx].terminal_error().await.unwrap_or(fallback) + } + + fn shutdown_on_failure(&self, trigger_index: usize, cause: &ZerobusError) { + if self.is_closed_fast() || self.failure.set(cause.clone()).is_err() { + return; + } + self.is_closed.store(true, Ordering::Relaxed); + self.closed_token.cancel(); + + error!( + trigger_stream_index = trigger_index, + cause = %cause, + num_streams = self.streams.len(), + "MultiplexedStream poisoned due to sub-stream failure" + ); + } + + // Keep the selected lane through backpressure; do not reroute. + fn pick_substream(&self) -> usize { + self.round_robin_counter.fetch_add(1, Ordering::Relaxed) % self.streams.len() + } + + async fn reserve_capacity(&self, stream: &S, idx: usize) -> ZerobusResult { + let started_at = tokio::time::Instant::now(); + let timeout_ms = CAPACITY_WAIT_TIMEOUT.as_millis(); + let table_name = stream.table_name(); + let (capacity_option, capacity) = stream.capacity(); + + self.check_closed()?; + + let wait_for_reservation = async { + let reservation = stream.reserve_capacity(); + tokio::pin!(reservation); + + if let Ok(result) = tokio::time::timeout(Duration::from_secs(1), &mut reservation).await + { + return result; + } + let waited_ms = started_at.elapsed().as_millis(); + warn!( + stream_index = idx, + table_name, + waited_ms, + timeout_ms, + capacity, + capacity_option, + "Backpressure: sub-stream at capacity, waiting for drain" + ); + reservation.await + }; + + let result = tokio::select! { + result = tokio::time::timeout(CAPACITY_WAIT_TIMEOUT, wait_for_reservation) => result, + _ = self.closed_token.cancelled() => { + let waited_ms = started_at.elapsed().as_millis(); + if let Some(failure) = self.failure.get() { + return Err(failure.clone()); + } + warn!( + stream_index = idx, + table_name, + waited_ms, + capacity, + capacity_option, + "Multiplexed capacity wait cancelled by shutdown" + ); + return Err(ZerobusError::InvalidStateError( + format!( + "MultiplexedStream closed after {waited_ms} ms while waiting for capacity on sub-stream {idx} for table {table_name} ({capacity_option}: {capacity})" + ), + )); + } + }; + + match result { + Ok(Ok(reservation)) => Ok(reservation), + Ok(Err(e)) => Err(self.handle_lane_error(idx, e).await), + Err(_) => { + self.check_closed()?; + if stream.is_closed() { + return Err(self + .handle_lane_error( + idx, + ZerobusError::ConnectionTimeout(format!( + "Timed out waiting for capacity on multiplexed sub-stream {idx}" + )), + ) + .await); + } + let waited_ms = started_at.elapsed().as_millis(); + warn!( + stream_index = idx, + table_name, + waited_ms, + timeout_ms, + capacity, + capacity_option, + "Timed out waiting for multiplexed sub-stream capacity" + ); + Err(ZerobusError::ConnectionTimeout(format!( + "Timed out after {waited_ms} ms waiting for capacity on multiplexed sub-stream {idx} for table {table_name} (configured timeout: {timeout_ms} ms, {capacity_option}: {capacity})" + ))) + } + } + } + + pub(super) fn first(&self) -> &S { + &self.streams[0] + } + + pub(super) async fn ingest(&self, prepare: F) -> ZerobusResult + where + F: FnOnce(&S) -> ZerobusResult, + { + self.check_closed()?; + let idx = self.pick_substream(); + let stream = &self.streams[idx]; + let encoded_batch = prepare(stream)?; + let reservation = self.reserve_capacity(stream, idx).await?; + let enqueue_result = stream + .enqueue_reserved_admitted(encoded_batch, reservation, || self.check_closed()) + .await; + + match enqueue_result { + Ok(off) => Ok(MessageId::new(idx, off)), + Err(e) => Err(self.handle_lane_error(idx, e).await), + } + } + + // Payload errors and wait timeouts leave the lane alive; only terminal + // lane errors poison the mux. + async fn handle_lane_error(&self, idx: usize, e: ZerobusError) -> ZerobusError { + if self.streams[idx].is_closed() { + let cause = self.lane_terminal_error(idx, e).await; + self.shutdown_on_failure(idx, &cause); + cause + } else { + warn!(stream_index = idx, error = %e, "Sub-stream operation errored but lane remains alive"); + e + } + } + + pub async fn flush(&self) -> ZerobusResult<()> { + if self.is_closed_fast() && self.failure.get().is_none() { + return Err(self.closed_error()); + } + if self.failure.get().is_none() { + if let Some(idx) = self.first_closed_stream() { + let error = self + .lane_terminal_error( + idx, + ZerobusError::InvalidStateError(format!( + "MultiplexedStream sub-stream {idx} is closed" + )), + ) + .await; + self.shutdown_on_failure(idx, &error); + } + } + + let results = join_all(self.streams.iter().map(MuxLane::flush)).await; + let mut first_error: Option = None; + let mut first_terminal: Option<(usize, ZerobusError)> = None; + for (i, result) in results.into_iter().enumerate() { + if let Err(e) = result { + if self.streams[i].is_closed() && first_terminal.is_none() { + let terminal_error = self.lane_terminal_error(i, e.clone()).await; + first_terminal = Some((i, terminal_error)); + } + if first_error.is_none() { + first_error = Some(e); + } else { + warn!( + stream_index = i, + error = %e, + "Additional sub-stream flush error" + ); + } + } + } + if let Some((i, error)) = first_terminal { + self.shutdown_on_failure(i, &error); + } else if let Some(error) = &first_error { + warn!(error = %error, "flush errored but sub-streams still alive"); + } + self.failure + .get() + .cloned() + .or(first_error) + .map_or(Ok(()), Err) + } + + pub async fn wait_for_message_id(&self, message_id: MessageId) -> ZerobusResult<()> { + let idx = message_id.stream_index(); + if idx >= self.streams.len() { + return Err(ZerobusError::InvalidArgument(format!( + "Invalid stream index {} in message id", + idx + ))); + } + match self.streams[idx] + .wait_for_offset(message_id.sub_offset()) + .await + { + Ok(()) => Ok(()), + Err(e) => Err(self.handle_lane_error(idx, e).await), + } + } + + pub async fn close(&mut self) -> ZerobusResult<()> { + info!("Closing MultiplexedStream"); + + self.is_closed.store(true, Ordering::Relaxed); + self.closed_token.cancel(); + + if self.close_flush_result.is_none() { + if self.failure.get().is_none() { + if let Some(idx) = self.first_closed_stream() { + if let Some(error) = self.streams[idx].terminal_error().await { + let _ = self.failure.set(error); + } + } + } + let mut first_error = None; + // All lanes stay live through this flush attempt. Cache its result + // before any supervisor is stopped or callback drain can be cancelled. + let results = join_all(self.streams.iter().map(MuxLane::flush_before_close)).await; + for (i, result) in results.into_iter().enumerate() { + if let Err(error) = result { + if self.failure.get().is_none() && self.streams[i].is_closed() { + if let Some(cause) = self.streams[i].terminal_error().await { + let _ = self.failure.set(cause); + } + } + if first_error.is_none() { + first_error = Some(error); + } else { + warn!(stream_index = i, error = %error, "Additional sub-stream flush error during close"); + } + } + } + self.close_flush_result = Some(first_error.map_or(Ok(()), Err)); + } + + let failure = &self.failure; + join_all(self.streams.iter_mut().enumerate().map(|(i, stream)| async move { + if let Some(error) = stream.close_after_flush().await { + // Publish immediately, not after join_all: another lane's + // callbacks may still be draining when this close is cancelled. + if let Err(error) = failure.set(error) { + warn!(stream_index = i, error = %error, "Additional terminal lane error during close"); + } + } + stream.shutdown_callbacks().await; + })) + .await; + + if let Some(error) = self.failure.get() { + Err(error.clone()) + } else { + self.close_flush_result + .as_ref() + .expect("close flush completed") + .clone() + } + } + + pub fn is_closed(&self) -> bool { + self.is_closed_fast() + } + + pub(super) async fn get_unacked_batches(&mut self) -> ZerobusResult> { + let _ = self.close().await; + let mut all_batches = Vec::new(); + for stream in &self.streams { + all_batches.extend(stream.get_unacked_batches().await?); + } + Ok(all_batches) + } +} + +impl Drop for MuxCore { + fn drop(&mut self) { + self.is_closed.store(true, Ordering::Relaxed); + self.closed_token.cancel(); + // Fire cancellation on every sub-stream in parallel so their + // background tasks can start unwinding concurrently. The Vec drop + // below then runs each lane's Drop, which aborts any + // JoinHandles that haven't already exited. + for stream in &self.streams { + stream.signal_shutdown(); + } + } +} diff --git a/rust/sdk/src/multiplexed_stream/lane.rs b/rust/sdk/src/multiplexed_stream/lane.rs new file mode 100644 index 00000000..f9684ad1 --- /dev/null +++ b/rust/sdk/src/multiplexed_stream/lane.rs @@ -0,0 +1,94 @@ +//! Private transport contract. `async_trait` keeps this compatible with the core MSRV. +//! +//! Reservations own capacity until admission or drop. Admission must check the mux +//! under the lane's lifecycle lock, before assigning an offset or retaining data. +//! A terminal admission error must await finalization so `is_closed` and the typed +//! terminal cause are visible together. Retryable recovery remains lane-owned. + +use crate::{EncodedBatch, OffsetId, ZerobusError, ZerobusResult, ZerobusStream}; +use async_trait::async_trait; + +#[async_trait] +pub(super) trait MuxLane: Send + Sync { + type Batch: Send; + type Reservation: Send; + + fn table_name(&self) -> &str; + fn capacity(&self) -> (&'static str, usize); + fn is_closed(&self) -> bool; + fn signal_shutdown(&self) {} + async fn terminal_error(&self) -> Option; + async fn reserve_capacity(&self) -> ZerobusResult; + async fn enqueue_reserved_admitted( + &self, + batch: Self::Batch, + reservation: Self::Reservation, + admit: F, + ) -> ZerobusResult + where + F: FnOnce() -> ZerobusResult<()> + Send; + async fn flush(&self) -> ZerobusResult<()>; + async fn wait_for_offset(&self, offset: OffsetId) -> ZerobusResult<()>; + // gRPC needs the shared pre-flush barrier. Arrow flushes within its close supervisor. + async fn flush_before_close(&self) -> ZerobusResult<()> { + Ok(()) + } + async fn close_after_flush(&mut self) -> Option; + async fn shutdown_callbacks(&mut self) {} + async fn get_unacked_batches(&self) -> ZerobusResult>; +} + +#[async_trait] +impl MuxLane for ZerobusStream { + type Batch = EncodedBatch; + type Reservation = crate::landing_zone::CapacityReservation; + + fn table_name(&self) -> &str { + &self.table_properties.table_name + } + fn capacity(&self) -> (&'static str, usize) { + ("max_inflight_requests", self.options.max_inflight_requests) + } + fn is_closed(&self) -> bool { + self.is_closed() + } + fn signal_shutdown(&self) { + self.signal_shutdown(); + } + async fn terminal_error(&self) -> Option { + self.terminal_error().await + } + async fn reserve_capacity(&self) -> ZerobusResult { + self.reserve_capacity().await + } + async fn enqueue_reserved_admitted( + &self, + batch: Self::Batch, + reservation: Self::Reservation, + admit: F, + ) -> ZerobusResult + where + F: FnOnce() -> ZerobusResult<()> + Send, + { + self.enqueue_reserved_admitted(batch, reservation, admit) + .await + } + async fn flush(&self) -> ZerobusResult<()> { + self.flush().await + } + async fn wait_for_offset(&self, offset: OffsetId) -> ZerobusResult<()> { + self.wait_for_offset(offset).await + } + async fn flush_before_close(&self) -> ZerobusResult<()> { + self.flush().await + } + async fn close_after_flush(&mut self) -> Option { + self.close_after_flush().await + } + async fn shutdown_callbacks(&mut self) { + self.shutdown_callbacks().await; + } + async fn get_unacked_batches(&self) -> ZerobusResult> { + self.get_unacked_batches().await + } +} From f55b30394ce68320bb3af4528e237614e059f4e8 Mon Sep 17 00:00:00 2001 From: danilo-najkov-db Date: Fri, 25 Sep 2026 15:58:06 +0000 Subject: [PATCH 2/2] [Rust] Simplify mux lane contract Signed-off-by: danilo-najkov-db --- rust/sdk/src/multiplexed_stream.rs | 2 +- rust/sdk/src/multiplexed_stream/core.rs | 37 ++++---- rust/sdk/src/multiplexed_stream/lane.rs | 85 ++++-------------- rust/sdk/src/stream/grpc/acks.rs | 10 +-- rust/sdk/src/stream/grpc/close.rs | 35 +------- rust/sdk/src/stream/grpc/ingest.rs | 29 ------ rust/sdk/src/stream/grpc/mod.rs | 2 + rust/sdk/src/stream/grpc/mux_lane.rs | 113 ++++++++++++++++++++++++ 8 files changed, 156 insertions(+), 157 deletions(-) create mode 100644 rust/sdk/src/stream/grpc/mux_lane.rs diff --git a/rust/sdk/src/multiplexed_stream.rs b/rust/sdk/src/multiplexed_stream.rs index 3c079d65..c247b287 100644 --- a/rust/sdk/src/multiplexed_stream.rs +++ b/rust/sdk/src/multiplexed_stream.rs @@ -21,7 +21,7 @@ use std::sync::Arc; mod core; -mod lane; +pub(crate) mod lane; use core::MuxCore; use crate::{ diff --git a/rust/sdk/src/multiplexed_stream/core.rs b/rust/sdk/src/multiplexed_stream/core.rs index 6a3f1e81..602494a3 100644 --- a/rust/sdk/src/multiplexed_stream/core.rs +++ b/rust/sdk/src/multiplexed_stream/core.rs @@ -1,5 +1,5 @@ //! Shared mux routing and lifecycle. Transport-specific invariants stay in `MuxLane`. -use super::lane::MuxLane; +use super::lane::{CapacityContext, MuxLane}; use super::{MessageId, STREAM_BITS}; use crate::{ZerobusError, ZerobusResult}; use futures::future::join_all; @@ -64,11 +64,11 @@ impl MuxCore { } fn first_closed_stream(&self) -> Option { - self.streams.iter().position(MuxLane::is_closed) + self.streams.iter().position(MuxLane::is_terminal) } async fn lane_terminal_error(&self, idx: usize, fallback: ZerobusError) -> ZerobusError { - self.streams[idx].terminal_error().await.unwrap_or(fallback) + self.streams[idx].terminal_cause().await.unwrap_or(fallback) } fn shutdown_on_failure(&self, trigger_index: usize, cause: &ZerobusError) { @@ -94,13 +94,16 @@ impl MuxCore { async fn reserve_capacity(&self, stream: &S, idx: usize) -> ZerobusResult { let started_at = tokio::time::Instant::now(); let timeout_ms = CAPACITY_WAIT_TIMEOUT.as_millis(); - let table_name = stream.table_name(); - let (capacity_option, capacity) = stream.capacity(); + let CapacityContext { + table_name, + capacity_option, + capacity, + } = stream.capacity_context(); self.check_closed()?; let wait_for_reservation = async { - let reservation = stream.reserve_capacity(); + let reservation = stream.reserve_slot(); tokio::pin!(reservation); if let Ok(result) = tokio::time::timeout(Duration::from_secs(1), &mut reservation).await @@ -148,7 +151,7 @@ impl MuxCore { Ok(Err(e)) => Err(self.handle_lane_error(idx, e).await), Err(_) => { self.check_closed()?; - if stream.is_closed() { + if stream.is_terminal() { return Err(self .handle_lane_error( idx, @@ -189,7 +192,7 @@ impl MuxCore { let encoded_batch = prepare(stream)?; let reservation = self.reserve_capacity(stream, idx).await?; let enqueue_result = stream - .enqueue_reserved_admitted(encoded_batch, reservation, || self.check_closed()) + .enqueue_admitted(encoded_batch, reservation, || self.check_closed()) .await; match enqueue_result { @@ -201,7 +204,7 @@ impl MuxCore { // Payload errors and wait timeouts leave the lane alive; only terminal // lane errors poison the mux. async fn handle_lane_error(&self, idx: usize, e: ZerobusError) -> ZerobusError { - if self.streams[idx].is_closed() { + if self.streams[idx].is_terminal() { let cause = self.lane_terminal_error(idx, e).await; self.shutdown_on_failure(idx, &cause); cause @@ -229,12 +232,12 @@ impl MuxCore { } } - let results = join_all(self.streams.iter().map(MuxLane::flush)).await; + let results = join_all(self.streams.iter().map(MuxLane::flush_lane)).await; let mut first_error: Option = None; let mut first_terminal: Option<(usize, ZerobusError)> = None; for (i, result) in results.into_iter().enumerate() { if let Err(e) = result { - if self.streams[i].is_closed() && first_terminal.is_none() { + if self.streams[i].is_terminal() && first_terminal.is_none() { let terminal_error = self.lane_terminal_error(i, e.clone()).await; first_terminal = Some((i, terminal_error)); } @@ -270,7 +273,7 @@ impl MuxCore { ))); } match self.streams[idx] - .wait_for_offset(message_id.sub_offset()) + .wait_for_local_offset(message_id.sub_offset()) .await { Ok(()) => Ok(()), @@ -287,7 +290,7 @@ impl MuxCore { if self.close_flush_result.is_none() { if self.failure.get().is_none() { if let Some(idx) = self.first_closed_stream() { - if let Some(error) = self.streams[idx].terminal_error().await { + if let Some(error) = self.streams[idx].terminal_cause().await { let _ = self.failure.set(error); } } @@ -298,8 +301,8 @@ impl MuxCore { let results = join_all(self.streams.iter().map(MuxLane::flush_before_close)).await; for (i, result) in results.into_iter().enumerate() { if let Err(error) = result { - if self.failure.get().is_none() && self.streams[i].is_closed() { - if let Some(cause) = self.streams[i].terminal_error().await { + if self.failure.get().is_none() && self.streams[i].is_terminal() { + if let Some(cause) = self.streams[i].terminal_cause().await { let _ = self.failure.set(cause); } } @@ -322,7 +325,7 @@ impl MuxCore { warn!(stream_index = i, error = %error, "Additional terminal lane error during close"); } } - stream.shutdown_callbacks().await; + stream.drain_callbacks().await; })) .await; @@ -344,7 +347,7 @@ impl MuxCore { let _ = self.close().await; let mut all_batches = Vec::new(); for stream in &self.streams { - all_batches.extend(stream.get_unacked_batches().await?); + all_batches.extend(stream.unacked_batches().await?); } Ok(all_batches) } diff --git a/rust/sdk/src/multiplexed_stream/lane.rs b/rust/sdk/src/multiplexed_stream/lane.rs index f9684ad1..ea4dc7dc 100644 --- a/rust/sdk/src/multiplexed_stream/lane.rs +++ b/rust/sdk/src/multiplexed_stream/lane.rs @@ -5,21 +5,27 @@ //! A terminal admission error must await finalization so `is_closed` and the typed //! terminal cause are visible together. Retryable recovery remains lane-owned. -use crate::{EncodedBatch, OffsetId, ZerobusError, ZerobusResult, ZerobusStream}; +use crate::{OffsetId, ZerobusError, ZerobusResult}; use async_trait::async_trait; +/// Diagnostic context for a selected lane's capacity wait. +pub(crate) struct CapacityContext<'a> { + pub(crate) table_name: &'a str, + pub(crate) capacity_option: &'static str, + pub(crate) capacity: usize, +} + #[async_trait] -pub(super) trait MuxLane: Send + Sync { +pub(crate) trait MuxLane: Send + Sync { type Batch: Send; type Reservation: Send; - fn table_name(&self) -> &str; - fn capacity(&self) -> (&'static str, usize); - fn is_closed(&self) -> bool; + fn capacity_context(&self) -> CapacityContext<'_>; + fn is_terminal(&self) -> bool; fn signal_shutdown(&self) {} - async fn terminal_error(&self) -> Option; - async fn reserve_capacity(&self) -> ZerobusResult; - async fn enqueue_reserved_admitted( + async fn terminal_cause(&self) -> Option; + async fn reserve_slot(&self) -> ZerobusResult; + async fn enqueue_admitted( &self, batch: Self::Batch, reservation: Self::Reservation, @@ -27,68 +33,13 @@ pub(super) trait MuxLane: Send + Sync { ) -> ZerobusResult where F: FnOnce() -> ZerobusResult<()> + Send; - async fn flush(&self) -> ZerobusResult<()>; - async fn wait_for_offset(&self, offset: OffsetId) -> ZerobusResult<()>; + async fn flush_lane(&self) -> ZerobusResult<()>; + async fn wait_for_local_offset(&self, offset: OffsetId) -> ZerobusResult<()>; // gRPC needs the shared pre-flush barrier. Arrow flushes within its close supervisor. async fn flush_before_close(&self) -> ZerobusResult<()> { Ok(()) } async fn close_after_flush(&mut self) -> Option; - async fn shutdown_callbacks(&mut self) {} - async fn get_unacked_batches(&self) -> ZerobusResult>; -} - -#[async_trait] -impl MuxLane for ZerobusStream { - type Batch = EncodedBatch; - type Reservation = crate::landing_zone::CapacityReservation; - - fn table_name(&self) -> &str { - &self.table_properties.table_name - } - fn capacity(&self) -> (&'static str, usize) { - ("max_inflight_requests", self.options.max_inflight_requests) - } - fn is_closed(&self) -> bool { - self.is_closed() - } - fn signal_shutdown(&self) { - self.signal_shutdown(); - } - async fn terminal_error(&self) -> Option { - self.terminal_error().await - } - async fn reserve_capacity(&self) -> ZerobusResult { - self.reserve_capacity().await - } - async fn enqueue_reserved_admitted( - &self, - batch: Self::Batch, - reservation: Self::Reservation, - admit: F, - ) -> ZerobusResult - where - F: FnOnce() -> ZerobusResult<()> + Send, - { - self.enqueue_reserved_admitted(batch, reservation, admit) - .await - } - async fn flush(&self) -> ZerobusResult<()> { - self.flush().await - } - async fn wait_for_offset(&self, offset: OffsetId) -> ZerobusResult<()> { - self.wait_for_offset(offset).await - } - async fn flush_before_close(&self) -> ZerobusResult<()> { - self.flush().await - } - async fn close_after_flush(&mut self) -> Option { - self.close_after_flush().await - } - async fn shutdown_callbacks(&mut self) { - self.shutdown_callbacks().await; - } - async fn get_unacked_batches(&self) -> ZerobusResult> { - self.get_unacked_batches().await - } + async fn drain_callbacks(&mut self) {} + async fn unacked_batches(&self) -> ZerobusResult>; } diff --git a/rust/sdk/src/stream/grpc/acks.rs b/rust/sdk/src/stream/grpc/acks.rs index 800aba25..f412f2d2 100644 --- a/rust/sdk/src/stream/grpc/acks.rs +++ b/rust/sdk/src/stream/grpc/acks.rs @@ -24,14 +24,6 @@ impl ZerobusStream { } } - /// Returns the final server error after this stream becomes terminal. - /// Multiplexed streams use this to preserve the lane's typed failure when - /// they discover an asynchronously closed lane. - pub(crate) async fn terminal_error(&self) -> Option { - self.terminal_token.cancelled().await; - self.server_error_rx.borrow().clone() - } - /// Internal method to wait for a specific offset to be acknowledged. /// Used by both `flush()` and `wait_for_offset()`. async fn wait_for_offset_internal( @@ -80,7 +72,7 @@ impl ZerobusStream { } // fail_stream sets is_closed, publishes the error, then cancels // terminal_token. This path can observe closure before the final - // watch value; mux terminal_error() waits on the token before reading it. + // watch value; the mux terminal-cause hook waits on the token before reading it. if let Some(server_error) = error_rx.borrow().clone() { return Err(Self::normalize_wait_error(server_error)); } diff --git a/rust/sdk/src/stream/grpc/close.rs b/rust/sdk/src/stream/grpc/close.rs index e85c5647..63702488 100644 --- a/rust/sdk/src/stream/grpc/close.rs +++ b/rust/sdk/src/stream/grpc/close.rs @@ -71,32 +71,10 @@ impl ZerobusStream { flush_result } - /// Stops a mux lane after its flush attempt. Cache the outcome and close - /// the lane before the caller starts cancellable callback draining. - pub(crate) async fn close_after_flush(&mut self) -> Option { - if let Some(result) = &self.supervisor_shutdown_result { - return result.as_ref().err().cloned(); - } - let task_error = self.shutdown_supervisor().await; - // The supervisor publishes the final error before cancelling this - // token. A clean shutdown's transient watch error must not be promoted - // to a terminal cause. - let lane_error = if self.terminal_token.is_cancelled() { - self.server_error_rx.borrow().clone() - } else { - None - }; - let result = lane_error.or(task_error).map_or(Ok(()), Err); - self.supervisor_shutdown_result = Some(result.clone()); - self.is_closed.store(true, Ordering::Relaxed); - self.terminal_token.cancel(); - result.err() - } - /// Waits up to one second for cooperative shutdown, then at most 100ms /// after abort. Synchronous user code may outlive that budget; the retained /// handle is aborted and the timeout outcome is cached for retries. - async fn shutdown_supervisor(&mut self) -> Option { + pub(super) async fn shutdown_supervisor(&mut self) -> Option { if let Some(result) = &self.supervisor_shutdown_result { return result.as_ref().err().cloned(); } @@ -168,17 +146,6 @@ impl ZerobusStream { let _ = task.await; } } - - // Signal the stream to stop accepting work and tear down its background - // tasks. Unlike `close`, this only needs `&self` — it relies on the - // cancellation token and `is_closed` flag, both of which are already - // interior-mutable. The `JoinHandle`s aren't reaped here; that happens in - // `close` or `Drop`. - pub(crate) fn signal_shutdown(&self) { - self.is_closed.store(true, Ordering::Relaxed); - self.terminal_token.cancel(); - self.cancellation_token.cancel(); - } } #[cfg(all(test, feature = "testing"))] diff --git a/rust/sdk/src/stream/grpc/ingest.rs b/rust/sdk/src/stream/grpc/ingest.rs index 17d86c2d..1ccbf2eb 100644 --- a/rust/sdk/src/stream/grpc/ingest.rs +++ b/rust/sdk/src/stream/grpc/ingest.rs @@ -297,33 +297,4 @@ impl ZerobusStream { } } } - - pub(crate) async fn enqueue_reserved_admitted( - &self, - encoded_batch: EncodedBatch, - reservation: crate::landing_zone::CapacityReservation, - admit: F, - ) -> ZerobusResult - where - F: FnOnce() -> ZerobusResult<()>, - { - let _guard = self.sync_mutex.lock().await; - admit()?; - self.check_open()?; - - let offset_id = self.logical_offset_id_generator.next(); - debug!( - offset_id, - record_count = encoded_batch.get_record_count(), - "Ingesting record(s)" - ); - self.landing_zone.enqueue_reserved( - Box::new(IngestRequest { - payload: encoded_batch, - offset_id, - }), - reservation, - ); - Ok(offset_id) - } } diff --git a/rust/sdk/src/stream/grpc/mod.rs b/rust/sdk/src/stream/grpc/mod.rs index e82a6c8a..96d574a1 100644 --- a/rust/sdk/src/stream/grpc/mod.rs +++ b/rust/sdk/src/stream/grpc/mod.rs @@ -11,6 +11,7 @@ //! |-----------------------|----------------------------------------------------|--------------------| //! | `types.rs` | Internal types (`IngestRequest`, channel messages) | Transport-agnostic | //! | `ingest.rs` | Public `ingest_*` methods | Transport-agnostic | +//! | `mux_lane.rs` | Private mux lane hooks and gRPC adapter | gRPC-specific | //! | `acks.rs` | `flush`, `wait_for_offset`, unacked queries | Transport-agnostic | //! | `close.rs` | `close`, `is_closed`, task shutdown | Transport-agnostic | //! | `callback_handler.rs` | User-callback dispatch task | Transport-agnostic | @@ -41,6 +42,7 @@ mod callback_handler; mod close; mod connection; mod ingest; +mod mux_lane; mod receiver; mod sender; mod supervisor; diff --git a/rust/sdk/src/stream/grpc/mux_lane.rs b/rust/sdk/src/stream/grpc/mux_lane.rs new file mode 100644 index 00000000..4cf52c13 --- /dev/null +++ b/rust/sdk/src/stream/grpc/mux_lane.rs @@ -0,0 +1,113 @@ +//! gRPC implementation of the private multiplexed-lane contract. +use std::sync::atomic::Ordering; + +use async_trait::async_trait; +use tracing::debug; + +use super::types::IngestRequest; +use super::ZerobusStream; +use crate::multiplexed_stream::lane::{CapacityContext, MuxLane}; +use crate::{EncodedBatch, OffsetId, ZerobusError, ZerobusResult}; + +#[async_trait] +impl MuxLane for ZerobusStream { + type Batch = EncodedBatch; + type Reservation = crate::landing_zone::CapacityReservation; + + fn capacity_context(&self) -> CapacityContext<'_> { + CapacityContext { + table_name: &self.table_properties.table_name, + capacity_option: "max_inflight_requests", + capacity: self.options.max_inflight_requests, + } + } + + fn is_terminal(&self) -> bool { + ZerobusStream::is_closed(self) + } + + fn signal_shutdown(&self) { + self.is_closed.store(true, Ordering::Relaxed); + self.terminal_token.cancel(); + self.cancellation_token.cancel(); + } + + async fn terminal_cause(&self) -> Option { + self.terminal_token.cancelled().await; + self.server_error_rx.borrow().clone() + } + + async fn reserve_slot(&self) -> ZerobusResult { + ZerobusStream::reserve_capacity(self).await + } + + async fn enqueue_admitted( + &self, + batch: Self::Batch, + reservation: Self::Reservation, + admit: F, + ) -> ZerobusResult + where + F: FnOnce() -> ZerobusResult<()> + Send, + { + let _guard = self.sync_mutex.lock().await; + admit()?; + self.check_open()?; + + let offset_id = self.logical_offset_id_generator.next(); + debug!( + offset_id, + record_count = batch.get_record_count(), + "Ingesting record(s)" + ); + self.landing_zone.enqueue_reserved( + Box::new(IngestRequest { + payload: batch, + offset_id, + }), + reservation, + ); + Ok(offset_id) + } + + async fn flush_lane(&self) -> ZerobusResult<()> { + ZerobusStream::flush(self).await + } + + async fn wait_for_local_offset(&self, offset: OffsetId) -> ZerobusResult<()> { + ZerobusStream::wait_for_offset(self, offset).await + } + + // Keep the gRPC pre-flush barrier and cache its outcome before teardown. + async fn flush_before_close(&self) -> ZerobusResult<()> { + ZerobusStream::flush(self).await + } + + async fn close_after_flush(&mut self) -> Option { + if let Some(result) = &self.supervisor_shutdown_result { + return result.as_ref().err().cloned(); + } + let task_error = ZerobusStream::shutdown_supervisor(self).await; + // The supervisor publishes the final error before cancelling this + // token. A clean shutdown's transient watch error must not be promoted + // to a terminal cause. + let lane_error = if self.terminal_token.is_cancelled() { + self.server_error_rx.borrow().clone() + } else { + None + }; + let result = lane_error.or(task_error).map_or(Ok(()), Err); + self.supervisor_shutdown_result = Some(result.clone()); + self.is_closed.store(true, Ordering::Relaxed); + self.terminal_token.cancel(); + result.err() + } + + async fn drain_callbacks(&mut self) { + ZerobusStream::shutdown_callbacks(self).await; + } + + async fn unacked_batches(&self) -> ZerobusResult> { + ZerobusStream::get_unacked_batches(self).await + } +}