[Rust] Multiplexed stream ingest followups

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<T: Clone> LandingZone<T> {
         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.

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<ZerobusStream>,
     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<crate::landing_zone::CapacityReservation> {
         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<MessageId> {
+        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