diff --git a/rust/sdk/src/landing_zone.rs b/rust/sdk/src/landing_zone.rs index f5e61daa..e97dbdcd 100644 --- a/rust/sdk/src/landing_zone.rs +++ b/rust/sdk/src/landing_zone.rs @@ -4,6 +4,8 @@ use std::sync::Arc; use thiserror::Error; use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore}; +pub(crate) struct CapacityReservation(OwnedSemaphorePermit); + #[derive(Debug, Error)] pub enum LandingZoneError { #[error("Attempted to remove non-observed element")] @@ -85,29 +87,32 @@ impl LandingZone { all_items } + pub(crate) async fn reserve_capacity(&self) -> CapacityReservation { + CapacityReservation( + self.semaphore + .clone() + .acquire_owned() + .await + .expect("Failed to acquire semaphore"), + ) + } + + pub(crate) fn enqueue_reserved(&self, request: T, reservation: CapacityReservation) { + let mut state = self.state.lock().expect("Lock poisoned"); + let mut permits = self.permits.lock().expect("Lock poisoned"); + state.queue.push_back(request); + permits.push_back(reservation.0); + // Unblock one of the waiting observe() calls. + self.new_item_notify.notify_one(); + } + /// Adds an item to the queue. /// /// This method will block if the maximum number of inflight requests has been reached, /// providing automatic backpressure control. - /// - /// # Arguments - /// - /// * `request` - The item to add to the queue pub async fn add(&self, request: T) { - let _permit = self - .semaphore - .clone() - .acquire_owned() - .await - .expect("Failed to acquire semaphore"); - let mut state = self.state.lock().expect("Lock poisoned"); - state.queue.push_back(request); - self.permits - .lock() - .expect("Lock poisoned") - .push_back(_permit); - // Unblock one of the waiting observe() calls. - self.new_item_notify.notify_one(); + let reservation = self.reserve_capacity().await; + self.enqueue_reserved(request, reservation); } /// Removes and returns the next observed item. diff --git a/rust/sdk/src/multiplexed_stream.rs b/rust/sdk/src/multiplexed_stream.rs index 95d169dd..8bd7a62f 100644 --- a/rust/sdk/src/multiplexed_stream.rs +++ b/rust/sdk/src/multiplexed_stream.rs @@ -25,6 +25,8 @@ use tracing::{error, info, warn}; use crate::{EncodedBatch, EncodedRecord, OffsetId, ZerobusError, ZerobusResult, ZerobusStream}; +const CAPACITY_WAIT_TIMEOUT: std::time::Duration = std::time::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; @@ -50,7 +52,7 @@ impl std::fmt::Display for MessageId { } impl MessageId { - fn new(stream_index: usize, sub_offset: OffsetId) -> Self { + pub(crate) fn new(stream_index: usize, sub_offset: OffsetId) -> Self { debug_assert!(stream_index < (1 << STREAM_BITS)); debug_assert!((0..=OFFSET_MASK).contains(&sub_offset)); Self(((stream_index as i64) << (64 - STREAM_BITS)) | (sub_offset & OFFSET_MASK)) @@ -89,11 +91,14 @@ pub struct MultiplexedStream { streams: Vec, round_robin_counter: AtomicUsize, is_closed: AtomicBool, + admission: tokio::sync::RwLock<()>, } impl MultiplexedStream { /// Creates a multiplexed stream over the given sub-streams. /// + /// Ingest waits up to 30 seconds for capacity on its selected sub-stream. + /// /// # Panics /// /// Panics if `streams` is empty or holds more than 64 sub-streams. @@ -111,6 +116,7 @@ impl MultiplexedStream { streams, round_robin_counter: AtomicUsize::new(0), is_closed: AtomicBool::new(false), + admission: tokio::sync::RwLock::new(()), } } @@ -140,6 +146,13 @@ impl MultiplexedStream { "MultiplexedStream poisoned due to sub-stream failure" ); + // Drain any readers already admitted before `is_closed` was set. The + // write lock is only a barrier: readers arriving after it is released + // will observe the closed state and reject the ingest. + { + let _admission = self.admission.write().await; + } + let flush_results = join_all(self.streams.iter().map(|s| s.flush())).await; for (i, result) in flush_results.into_iter().enumerate() { if let Err(e) = result { @@ -161,10 +174,15 @@ impl MultiplexedStream { self.round_robin_counter.fetch_add(1, Ordering::Relaxed) % self.streams.len() } - async fn wait_for_capacity(&self, stream: &ZerobusStream, idx: usize) -> ZerobusResult<()> { + async fn reserve_capacity( + &self, + stream: &ZerobusStream, + idx: usize, + ) -> ZerobusResult { let mut backoff_ms = 1u64; - let mut total_wait_ms = 0u64; let mut logged_backpressure = false; + let started_at = tokio::time::Instant::now(); + let deadline = started_at + CAPACITY_WAIT_TIMEOUT; loop { self.check_closed()?; @@ -178,14 +196,28 @@ impl MultiplexedStream { return Err(err); } - if stream.has_capacity() { - return Ok(()); + let now = tokio::time::Instant::now(); + if now >= deadline { + return Err(ZerobusError::ConnectionTimeout(format!( + "Timed out waiting for capacity on multiplexed sub-stream {}", + idx + ))); + } + + let wait_duration = std::time::Duration::from_millis(backoff_ms) + .min(deadline.saturating_duration_since(now)); + + match tokio::time::timeout(wait_duration, stream.reserve_capacity()).await { + Ok(Ok(reservation)) => return Ok(reservation), + Ok(Err(e)) => return Err(self.handle_ingest_error(e, stream, idx).await), + // Timed out waiting for a permit: the sub-stream is still at + // capacity. Loop to re-check liveness and keep waiting. + Err(_elapsed) => {} } - tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await; - total_wait_ms += backoff_ms; backoff_ms = (backoff_ms * 2).min(50); + let total_wait_ms = started_at.elapsed().as_millis(); if !logged_backpressure && total_wait_ms >= 1000 { warn!( stream_index = idx, @@ -196,6 +228,27 @@ impl MultiplexedStream { } } + 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, || async { + let admission = self.admission.read().await; + self.check_closed()?; + Ok(admission) + }) + .await; + + match enqueue_result { + Ok(off) => Ok(MessageId::new(idx, off)), + Err(e) => Err(self.handle_ingest_error(e, stream, idx).await), + } + } + // Only poison the mux when the sub-stream itself has reached a terminal // state (`is_closed`): recovery is exhausted or a non-retryable server // error fired, so its offsets/pending records are unrecoverable. Other @@ -230,13 +283,8 @@ impl MultiplexedStream { let record = payload.into(); let idx = self.pick_substream(); let stream = &self.streams[idx]; - self.wait_for_capacity(stream, idx).await?; - self.check_closed()?; - - match stream.ingest_record_offset(record).await { - Ok(off) => Ok(MessageId::new(idx, off)), - Err(e) => Err(self.handle_ingest_error(e, stream, idx).await), - } + let encoded_batch = stream.prepare_record_batch(record)?; + self.enqueue_reserved(stream, idx, encoded_batch).await } /// Ingests a batch of records into a single sub-stream (round-robin). @@ -256,13 +304,10 @@ impl MultiplexedStream { } let idx = self.pick_substream(); let stream = &self.streams[idx]; - self.wait_for_capacity(stream, idx).await?; - self.check_closed()?; - - match stream.ingest_records_offset(records).await { - Ok(sub_offset) => Ok(sub_offset.map(|off| MessageId::new(idx, off))), - Err(e) => Err(self.handle_ingest_error(e, stream, idx).await), - } + let encoded_batch = stream.prepare_records_batch(records)?; + self.enqueue_reserved(stream, idx, encoded_batch) + .await + .map(Some) } /// Waits until every record already queued on every sub-stream is diff --git a/rust/sdk/src/stream/grpc/ingest.rs b/rust/sdk/src/stream/grpc/ingest.rs index ca4c9e89..eeb63e4b 100644 --- a/rust/sdk/src/stream/grpc/ingest.rs +++ b/rust/sdk/src/stream/grpc/ingest.rs @@ -53,14 +53,8 @@ impl ZerobusStream { &self, payload: impl Into, ) -> ZerobusResult { - let encoded_batch = EncodedBatch::try_from_record(payload, self.options.record_type) - .ok_or_else(|| { - ZerobusError::InvalidArgument( - "Record type does not match stream configuration".to_string(), - ) - })?; - - self.ingest_internal_v2(encoded_batch).await + let encoded_batch = self.prepare_record_batch(payload)?; + self.enqueue_prepared_batch(encoded_batch).await } /// Ingests a batch of records and returns the logical offset directly. @@ -100,6 +94,38 @@ impl ZerobusStream { /// # } /// ``` pub async fn ingest_records_offset(&self, payload: I) -> ZerobusResult> + where + I: IntoIterator, + T: Into, + { + let encoded_batch = self.prepare_records_batch(payload)?; + + if encoded_batch.is_empty() { + Ok(None) + } else { + self.enqueue_prepared_batch(encoded_batch) + .await + .map(Option::Some) + } + } + + #[allow(clippy::result_large_err)] + pub(crate) fn prepare_record_batch( + &self, + payload: impl Into, + ) -> ZerobusResult { + let encoded_batch = EncodedBatch::try_from_record(payload, self.options.record_type) + .ok_or_else(|| { + ZerobusError::InvalidArgument( + "Record type does not match stream configuration".to_string(), + ) + })?; + self.validate_ingest_payload(&encoded_batch)?; + Ok(encoded_batch) + } + + #[allow(clippy::result_large_err)] + pub(crate) fn prepare_records_batch(&self, payload: I) -> ZerobusResult where I: IntoIterator, T: Into, @@ -110,14 +136,31 @@ impl ZerobusStream { "Record type does not match stream configuration".to_string(), ) })?; + self.validate_ingest_payload(&encoded_batch)?; + Ok(encoded_batch) + } - if encoded_batch.is_empty() { - Ok(None) - } else { - self.ingest_internal_v2(encoded_batch) - .await - .map(Option::Some) + #[allow(clippy::result_large_err)] + fn validate_ingest_payload(&self, encoded_batch: &EncodedBatch) -> ZerobusResult<()> { + let byte_size = encoded_batch.total_byte_size(); + let max_payload_bytes = self.options.max_ingest_payload_bytes; + if byte_size > max_payload_bytes { + return Err(ZerobusError::InvalidArgument(format!( + "Ingest payload too large: {byte_size} bytes exceeds the configured limit of {max_payload_bytes} bytes" + ))); } + Ok(()) + } + + #[allow(clippy::result_large_err)] + pub(crate) fn check_open(&self) -> ZerobusResult<()> { + if self.is_closed.load(Ordering::Relaxed) { + error!(table_name = %self.table_properties.table_name, "Stream closed"); + return Err(ZerobusError::StreamClosedError(tonic::Status::internal( + "Stream closed", + ))); + } + Ok(()) } /// Internal unified method for ingesting records and batches. @@ -128,12 +171,7 @@ impl ZerobusStream { &self, encoded_batch: EncodedBatch, ) -> ZerobusResult>> { - if self.is_closed.load(Ordering::Relaxed) { - error!(table_name = %self.table_properties.table_name, "Stream closed"); - return Err(ZerobusError::StreamClosedError(tonic::Status::internal( - "Stream closed", - ))); - } + self.check_open()?; let _guard = self.sync_mutex.lock().await; @@ -177,21 +215,8 @@ impl ZerobusStream { /// /// Returns the logical offset directly without waiting for acknowledgment. /// Used by the public `ingest_*_offset` methods. - async fn ingest_internal_v2(&self, encoded_batch: EncodedBatch) -> ZerobusResult { - let byte_size = encoded_batch.total_byte_size(); - let max_payload_bytes = self.options.max_ingest_payload_bytes; - if byte_size > max_payload_bytes { - return Err(ZerobusError::InvalidArgument(format!( - "Ingest payload too large: {byte_size} bytes exceeds the configured limit of {max_payload_bytes} bytes" - ))); - } - - if self.is_closed.load(Ordering::Relaxed) { - error!(table_name = %self.table_properties.table_name, "Stream closed"); - return Err(ZerobusError::StreamClosedError(tonic::Status::internal( - "Stream closed", - ))); - } + async fn enqueue_prepared_batch(&self, encoded_batch: EncodedBatch) -> ZerobusResult { + self.check_open()?; let _guard = self.sync_mutex.lock().await; @@ -211,7 +236,42 @@ impl ZerobusStream { } #[cfg(feature = "testing")] - pub(crate) fn has_capacity(&self) -> bool { - self.landing_zone.len() < self.options.max_inflight_requests + pub(crate) async fn reserve_capacity( + &self, + ) -> ZerobusResult { + self.check_open()?; + Ok(self.landing_zone.reserve_capacity().await) + } + + #[cfg(feature = "testing")] + pub(crate) async fn enqueue_reserved_admitted( + &self, + encoded_batch: EncodedBatch, + reservation: crate::landing_zone::CapacityReservation, + admit: F, + ) -> ZerobusResult + where + F: FnOnce() -> Fut, + Fut: Future>, + { + let _guard = self.sync_mutex.lock().await; + let admission_guard = admit().await?; + 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, + ); + drop(admission_guard); + Ok(offset_id) } } diff --git a/rust/tests/src/multiplexed_stream_tests.rs b/rust/tests/src/multiplexed_stream_tests.rs index c8cc8538..82bd7bf2 100644 --- a/rust/tests/src/multiplexed_stream_tests.rs +++ b/rust/tests/src/multiplexed_stream_tests.rs @@ -986,6 +986,122 @@ mod failure_tests { Ok(()) } + #[tokio::test] + async fn test_concurrent_ingest_waiting_for_capacity_fails_after_mux_poison( + ) -> Result<(), Box> { + setup_tracing(); + info!("Starting test_concurrent_ingest_waiting_for_capacity_fails_after_mux_poison"); + + let (mock_server, server_url) = start_mock_server().await?; + mock_server + .inject_responses( + TABLE_FAIL, + vec![ + MockResponse::CreateStream { + stream_id: "poisoned".to_string(), + delay_ms: 0, + }, + MockResponse::Error { + status: tonic::Status::permission_denied("sub-stream failure"), + delay_ms: 0, + }, + ], + ) + .await; + + let sdk = create_test_sdk(&server_url).await?; + let opts = TestOpts { + max_inflight_requests: 1, + ..default_options() + }; + let stream = create_test_stream(&sdk, TABLE_FAIL, opts).await?; + let mux = Arc::new(MultiplexedStream::new(vec![stream])); + + const TASKS: usize = 16; + let barrier = Arc::new(tokio::sync::Barrier::new(TASKS)); + let mut handles = Vec::with_capacity(TASKS); + + for i in 0..TASKS { + let mux = Arc::clone(&mux); + let barrier = Arc::clone(&barrier); + handles.push(tokio::spawn(async move { + barrier.wait().await; + mux.ingest_record(format!("record-{i}").into_bytes()).await + })); + } + + let mut successes = Vec::new(); + let mut errors = 0; + for handle in handles { + match tokio::time::timeout(std::time::Duration::from_secs(2), handle) + .await + .expect("ingest task should finish after mux poison")? + { + Ok(message_id) => successes.push(message_id), + Err(ZerobusError::InvalidStateError(_)) + | Err(ZerobusError::StreamClosedError(_)) => errors += 1, + Err(e) => panic!("unexpected ingest error: {e:?}"), + } + } + + assert_eq!( + successes.len(), + 1, + "Only the first record should be admitted before poison" + ); + assert_eq!(successes[0].stream_index(), 0); + assert_eq!(successes[0].sub_offset(), 0); + assert_eq!(errors, TASKS - 1); + assert!(mux.is_closed(), "Mux should report the failed sub-stream"); + assert_eq!(mock_server.get_write_count().await, 1); + + Ok(()) + } + + #[tokio::test] + async fn test_ingest_times_out_when_capacity_never_recovers( + ) -> Result<(), Box> { + setup_tracing(); + + let (mock_server, server_url) = start_mock_server().await?; + mock_server + .inject_responses( + TABLE_OK, + vec![MockResponse::CreateStream { + stream_id: "stalled".to_string(), + delay_ms: 0, + }], + ) + .await; + + let sdk = create_test_sdk(&server_url).await?; + let stream = create_test_stream( + &sdk, + TABLE_OK, + TestOpts { + max_inflight_requests: 1, + flush_timeout_ms: None, + }, + ) + .await?; + let mux = MultiplexedStream::new(vec![stream]); + + mux.ingest_record(b"fills-capacity".to_vec()).await?; + let result = mux.ingest_record(b"times-out".to_vec()).await; + + assert!(matches!( + result, + Err(ZerobusError::ConnectionTimeout(message)) + if message.contains("sub-stream 0") + )); + assert!( + !mux.is_closed(), + "capacity timeout should not poison the mux" + ); + + Ok(()) + } + #[tokio::test] async fn test_public_is_closed_reports_closed_substream( ) -> Result<(), Box> {