[Rust] Multiplexed stream ingest followups

[danilonajkov-db]

What changes are proposed in this pull request?

This fixes the MultiplexedStream poison/admission race and makes capacity waits bounded.

The new ingest flow is:

1. Reserve real landing-zone capacity on the selected sub-stream.
2. Enter the sub-stream's private ingest-ordering section.
3. Acquire the mux admission read lock and re-check mux state.
4. Re-check sub-stream state.
5. Allocate the offset and enqueue using the reserved capacity.

The poison path acquires the mux admission write lock before flushing sub-streams, so no record can be admitted between the final closed check and the poison flush.

ZerobusStream retains ownership of its internal ordering lock through enqueue_reserved_admitted; the mux supplies only the admission check. Capacity reservations are opaque and cannot expose the underlying semaphore implementation.

Capacity waits are bounded by the existing flush_timeout_ms option. If acknowledgments never free capacity, ingest_record / ingest_records return ZerobusError::ConnectionTimeout instead of waiting indefinitely. A capacity timeout does not poison an otherwise healthy mux.

How is this tested?

• Regression test for concurrent ingests waiting for capacity while a sub-stream failure poisons the mux.
• Regression test for a permanently full sub-stream returning ConnectionTimeout after the configured timeout.
• Full multiplexed stream integration suite: 23 tests passed.
• Rust SDK unit suite: 141 tests passed.
• Doc tests passed.
• cargo fmt --all --check
• cargo clippy -p databricks-zerobus-ingest-sdk --all-targets -- -D warnings

[danilonajkov-db]

@elenagaljak-db thanks for the review, i've decided to move the ack callback to another PR to not mix contexts here. I've pasted your comments regarding the ack callback on that PR and resolved the ones here

[teodordelibasic-db]

Not sure if we discussed this before.

Since we hold this lock across join_all(...flush()), an ingest already parked at admission.read().await in enqueue_reserved_admitted cannot acquire the read lock until the whole flush finishes, so it stays blocked for that window before it sees is_closed and errors, and because it also holds sync_mutex, other ingests on that sub-stream queue behind it.

The write lock maybe only needs to be held long enough to drain readers that are mid-enqueue, not for the whole flush. Could we take and immediately drop it as a barrier after setting is_closed, then run the flush and signal_shutdown() unlocked? Any reader that acquires the read lock afterward fails check_closed() anyway, so nothing new is admitted.

[danilonajkov-db]

make sense to not block ingests if stream is mid close/failure and return error. Fixed

[teodordelibasic-db]

Do you think it would make sense to have a separate knob for this timeout other than the flush one?

[danilonajkov-db]

I added a seperate timeout variable, but didn't make it configureable for the clients. My thinking is that in the future we want the mux stream automatically try the next stream if the first one is full, so it will become an implementation detail the customer doesn't need to worry about. WDYT?

[danilonajkov-db]

ill wait for #523 to be merged so I can fix the tests

[teodordelibasic-db]

nit: What do you think about renaming to prepare_encoded_batch since that's the actual name of the returned structure and also RecordBatch is an Arrow concept?

[danilonajkov-db]

yep makes sense

[teodordelibasic-db]

Ah the comment above kind of gets debunked by this method. Still, maybe some alternative could be applied to both methods, not sure...

[danilonajkov-db]

Still I can rename to something like prepare_encoded_batch

[teodordelibasic-db]

The capacity deadline is always 30 seconds, so a stream configured with .flush_timeout_ms(100) still blocks for 30 seconds, while one using the default five-minute timeout starts failing much earlier. This also differs from the intended contract that the capacity wait uses the existing flush_timeout_ms option. Could we derive the deadline from the selected stream instead?

let capacity_wait_timeout =
    std::time::Duration::from_millis(stream.options.flush_timeout_ms);
let deadline = started_at + capacity_wait_timeout;

Could we also set flush_timeout_ms: Some(100) in test_ingest_times_out_when_capacity_never_recovers and assert the timeout after advancing past 100 ms? That would cover the configured contract rather than the fixed constant.

[danilonajkov-db]

flush() still behaves correctly.
If capacity is full because record A is unacknowledged:

1. Record B waits for capacity.
2. flush() snapshots record A.
3. If A remains unacknowledged past flush_timeout_ms, flush() fails.
4. If A is acknowledged, flush() may succeed and B may enqueue afterward.

That last case is intentional: records ingested concurrently with flush() are explicitly excluded. A capacity-waiting ingest has no offset yet, so it has not been admitted and cannot be part of the flush snapshot.

[teodordelibasic-db]

tokio::time::timeout cancels stream.reserve_capacity() on every 1 to 50 ms poll. Tokio's semaphore is FIFO and documents that cancelling acquire_owned loses the waiter's place, so concurrent ingests repeatedly move to the back of the queue and can reach the capacity timeout even while permits are being released. Could we keep one reservation future alive for the whole wait and use separate timer branches only for liveness checks and the final deadline?

let reservation = stream.reserve_capacity();
tokio::pin!(reservation);
let mut health_check = tokio::time::interval(std::time::Duration::from_millis(50));

loop {
    tokio::select! {
        result = &mut reservation => {
            return match result {
                Ok(reservation) => Ok(reservation),
                Err(error) => Err(self.handle_ingest_error(error, stream, idx).await),
            };
        }
        _ = tokio::time::sleep_until(deadline) => {
            return Err(ZerobusError::ConnectionTimeout(format!(
                "Timed out waiting for capacity on multiplexed sub-stream {idx}",
            )));
        }
        _ = health_check.tick() => {
            self.check_closed()?;
            // Keep the existing sub-stream liveness check here.
        }
    }
}

[danilonajkov-db]

I fixed it by pinning one reservation future for the entire capacity wait. The diagnostic timeout only borrows that future and does not recreate the semaphore acquisition, so its FIFO position is preserved. Stream and mux cancellation tokens now handle liveness without periodic polling.