From 3aba69fbdad9b734df3d746fd45fce204fb52b5c Mon Sep 17 00:00:00 2001 From: elenagaljak-db Date: Thu, 20 Aug 2026 08:16:49 +0000 Subject: [PATCH 1/7] [Rust] Add persistent gRPC transport Signed-off-by: elenagaljak-db --- rust/NEXT_CHANGELOG.md | 3 + rust/sdk/Cargo.toml | 2 + rust/sdk/src/landing_zone.rs | 25 ++ rust/sdk/src/stream/grpc/connection.rs | 177 ++++++------ rust/sdk/src/stream/grpc/mod.rs | 88 +++++- rust/sdk/src/stream/grpc/receiver.rs | 197 ++++++------- rust/sdk/src/stream/grpc/sender.rs | 34 ++- rust/sdk/src/stream/grpc/supervisor.rs | 125 +++++++- rust/sdk/src/stream/grpc/transport.rs | 386 +++++++++++++++++++++++++ 9 files changed, 813 insertions(+), 224 deletions(-) create mode 100644 rust/sdk/src/stream/grpc/transport.rs diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index 6fc696f3..3e84f409 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -47,6 +47,9 @@ ### Internal Changes +- Added the feature-gated persistent gRPC transport, durable wire offsets, and + resume-watermark reconciliation for recovery after a lost acknowledgment. + ### Breaking Changes ### Deprecations diff --git a/rust/sdk/Cargo.toml b/rust/sdk/Cargo.toml index 1de5b671..fcd7c12a 100644 --- a/rust/sdk/Cargo.toml +++ b/rust/sdk/Cargo.toml @@ -77,6 +77,8 @@ internal-arrow-c-data = [ ] # Zero-copy protobuf parser. zeroparser = ["dep:self_cell", "dep:prost-build"] +# Persistent (Eos) streams; in development, API is unstable +eos = [] testing = ["dep:futures"] # Test-only deterministic seams (barriers/notifies) for the Arrow stream. Zero footprint # unless enabled; never enabled by FFI/production builds. diff --git a/rust/sdk/src/landing_zone.rs b/rust/sdk/src/landing_zone.rs index f5e61daa..8bd44aad 100644 --- a/rust/sdk/src/landing_zone.rs +++ b/rust/sdk/src/landing_zone.rs @@ -85,6 +85,31 @@ impl LandingZone { all_items } + /// Removes and returns items from the front of the unobserved queue while + /// `pred` holds, stopping at the first item that fails it. + /// + /// Used by the persistent-stream resume path to reconcile the retained tail + /// against the server's committed watermark: it removes the prefix of + /// records the server has already durably stored (offset ≤ resume watermark) + /// before re-sending the rest. Only the unobserved queue is consulted: + /// callers reset observation (`reset_observe`) first, so every retained item + /// lives in the queue in offset order. One semaphore permit is released per + /// removed item, mirroring `remove_observed`. + pub fn remove_front_while bool>(&self, mut pred: F) -> Vec { + let mut state = self.state.lock().expect("Lock poisoned"); + let mut permits = self.permits.lock().expect("Lock poisoned"); + let mut removed = Vec::new(); + while let Some(front) = state.queue.front() { + if !pred(front) { + break; + } + let item = state.queue.pop_front().expect("front just checked"); + permits.pop_front(); + removed.push(item); + } + removed + } + /// Adds an item to the queue. /// /// This method will block if the maximum number of inflight requests has been reached, diff --git a/rust/sdk/src/stream/grpc/connection.rs b/rust/sdk/src/stream/grpc/connection.rs index e1b1ed99..a7a34420 100644 --- a/rust/sdk/src/stream/grpc/connection.rs +++ b/rust/sdk/src/stream/grpc/connection.rs @@ -1,51 +1,55 @@ //! gRPC stream connection setup. //! -//! This module is transport-specific: it builds the bidirectional gRPC stream -//! used by `ZerobusStream` to talk to the Zerobus service. The Arrow Flight -//! transport has its own equivalent in `stream/arrow/connection.rs`. +//! This module is transport-specific: it opens the bidirectional gRPC stream +//! used by `ZerobusStream`. It handles both stream kinds through the transport +//! seam (`super::transport`): ephemeral streams over the `EphemeralStream` RPC +//! and, behind the `eos` feature, persistent streams over `PersistentStream`. +//! The Arrow Flight transport has its own equivalent under `stream/arrow/`. use std::sync::Arc; use prost::Message; use tokio::time::Duration; -use tokio_stream::wrappers::ReceiverStream; use tonic::metadata::MetadataValue; use tonic::transport::Channel; use tracing::{debug, error, info, instrument, warn}; -use super::ZerobusStream; -use crate::databricks::zerobus::ephemeral_stream_request::Payload as RequestPayload; -use crate::databricks::zerobus::ephemeral_stream_response::Payload as ResponsePayload; +use super::transport::{self, InboundStream, OutboundSink, TransportKind}; use crate::databricks::zerobus::zerobus_client::ZerobusClient; -use crate::databricks::zerobus::{ - CreateIngestStreamRequest, EphemeralStreamRequest, EphemeralStreamResponse, RecordType, -}; -use crate::{HeadersProvider, TableProperties, ZerobusError, ZerobusResult}; - -impl ZerobusStream { - /// Creates a stream connection to the Zerobus API. - /// Returns a tuple containing the sender, response gRPC stream, and stream ID. - /// If the stream creation fails, it returns an error. +use crate::databricks::zerobus::{CreateIngestStreamRequest, RecordType}; +use crate::{HeadersProvider, OffsetId, TableProperties, ZerobusError, ZerobusResult}; + +/// A freshly opened stream connection: the outbound sink, the inbound response +/// stream, the server-assigned `stream_id`, and (persistent resume only) the +/// committed-offset watermark to resume from. +pub(super) struct StreamConnection { + pub(super) sink: OutboundSink, + pub(super) inbound: InboundStream, + pub(super) stream_id: String, + pub(super) last_committed_offset: Option, +} + +impl super::ZerobusStream { + /// Opens a stream connection to the Zerobus API for the given transport + /// kind. Returns the sink, inbound stream, stream id, and resume watermark. /// - /// On a server-side authentication rejection it asks the headers provider to - /// invalidate cached credentials so the next attempt re-derives them. This - /// covers IdP-revoked tokens, not a same-named table recreated within the - /// token's lifetime, which the server accepts. + /// On a server-side authentication rejection it asks the headers provider + /// to invalidate cached credentials so the next attempt re-derives them. + /// This covers IdP-revoked tokens, not a same-named table recreated within + /// the token's lifetime, which the server accepts. pub(super) async fn create_stream_connection( channel: ZerobusClient, table_properties: &TableProperties, headers_provider: &Arc, record_type: RecordType, - ) -> ZerobusResult<( - tokio::sync::mpsc::Sender, - tonic::Streaming, - String, - )> { + kind: &TransportKind, + ) -> ZerobusResult { let result = Self::create_stream_connection_inner( channel, table_properties, headers_provider, record_type, + kind, ) .await; if let Err(err) = &result { @@ -69,12 +73,9 @@ impl ZerobusStream { table_properties: &TableProperties, headers_provider: &Arc, record_type: RecordType, + kind: &TransportKind, recovery_timeout_ms: u64, - ) -> ZerobusResult<( - tokio::sync::mpsc::Sender, - tonic::Streaming, - String, - )> { + ) -> ZerobusResult { let attempt_timeout = Duration::from_millis(recovery_timeout_ms); let attempt_started = tokio::time::Instant::now(); let result = tokio::time::timeout( @@ -84,6 +85,7 @@ impl ZerobusStream { table_properties, headers_provider, record_type, + kind, ), ) .await @@ -115,18 +117,13 @@ impl ZerobusStream { table_properties: &TableProperties, headers_provider: &Arc, record_type: RecordType, - ) -> ZerobusResult<( - tokio::sync::mpsc::Sender, - tonic::Streaming, - String, - )> { - const CHANNEL_BUFFER_SIZE: usize = 2048; - let (tx, rx) = tokio::sync::mpsc::channel(CHANNEL_BUFFER_SIZE); - let mut request_stream = tonic::Request::new(ReceiverStream::new(rx)); - - let stream_metadata = request_stream.metadata_mut(); - let headers = headers_provider.get_headers().await?; + kind: &TransportKind, + ) -> ZerobusResult { + let (sink, request_body) = transport::make_outbound(kind); + let mut request = tonic::Request::new(request_body); + let stream_metadata = request.metadata_mut(); + let headers = headers_provider.get_headers().await?; for (key, value) in headers { match key { "x-databricks-zerobus-table-name" => { @@ -152,12 +149,45 @@ impl ZerobusStream { } } - let mut response_grpc_stream = channel - .ephemeral_stream(request_stream) - .await - .map_err(ZerobusError::CreateStreamError)? - .into_inner(); + let mut inbound = transport::open_rpc(&mut channel, request).await?; + + let create_request = Self::build_create_request(table_properties, record_type)?; + + debug!("Sending stream-open request."); + sink.send_open(kind, create_request).await.map_err(|_| { + error!(table_name = %table_properties.table_name, "Failed to send stream-open request"); + ZerobusError::StreamClosedError(tonic::Status::internal( + "Failed to send stream-open request", + )) + })?; + + debug!("Waiting for stream-open response."); + let opened = inbound.recv_open().await?; + + // On a persistent resume the server does not re-send the stream_id (the + // client supplied it), so fall back to the id being resumed. + let stream_id = if opened.stream_id.is_empty() { + Self::resume_stream_id(kind).unwrap_or_default() + } else { + opened.stream_id + }; + info!(stream_id = %stream_id, last_committed_offset = ?opened.last_committed_offset, "Successfully opened stream"); + Ok(StreamConnection { + sink, + inbound, + stream_id, + last_committed_offset: opened.last_committed_offset, + }) + } + + /// Builds the `CreateIngestStreamRequest` sent when opening (or, on the + /// persistent path, creating) a stream. Encodes the descriptor for proto + /// streams and validates its presence. + fn build_create_request( + table_properties: &TableProperties, + record_type: RecordType, + ) -> ZerobusResult { let descriptor_proto = if record_type == RecordType::Proto { Some( table_properties @@ -174,56 +204,19 @@ impl ZerobusStream { None }; - let create_stream_request = RequestPayload::CreateStream(CreateIngestStreamRequest { + Ok(CreateIngestStreamRequest { table_name: Some(table_properties.table_name.to_string()), descriptor_proto, record_type: Some(record_type.into()), - }); - - debug!("Sending CreateStream request."); - tx.send(EphemeralStreamRequest { - payload: Some(create_stream_request), }) - .await - .map_err(|_| { - error!(table_name = %table_properties.table_name, "Failed to send CreateStream request"); - ZerobusError::StreamClosedError(tonic::Status::internal( - "Failed to send CreateStream request", - )) - })?; - debug!("Waiting for CreateStream response."); - let create_stream_response = response_grpc_stream.message().await; - - match create_stream_response { - Ok(Some(create_stream_response)) => match create_stream_response.payload { - Some(ResponsePayload::CreateStreamResponse(resp)) => { - if let Some(stream_id) = resp.stream_id { - info!(stream_id = %stream_id, "Successfully created stream"); - Ok((tx, response_grpc_stream, stream_id)) - } else { - error!("Successfully created a stream but stream_id is None"); - Err(ZerobusError::CreateStreamError(tonic::Status::internal( - "Successfully created a stream but stream_id is None", - ))) - } - } - unexpected_message => { - error!("Unexpected response from server {unexpected_message:?}"); - Err(ZerobusError::CreateStreamError(tonic::Status::internal( - "Unexpected response from server", - ))) - } - }, - Ok(None) => { - info!("Server closed the stream gracefully before sending CreateStream response"); - Err(ZerobusError::CreateStreamError(tonic::Status::ok( - "Stream closed gracefully by server", - ))) - } - Err(status) => { - error!("CreateStream RPC failed: {status:?}"); - Err(ZerobusError::CreateStreamError(status)) - } + } + + /// The stream id being resumed, if this is a persistent resume. + fn resume_stream_id(kind: &TransportKind) -> Option { + match kind { + TransportKind::Ephemeral => None, + #[cfg(feature = "eos")] + TransportKind::Persistent { resume_stream_id } => resume_stream_id.clone(), } } } diff --git a/rust/sdk/src/stream/grpc/mod.rs b/rust/sdk/src/stream/grpc/mod.rs index 8caac1bb..e1e41205 100644 --- a/rust/sdk/src/stream/grpc/mod.rs +++ b/rust/sdk/src/stream/grpc/mod.rs @@ -15,6 +15,7 @@ //! | `close.rs` | `close`, `is_closed`, task shutdown | Transport-agnostic | //! | `callback_handler.rs` | User-callback dispatch task | Transport-agnostic | //! | `connection.rs` | gRPC bidirectional stream setup | gRPC-specific | +//! | `transport.rs` | Ephemeral/persistent RPC seam (`eos`) | gRPC-specific | //! | `sender.rs` | Outbound gRPC sender task | gRPC-specific | //! | `receiver.rs` | Inbound gRPC receiver task | gRPC-specific | //! | `supervisor.rs` | Create → spawn → recover loop | gRPC-specific | @@ -43,8 +44,10 @@ mod ingest; mod receiver; mod sender; mod supervisor; +mod transport; mod types; +use transport::TransportKind; use types::{IngestRequest, OneshotMap, RecordLandingZone}; #[cfg(feature = "testing")] @@ -93,6 +96,13 @@ pub struct ZerobusStream { pub(crate) stream_id: Option, /// Type of gRPC stream that is used when sending records. pub stream_type: StreamType, + /// For a persistent (Eos) stream resumed from a prior session, the offset + /// the server had durably committed at resume time (`last_committed_offset` + /// from the resume response). `None` for ephemeral streams and for a + /// freshly created persistent stream (nothing committed yet). Only read + /// through the `eos`-gated accessor. + #[cfg_attr(not(feature = "eos"), allow(dead_code))] + pub(crate) last_committed_offset: Option, /// Gets headers which are used in the first request to establish connection with the server. pub headers_provider: Arc, /// The stream configuration options related to recovery, fetching OAuth tokens, etc. @@ -141,9 +151,59 @@ impl ZerobusStream { table_properties: TableProperties, headers_provider: Arc, options: StreamConfigurationOptions, + ) -> ZerobusResult { + Self::new_with_kind( + channel, + table_properties, + headers_provider, + options, + StreamType::Ephemeral, + TransportKind::Ephemeral, + ) + .await + } + + /// Creates or resumes a persistent (Eos) stream. + /// + /// With `resume_stream_id = None` the server mints a new stream and its id + /// is available via [`stream_id`](Self::stream_id). With `Some(id)` the SDK + /// reconnects to an existing persistent stream, reseeds its offset generator + /// to continue after the server's committed offset, and re-sends only the + /// records the server has not yet durably stored. + #[cfg(feature = "eos")] + #[instrument(level = "debug", skip_all)] + pub(crate) async fn new_persistent_stream( + channel: ZerobusClient, + table_properties: TableProperties, + headers_provider: Arc, + options: StreamConfigurationOptions, + resume_stream_id: Option, + ) -> ZerobusResult { + Self::new_with_kind( + channel, + table_properties, + headers_provider, + options, + StreamType::Persistent, + TransportKind::Persistent { resume_stream_id }, + ) + .await + } + + /// Shared constructor for both stream kinds. Wires the supervisor + IO + /// tasks, waits for the first open to complete, and — for a persistent + /// resume — reseeds the logical offset generator so records ingested after + /// resume continue past the server's committed offset. + async fn new_with_kind( + channel: ZerobusClient, + table_properties: TableProperties, + headers_provider: Arc, + options: StreamConfigurationOptions, + stream_type: StreamType, + kind: TransportKind, ) -> ZerobusResult { let (stream_init_result_tx, stream_init_result_rx) = - tokio::sync::oneshot::channel::>(); + tokio::sync::oneshot::channel::>(); let (logical_last_received_offset_id_tx, _logical_last_received_offset_id_rx) = tokio::sync::watch::channel(None); @@ -176,6 +236,7 @@ impl ZerobusStream { table_properties.clone(), Arc::clone(&headers_provider), options.clone(), + kind, Arc::clone(&landing_zone), Arc::clone(&oneshot_map), logical_last_received_offset_id_tx.clone(), @@ -186,20 +247,37 @@ impl ZerobusStream { cancellation_token.clone(), callback_tx.clone(), )); - let stream_id = Some(stream_init_result_rx.await.map_err(|_| { + let init_info = stream_init_result_rx.await.map_err(|_| { ZerobusError::UnexpectedStreamResponseError( "Supervisor task died before stream creation".to_string(), ) - })??); + })??; + + // On a persistent resume, continue offset generation past the server's + // committed offset so newly ingested records get the next durable + // offsets rather than restarting at 0. Also seed the last-received-offset + // watch to the watermark: everything up to it is already durable, so a + // `flush()` / `wait_for_offset()` targeting an already-committed offset + // must resolve immediately. Without this the watch stays `None` and, on a + // fresh-process resume with no new ingests, `flush()` (and the `close()` + // that calls it) would block until the flush timeout and then error, + // since the server never re-acks offsets it committed in a prior session. + // Safe because the constructor returns before any user ingest, so no real + // ack can race this initial value. + if let Some(watermark) = init_info.last_committed_offset { + logical_offset_id_generator.set_next(watermark + 1); + let _ = logical_last_received_offset_id_tx.send(Some(watermark)); + } // Cloned out before `table_properties` is moved into the struct below. let dynamic_message_descriptor = table_properties.message_descriptor.clone(); let stream = Self { - stream_type: StreamType::Ephemeral, + stream_type, headers_provider, options: options.clone(), table_properties, - stream_id, + stream_id: Some(init_info.stream_id), + last_committed_offset: init_info.last_committed_offset, landing_zone, oneshot_map, supervisor_task, diff --git a/rust/sdk/src/stream/grpc/receiver.rs b/rust/sdk/src/stream/grpc/receiver.rs index 206f99cc..3c5be9a3 100644 --- a/rust/sdk/src/stream/grpc/receiver.rs +++ b/rust/sdk/src/stream/grpc/receiver.rs @@ -1,9 +1,9 @@ //! Inbound gRPC receiver task. //! -//! Transport-specific: reads `EphemeralStreamResponse` messages from the -//! gRPC inbound stream, dispatches durability acks to oneshot senders / -//! callbacks, and signals the supervisor via `server_error_tx` / pause -//! deadlines. +//! Transport-specific: reads server messages from the `InboundStream` (which +//! hides whether the RPC is `EphemeralStream` or `PersistentStream`), dispatches +//! durability acks to oneshot senders / callbacks, and signals the supervisor +//! via `server_error_tx` / pause deadlines. use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; @@ -12,22 +12,26 @@ use tokio::time::Duration; use tokio_util::sync::CancellationToken; use tracing::{error, info, instrument, span, Level}; +use super::transport::{InboundMessage, InboundStream}; use super::types::{CallbackMessage, OneshotMap, RecordLandingZone}; use super::{ZerobusStream, STREAM_TEARDOWN_DRAIN_TIMEOUT_MS}; -use crate::databricks::zerobus::ephemeral_stream_response::Payload as ResponsePayload; -use crate::databricks::zerobus::{ - CloseStreamSignal, EphemeralStreamResponse, IngestRecordResponse, -}; +use crate::databricks::zerobus::{CloseStreamSignal, IngestRecordResponse}; use crate::{OffsetId, StreamConfigurationOptions, ZerobusError, ZerobusResult}; impl ZerobusStream { - /// Spawns a task that continuously reads from `response_grpc_stream` - /// and propagates the received durability acknowledgements to the - /// corresponding pending acks promises. + /// Spawns a task that continuously reads server messages and propagates the + /// received durability acknowledgements to the corresponding pending acks + /// promises. + /// + /// `initial_last_acked_offset` is the wire offset already acknowledged when + /// this session starts: `-1` for a fresh stream (ephemeral, or a new + /// persistent create), or the resume watermark for a persistent resume so + /// the ack-counting below aligns with the durable wire offsets the server + /// reports. #[instrument(level = "debug", skip_all)] #[allow(clippy::too_many_arguments)] pub(super) fn spawn_receiver_task( - mut response_grpc_stream: tonic::Streaming, + mut inbound: InboundStream, last_received_offset_id_tx: tokio::sync::watch::Sender>, landing_zone: RecordLandingZone, oneshot_map: Arc>, @@ -36,11 +40,12 @@ impl ZerobusStream { server_error_tx: tokio::sync::watch::Sender>, recv_drain_token: CancellationToken, callback_tx: Option>, + initial_last_acked_offset: OffsetId, ) -> tokio::task::JoinHandle> { tokio::spawn(async move { let span = span!(Level::DEBUG, "inbound_stream_processor"); let _guard = span.enter(); - let mut last_acked_offset = -1; + let mut last_acked_offset = initial_last_acked_offset; let mut pause_deadline: Option = None; // Set when we exit because the supervisor signalled close (`recv_drain_token`). // On that path we drain the response stream inline so the server sees END_STREAM @@ -74,7 +79,7 @@ impl ZerobusStream { } res = tokio::time::timeout( Duration::from_millis(options.server_lack_of_ack_timeout_ms), - response_grpc_stream.message(), + inbound.message(), ) => res, } } else { @@ -86,98 +91,94 @@ impl ZerobusStream { } res = tokio::time::timeout( Duration::from_millis(options.server_lack_of_ack_timeout_ms), - response_grpc_stream.message(), + inbound.message(), ) => res, } }; match message_result { - Ok(Ok(Some(ingest_record_response))) => match ingest_record_response.payload { - Some(ResponsePayload::IngestRecordResponse(IngestRecordResponse { - durability_ack_up_to_offset, - })) => { - let durability_ack_up_to_offset = match durability_ack_up_to_offset { - Some(offset) => offset, - None => { - error!("Missing ack offset in server response"); - let error = - ZerobusError::StreamClosedError(tonic::Status::internal( - "Missing ack offset in server response", - )); - let _ = server_error_tx.send(Some(error.clone())); - return Err(error); - } - }; - let mut last_logical_acked_offset = -2; - let mut map = oneshot_map.lock().await; - for _offset_to_ack in - (last_acked_offset + 1)..=durability_ack_up_to_offset - { - if let Ok(record) = landing_zone.remove_observed() { - let logical_offset = record.offset_id; - last_logical_acked_offset = logical_offset; + Ok(Ok(Some(InboundMessage::Ack(IngestRecordResponse { + durability_ack_up_to_offset, + })))) => { + let durability_ack_up_to_offset = match durability_ack_up_to_offset { + Some(offset) => offset, + None => { + error!("Missing ack offset in server response"); + let error = + ZerobusError::StreamClosedError(tonic::Status::internal( + "Missing ack offset in server response", + )); + let _ = server_error_tx.send(Some(error.clone())); + return Err(error); + } + }; + let mut last_logical_acked_offset = -2; + let mut map = oneshot_map.lock().await; + for _offset_to_ack in (last_acked_offset + 1)..=durability_ack_up_to_offset + { + if let Ok(record) = landing_zone.remove_observed() { + let logical_offset = record.offset_id; + last_logical_acked_offset = logical_offset; - if let Some(sender) = map.remove(&logical_offset) { - let _ = sender.send(Ok(logical_offset)); - } + if let Some(sender) = map.remove(&logical_offset) { + let _ = sender.send(Ok(logical_offset)); + } - if let Some(ref tx) = callback_tx { - let _ = tx.send(CallbackMessage::Ack(logical_offset)); - } + if let Some(ref tx) = callback_tx { + let _ = tx.send(CallbackMessage::Ack(logical_offset)); } } - drop(map); - last_acked_offset = durability_ack_up_to_offset; - if last_logical_acked_offset != -2 { - let _ignore_on_channel_break = last_received_offset_id_tx - .send(Some(last_logical_acked_offset)); - } } - Some(ResponsePayload::CloseStreamSignal(CloseStreamSignal { - duration, - })) => { - if options.recovery { - let server_duration_ms = duration - .as_ref() - .map(|d| d.seconds as u64 * 1000 + d.nanos as u64 / 1_000_000) - .unwrap_or(0); - - let wait_duration_ms = match options.stream_paused_max_wait_time_ms - { - None => server_duration_ms, - Some(0) => { - // Immediate recovery - info!("Server will close the stream in {}ms. Triggering stream recovery.", server_duration_ms); - break 'recv_loop; - } - Some(max_wait) => std::cmp::min(max_wait, server_duration_ms), - }; + drop(map); + last_acked_offset = durability_ack_up_to_offset; + if last_logical_acked_offset != -2 { + let _ignore_on_channel_break = + last_received_offset_id_tx.send(Some(last_logical_acked_offset)); + } + } + Ok(Ok(Some(InboundMessage::Close(CloseStreamSignal { duration })))) => { + if options.recovery { + let server_duration_ms = duration + .as_ref() + .map(|d| d.seconds as u64 * 1000 + d.nanos as u64 / 1_000_000) + .unwrap_or(0); - if wait_duration_ms == 0 { - info!("Server will close the stream. Triggering immediate recovery."); + let wait_duration_ms = match options.stream_paused_max_wait_time_ms { + None => server_duration_ms, + Some(0) => { + // Immediate recovery + info!("Server will close the stream in {}ms. Triggering stream recovery.", server_duration_ms); break 'recv_loop; } + Some(max_wait) => std::cmp::min(max_wait, server_duration_ms), + }; - is_paused.store(true, Ordering::Relaxed); - pause_deadline = Some( - tokio::time::Instant::now() - + Duration::from_millis(wait_duration_ms), - ); + if wait_duration_ms == 0 { info!( - "Server will close the stream in {}ms. Entering graceful close period (waiting up to {}ms for in-flight acks).", - server_duration_ms, wait_duration_ms + "Server will close the stream. Triggering immediate recovery." ); + break 'recv_loop; } + + is_paused.store(true, Ordering::Relaxed); + pause_deadline = Some( + tokio::time::Instant::now() + + Duration::from_millis(wait_duration_ms), + ); + info!( + "Server will close the stream in {}ms. Entering graceful close period (waiting up to {}ms for in-flight acks).", + server_duration_ms, wait_duration_ms + ); } - unexpected_message => { - error!("Unexpected response from server {unexpected_message:?}"); - let error = ZerobusError::StreamClosedError(tonic::Status::internal( - "Unexpected response from server", - )); - let _ = server_error_tx.send(Some(error.clone())); - return Err(error); - } - }, + } + Ok(Ok(Some(InboundMessage::Other))) => { + error!("Unexpected response from server"); + let error = ZerobusError::StreamClosedError(tonic::Status::internal( + "Unexpected response from server", + )); + let _ = server_error_tx.send(Some(error.clone())); + return Err(error); + } Ok(Ok(None)) => { info!("Server closed the stream without errors."); let error = ZerobusError::StreamClosedError(tonic::Status::ok( @@ -215,30 +216,14 @@ impl ZerobusStream { if close_initiated { let _ = tokio::time::timeout( Duration::from_millis(STREAM_TEARDOWN_DRAIN_TIMEOUT_MS), - async { - while response_grpc_stream - .message() - .await - .ok() - .flatten() - .is_some() - {} - }, + inbound.drain(), ) .await; } else { tokio::spawn(async move { let _ = tokio::time::timeout( Duration::from_millis(STREAM_TEARDOWN_DRAIN_TIMEOUT_MS), - async move { - while response_grpc_stream - .message() - .await - .ok() - .flatten() - .is_some() - {} - }, + inbound.drain(), ) .await; }); diff --git a/rust/sdk/src/stream/grpc/sender.rs b/rust/sdk/src/stream/grpc/sender.rs index 540d1491..246da3ec 100644 --- a/rust/sdk/src/stream/grpc/sender.rs +++ b/rust/sdk/src/stream/grpc/sender.rs @@ -1,7 +1,8 @@ //! Outbound gRPC sender task. //! //! Transport-specific: reads from the landing zone (transport-agnostic) and -//! writes `EphemeralStreamRequest` messages over the gRPC outbound channel. +//! writes ingest messages through the `OutboundSink`, which hides whether the +//! underlying RPC is `EphemeralStream` or `PersistentStream`. use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; @@ -9,21 +10,28 @@ use std::sync::Arc; use tokio_util::sync::CancellationToken; use tracing::error; +use super::transport::OutboundSink; use super::types::RecordLandingZone; use super::ZerobusStream; -use crate::databricks::zerobus::EphemeralStreamRequest; use crate::offset_generator::OffsetIdGenerator; use crate::{ZerobusError, ZerobusResult}; impl ZerobusStream { /// Spawns a task that continuously sends records to the Zerobus API by observing the landing zone /// to get records and sending them through the outbound stream to the gRPC stream. + /// + /// `durable_wire_offset` selects the offset policy. Ephemeral streams number + /// records with a fresh 0-based physical counter each session (the server + /// tracks nothing across reconnects). Persistent streams put the record's + /// durable logical offset on the wire so the server can dedup and resume by + /// it — that offset already lives on the landing-zone item. pub(super) fn spawn_sender_task( - outbound_stream: tokio::sync::mpsc::Sender, + sink: OutboundSink, landing_zone: RecordLandingZone, is_paused: Arc, server_error_tx: tokio::sync::watch::Sender>, cancellation_token: CancellationToken, + durable_wire_offset: bool, ) -> tokio::task::JoinHandle> { tokio::spawn(async move { let physical_offset_id_generator = OffsetIdGenerator::default(); @@ -39,22 +47,18 @@ impl ZerobusStream { } } => item.clone(), }; - let offset_id = physical_offset_id_generator.next(); - let request_payload = item.payload.into_request_payload(offset_id); + let wire_offset = if durable_wire_offset { + item.offset_id + } else { + physical_offset_id_generator.next() + }; - let send_result = outbound_stream - .send(EphemeralStreamRequest { - payload: Some(request_payload), - }) - .await; + let send_result = sink.send_ingest(item.payload, wire_offset).await; if let Err(err) = send_result { error!("Failed to send record: {}", err); - let error = ZerobusError::StreamClosedError(tonic::Status::internal( - "Failed to send record", - )); - let _ = server_error_tx.send(Some(error.clone())); - return Err(error); + let _ = server_error_tx.send(Some(err.clone())); + return Err(err); } } }) diff --git a/rust/sdk/src/stream/grpc/supervisor.rs b/rust/sdk/src/stream/grpc/supervisor.rs index 74b03cad..e11bf43c 100644 --- a/rust/sdk/src/stream/grpc/supervisor.rs +++ b/rust/sdk/src/stream/grpc/supervisor.rs @@ -18,6 +18,7 @@ use tokio_util::sync::CancellationToken; use tonic::transport::Channel; use tracing::{debug, error, info, instrument, warn}; +use super::transport::TransportKind; use super::types::{CallbackMessage, OneshotMap, RecordLandingZone}; use super::{ZerobusStream, STREAM_TEARDOWN_DRAIN_TIMEOUT_MS}; use crate::databricks::zerobus::zerobus_client::ZerobusClient; @@ -27,9 +28,23 @@ use crate::{ ZerobusError, ZerobusResult, }; +/// What the supervisor reports back to the constructor once the stream is +/// first opened: the server-assigned identity and, for a persistent resume, the +/// committed-offset watermark the caller uses to reseed its offset generator. +pub(super) struct StreamInitInfo { + pub(super) stream_id: String, + pub(super) last_committed_offset: Option, +} + impl ZerobusStream { /// Supervisor task is responsible for managing the stream lifecycle. /// It handles stream creation, recovery, and error handling. + /// + /// `kind` selects the transport (ephemeral vs persistent). For a persistent + /// stream the supervisor opens with a create (or a resume, if the caller + /// supplied a stream id) on the first iteration, then resumes by the minted + /// id on every subsequent recovery — an ephemeral stream re-creates each + /// time, as before. #[allow(clippy::too_many_arguments)] #[instrument(level = "debug", skip_all, fields(table_name = %table_properties.table_name))] pub(super) async fn supervisor_task( @@ -37,16 +52,20 @@ impl ZerobusStream { table_properties: TableProperties, headers_provider: Arc, options: StreamConfigurationOptions, + // Mutated only on the persistent recovery path (flip create → resume); + // that mutation is `eos`-gated, so `mut` is otherwise unused. + #[cfg_attr(not(feature = "eos"), allow(unused_mut))] mut kind: TransportKind, landing_zone: RecordLandingZone, oneshot_map: Arc>, logical_last_received_offset_id_tx: tokio::sync::watch::Sender>, is_closed: Arc, failed_records: Arc>>, - stream_init_result_tx: tokio::sync::oneshot::Sender>, + stream_init_result_tx: tokio::sync::oneshot::Sender>, server_error_tx: tokio::sync::watch::Sender>, cancellation_token: CancellationToken, callback_tx: Option>, ) -> ZerobusResult<()> { + let durable_wire_offset = kind.uses_durable_wire_offset(); let mut initial_stream_creation = true; let mut stream_init_result_tx = Some(stream_init_result_tx); // One-shot budget: initial setup may spend a single recovery retry to refresh a @@ -92,6 +111,7 @@ impl ZerobusStream { let headers_provider = Arc::clone(&headers_provider); let record_type = options.record_type; let attempt = &attempt; + let kind = &kind; async move { let attempt_no = attempt.fetch_add(1, Ordering::Relaxed) + 1; @@ -101,6 +121,7 @@ impl ZerobusStream { &table_properties, &headers_provider, record_type, + kind, options.recovery_timeout_ms, ) .await @@ -112,6 +133,7 @@ impl ZerobusStream { &table_properties, &headers_provider, record_type, + kind, ), ) .await @@ -148,8 +170,8 @@ impl ZerobusStream { }; let creation = RetryIf::spawn(strategy, create_attempt, should_retry).await; - let (tx, response_grpc_stream, stream_id) = match creation { - Ok((tx, response_grpc_stream, stream_id)) => (tx, response_grpc_stream, stream_id), + let connection = match creation { + Ok(connection) => connection, Err(e) => { if initial_stream_creation { if let Some(tx) = stream_init_result_tx.take() { @@ -176,11 +198,22 @@ impl ZerobusStream { return Err(e); } }; + let stream_id = connection.stream_id; + let last_committed_offset = connection.last_committed_offset; if initial_stream_creation { if let Some(stream_init_result_tx_inner) = stream_init_result_tx.take() { - let _ = stream_init_result_tx_inner.send(Ok(stream_id.clone())); + let _ = stream_init_result_tx_inner.send(Ok(StreamInitInfo { + stream_id: stream_id.clone(), + last_committed_offset, + })); } initial_stream_creation = false; + // A persistent stream recovers by resuming its now-known id, not + // by creating a fresh stream on every reconnect. + #[cfg(feature = "eos")] + if let TransportKind::Persistent { resume_stream_id } = &mut kind { + *resume_stream_id = Some(stream_id.clone()); + } info!(stream_id = %stream_id, "Successfully created stream"); } else { info!(stream_id = %stream_id, "Successfully recovered stream"); @@ -190,6 +223,31 @@ impl ZerobusStream { // 2. Reset landing zone. let resent_records = landing_zone_recovery.observed_count(); let resent_batches = landing_zone_recovery.reset_observe(); + + // Resume alignment (persistent streams). The server reports how far + // it has durably committed. A dropped connection can lose the ack + // for records the server did commit, so the retained tail can start + // at or below that watermark. Re-sending an already-committed offset + // would be rejected by the server and kill the stream (Decision 6: + // server rejects), so reconcile first: resolve those records' + // pending acks locally and re-send only the offsets above the + // watermark. This is not user-duplicate dedup — it realigns the + // SDK's retained tail with the server's durable position on resume. + let mut initial_last_acked_offset: OffsetId = -1; + if durable_wire_offset { + if let Some(watermark) = last_committed_offset { + initial_last_acked_offset = watermark; + Self::reconcile_committed_on_resume( + &landing_zone_recovery, + &oneshot_map, + &logical_last_received_offset_id_tx, + &callback_tx, + watermark, + ) + .await; + } + } + if resent_batches > 0 { info!( stream_id = %stream_id, @@ -210,7 +268,7 @@ impl ZerobusStream { let recv_drain_token = CancellationToken::new(); let mut recv_task = Self::spawn_receiver_task( - response_grpc_stream, + connection.inbound, logical_last_received_offset_id_tx.clone(), landing_zone_receiver, oneshot_map.clone(), @@ -219,13 +277,15 @@ impl ZerobusStream { server_error_tx.clone(), recv_drain_token.clone(), callback_tx.clone(), + initial_last_acked_offset, ); let mut send_task = Self::spawn_sender_task( - tx, + connection.sink, landing_zone_sender, Arc::clone(&is_paused), server_error_tx.clone(), per_stream_token.clone(), + durable_wire_offset, ); // 4. Wait for any of the two tasks to end. @@ -307,6 +367,59 @@ impl ZerobusStream { } } + /// Reconciles the retained landing-zone tail with the server's durable + /// position when resuming a persistent stream: removes the prefix of records + /// the server has already committed (logical offset ≤ `watermark`) and + /// resolves their pending acks locally, so they are never re-sent. + /// + /// Re-sending an already-committed offset would be rejected by the server + /// and kill the stream (Decision 6: server rejects), so this realignment is + /// what makes resume survive a lost-ack race. It is not user-duplicate + /// dedup — the records removed here are ones the SDK itself queued and the + /// server durably stored; only their acknowledgement was lost with the + /// dropped connection. + /// + /// The caller must have reset observation first, so every retained record + /// sits in the landing-zone queue in ascending offset order — the committed + /// ones therefore form a contiguous front prefix. Each removed offset is + /// acked to its waiter, reported to the ack callback, and advances the + /// last-received-offset watermark, matching what a real server ack would do. + async fn reconcile_committed_on_resume( + landing_zone: &RecordLandingZone, + oneshot_map: &Arc>, + logical_last_received_offset_id_tx: &tokio::sync::watch::Sender>, + callback_tx: &Option>, + watermark: OffsetId, + ) { + let committed = landing_zone.remove_front_while(|item| item.offset_id <= watermark); + if committed.is_empty() { + return; + } + let reconciled_records: usize = + committed.iter().map(|r| r.payload.get_record_count()).sum(); + let mut map = oneshot_map.lock().await; + let mut highest = None; + for item in &committed { + if let Some(sender) = map.remove(&item.offset_id) { + let _ = sender.send(Ok(item.offset_id)); + } + if let Some(tx) = callback_tx { + let _ = tx.send(CallbackMessage::Ack(item.offset_id)); + } + highest = Some(item.offset_id); + } + drop(map); + if let Some(highest) = highest { + let _ = logical_last_received_offset_id_tx.send(Some(highest)); + } + info!( + watermark, + reconciled_batches = committed.len(), + reconciled_records, + "Persistent resume: reconciled already-committed records at/below the watermark" + ); + } + /// Fails all pending records by removing them from the landing zone and sending error to all pending acks promises. pub(super) async fn fail_all_pending_records( landing_zone: RecordLandingZone, diff --git a/rust/sdk/src/stream/grpc/transport.rs b/rust/sdk/src/stream/grpc/transport.rs new file mode 100644 index 00000000..3f16dd89 --- /dev/null +++ b/rust/sdk/src/stream/grpc/transport.rs @@ -0,0 +1,386 @@ +//! Transport seam between the ephemeral and persistent gRPC stream kinds. +//! +//! The two stream kinds ride different RPCs (`EphemeralStream` vs +//! `PersistentStream`) with different request/response envelopes, and differ +//! in exactly three ways: +//! +//! 1. **Opening** — ephemeral always sends `create_stream`; persistent sends +//! `create_stream` (new) or `resume_stream` (reconnect), and a resume comes +//! back with a committed-offset watermark. +//! 2. **Offset on the wire** — ephemeral numbers records with a fresh 0-based +//! counter each session (the server tracks nothing across reconnects); +//! persistent puts the record's durable logical offset on the wire so the +//! server can dedup and resume by it. +//! 3. **Response parsing** — the oneof variants live in different generated +//! enums. +//! +//! Everything else — the landing zone, backpressure, ack tracking, flush, +//! close, callbacks, and the create → spawn → recover supervisor loop — is +//! identical and stays transport-agnostic. This module normalizes the three +//! differences behind two small enums (`OutboundSink`, `InboundStream`) plus a +//! neutral `InboundMessage`, so the sender/receiver tasks never mention a +//! concrete RPC. The enums carry no `dyn` and compile away to the ephemeral +//! arm when the `eos` feature is off. + +use tokio_stream::wrappers::ReceiverStream; +use tonic::transport::Channel; + +use crate::databricks::zerobus::ephemeral_stream_request::Payload as EphemeralRequestPayload; +use crate::databricks::zerobus::ephemeral_stream_response::Payload as EphemeralResponsePayload; +use crate::databricks::zerobus::zerobus_client::ZerobusClient; +use crate::databricks::zerobus::{ + CloseStreamSignal, CreateIngestStreamRequest, EphemeralStreamRequest, EphemeralStreamResponse, + IngestRecordResponse, +}; +use crate::{EncodedBatch, OffsetId, ZerobusError, ZerobusResult}; + +#[cfg(feature = "eos")] +use crate::databricks::zerobus::persistent_stream_request::Payload as PersistentRequestPayload; +#[cfg(feature = "eos")] +use crate::databricks::zerobus::persistent_stream_response::Payload as PersistentResponsePayload; +#[cfg(feature = "eos")] +use crate::databricks::zerobus::resume_ingest_stream_request::Identifier as ResumeIdentifier; +#[cfg(feature = "eos")] +use crate::databricks::zerobus::{ + CreatePersistentStreamRequest, PersistentStreamRequest, PersistentStreamResponse, + ResumeIngestStreamRequest, +}; + +/// Buffer size for the outbound request channel handed to tonic. +pub(super) const CHANNEL_BUFFER_SIZE: usize = 2048; + +/// Which kind of stream a connection serves. Selects the RPC, the offset +/// policy on the wire, and (for persistent) whether the first message is a +/// create or a resume. +#[derive(Clone)] +pub(super) enum TransportKind { + /// Ephemeral stream: `EphemeralStream` RPC, per-session 0-based wire offsets. + Ephemeral, + /// Persistent (Eos) stream: `PersistentStream` RPC, durable wire offsets. + /// `resume_stream_id` is `Some` when reconnecting to an existing stream and + /// `None` when creating a new one. + #[cfg(feature = "eos")] + Persistent { resume_stream_id: Option }, +} + +impl TransportKind { + /// Whether records are numbered on the wire by their durable logical offset + /// (persistent) rather than a fresh per-session physical counter (ephemeral). + pub(super) fn uses_durable_wire_offset(&self) -> bool { + match self { + TransportKind::Ephemeral => false, + #[cfg(feature = "eos")] + TransportKind::Persistent { .. } => true, + } + } +} + +/// Outcome of opening a connection: the first server message, normalized. +pub(super) struct Opened { + /// The stream identity the server assigned (create) or echoed (resume). + pub(super) stream_id: String, + /// Resume watermark: the highest offset the server has durably committed + /// for this stream. `None` for a freshly created stream (nothing committed + /// yet) and always `None` for ephemeral streams. + pub(super) last_committed_offset: Option, +} + +/// The outbound half of a stream: wraps the concrete tonic request sender and +/// exposes a single neutral `send` that both IO tasks use. +pub(super) enum OutboundSink { + Ephemeral(tokio::sync::mpsc::Sender), + #[cfg(feature = "eos")] + Persistent(tokio::sync::mpsc::Sender), +} + +impl OutboundSink { + /// Sends the opening message: `create_stream` for a new stream (ephemeral + /// or persistent), or `resume_stream` when reconnecting to a persistent + /// one. `create` carries the table name, descriptor, and record type; it is + /// ignored on the resume path (the server keeps that state from creation). + pub(super) async fn send_open( + &self, + kind: &TransportKind, + create: CreateIngestStreamRequest, + ) -> ZerobusResult<()> { + match (self, kind) { + (OutboundSink::Ephemeral(tx), TransportKind::Ephemeral) => tx + .send(EphemeralStreamRequest { + payload: Some(EphemeralRequestPayload::CreateStream(create)), + }) + .await + .map_err(|_| Self::open_failed()), + #[cfg(feature = "eos")] + (OutboundSink::Persistent(tx), TransportKind::Persistent { resume_stream_id }) => { + let payload = match resume_stream_id { + None => PersistentRequestPayload::CreateStream(CreatePersistentStreamRequest { + create_stream: Some(create), + }), + Some(stream_id) => { + // The descriptor is fixed at creation but re-sent on + // resume so the server can re-validate it against the + // table schema (required for PROTO streams; absent for + // JSON / Arrow, matching `create`). + PersistentRequestPayload::ResumeStream(ResumeIngestStreamRequest { + identifier: Some(ResumeIdentifier::StreamId(stream_id.clone())), + descriptor_proto: create.descriptor_proto, + record_type: create.record_type, + }) + } + }; + tx.send(PersistentStreamRequest { + payload: Some(payload), + }) + .await + .map_err(|_| Self::open_failed()) + } + // The sink and kind are always constructed together by + // `make_outbound`, so a mismatch is unreachable. Only present when + // more than one variant exists (i.e. `eos` is enabled). + #[cfg(feature = "eos")] + _ => Err(Self::open_failed()), + } + } + + fn open_failed() -> ZerobusError { + ZerobusError::StreamClosedError(tonic::Status::internal( + "Failed to send stream-open request", + )) + } + + /// Sends one ingest batch on the wire, numbered with `offset_id`. + pub(super) async fn send_ingest( + &self, + batch: EncodedBatch, + offset_id: OffsetId, + ) -> ZerobusResult<()> { + match self { + OutboundSink::Ephemeral(tx) => { + let payload = batch.into_request_payload(offset_id); + tx.send(EphemeralStreamRequest { + payload: Some(payload), + }) + .await + .map_err(|_| Self::send_failed()) + } + #[cfg(feature = "eos")] + OutboundSink::Persistent(tx) => { + let payload = ingest_payload_to_persistent(batch.into_request_payload(offset_id)); + tx.send(PersistentStreamRequest { + payload: Some(payload), + }) + .await + .map_err(|_| Self::send_failed()) + } + } + } + + fn send_failed() -> ZerobusError { + ZerobusError::StreamClosedError(tonic::Status::internal("Failed to send record")) + } +} + +/// A server message, normalized across the two response envelopes. Both IO +/// tasks match on this instead of the concrete generated oneof. +pub(super) enum InboundMessage { + /// A durability acknowledgement carrying the highest acked wire offset. + Ack(IngestRecordResponse), + /// The server will close the current session after the given duration. + Close(CloseStreamSignal), + /// Any other payload (create/resume responses arrive only while opening and + /// are handled there; seeing one mid-stream is unexpected). + Other, +} + +/// The inbound half of a stream: wraps the concrete tonic response stream and +/// yields normalized `InboundMessage`s. +pub(super) enum InboundStream { + Ephemeral(tonic::Streaming), + #[cfg(feature = "eos")] + Persistent(tonic::Streaming), +} + +impl InboundStream { + /// Reads the next server message. `Ok(None)` means the server closed the + /// stream gracefully; `Err` is a transport error. + pub(super) async fn message(&mut self) -> Result, tonic::Status> { + match self { + InboundStream::Ephemeral(s) => Ok(s.message().await?.map(|resp| match resp.payload { + Some(EphemeralResponsePayload::IngestRecordResponse(ack)) => { + InboundMessage::Ack(ack) + } + Some(EphemeralResponsePayload::CloseStreamSignal(sig)) => { + InboundMessage::Close(sig) + } + _ => InboundMessage::Other, + })), + #[cfg(feature = "eos")] + InboundStream::Persistent(s) => Ok(s.message().await?.map(|resp| match resp.payload { + Some(PersistentResponsePayload::IngestRecordResponse(ack)) => { + InboundMessage::Ack(ack) + } + Some(PersistentResponsePayload::CloseStreamSignal(sig)) => { + InboundMessage::Close(sig) + } + _ => InboundMessage::Other, + })), + } + } + + /// Reads and interprets the first server message after opening. + /// + /// For a create (ephemeral or persistent) this is a `CreateStreamResponse` + /// carrying the minted `stream_id`. For a persistent resume it is a + /// `ResumeStreamResponse` carrying the committed-offset watermark; the + /// `stream_id` is already known to the caller, so `Opened.stream_id` is + /// left empty on that path and filled in by the caller. + pub(super) async fn recv_open(&mut self) -> ZerobusResult { + match self { + InboundStream::Ephemeral(s) => { + let msg = Self::first_message(s.message().await)?; + match msg.payload { + Some(EphemeralResponsePayload::CreateStreamResponse(resp)) => { + let stream_id = resp.stream_id.ok_or_else(Self::missing_stream_id)?; + Ok(Opened { + stream_id, + last_committed_offset: None, + }) + } + other => Err(Self::unexpected_open(&other)), + } + } + #[cfg(feature = "eos")] + InboundStream::Persistent(s) => { + let msg = Self::first_message(s.message().await)?; + match msg.payload { + Some(PersistentResponsePayload::CreateStreamResponse(resp)) => { + let stream_id = resp.stream_id.ok_or_else(Self::missing_stream_id)?; + Ok(Opened { + stream_id, + last_committed_offset: None, + }) + } + Some(PersistentResponsePayload::ResumeStreamResponse(resp)) => Ok(Opened { + // The caller supplied the id it asked to resume. + stream_id: String::new(), + last_committed_offset: resp.last_committed_offset, + }), + other => Err(Self::unexpected_open(&other)), + } + } + } + } + + fn first_message(result: Result, tonic::Status>) -> ZerobusResult { + match result { + Ok(Some(msg)) => Ok(msg), + Ok(None) => Err(ZerobusError::CreateStreamError(tonic::Status::ok( + "Stream closed gracefully by server", + ))), + Err(status) => Err(ZerobusError::CreateStreamError(status)), + } + } + + fn missing_stream_id() -> ZerobusError { + ZerobusError::CreateStreamError(tonic::Status::internal( + "Successfully opened a stream but stream_id is None", + )) + } + + fn unexpected_open(payload: &T) -> ZerobusError { + ZerobusError::CreateStreamError(tonic::Status::internal(format!( + "Unexpected response from server while opening stream: {payload:?}" + ))) + } + + /// Drains and discards remaining server messages during teardown so the + /// server observes END_STREAM instead of a client RST_STREAM. + pub(super) async fn drain(&mut self) { + match self { + InboundStream::Ephemeral(s) => while matches!(s.message().await, Ok(Some(_))) {}, + #[cfg(feature = "eos")] + InboundStream::Persistent(s) => while matches!(s.message().await, Ok(Some(_))) {}, + } + } +} + +/// Opens the outbound request channel for the given transport kind, returning +/// the neutral sink and the raw tonic request stream to hand to the RPC. +/// +/// Split from the RPC call so the connection module can attach metadata to the +/// request before dispatching. +pub(super) fn make_outbound(kind: &TransportKind) -> (OutboundSink, OutboundRequestStream) { + match kind { + TransportKind::Ephemeral => { + let (tx, rx) = tokio::sync::mpsc::channel(CHANNEL_BUFFER_SIZE); + ( + OutboundSink::Ephemeral(tx), + OutboundRequestStream::Ephemeral(ReceiverStream::new(rx)), + ) + } + #[cfg(feature = "eos")] + TransportKind::Persistent { .. } => { + let (tx, rx) = tokio::sync::mpsc::channel(CHANNEL_BUFFER_SIZE); + ( + OutboundSink::Persistent(tx), + OutboundRequestStream::Persistent(ReceiverStream::new(rx)), + ) + } + } +} + +/// The raw request stream handed to the tonic RPC method. Kept concrete because +/// `ZerobusClient::ephemeral_stream` / `persistent_stream` want the exact type. +pub(super) enum OutboundRequestStream { + Ephemeral(ReceiverStream), + #[cfg(feature = "eos")] + Persistent(ReceiverStream), +} + +/// Dispatches the opening RPC on `channel` with `request` (metadata already +/// attached) and returns the inbound stream. The first server message is read +/// by the caller (`connection.rs`) so it can extract the stream_id / watermark. +pub(super) async fn open_rpc( + channel: &mut ZerobusClient, + request: tonic::Request, +) -> ZerobusResult { + let (metadata, extensions, body) = request.into_parts(); + match body { + OutboundRequestStream::Ephemeral(stream) => { + let req = tonic::Request::from_parts(metadata, extensions, stream); + let resp = channel + .ephemeral_stream(req) + .await + .map_err(ZerobusError::CreateStreamError)?; + Ok(InboundStream::Ephemeral(resp.into_inner())) + } + #[cfg(feature = "eos")] + OutboundRequestStream::Persistent(stream) => { + let req = tonic::Request::from_parts(metadata, extensions, stream); + let resp = channel + .persistent_stream(req) + .await + .map_err(ZerobusError::CreateStreamError)?; + Ok(InboundStream::Persistent(resp.into_inner())) + } + } +} + +/// Translates a shared ingest request payload (built by `EncodedBatch`) into +/// the persistent envelope's oneof. The two enums are structurally identical +/// for the ingest variants; only the create/resume-specific variants differ, +/// and those never flow through here. +#[cfg(feature = "eos")] +fn ingest_payload_to_persistent(payload: EphemeralRequestPayload) -> PersistentRequestPayload { + match payload { + EphemeralRequestPayload::IngestRecord(r) => PersistentRequestPayload::IngestRecord(r), + EphemeralRequestPayload::IngestRecordBatch(b) => { + PersistentRequestPayload::IngestRecordBatch(b) + } + // `into_request_payload` only ever produces ingest variants; a create + // payload here would be a programming error. + EphemeralRequestPayload::CreateStream(_) => unreachable!( + "into_request_payload never yields CreateStream; create is sent by connection.rs" + ), + } +} From f28ce736f31fbde93f4a0bf0d8703cbef2509576 Mon Sep 17 00:00:00 2001 From: elenagaljak-db Date: Tue, 25 Aug 2026 12:12:45 +0000 Subject: [PATCH 2/7] [Rust] Address transport review Signed-off-by: elenagaljak-db --- rust/sdk/Cargo.toml | 2 +- rust/sdk/src/landing_zone.rs | 2 +- rust/sdk/src/record_types.rs | 34 ++++++++++++++++++++++++++ rust/sdk/src/stream/grpc/supervisor.rs | 22 +++++------------ rust/sdk/src/stream/grpc/transport.rs | 29 ++++------------------ 5 files changed, 47 insertions(+), 42 deletions(-) diff --git a/rust/sdk/Cargo.toml b/rust/sdk/Cargo.toml index fcd7c12a..9c38359d 100644 --- a/rust/sdk/Cargo.toml +++ b/rust/sdk/Cargo.toml @@ -77,7 +77,7 @@ internal-arrow-c-data = [ ] # Zero-copy protobuf parser. zeroparser = ["dep:self_cell", "dep:prost-build"] -# Persistent (Eos) streams; in development, API is unstable +# Persistent streams with exactly-once semantics (EoS); API is in development eos = [] testing = ["dep:futures"] # Test-only deterministic seams (barriers/notifies) for the Arrow stream. Zero footprint diff --git a/rust/sdk/src/landing_zone.rs b/rust/sdk/src/landing_zone.rs index 8bd44aad..e80af7f9 100644 --- a/rust/sdk/src/landing_zone.rs +++ b/rust/sdk/src/landing_zone.rs @@ -90,7 +90,7 @@ impl LandingZone { /// /// Used by the persistent-stream resume path to reconcile the retained tail /// against the server's committed watermark: it removes the prefix of - /// records the server has already durably stored (offset ≤ resume watermark) + /// records the server has already durably stored (offset <= resume watermark) /// before re-sending the rest. Only the unobserved queue is consulted: /// callers reset observation (`reset_observe`) first, so every retained item /// lives in the queue in offset order. One semaphore permit is released per diff --git a/rust/sdk/src/record_types.rs b/rust/sdk/src/record_types.rs index e805af37..4937f99d 100644 --- a/rust/sdk/src/record_types.rs +++ b/rust/sdk/src/record_types.rs @@ -20,6 +20,9 @@ use crate::databricks::zerobus::{ }; use crate::OffsetId; +#[cfg(feature = "eos")] +use crate::databricks::zerobus::persistent_stream_request::Payload as PersistentRequestPayload; + /// A type alias for a protobuf-encoded record. pub type ProtoEncodedRecord = Vec; @@ -257,6 +260,22 @@ impl EncodedBatch { } } + #[cfg(feature = "eos")] + pub(crate) fn into_persistent_request_payload( + self, + offset_id: OffsetId, + ) -> PersistentRequestPayload { + match self.into_request_payload(offset_id) { + RequestPayload::IngestRecord(record) => PersistentRequestPayload::IngestRecord(record), + RequestPayload::IngestRecordBatch(batch) => { + PersistentRequestPayload::IngestRecordBatch(batch) + } + RequestPayload::CreateStream(_) => { + unreachable!("encoded batches only produce ingest payloads") + } + } + } + /// Returns the number of records in this batch. pub fn get_record_count(&self) -> usize { match self { @@ -778,6 +797,21 @@ mod tests { _ => panic!("Expected IngestRecordBatch payload"), } } + + #[cfg(feature = "eos")] + #[test] + fn test_into_persistent_request_payload() { + let record = r#"{"id": 1}"#.to_string(); + let batch = EncodedBatch::Json(smallvec![record.clone()]); + + match batch.into_persistent_request_payload(42) { + PersistentRequestPayload::IngestRecord(req) => { + assert_eq!(req.offset_id, Some(42)); + assert_eq!(req.record, Some(IngestRequestRecord::JsonRecord(record))); + } + _ => panic!("Expected persistent IngestRecord payload"), + } + } } mod encoded_batch_iter { diff --git a/rust/sdk/src/stream/grpc/supervisor.rs b/rust/sdk/src/stream/grpc/supervisor.rs index e11bf43c..3dbcc156 100644 --- a/rust/sdk/src/stream/grpc/supervisor.rs +++ b/rust/sdk/src/stream/grpc/supervisor.rs @@ -224,15 +224,9 @@ impl ZerobusStream { let resent_records = landing_zone_recovery.observed_count(); let resent_batches = landing_zone_recovery.reset_observe(); - // Resume alignment (persistent streams). The server reports how far - // it has durably committed. A dropped connection can lose the ack - // for records the server did commit, so the retained tail can start - // at or below that watermark. Re-sending an already-committed offset - // would be rejected by the server and kill the stream (Decision 6: - // server rejects), so reconcile first: resolve those records' - // pending acks locally and re-send only the offsets above the - // watermark. This is not user-duplicate dedup — it realigns the - // SDK's retained tail with the server's durable position on resume. + // A dropped connection can lose ACKs for records the server already + // committed. Resolve retained records through the resume watermark + // locally, then resend only records above it. let mut initial_last_acked_offset: OffsetId = -1; if durable_wire_offset { if let Some(watermark) = last_committed_offset { @@ -369,15 +363,11 @@ impl ZerobusStream { /// Reconciles the retained landing-zone tail with the server's durable /// position when resuming a persistent stream: removes the prefix of records - /// the server has already committed (logical offset ≤ `watermark`) and + /// the server has already committed (logical offset <= `watermark`) and /// resolves their pending acks locally, so they are never re-sent. /// - /// Re-sending an already-committed offset would be rejected by the server - /// and kill the stream (Decision 6: server rejects), so this realignment is - /// what makes resume survive a lost-ack race. It is not user-duplicate - /// dedup — the records removed here are ones the SDK itself queued and the - /// server durably stored; only their acknowledgement was lost with the - /// dropped connection. + /// This realignment makes resume survive a lost-ACK race: the records were + /// durably stored, but their acknowledgements were lost with the connection. /// /// The caller must have reset observation first, so every retained record /// sits in the landing-zone queue in ascending offset order — the committed diff --git a/rust/sdk/src/stream/grpc/transport.rs b/rust/sdk/src/stream/grpc/transport.rs index 3f16dd89..725985e7 100644 --- a/rust/sdk/src/stream/grpc/transport.rs +++ b/rust/sdk/src/stream/grpc/transport.rs @@ -161,22 +161,22 @@ impl OutboundSink { payload: Some(payload), }) .await - .map_err(|_| Self::send_failed()) + .map_err(|_| Self::ingest_failed()) } #[cfg(feature = "eos")] OutboundSink::Persistent(tx) => { - let payload = ingest_payload_to_persistent(batch.into_request_payload(offset_id)); + let payload = batch.into_persistent_request_payload(offset_id); tx.send(PersistentStreamRequest { payload: Some(payload), }) .await - .map_err(|_| Self::send_failed()) + .map_err(|_| Self::ingest_failed()) } } } - fn send_failed() -> ZerobusError { - ZerobusError::StreamClosedError(tonic::Status::internal("Failed to send record")) + fn ingest_failed() -> ZerobusError { + ZerobusError::StreamClosedError(tonic::Status::internal("Failed to send batch")) } } @@ -365,22 +365,3 @@ pub(super) async fn open_rpc( } } } - -/// Translates a shared ingest request payload (built by `EncodedBatch`) into -/// the persistent envelope's oneof. The two enums are structurally identical -/// for the ingest variants; only the create/resume-specific variants differ, -/// and those never flow through here. -#[cfg(feature = "eos")] -fn ingest_payload_to_persistent(payload: EphemeralRequestPayload) -> PersistentRequestPayload { - match payload { - EphemeralRequestPayload::IngestRecord(r) => PersistentRequestPayload::IngestRecord(r), - EphemeralRequestPayload::IngestRecordBatch(b) => { - PersistentRequestPayload::IngestRecordBatch(b) - } - // `into_request_payload` only ever produces ingest variants; a create - // payload here would be a programming error. - EphemeralRequestPayload::CreateStream(_) => unreachable!( - "into_request_payload never yields CreateStream; create is sent by connection.rs" - ), - } -} From 973c7aa55b83926d8f95498171baa47d232ebdb1 Mon Sep 17 00:00:00 2001 From: elenagaljak-db Date: Tue, 25 Aug 2026 13:11:46 +0000 Subject: [PATCH 3/7] [Rust] Validate persistent protocol Signed-off-by: elenagaljak-db --- rust/NEXT_CHANGELOG.md | 5 +- rust/sdk/src/stream/grpc/connection.rs | 96 ++++++++++++++++++++------ rust/sdk/src/stream/grpc/mod.rs | 7 +- rust/sdk/src/stream/grpc/receiver.rs | 58 ++++++++++++++++ rust/sdk/src/stream/grpc/supervisor.rs | 9 ++- rust/sdk/src/stream/grpc/transport.rs | 36 +++++----- 6 files changed, 165 insertions(+), 46 deletions(-) diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index 3e84f409..850139f9 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -47,8 +47,9 @@ ### Internal Changes -- Added the feature-gated persistent gRPC transport, durable wire offsets, and - resume-watermark reconciliation for recovery after a lost acknowledgment. +- Added the feature-gated persistent gRPC transport, durable wire offsets, + resume-watermark reconciliation after a lost acknowledgment, and validation + for setup responses, acknowledgment bounds, and offset overflow. ### Breaking Changes diff --git a/rust/sdk/src/stream/grpc/connection.rs b/rust/sdk/src/stream/grpc/connection.rs index a7a34420..6abe2d3a 100644 --- a/rust/sdk/src/stream/grpc/connection.rs +++ b/rust/sdk/src/stream/grpc/connection.rs @@ -14,10 +14,11 @@ use tonic::metadata::MetadataValue; use tonic::transport::Channel; use tracing::{debug, error, info, instrument, warn}; -use super::transport::{self, InboundStream, OutboundSink, TransportKind}; +use super::supervisor::StreamInitInfo; +use super::transport::{self, InboundStream, Opened, OutboundSink, TransportKind}; use crate::databricks::zerobus::zerobus_client::ZerobusClient; use crate::databricks::zerobus::{CreateIngestStreamRequest, RecordType}; -use crate::{HeadersProvider, OffsetId, TableProperties, ZerobusError, ZerobusResult}; +use crate::{HeadersProvider, TableProperties, ZerobusError, ZerobusResult}; /// A freshly opened stream connection: the outbound sink, the inbound response /// stream, the server-assigned `stream_id`, and (persistent resume only) the @@ -25,8 +26,7 @@ use crate::{HeadersProvider, OffsetId, TableProperties, ZerobusError, ZerobusRes pub(super) struct StreamConnection { pub(super) sink: OutboundSink, pub(super) inbound: InboundStream, - pub(super) stream_id: String, - pub(super) last_committed_offset: Option, + pub(super) init_info: StreamInitInfo, } impl super::ZerobusStream { @@ -164,23 +164,58 @@ impl super::ZerobusStream { debug!("Waiting for stream-open response."); let opened = inbound.recv_open().await?; - // On a persistent resume the server does not re-send the stream_id (the - // client supplied it), so fall back to the id being resumed. - let stream_id = if opened.stream_id.is_empty() { - Self::resume_stream_id(kind).unwrap_or_default() - } else { - opened.stream_id - }; - info!(stream_id = %stream_id, last_committed_offset = ?opened.last_committed_offset, "Successfully opened stream"); + let init_info = Self::validate_open_response(kind, opened)?; + info!(stream_id = %init_info.stream_id, last_committed_offset = ?init_info.last_committed_offset, "Successfully opened stream"); Ok(StreamConnection { sink, inbound, - stream_id, - last_committed_offset: opened.last_committed_offset, + init_info, }) } + fn validate_open_response( + kind: &TransportKind, + opened: Opened, + ) -> ZerobusResult { + let init_info = match (kind, opened) { + (TransportKind::Ephemeral, Opened::Created { stream_id }) => StreamInitInfo { + stream_id, + last_committed_offset: None, + }, + #[cfg(feature = "eos")] + ( + TransportKind::Persistent { + resume_stream_id: None, + }, + Opened::Created { stream_id }, + ) => StreamInitInfo { + stream_id, + last_committed_offset: None, + }, + #[cfg(feature = "eos")] + ( + TransportKind::Persistent { + resume_stream_id: Some(stream_id), + }, + Opened::Resumed { + last_committed_offset, + }, + ) => StreamInitInfo { + stream_id: stream_id.clone(), + last_committed_offset, + }, + #[cfg(feature = "eos")] + _ => { + return Err(ZerobusError::UnexpectedStreamResponseError( + "Persistent stream setup response did not match the requested operation" + .to_string(), + )); + } + }; + Ok(init_info) + } + /// Builds the `CreateIngestStreamRequest` sent when opening (or, on the /// persistent path, creating) a stream. Encodes the descriptor for proto /// streams and validates its presence. @@ -210,13 +245,32 @@ impl super::ZerobusStream { record_type: Some(record_type.into()), }) } +} - /// The stream id being resumed, if this is a persistent resume. - fn resume_stream_id(kind: &TransportKind) -> Option { - match kind { - TransportKind::Ephemeral => None, - #[cfg(feature = "eos")] - TransportKind::Persistent { resume_stream_id } => resume_stream_id.clone(), - } +#[cfg(all(test, feature = "eos"))] +mod tests { + use super::*; + use crate::ZerobusStream; + + #[test] + fn persistent_create_rejects_resume_response() { + let kind = TransportKind::Persistent { + resume_stream_id: None, + }; + let opened = Opened::Resumed { + last_committed_offset: Some(3), + }; + assert!(ZerobusStream::validate_open_response(&kind, opened).is_err()); + } + + #[test] + fn persistent_resume_rejects_create_response() { + let kind = TransportKind::Persistent { + resume_stream_id: Some("stream-id".to_string()), + }; + let opened = Opened::Created { + stream_id: "other-id".to_string(), + }; + assert!(ZerobusStream::validate_open_response(&kind, opened).is_err()); } } diff --git a/rust/sdk/src/stream/grpc/mod.rs b/rust/sdk/src/stream/grpc/mod.rs index e1e41205..6ed86380 100644 --- a/rust/sdk/src/stream/grpc/mod.rs +++ b/rust/sdk/src/stream/grpc/mod.rs @@ -265,7 +265,12 @@ impl ZerobusStream { // Safe because the constructor returns before any user ingest, so no real // ack can race this initial value. if let Some(watermark) = init_info.last_committed_offset { - logical_offset_id_generator.set_next(watermark + 1); + let next_offset = watermark.checked_add(1).ok_or_else(|| { + ZerobusError::UnexpectedStreamResponseError( + "Persistent stream offset space is exhausted".to_string(), + ) + })?; + logical_offset_id_generator.set_next(next_offset); let _ = logical_last_received_offset_id_tx.send(Some(watermark)); } diff --git a/rust/sdk/src/stream/grpc/receiver.rs b/rust/sdk/src/stream/grpc/receiver.rs index 3c5be9a3..7d9eb1be 100644 --- a/rust/sdk/src/stream/grpc/receiver.rs +++ b/rust/sdk/src/stream/grpc/receiver.rs @@ -112,6 +112,14 @@ impl ZerobusStream { return Err(error); } }; + if let Err(error) = Self::validate_ack_offset( + last_acked_offset, + durability_ack_up_to_offset, + landing_zone.observed_count(), + ) { + let _ = server_error_tx.send(Some(error.clone())); + return Err(error); + } let mut last_logical_acked_offset = -2; let mut map = oneshot_map.lock().await; for _offset_to_ack in (last_acked_offset + 1)..=durability_ack_up_to_offset @@ -231,4 +239,54 @@ impl ZerobusStream { Ok(()) }) } + + fn validate_ack_offset( + last_acked_offset: OffsetId, + ack_offset: OffsetId, + pending_batches: usize, + ) -> ZerobusResult<()> { + let pending_batches = OffsetId::try_from(pending_batches).map_err(|_| { + ZerobusError::StreamClosedError(tonic::Status::invalid_argument( + "Pending batch count exceeds the supported offset range", + )) + })?; + let greatest_sent_offset = + last_acked_offset + .checked_add(pending_batches) + .ok_or_else(|| { + ZerobusError::StreamClosedError(tonic::Status::invalid_argument( + "Greatest sent offset exceeds the supported offset range", + )) + })?; + + if ack_offset < last_acked_offset || ack_offset > greatest_sent_offset { + return Err(ZerobusError::StreamClosedError( + tonic::Status::invalid_argument(format!( + "Invalid durability ACK offset {ack_offset}; expected {last_acked_offset}..={greatest_sent_offset}" + )), + )); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::ZerobusStream; + + #[test] + fn ack_offset_must_not_regress() { + assert!(ZerobusStream::validate_ack_offset(4, 3, 2).is_err()); + } + + #[test] + fn ack_offset_must_not_exceed_greatest_sent() { + assert!(ZerobusStream::validate_ack_offset(4, 7, 2).is_err()); + } + + #[test] + fn duplicate_and_cumulative_ack_offsets_are_valid() { + assert!(ZerobusStream::validate_ack_offset(4, 4, 2).is_ok()); + assert!(ZerobusStream::validate_ack_offset(4, 6, 2).is_ok()); + } } diff --git a/rust/sdk/src/stream/grpc/supervisor.rs b/rust/sdk/src/stream/grpc/supervisor.rs index 3dbcc156..4fd0bb11 100644 --- a/rust/sdk/src/stream/grpc/supervisor.rs +++ b/rust/sdk/src/stream/grpc/supervisor.rs @@ -103,6 +103,9 @@ impl ZerobusStream { // rejection, so the retry re-mints via the headers provider. Initial setup bounds // that invalidation under the setup deadline so a stalled provider cannot bypass // the one-shot limit; reconnect keeps its existing timeout-wrapped path unchanged. + // A persistent create has no client-supplied id yet, so if the server creates the + // stream but its response is lost, retrying can create an orphaned stream. Resumes + // are idempotent because they use the already-minted stream id. let is_initial = initial_stream_creation; let attempt = AtomicUsize::new(0); let create_attempt = || { @@ -198,8 +201,10 @@ impl ZerobusStream { return Err(e); } }; - let stream_id = connection.stream_id; - let last_committed_offset = connection.last_committed_offset; + let StreamInitInfo { + stream_id, + last_committed_offset, + } = connection.init_info; if initial_stream_creation { if let Some(stream_init_result_tx_inner) = stream_init_result_tx.take() { let _ = stream_init_result_tx_inner.send(Ok(StreamInitInfo { diff --git a/rust/sdk/src/stream/grpc/transport.rs b/rust/sdk/src/stream/grpc/transport.rs index 725985e7..c5f14d49 100644 --- a/rust/sdk/src/stream/grpc/transport.rs +++ b/rust/sdk/src/stream/grpc/transport.rs @@ -75,14 +75,16 @@ impl TransportKind { } } -/// Outcome of opening a connection: the first server message, normalized. -pub(super) struct Opened { - /// The stream identity the server assigned (create) or echoed (resume). - pub(super) stream_id: String, - /// Resume watermark: the highest offset the server has durably committed - /// for this stream. `None` for a freshly created stream (nothing committed - /// yet) and always `None` for ephemeral streams. - pub(super) last_committed_offset: Option, +/// Outcome of opening a connection, preserving which operation the server +/// acknowledged so the caller can reject mismatched setup responses. +pub(super) enum Opened { + Created { + stream_id: String, + }, + #[cfg(feature = "eos")] + Resumed { + last_committed_offset: Option, + }, } /// The outbound half of a stream: wraps the concrete tonic request sender and @@ -241,10 +243,7 @@ impl InboundStream { match msg.payload { Some(EphemeralResponsePayload::CreateStreamResponse(resp)) => { let stream_id = resp.stream_id.ok_or_else(Self::missing_stream_id)?; - Ok(Opened { - stream_id, - last_committed_offset: None, - }) + Ok(Opened::Created { stream_id }) } other => Err(Self::unexpected_open(&other)), } @@ -255,16 +254,13 @@ impl InboundStream { match msg.payload { Some(PersistentResponsePayload::CreateStreamResponse(resp)) => { let stream_id = resp.stream_id.ok_or_else(Self::missing_stream_id)?; - Ok(Opened { - stream_id, - last_committed_offset: None, + Ok(Opened::Created { stream_id }) + } + Some(PersistentResponsePayload::ResumeStreamResponse(resp)) => { + Ok(Opened::Resumed { + last_committed_offset: resp.last_committed_offset, }) } - Some(PersistentResponsePayload::ResumeStreamResponse(resp)) => Ok(Opened { - // The caller supplied the id it asked to resume. - stream_id: String::new(), - last_committed_offset: resp.last_committed_offset, - }), other => Err(Self::unexpected_open(&other)), } } From 4c0c9e8c06c24d55749f9e1c3f6e2a916b45cddc Mon Sep 17 00:00:00 2001 From: elenagaljak-db Date: Tue, 25 Aug 2026 16:51:29 +0000 Subject: [PATCH 4/7] [Rust] Clarify persistent recovery Signed-off-by: elenagaljak-db --- rust/sdk/src/stream/grpc/connection.rs | 102 ++++++++++++++++--------- rust/sdk/src/stream/grpc/mod.rs | 8 +- rust/sdk/src/stream/grpc/sender.rs | 12 ++- rust/sdk/src/stream/grpc/supervisor.rs | 14 ++-- rust/sdk/src/stream/grpc/transport.rs | 33 ++++---- 5 files changed, 100 insertions(+), 69 deletions(-) diff --git a/rust/sdk/src/stream/grpc/connection.rs b/rust/sdk/src/stream/grpc/connection.rs index 6abe2d3a..892b4d5d 100644 --- a/rust/sdk/src/stream/grpc/connection.rs +++ b/rust/sdk/src/stream/grpc/connection.rs @@ -15,7 +15,7 @@ use tonic::transport::Channel; use tracing::{debug, error, info, instrument, warn}; use super::supervisor::StreamInitInfo; -use super::transport::{self, InboundStream, Opened, OutboundSink, TransportKind}; +use super::transport::{self, GrpcConnectionMode, InboundStream, Opened, OutboundSink}; use crate::databricks::zerobus::zerobus_client::ZerobusClient; use crate::databricks::zerobus::{CreateIngestStreamRequest, RecordType}; use crate::{HeadersProvider, TableProperties, ZerobusError, ZerobusResult}; @@ -42,7 +42,7 @@ impl super::ZerobusStream { table_properties: &TableProperties, headers_provider: &Arc, record_type: RecordType, - kind: &TransportKind, + kind: &GrpcConnectionMode, ) -> ZerobusResult { let result = Self::create_stream_connection_inner( channel, @@ -73,7 +73,7 @@ impl super::ZerobusStream { table_properties: &TableProperties, headers_provider: &Arc, record_type: RecordType, - kind: &TransportKind, + kind: &GrpcConnectionMode, recovery_timeout_ms: u64, ) -> ZerobusResult { let attempt_timeout = Duration::from_millis(recovery_timeout_ms); @@ -117,7 +117,7 @@ impl super::ZerobusStream { table_properties: &TableProperties, headers_provider: &Arc, record_type: RecordType, - kind: &TransportKind, + kind: &GrpcConnectionMode, ) -> ZerobusResult { let (sink, request_body) = transport::make_outbound(kind); let mut request = tonic::Request::new(request_body); @@ -175,45 +175,53 @@ impl super::ZerobusStream { } fn validate_open_response( - kind: &TransportKind, + kind: &GrpcConnectionMode, opened: Opened, ) -> ZerobusResult { - let init_info = match (kind, opened) { - (TransportKind::Ephemeral, Opened::Created { stream_id }) => StreamInitInfo { - stream_id, - last_committed_offset: None, + match kind { + GrpcConnectionMode::Ephemeral => match opened { + Opened::Created { stream_id } => Ok(StreamInitInfo { + stream_id, + last_committed_offset: None, + }), + #[cfg(feature = "eos")] + _ => Err(Self::mismatched_open_response()), }, #[cfg(feature = "eos")] - ( - TransportKind::Persistent { - resume_stream_id: None, - }, - Opened::Created { stream_id }, - ) => StreamInitInfo { + GrpcConnectionMode::Persistent { resume_stream_id } => { + Self::validate_persistent_open_response(resume_stream_id.as_deref(), opened) + } + } + } + + #[cfg(feature = "eos")] + fn validate_persistent_open_response( + resume_stream_id: Option<&str>, + opened: Opened, + ) -> ZerobusResult { + match (resume_stream_id, opened) { + (None, Opened::Created { stream_id }) => Ok(StreamInitInfo { stream_id, last_committed_offset: None, - }, - #[cfg(feature = "eos")] + }), ( - TransportKind::Persistent { - resume_stream_id: Some(stream_id), - }, + Some(stream_id), Opened::Resumed { last_committed_offset, }, - ) => StreamInitInfo { - stream_id: stream_id.clone(), + ) => Ok(StreamInitInfo { + stream_id: stream_id.to_string(), last_committed_offset, - }, - #[cfg(feature = "eos")] - _ => { - return Err(ZerobusError::UnexpectedStreamResponseError( - "Persistent stream setup response did not match the requested operation" - .to_string(), - )); - } - }; - Ok(init_info) + }), + _ => Err(Self::mismatched_open_response()), + } + } + + #[cfg(feature = "eos")] + fn mismatched_open_response() -> ZerobusError { + ZerobusError::UnexpectedStreamResponseError( + "Persistent stream setup response did not match the requested operation".to_string(), + ) } /// Builds the `CreateIngestStreamRequest` sent when opening (or, on the @@ -252,9 +260,35 @@ mod tests { use super::*; use crate::ZerobusStream; + #[test] + fn persistent_create_accepts_create_response() { + let kind = GrpcConnectionMode::Persistent { + resume_stream_id: None, + }; + let opened = Opened::Created { + stream_id: "created-id".to_string(), + }; + let init = ZerobusStream::validate_open_response(&kind, opened).unwrap(); + assert_eq!(init.stream_id, "created-id"); + assert_eq!(init.last_committed_offset, None); + } + + #[test] + fn persistent_resume_accepts_resume_response() { + let kind = GrpcConnectionMode::Persistent { + resume_stream_id: Some("stream-id".to_string()), + }; + let opened = Opened::Resumed { + last_committed_offset: Some(3), + }; + let init = ZerobusStream::validate_open_response(&kind, opened).unwrap(); + assert_eq!(init.stream_id, "stream-id"); + assert_eq!(init.last_committed_offset, Some(3)); + } + #[test] fn persistent_create_rejects_resume_response() { - let kind = TransportKind::Persistent { + let kind = GrpcConnectionMode::Persistent { resume_stream_id: None, }; let opened = Opened::Resumed { @@ -265,7 +299,7 @@ mod tests { #[test] fn persistent_resume_rejects_create_response() { - let kind = TransportKind::Persistent { + let kind = GrpcConnectionMode::Persistent { resume_stream_id: Some("stream-id".to_string()), }; let opened = Opened::Created { diff --git a/rust/sdk/src/stream/grpc/mod.rs b/rust/sdk/src/stream/grpc/mod.rs index 6ed86380..b996b150 100644 --- a/rust/sdk/src/stream/grpc/mod.rs +++ b/rust/sdk/src/stream/grpc/mod.rs @@ -47,7 +47,7 @@ mod supervisor; mod transport; mod types; -use transport::TransportKind; +use transport::GrpcConnectionMode; use types::{IngestRequest, OneshotMap, RecordLandingZone}; #[cfg(feature = "testing")] @@ -158,7 +158,7 @@ impl ZerobusStream { headers_provider, options, StreamType::Ephemeral, - TransportKind::Ephemeral, + GrpcConnectionMode::Ephemeral, ) .await } @@ -185,7 +185,7 @@ impl ZerobusStream { headers_provider, options, StreamType::Persistent, - TransportKind::Persistent { resume_stream_id }, + GrpcConnectionMode::Persistent { resume_stream_id }, ) .await } @@ -200,7 +200,7 @@ impl ZerobusStream { headers_provider: Arc, options: StreamConfigurationOptions, stream_type: StreamType, - kind: TransportKind, + kind: GrpcConnectionMode, ) -> ZerobusResult { let (stream_init_result_tx, stream_init_result_rx) = tokio::sync::oneshot::channel::>(); diff --git a/rust/sdk/src/stream/grpc/sender.rs b/rust/sdk/src/stream/grpc/sender.rs index 246da3ec..7dd8b50c 100644 --- a/rust/sdk/src/stream/grpc/sender.rs +++ b/rust/sdk/src/stream/grpc/sender.rs @@ -20,18 +20,16 @@ impl ZerobusStream { /// Spawns a task that continuously sends records to the Zerobus API by observing the landing zone /// to get records and sending them through the outbound stream to the gRPC stream. /// - /// `durable_wire_offset` selects the offset policy. Ephemeral streams number - /// records with a fresh 0-based physical counter each session (the server - /// tracks nothing across reconnects). Persistent streams put the record's - /// durable logical offset on the wire so the server can dedup and resume by - /// it — that offset already lives on the landing-zone item. + /// Persistent SDK and wire offsets match exactly. Ephemeral streams use a + /// fresh 0-based wire sequence for each server stream so their SDK offsets + /// can preserve continuity across recovery. pub(super) fn spawn_sender_task( sink: OutboundSink, landing_zone: RecordLandingZone, is_paused: Arc, server_error_tx: tokio::sync::watch::Sender>, cancellation_token: CancellationToken, - durable_wire_offset: bool, + wire_offsets_match: bool, ) -> tokio::task::JoinHandle> { tokio::spawn(async move { let physical_offset_id_generator = OffsetIdGenerator::default(); @@ -47,7 +45,7 @@ impl ZerobusStream { } } => item.clone(), }; - let wire_offset = if durable_wire_offset { + let wire_offset = if wire_offsets_match { item.offset_id } else { physical_offset_id_generator.next() diff --git a/rust/sdk/src/stream/grpc/supervisor.rs b/rust/sdk/src/stream/grpc/supervisor.rs index 4fd0bb11..97c667f7 100644 --- a/rust/sdk/src/stream/grpc/supervisor.rs +++ b/rust/sdk/src/stream/grpc/supervisor.rs @@ -18,7 +18,7 @@ use tokio_util::sync::CancellationToken; use tonic::transport::Channel; use tracing::{debug, error, info, instrument, warn}; -use super::transport::TransportKind; +use super::transport::GrpcConnectionMode; use super::types::{CallbackMessage, OneshotMap, RecordLandingZone}; use super::{ZerobusStream, STREAM_TEARDOWN_DRAIN_TIMEOUT_MS}; use crate::databricks::zerobus::zerobus_client::ZerobusClient; @@ -54,7 +54,7 @@ impl ZerobusStream { options: StreamConfigurationOptions, // Mutated only on the persistent recovery path (flip create → resume); // that mutation is `eos`-gated, so `mut` is otherwise unused. - #[cfg_attr(not(feature = "eos"), allow(unused_mut))] mut kind: TransportKind, + #[cfg_attr(not(feature = "eos"), allow(unused_mut))] mut kind: GrpcConnectionMode, landing_zone: RecordLandingZone, oneshot_map: Arc>, logical_last_received_offset_id_tx: tokio::sync::watch::Sender>, @@ -65,7 +65,7 @@ impl ZerobusStream { cancellation_token: CancellationToken, callback_tx: Option>, ) -> ZerobusResult<()> { - let durable_wire_offset = kind.uses_durable_wire_offset(); + let wire_offsets_match = kind.wire_offsets_match(); let mut initial_stream_creation = true; let mut stream_init_result_tx = Some(stream_init_result_tx); // One-shot budget: initial setup may spend a single recovery retry to refresh a @@ -216,7 +216,7 @@ impl ZerobusStream { // A persistent stream recovers by resuming its now-known id, not // by creating a fresh stream on every reconnect. #[cfg(feature = "eos")] - if let TransportKind::Persistent { resume_stream_id } = &mut kind { + if let GrpcConnectionMode::Persistent { resume_stream_id } = &mut kind { *resume_stream_id = Some(stream_id.clone()); } info!(stream_id = %stream_id, "Successfully created stream"); @@ -233,7 +233,7 @@ impl ZerobusStream { // committed. Resolve retained records through the resume watermark // locally, then resend only records above it. let mut initial_last_acked_offset: OffsetId = -1; - if durable_wire_offset { + if wire_offsets_match { if let Some(watermark) = last_committed_offset { initial_last_acked_offset = watermark; Self::reconcile_committed_on_resume( @@ -284,7 +284,7 @@ impl ZerobusStream { Arc::clone(&is_paused), server_error_tx.clone(), per_stream_token.clone(), - durable_wire_offset, + wire_offsets_match, ); // 4. Wait for any of the two tasks to end. @@ -368,7 +368,7 @@ impl ZerobusStream { /// Reconciles the retained landing-zone tail with the server's durable /// position when resuming a persistent stream: removes the prefix of records - /// the server has already committed (logical offset <= `watermark`) and + /// the server has already committed (SDK/wire offset <= `watermark`) and /// resolves their pending acks locally, so they are never re-sent. /// /// This realignment makes resume survive a lost-ACK race: the records were diff --git a/rust/sdk/src/stream/grpc/transport.rs b/rust/sdk/src/stream/grpc/transport.rs index c5f14d49..747a12a5 100644 --- a/rust/sdk/src/stream/grpc/transport.rs +++ b/rust/sdk/src/stream/grpc/transport.rs @@ -7,10 +7,9 @@ //! 1. **Opening** — ephemeral always sends `create_stream`; persistent sends //! `create_stream` (new) or `resume_stream` (reconnect), and a resume comes //! back with a committed-offset watermark. -//! 2. **Offset on the wire** — ephemeral numbers records with a fresh 0-based -//! counter each session (the server tracks nothing across reconnects); -//! persistent puts the record's durable logical offset on the wire so the -//! server can dedup and resume by it. +//! 2. **Offset on the wire** — persistent SDK and wire offsets match exactly. +//! Ephemeral uses a fresh 0-based wire offset for each server stream while +//! its SDK offset preserves continuity across recovery. //! 3. **Response parsing** — the oneof variants live in different generated //! enums. //! @@ -53,7 +52,7 @@ pub(super) const CHANNEL_BUFFER_SIZE: usize = 2048; /// policy on the wire, and (for persistent) whether the first message is a /// create or a resume. #[derive(Clone)] -pub(super) enum TransportKind { +pub(super) enum GrpcConnectionMode { /// Ephemeral stream: `EphemeralStream` RPC, per-session 0-based wire offsets. Ephemeral, /// Persistent (Eos) stream: `PersistentStream` RPC, durable wire offsets. @@ -63,14 +62,14 @@ pub(super) enum TransportKind { Persistent { resume_stream_id: Option }, } -impl TransportKind { - /// Whether records are numbered on the wire by their durable logical offset - /// (persistent) rather than a fresh per-session physical counter (ephemeral). - pub(super) fn uses_durable_wire_offset(&self) -> bool { +impl GrpcConnectionMode { + /// Whether SDK and wire offsets match. They match for persistent streams; + /// ephemeral streams use a fresh wire sequence for each server stream. + pub(super) fn wire_offsets_match(&self) -> bool { match self { - TransportKind::Ephemeral => false, + GrpcConnectionMode::Ephemeral => false, #[cfg(feature = "eos")] - TransportKind::Persistent { .. } => true, + GrpcConnectionMode::Persistent { .. } => true, } } } @@ -102,18 +101,18 @@ impl OutboundSink { /// ignored on the resume path (the server keeps that state from creation). pub(super) async fn send_open( &self, - kind: &TransportKind, + kind: &GrpcConnectionMode, create: CreateIngestStreamRequest, ) -> ZerobusResult<()> { match (self, kind) { - (OutboundSink::Ephemeral(tx), TransportKind::Ephemeral) => tx + (OutboundSink::Ephemeral(tx), GrpcConnectionMode::Ephemeral) => tx .send(EphemeralStreamRequest { payload: Some(EphemeralRequestPayload::CreateStream(create)), }) .await .map_err(|_| Self::open_failed()), #[cfg(feature = "eos")] - (OutboundSink::Persistent(tx), TransportKind::Persistent { resume_stream_id }) => { + (OutboundSink::Persistent(tx), GrpcConnectionMode::Persistent { resume_stream_id }) => { let payload = match resume_stream_id { None => PersistentRequestPayload::CreateStream(CreatePersistentStreamRequest { create_stream: Some(create), @@ -305,9 +304,9 @@ impl InboundStream { /// /// Split from the RPC call so the connection module can attach metadata to the /// request before dispatching. -pub(super) fn make_outbound(kind: &TransportKind) -> (OutboundSink, OutboundRequestStream) { +pub(super) fn make_outbound(kind: &GrpcConnectionMode) -> (OutboundSink, OutboundRequestStream) { match kind { - TransportKind::Ephemeral => { + GrpcConnectionMode::Ephemeral => { let (tx, rx) = tokio::sync::mpsc::channel(CHANNEL_BUFFER_SIZE); ( OutboundSink::Ephemeral(tx), @@ -315,7 +314,7 @@ pub(super) fn make_outbound(kind: &TransportKind) -> (OutboundSink, OutboundRequ ) } #[cfg(feature = "eos")] - TransportKind::Persistent { .. } => { + GrpcConnectionMode::Persistent { .. } => { let (tx, rx) = tokio::sync::mpsc::channel(CHANNEL_BUFFER_SIZE); ( OutboundSink::Persistent(tx), From 4a23f77facd36106995a51ef83b8eecf5957b048 Mon Sep 17 00:00:00 2001 From: elenagaljak-db Date: Wed, 26 Aug 2026 15:38:37 +0000 Subject: [PATCH 5/7] [Rust] Simplify persistent recovery Signed-off-by: elenagaljak-db --- rust/sdk/src/landing_zone.rs | 51 +++-- rust/sdk/src/record_types.rs | 3 - rust/sdk/src/stream/grpc/connection.rs | 42 ++--- rust/sdk/src/stream/grpc/mod.rs | 2 - rust/sdk/src/stream/grpc/sender.rs | 15 +- rust/sdk/src/stream/grpc/supervisor.rs | 31 +-- rust/sdk/src/stream/grpc/transport.rs | 250 +++++++++++++++---------- 7 files changed, 224 insertions(+), 170 deletions(-) diff --git a/rust/sdk/src/landing_zone.rs b/rust/sdk/src/landing_zone.rs index e80af7f9..6c369773 100644 --- a/rust/sdk/src/landing_zone.rs +++ b/rust/sdk/src/landing_zone.rs @@ -85,25 +85,22 @@ impl LandingZone { all_items } - /// Removes and returns items from the front of the unobserved queue while - /// `pred` holds, stopping at the first item that fails it. + /// Removes a prefix of observed items while `should_remove` returns true. /// - /// Used by the persistent-stream resume path to reconcile the retained tail - /// against the server's committed watermark: it removes the prefix of - /// records the server has already durably stored (offset <= resume watermark) - /// before re-sending the rest. Only the unobserved queue is consulted: - /// callers reset observation (`reset_observe`) first, so every retained item - /// lives in the queue in offset order. One semaphore permit is released per - /// removed item, mirroring `remove_observed`. - pub fn remove_front_while bool>(&self, mut pred: F) -> Vec { + /// Observed items stay in FIFO order. Stops at the first item that should + /// remain. Releases the corresponding backpressure permits. + pub fn remove_observed_prefix(&self, mut should_remove: impl FnMut(&T) -> bool) -> Vec { let mut state = self.state.lock().expect("Lock poisoned"); let mut permits = self.permits.lock().expect("Lock poisoned"); let mut removed = Vec::new(); - while let Some(front) = state.queue.front() { - if !pred(front) { + while let Some(front) = state.observed_items.front() { + if !should_remove(front) { break; } - let item = state.queue.pop_front().expect("front just checked"); + let item = state + .observed_items + .pop_front() + .expect("front existed before pop"); permits.pop_front(); removed.push(item); } @@ -307,6 +304,34 @@ mod tests { )); } + #[tokio::test] + async fn test_remove_observed_prefix_preserves_fifo_suffix() { + let lz = LandingZone::new(4); + for value in 1..=4 { + lz.add(value).await; + } + for _ in 0..3 { + lz.observe().await; + } + + assert_eq!(lz.remove_observed_prefix(|value| *value <= 2), vec![1, 2]); + assert_eq!(lz.reset_observe(), 1); + assert_eq!(lz.observe().await, 3); + assert_eq!(lz.observe().await, 4); + } + + #[tokio::test] + async fn test_remove_observed_prefix_never_removes_unsent_items() { + let lz = LandingZone::new(2); + lz.add(1).await; + lz.add(2).await; + lz.observe().await; + + assert_eq!(lz.remove_observed_prefix(|_| true), vec![1]); + assert_eq!(lz.len(), 1); + assert_eq!(lz.observe().await, 2); + } + #[tokio::test] async fn test_remove_all() { let lz = Arc::new(LandingZone::new(10)); diff --git a/rust/sdk/src/record_types.rs b/rust/sdk/src/record_types.rs index 4937f99d..d6cf1ec3 100644 --- a/rust/sdk/src/record_types.rs +++ b/rust/sdk/src/record_types.rs @@ -20,7 +20,6 @@ use crate::databricks::zerobus::{ }; use crate::OffsetId; -#[cfg(feature = "eos")] use crate::databricks::zerobus::persistent_stream_request::Payload as PersistentRequestPayload; /// A type alias for a protobuf-encoded record. @@ -260,7 +259,6 @@ impl EncodedBatch { } } - #[cfg(feature = "eos")] pub(crate) fn into_persistent_request_payload( self, offset_id: OffsetId, @@ -798,7 +796,6 @@ mod tests { } } - #[cfg(feature = "eos")] #[test] fn test_into_persistent_request_payload() { let record = r#"{"id": 1}"#.to_string(); diff --git a/rust/sdk/src/stream/grpc/connection.rs b/rust/sdk/src/stream/grpc/connection.rs index 892b4d5d..2e7a58bb 100644 --- a/rust/sdk/src/stream/grpc/connection.rs +++ b/rust/sdk/src/stream/grpc/connection.rs @@ -3,7 +3,7 @@ //! This module is transport-specific: it opens the bidirectional gRPC stream //! used by `ZerobusStream`. It handles both stream kinds through the transport //! seam (`super::transport`): ephemeral streams over the `EphemeralStream` RPC -//! and, behind the `eos` feature, persistent streams over `PersistentStream`. +//! and persistent streams over `PersistentStream`. //! The Arrow Flight transport has its own equivalent under `stream/arrow/`. use std::sync::Arc; @@ -15,9 +15,11 @@ use tonic::transport::Channel; use tracing::{debug, error, info, instrument, warn}; use super::supervisor::StreamInitInfo; -use super::transport::{self, GrpcConnectionMode, InboundStream, Opened, OutboundSink}; +use super::transport::{ + self, GrpcConnectionMode, InboundStream, Opened, OutboundSink, StreamOpenParams, +}; use crate::databricks::zerobus::zerobus_client::ZerobusClient; -use crate::databricks::zerobus::{CreateIngestStreamRequest, RecordType}; +use crate::databricks::zerobus::RecordType; use crate::{HeadersProvider, TableProperties, ZerobusError, ZerobusResult}; /// A freshly opened stream connection: the outbound sink, the inbound response @@ -119,10 +121,8 @@ impl super::ZerobusStream { record_type: RecordType, kind: &GrpcConnectionMode, ) -> ZerobusResult { - let (sink, request_body) = transport::make_outbound(kind); - let mut request = tonic::Request::new(request_body); - - let stream_metadata = request.metadata_mut(); + let outbound = transport::make_outbound(kind); + let mut stream_metadata = tonic::metadata::MetadataMap::new(); let headers = headers_provider.get_headers().await?; for (key, value) in headers { match key { @@ -149,12 +149,13 @@ impl super::ZerobusStream { } } - let mut inbound = transport::open_rpc(&mut channel, request).await?; + let (sink, mut inbound) = + transport::open_rpc(outbound, &mut channel, stream_metadata).await?; - let create_request = Self::build_create_request(table_properties, record_type)?; + let open_params = Self::build_open_params(table_properties, record_type)?; debug!("Sending stream-open request."); - sink.send_open(kind, create_request).await.map_err(|_| { + sink.send_open(open_params).await.map_err(|_| { error!(table_name = %table_properties.table_name, "Failed to send stream-open request"); ZerobusError::StreamClosedError(tonic::Status::internal( "Failed to send stream-open request", @@ -184,17 +185,14 @@ impl super::ZerobusStream { stream_id, last_committed_offset: None, }), - #[cfg(feature = "eos")] _ => Err(Self::mismatched_open_response()), }, - #[cfg(feature = "eos")] GrpcConnectionMode::Persistent { resume_stream_id } => { Self::validate_persistent_open_response(resume_stream_id.as_deref(), opened) } } } - #[cfg(feature = "eos")] fn validate_persistent_open_response( resume_stream_id: Option<&str>, opened: Opened, @@ -217,20 +215,18 @@ impl super::ZerobusStream { } } - #[cfg(feature = "eos")] fn mismatched_open_response() -> ZerobusError { ZerobusError::UnexpectedStreamResponseError( "Persistent stream setup response did not match the requested operation".to_string(), ) } - /// Builds the `CreateIngestStreamRequest` sent when opening (or, on the - /// persistent path, creating) a stream. Encodes the descriptor for proto - /// streams and validates its presence. - fn build_create_request( + /// Resolves the schema inputs used to construct either a create or resume + /// opening message. Encodes and validates the descriptor for proto streams. + fn build_open_params( table_properties: &TableProperties, record_type: RecordType, - ) -> ZerobusResult { + ) -> ZerobusResult { let descriptor_proto = if record_type == RecordType::Proto { Some( table_properties @@ -247,15 +243,15 @@ impl super::ZerobusStream { None }; - Ok(CreateIngestStreamRequest { - table_name: Some(table_properties.table_name.to_string()), + Ok(StreamOpenParams { + table_name: table_properties.table_name.to_string(), descriptor_proto, - record_type: Some(record_type.into()), + record_type, }) } } -#[cfg(all(test, feature = "eos"))] +#[cfg(test)] mod tests { use super::*; use crate::ZerobusStream; diff --git a/rust/sdk/src/stream/grpc/mod.rs b/rust/sdk/src/stream/grpc/mod.rs index b996b150..ed959ac8 100644 --- a/rust/sdk/src/stream/grpc/mod.rs +++ b/rust/sdk/src/stream/grpc/mod.rs @@ -101,7 +101,6 @@ pub struct ZerobusStream { /// from the resume response). `None` for ephemeral streams and for a /// freshly created persistent stream (nothing committed yet). Only read /// through the `eos`-gated accessor. - #[cfg_attr(not(feature = "eos"), allow(dead_code))] pub(crate) last_committed_offset: Option, /// Gets headers which are used in the first request to establish connection with the server. pub headers_provider: Arc, @@ -170,7 +169,6 @@ impl ZerobusStream { /// reconnects to an existing persistent stream, reseeds its offset generator /// to continue after the server's committed offset, and re-sends only the /// records the server has not yet durably stored. - #[cfg(feature = "eos")] #[instrument(level = "debug", skip_all)] pub(crate) async fn new_persistent_stream( channel: ZerobusClient, diff --git a/rust/sdk/src/stream/grpc/sender.rs b/rust/sdk/src/stream/grpc/sender.rs index 7dd8b50c..ad89c4fb 100644 --- a/rust/sdk/src/stream/grpc/sender.rs +++ b/rust/sdk/src/stream/grpc/sender.rs @@ -13,26 +13,19 @@ use tracing::error; use super::transport::OutboundSink; use super::types::RecordLandingZone; use super::ZerobusStream; -use crate::offset_generator::OffsetIdGenerator; use crate::{ZerobusError, ZerobusResult}; impl ZerobusStream { /// Spawns a task that continuously sends records to the Zerobus API by observing the landing zone /// to get records and sending them through the outbound stream to the gRPC stream. - /// - /// Persistent SDK and wire offsets match exactly. Ephemeral streams use a - /// fresh 0-based wire sequence for each server stream so their SDK offsets - /// can preserve continuity across recovery. pub(super) fn spawn_sender_task( sink: OutboundSink, landing_zone: RecordLandingZone, is_paused: Arc, server_error_tx: tokio::sync::watch::Sender>, cancellation_token: CancellationToken, - wire_offsets_match: bool, ) -> tokio::task::JoinHandle> { tokio::spawn(async move { - let physical_offset_id_generator = OffsetIdGenerator::default(); loop { let item = tokio::select! { biased; @@ -45,13 +38,7 @@ impl ZerobusStream { } } => item.clone(), }; - let wire_offset = if wire_offsets_match { - item.offset_id - } else { - physical_offset_id_generator.next() - }; - - let send_result = sink.send_ingest(item.payload, wire_offset).await; + let send_result = sink.send_ingest(item.payload, item.offset_id).await; if let Err(err) = send_result { error!("Failed to send record: {}", err); diff --git a/rust/sdk/src/stream/grpc/supervisor.rs b/rust/sdk/src/stream/grpc/supervisor.rs index 97c667f7..714c83e4 100644 --- a/rust/sdk/src/stream/grpc/supervisor.rs +++ b/rust/sdk/src/stream/grpc/supervisor.rs @@ -54,7 +54,7 @@ impl ZerobusStream { options: StreamConfigurationOptions, // Mutated only on the persistent recovery path (flip create → resume); // that mutation is `eos`-gated, so `mut` is otherwise unused. - #[cfg_attr(not(feature = "eos"), allow(unused_mut))] mut kind: GrpcConnectionMode, + mut kind: GrpcConnectionMode, landing_zone: RecordLandingZone, oneshot_map: Arc>, logical_last_received_offset_id_tx: tokio::sync::watch::Sender>, @@ -65,7 +65,6 @@ impl ZerobusStream { cancellation_token: CancellationToken, callback_tx: Option>, ) -> ZerobusResult<()> { - let wire_offsets_match = kind.wire_offsets_match(); let mut initial_stream_creation = true; let mut stream_init_result_tx = Some(stream_init_result_tx); // One-shot budget: initial setup may spend a single recovery retry to refresh a @@ -215,7 +214,6 @@ impl ZerobusStream { initial_stream_creation = false; // A persistent stream recovers by resuming its now-known id, not // by creating a fresh stream on every reconnect. - #[cfg(feature = "eos")] if let GrpcConnectionMode::Persistent { resume_stream_id } = &mut kind { *resume_stream_id = Some(stream_id.clone()); } @@ -225,17 +223,15 @@ impl ZerobusStream { let _ = server_error_tx.send(None); } - // 2. Reset landing zone. - let resent_records = landing_zone_recovery.observed_count(); - let resent_batches = landing_zone_recovery.reset_observe(); - - // A dropped connection can lose ACKs for records the server already - // committed. Resolve retained records through the resume watermark - // locally, then resend only records above it. + // A dropped persistent connection can lose ACKs for records the + // server already committed. During automatic recovery only, resolve + // the committed observed prefix locally before resetting the + // remaining observed records for resend. A first-process resume has + // an empty landing zone and only seeds the receiver watermark. let mut initial_last_acked_offset: OffsetId = -1; - if wire_offsets_match { - if let Some(watermark) = last_committed_offset { - initial_last_acked_offset = watermark; + if let Some(watermark) = last_committed_offset { + initial_last_acked_offset = watermark; + if !is_initial { Self::reconcile_committed_on_resume( &landing_zone_recovery, &oneshot_map, @@ -247,6 +243,12 @@ impl ZerobusStream { } } + // Move only the still-unacknowledged observed records back for + // resend. Capture counts after reconciliation so the log describes + // what will actually be sent again. + let resent_records = landing_zone_recovery.observed_count(); + let resent_batches = landing_zone_recovery.reset_observe(); + if resent_batches > 0 { info!( stream_id = %stream_id, @@ -284,7 +286,6 @@ impl ZerobusStream { Arc::clone(&is_paused), server_error_tx.clone(), per_stream_token.clone(), - wire_offsets_match, ); // 4. Wait for any of the two tasks to end. @@ -386,7 +387,7 @@ impl ZerobusStream { callback_tx: &Option>, watermark: OffsetId, ) { - let committed = landing_zone.remove_front_while(|item| item.offset_id <= watermark); + let committed = landing_zone.remove_observed_prefix(|item| item.offset_id <= watermark); if committed.is_empty() { return; } diff --git a/rust/sdk/src/stream/grpc/transport.rs b/rust/sdk/src/stream/grpc/transport.rs index 747a12a5..9f32dcb1 100644 --- a/rust/sdk/src/stream/grpc/transport.rs +++ b/rust/sdk/src/stream/grpc/transport.rs @@ -16,10 +16,7 @@ //! Everything else — the landing zone, backpressure, ack tracking, flush, //! close, callbacks, and the create → spawn → recover supervisor loop — is //! identical and stays transport-agnostic. This module normalizes the three -//! differences behind two small enums (`OutboundSink`, `InboundStream`) plus a -//! neutral `InboundMessage`, so the sender/receiver tasks never mention a -//! concrete RPC. The enums carry no `dyn` and compile away to the ephemeral -//! arm when the `eos` feature is off. +//! differences so the sender/receiver tasks never mention a concrete RPC. use tokio_stream::wrappers::ReceiverStream; use tonic::transport::Channel; @@ -29,17 +26,13 @@ use crate::databricks::zerobus::ephemeral_stream_response::Payload as EphemeralR use crate::databricks::zerobus::zerobus_client::ZerobusClient; use crate::databricks::zerobus::{ CloseStreamSignal, CreateIngestStreamRequest, EphemeralStreamRequest, EphemeralStreamResponse, - IngestRecordResponse, + IngestRecordResponse, RecordType, }; -use crate::{EncodedBatch, OffsetId, ZerobusError, ZerobusResult}; +use crate::{EncodedBatch, OffsetId, OffsetIdGenerator, ZerobusError, ZerobusResult}; -#[cfg(feature = "eos")] use crate::databricks::zerobus::persistent_stream_request::Payload as PersistentRequestPayload; -#[cfg(feature = "eos")] use crate::databricks::zerobus::persistent_stream_response::Payload as PersistentResponsePayload; -#[cfg(feature = "eos")] use crate::databricks::zerobus::resume_ingest_stream_request::Identifier as ResumeIdentifier; -#[cfg(feature = "eos")] use crate::databricks::zerobus::{ CreatePersistentStreamRequest, PersistentStreamRequest, PersistentStreamResponse, ResumeIngestStreamRequest, @@ -58,74 +51,79 @@ pub(super) enum GrpcConnectionMode { /// Persistent (Eos) stream: `PersistentStream` RPC, durable wire offsets. /// `resume_stream_id` is `Some` when reconnecting to an existing stream and /// `None` when creating a new one. - #[cfg(feature = "eos")] Persistent { resume_stream_id: Option }, } -impl GrpcConnectionMode { - /// Whether SDK and wire offsets match. They match for persistent streams; - /// ephemeral streams use a fresh wire sequence for each server stream. - pub(super) fn wire_offsets_match(&self) -> bool { - match self { - GrpcConnectionMode::Ephemeral => false, - #[cfg(feature = "eos")] - GrpcConnectionMode::Persistent { .. } => true, - } - } -} - /// Outcome of opening a connection, preserving which operation the server /// acknowledged so the caller can reject mismatched setup responses. pub(super) enum Opened { Created { stream_id: String, }, - #[cfg(feature = "eos")] Resumed { last_committed_offset: Option, }, } +/// Schema and destination inputs needed to construct the protocol-specific +/// opening message. +pub(super) struct StreamOpenParams { + pub(super) table_name: String, + pub(super) descriptor_proto: Option>, + pub(super) record_type: RecordType, +} + +impl StreamOpenParams { + fn into_create_request(self) -> CreateIngestStreamRequest { + CreateIngestStreamRequest { + table_name: Some(self.table_name), + descriptor_proto: self.descriptor_proto, + record_type: Some(self.record_type.into()), + } + } +} + /// The outbound half of a stream: wraps the concrete tonic request sender and /// exposes a single neutral `send` that both IO tasks use. pub(super) enum OutboundSink { - Ephemeral(tokio::sync::mpsc::Sender), - #[cfg(feature = "eos")] - Persistent(tokio::sync::mpsc::Sender), + Ephemeral { + tx: tokio::sync::mpsc::Sender, + wire_offsets: OffsetIdGenerator, + }, + Persistent { + tx: tokio::sync::mpsc::Sender, + resume_stream_id: Option, + }, } impl OutboundSink { /// Sends the opening message: `create_stream` for a new stream (ephemeral /// or persistent), or `resume_stream` when reconnecting to a persistent - /// one. `create` carries the table name, descriptor, and record type; it is - /// ignored on the resume path (the server keeps that state from creation). - pub(super) async fn send_open( - &self, - kind: &GrpcConnectionMode, - create: CreateIngestStreamRequest, - ) -> ZerobusResult<()> { - match (self, kind) { - (OutboundSink::Ephemeral(tx), GrpcConnectionMode::Ephemeral) => tx + /// one. Create carries the destination table; resume carries only the + /// descriptor and record type needed to validate the reopened stream. + pub(super) async fn send_open(&self, params: StreamOpenParams) -> ZerobusResult<()> { + match self { + OutboundSink::Ephemeral { tx, .. } => tx .send(EphemeralStreamRequest { - payload: Some(EphemeralRequestPayload::CreateStream(create)), + payload: Some(EphemeralRequestPayload::CreateStream( + params.into_create_request(), + )), }) .await .map_err(|_| Self::open_failed()), - #[cfg(feature = "eos")] - (OutboundSink::Persistent(tx), GrpcConnectionMode::Persistent { resume_stream_id }) => { + OutboundSink::Persistent { + tx, + resume_stream_id, + } => { let payload = match resume_stream_id { None => PersistentRequestPayload::CreateStream(CreatePersistentStreamRequest { - create_stream: Some(create), + create_stream: Some(params.into_create_request()), }), Some(stream_id) => { - // The descriptor is fixed at creation but re-sent on - // resume so the server can re-validate it against the - // table schema (required for PROTO streams; absent for - // JSON / Arrow, matching `create`). PersistentRequestPayload::ResumeStream(ResumeIngestStreamRequest { identifier: Some(ResumeIdentifier::StreamId(stream_id.clone())), - descriptor_proto: create.descriptor_proto, - record_type: create.record_type, + descriptor_proto: params.descriptor_proto, + record_type: Some(params.record_type.into()), }) } }; @@ -135,11 +133,6 @@ impl OutboundSink { .await .map_err(|_| Self::open_failed()) } - // The sink and kind are always constructed together by - // `make_outbound`, so a mismatch is unreachable. Only present when - // more than one variant exists (i.e. `eos` is enabled). - #[cfg(feature = "eos")] - _ => Err(Self::open_failed()), } } @@ -149,24 +142,26 @@ impl OutboundSink { )) } - /// Sends one ingest batch on the wire, numbered with `offset_id`. + /// Sends one ingest batch, mapping the SDK offset to the wire sequence for + /// this connection. Persistent offsets pass through unchanged; ephemeral + /// connections use a fresh zero-based sequence. pub(super) async fn send_ingest( &self, batch: EncodedBatch, - offset_id: OffsetId, + sdk_offset: OffsetId, ) -> ZerobusResult<()> { + let wire_offset = self.next_wire_offset(sdk_offset); match self { - OutboundSink::Ephemeral(tx) => { - let payload = batch.into_request_payload(offset_id); + OutboundSink::Ephemeral { tx, .. } => { + let payload = batch.into_request_payload(wire_offset); tx.send(EphemeralStreamRequest { payload: Some(payload), }) .await .map_err(|_| Self::ingest_failed()) } - #[cfg(feature = "eos")] - OutboundSink::Persistent(tx) => { - let payload = batch.into_persistent_request_payload(offset_id); + OutboundSink::Persistent { tx, .. } => { + let payload = batch.into_persistent_request_payload(wire_offset); tx.send(PersistentStreamRequest { payload: Some(payload), }) @@ -179,6 +174,13 @@ impl OutboundSink { fn ingest_failed() -> ZerobusError { ZerobusError::StreamClosedError(tonic::Status::internal("Failed to send batch")) } + + fn next_wire_offset(&self, sdk_offset: OffsetId) -> OffsetId { + match self { + OutboundSink::Ephemeral { wire_offsets, .. } => wire_offsets.next(), + OutboundSink::Persistent { .. } => sdk_offset, + } + } } /// A server message, normalized across the two response envelopes. Both IO @@ -197,7 +199,6 @@ pub(super) enum InboundMessage { /// yields normalized `InboundMessage`s. pub(super) enum InboundStream { Ephemeral(tonic::Streaming), - #[cfg(feature = "eos")] Persistent(tonic::Streaming), } @@ -215,7 +216,6 @@ impl InboundStream { } _ => InboundMessage::Other, })), - #[cfg(feature = "eos")] InboundStream::Persistent(s) => Ok(s.message().await?.map(|resp| match resp.payload { Some(PersistentResponsePayload::IngestRecordResponse(ack)) => { InboundMessage::Ack(ack) @@ -247,7 +247,6 @@ impl InboundStream { other => Err(Self::unexpected_open(&other)), } } - #[cfg(feature = "eos")] InboundStream::Persistent(s) => { let msg = Self::first_message(s.message().await)?; match msg.payload { @@ -293,70 +292,121 @@ impl InboundStream { pub(super) async fn drain(&mut self) { match self { InboundStream::Ephemeral(s) => while matches!(s.message().await, Ok(Some(_))) {}, - #[cfg(feature = "eos")] InboundStream::Persistent(s) => while matches!(s.message().await, Ok(Some(_))) {}, } } } -/// Opens the outbound request channel for the given transport kind, returning -/// the neutral sink and the raw tonic request stream to hand to the RPC. -/// -/// Split from the RPC call so the connection module can attach metadata to the -/// request before dispatching. -pub(super) fn make_outbound(kind: &GrpcConnectionMode) -> (OutboundSink, OutboundRequestStream) { +/// Both halves of a newly allocated outbound channel. Keeping them in the same +/// variant guarantees the request stream, retained sender, and opening mode +/// cannot disagree. +pub(super) enum OutboundConnection { + Ephemeral { + tx: tokio::sync::mpsc::Sender, + requests: ReceiverStream, + }, + Persistent { + tx: tokio::sync::mpsc::Sender, + requests: ReceiverStream, + resume_stream_id: Option, + }, +} + +pub(super) fn make_outbound(kind: &GrpcConnectionMode) -> OutboundConnection { match kind { GrpcConnectionMode::Ephemeral => { let (tx, rx) = tokio::sync::mpsc::channel(CHANNEL_BUFFER_SIZE); - ( - OutboundSink::Ephemeral(tx), - OutboundRequestStream::Ephemeral(ReceiverStream::new(rx)), - ) + OutboundConnection::Ephemeral { + tx, + requests: ReceiverStream::new(rx), + } } - #[cfg(feature = "eos")] - GrpcConnectionMode::Persistent { .. } => { + GrpcConnectionMode::Persistent { resume_stream_id } => { let (tx, rx) = tokio::sync::mpsc::channel(CHANNEL_BUFFER_SIZE); - ( - OutboundSink::Persistent(tx), - OutboundRequestStream::Persistent(ReceiverStream::new(rx)), - ) + OutboundConnection::Persistent { + tx, + requests: ReceiverStream::new(rx), + resume_stream_id: resume_stream_id.clone(), + } } } } -/// The raw request stream handed to the tonic RPC method. Kept concrete because -/// `ZerobusClient::ephemeral_stream` / `persistent_stream` want the exact type. -pub(super) enum OutboundRequestStream { - Ephemeral(ReceiverStream), - #[cfg(feature = "eos")] - Persistent(ReceiverStream), -} - -/// Dispatches the opening RPC on `channel` with `request` (metadata already -/// attached) and returns the inbound stream. The first server message is read -/// by the caller (`connection.rs`) so it can extract the stream_id / watermark. +/// Dispatches the opening RPC with metadata prepared by the caller and returns +/// both normalized connection halves. pub(super) async fn open_rpc( + outbound: OutboundConnection, channel: &mut ZerobusClient, - request: tonic::Request, -) -> ZerobusResult { - let (metadata, extensions, body) = request.into_parts(); - match body { - OutboundRequestStream::Ephemeral(stream) => { - let req = tonic::Request::from_parts(metadata, extensions, stream); + metadata: tonic::metadata::MetadataMap, +) -> ZerobusResult<(OutboundSink, InboundStream)> { + match outbound { + OutboundConnection::Ephemeral { tx, requests } => { + let mut req = tonic::Request::new(requests); + *req.metadata_mut() = metadata; let resp = channel .ephemeral_stream(req) .await .map_err(ZerobusError::CreateStreamError)?; - Ok(InboundStream::Ephemeral(resp.into_inner())) + Ok(( + OutboundSink::Ephemeral { + tx, + wire_offsets: OffsetIdGenerator::default(), + }, + InboundStream::Ephemeral(resp.into_inner()), + )) } - #[cfg(feature = "eos")] - OutboundRequestStream::Persistent(stream) => { - let req = tonic::Request::from_parts(metadata, extensions, stream); + OutboundConnection::Persistent { + tx, + requests, + resume_stream_id, + } => { + let mut req = tonic::Request::new(requests); + *req.metadata_mut() = metadata; let resp = channel .persistent_stream(req) .await .map_err(ZerobusError::CreateStreamError)?; - Ok(InboundStream::Persistent(resp.into_inner())) + Ok(( + OutboundSink::Persistent { + tx, + resume_stream_id, + }, + InboundStream::Persistent(resp.into_inner()), + )) } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ephemeral_wire_offsets_restart_for_each_connection() { + let (tx, _rx) = tokio::sync::mpsc::channel(CHANNEL_BUFFER_SIZE); + let sink = OutboundSink::Ephemeral { + tx, + wire_offsets: OffsetIdGenerator::default(), + }; + assert_eq!(sink.next_wire_offset(40), 0); + assert_eq!(sink.next_wire_offset(41), 1); + + let (tx, _rx) = tokio::sync::mpsc::channel(CHANNEL_BUFFER_SIZE); + let recovered = OutboundSink::Ephemeral { + tx, + wire_offsets: OffsetIdGenerator::default(), + }; + assert_eq!(recovered.next_wire_offset(42), 0); + } + + #[test] + fn persistent_wire_offsets_match_sdk_offsets() { + let (tx, _rx) = tokio::sync::mpsc::channel(CHANNEL_BUFFER_SIZE); + let sink = OutboundSink::Persistent { + tx, + resume_stream_id: Some("stream-id".to_string()), + }; + assert_eq!(sink.next_wire_offset(40), 40); + assert_eq!(sink.next_wire_offset(41), 41); + } +} From 9e4348e95c1b096a5cc01cf9a4f81ba6c6599e63 Mon Sep 17 00:00:00 2001 From: elenagaljak-db Date: Fri, 28 Aug 2026 16:34:58 +0000 Subject: [PATCH 6/7] [Rust] Note persistent hardening gaps Signed-off-by: elenagaljak-db --- rust/sdk/Cargo.toml | 2 +- rust/sdk/src/stream/grpc/mod.rs | 2 ++ rust/sdk/src/stream/grpc/receiver.rs | 3 +++ rust/sdk/src/stream/grpc/supervisor.rs | 5 +++-- 4 files changed, 9 insertions(+), 3 deletions(-) diff --git a/rust/sdk/Cargo.toml b/rust/sdk/Cargo.toml index 9c38359d..4bb89e27 100644 --- a/rust/sdk/Cargo.toml +++ b/rust/sdk/Cargo.toml @@ -77,7 +77,7 @@ internal-arrow-c-data = [ ] # Zero-copy protobuf parser. zeroparser = ["dep:self_cell", "dep:prost-build"] -# Persistent streams with exactly-once semantics (EoS); API is in development +# Persistent streams with exactly-once semantics (EoS); API is in development. eos = [] testing = ["dep:futures"] # Test-only deterministic seams (barriers/notifies) for the Arrow stream. Zero footprint diff --git a/rust/sdk/src/stream/grpc/mod.rs b/rust/sdk/src/stream/grpc/mod.rs index ed959ac8..e80d79bb 100644 --- a/rust/sdk/src/stream/grpc/mod.rs +++ b/rust/sdk/src/stream/grpc/mod.rs @@ -262,6 +262,8 @@ impl ZerobusStream { // since the server never re-acks offsets it committed in a prior session. // Safe because the constructor returns before any user ingest, so no real // ack can race this initial value. + // TODO: Validate the resume watermark before spawning tasks, or cancel and abort + // the spawned tasks on post-initialization errors so their handles are not detached. if let Some(watermark) = init_info.last_committed_offset { let next_offset = watermark.checked_add(1).ok_or_else(|| { ZerobusError::UnexpectedStreamResponseError( diff --git a/rust/sdk/src/stream/grpc/receiver.rs b/rust/sdk/src/stream/grpc/receiver.rs index 7d9eb1be..2431bb92 100644 --- a/rust/sdk/src/stream/grpc/receiver.rs +++ b/rust/sdk/src/stream/grpc/receiver.rs @@ -112,6 +112,9 @@ impl ZerobusStream { return Err(error); } }; + // TODO: Bound ACKs by the highest successfully sent wire offset, not + // landing-zone counts; batches use one offset and observation precedes + // a successful send. if let Err(error) = Self::validate_ack_offset( last_acked_offset, durability_ack_up_to_offset, diff --git a/rust/sdk/src/stream/grpc/supervisor.rs b/rust/sdk/src/stream/grpc/supervisor.rs index 714c83e4..bf43b159 100644 --- a/rust/sdk/src/stream/grpc/supervisor.rs +++ b/rust/sdk/src/stream/grpc/supervisor.rs @@ -52,8 +52,7 @@ impl ZerobusStream { table_properties: TableProperties, headers_provider: Arc, options: StreamConfigurationOptions, - // Mutated only on the persistent recovery path (flip create → resume); - // that mutation is `eos`-gated, so `mut` is otherwise unused. + // Mutated on the persistent recovery path to flip create → resume. mut kind: GrpcConnectionMode, landing_zone: RecordLandingZone, oneshot_map: Arc>, @@ -228,6 +227,8 @@ impl ZerobusStream { // the committed observed prefix locally before resetting the // remaining observed records for resend. A first-process resume has // an empty landing zone and only seeds the receiver watermark. + // TODO: Validate the resumed watermark against retained ACK and sent-offset + // bounds before reconciling; treat None as -1 and fail closed if out of range. let mut initial_last_acked_offset: OffsetId = -1; if let Some(watermark) = last_committed_offset { initial_last_acked_offset = watermark; From 7a5f0265fc660ef67aaa64ac6b1f51f654b05bc2 Mon Sep 17 00:00:00 2001 From: elenagaljak-db Date: Fri, 28 Aug 2026 17:15:20 +0000 Subject: [PATCH 7/7] [Rust] Allow staged persistent internals Signed-off-by: elenagaljak-db --- rust/sdk/src/stream/grpc/mod.rs | 4 ++++ rust/sdk/src/stream/grpc/transport.rs | 2 ++ 2 files changed, 6 insertions(+) diff --git a/rust/sdk/src/stream/grpc/mod.rs b/rust/sdk/src/stream/grpc/mod.rs index e80d79bb..55b11ffc 100644 --- a/rust/sdk/src/stream/grpc/mod.rs +++ b/rust/sdk/src/stream/grpc/mod.rs @@ -101,6 +101,8 @@ pub struct ZerobusStream { /// from the resume response). `None` for ephemeral streams and for a /// freshly created persistent stream (nothing committed yet). Only read /// through the `eos`-gated accessor. + // Used by the feature-gated public API introduced in the downstream PR. + #[allow(dead_code)] pub(crate) last_committed_offset: Option, /// Gets headers which are used in the first request to establish connection with the server. pub headers_provider: Arc, @@ -169,6 +171,8 @@ impl ZerobusStream { /// reconnects to an existing persistent stream, reseeds its offset generator /// to continue after the server's committed offset, and re-sends only the /// records the server has not yet durably stored. + // Used by the feature-gated public API introduced in the downstream PR. + #[allow(dead_code)] #[instrument(level = "debug", skip_all)] pub(crate) async fn new_persistent_stream( channel: ZerobusClient, diff --git a/rust/sdk/src/stream/grpc/transport.rs b/rust/sdk/src/stream/grpc/transport.rs index 9f32dcb1..85733393 100644 --- a/rust/sdk/src/stream/grpc/transport.rs +++ b/rust/sdk/src/stream/grpc/transport.rs @@ -51,6 +51,8 @@ pub(super) enum GrpcConnectionMode { /// Persistent (Eos) stream: `PersistentStream` RPC, durable wire offsets. /// `resume_stream_id` is `Some` when reconnecting to an existing stream and /// `None` when creating a new one. + // Constructed by the feature-gated public API introduced in the downstream PR. + #[allow(dead_code)] Persistent { resume_stream_id: Option }, }