diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index 6fc696f3..850139f9 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -47,6 +47,10 @@ ### Internal Changes +- 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 ### Deprecations diff --git a/rust/sdk/Cargo.toml b/rust/sdk/Cargo.toml index 1de5b671..4bb89e27 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 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 # 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..6c369773 100644 --- a/rust/sdk/src/landing_zone.rs +++ b/rust/sdk/src/landing_zone.rs @@ -85,6 +85,28 @@ impl LandingZone { all_items } + /// Removes a prefix of observed items while `should_remove` returns true. + /// + /// 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.observed_items.front() { + if !should_remove(front) { + break; + } + let item = state + .observed_items + .pop_front() + .expect("front existed before pop"); + 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, @@ -282,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 e805af37..d6cf1ec3 100644 --- a/rust/sdk/src/record_types.rs +++ b/rust/sdk/src/record_types.rs @@ -20,6 +20,8 @@ use crate::databricks::zerobus::{ }; use crate::OffsetId; +use crate::databricks::zerobus::persistent_stream_request::Payload as PersistentRequestPayload; + /// A type alias for a protobuf-encoded record. pub type ProtoEncodedRecord = Vec; @@ -257,6 +259,21 @@ impl EncodedBatch { } } + 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 +795,20 @@ mod tests { _ => panic!("Expected IngestRecordBatch payload"), } } + + #[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/connection.rs b/rust/sdk/src/stream/grpc/connection.rs index e1b1ed99..2e7a58bb 100644 --- a/rust/sdk/src/stream/grpc/connection.rs +++ b/rust/sdk/src/stream/grpc/connection.rs @@ -1,51 +1,57 @@ //! 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 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 crate::databricks::zerobus::zerobus_client::ZerobusClient; -use crate::databricks::zerobus::{ - CreateIngestStreamRequest, EphemeralStreamRequest, EphemeralStreamResponse, RecordType, +use super::supervisor::StreamInitInfo; +use super::transport::{ + self, GrpcConnectionMode, InboundStream, Opened, OutboundSink, StreamOpenParams, }; +use crate::databricks::zerobus::zerobus_client::ZerobusClient; +use crate::databricks::zerobus::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. +/// 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) init_info: StreamInitInfo, +} + +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: &GrpcConnectionMode, + ) -> ZerobusResult { let result = Self::create_stream_connection_inner( channel, table_properties, headers_provider, record_type, + kind, ) .await; if let Err(err) = &result { @@ -69,12 +75,9 @@ impl ZerobusStream { table_properties: &TableProperties, headers_provider: &Arc, record_type: RecordType, + kind: &GrpcConnectionMode, 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 +87,7 @@ impl ZerobusStream { table_properties, headers_provider, record_type, + kind, ), ) .await @@ -115,18 +119,11 @@ 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(); + kind: &GrpcConnectionMode, + ) -> ZerobusResult { + 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 { "x-databricks-zerobus-table-name" => { @@ -152,12 +149,84 @@ impl ZerobusStream { } } - let mut response_grpc_stream = channel - .ephemeral_stream(request_stream) - .await - .map_err(ZerobusError::CreateStreamError)? - .into_inner(); + let (sink, mut inbound) = + transport::open_rpc(outbound, &mut channel, stream_metadata).await?; + + let open_params = Self::build_open_params(table_properties, record_type)?; + + debug!("Sending stream-open request."); + 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", + )) + })?; + + debug!("Waiting for stream-open response."); + let opened = inbound.recv_open().await?; + + 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, + init_info, + }) + } + + fn validate_open_response( + kind: &GrpcConnectionMode, + opened: Opened, + ) -> ZerobusResult { + match kind { + GrpcConnectionMode::Ephemeral => match opened { + Opened::Created { stream_id } => Ok(StreamInitInfo { + stream_id, + last_committed_offset: None, + }), + _ => Err(Self::mismatched_open_response()), + }, + GrpcConnectionMode::Persistent { resume_stream_id } => { + Self::validate_persistent_open_response(resume_stream_id.as_deref(), opened) + } + } + } + + 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, + }), + ( + Some(stream_id), + Opened::Resumed { + last_committed_offset, + }, + ) => Ok(StreamInitInfo { + stream_id: stream_id.to_string(), + last_committed_offset, + }), + _ => Err(Self::mismatched_open_response()), + } + } + fn mismatched_open_response() -> ZerobusError { + ZerobusError::UnexpectedStreamResponseError( + "Persistent stream setup response did not match the requested operation".to_string(), + ) + } + + /// 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 { let descriptor_proto = if record_type == RecordType::Proto { Some( table_properties @@ -174,56 +243,64 @@ impl ZerobusStream { None }; - let create_stream_request = RequestPayload::CreateStream(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()), - }); - - debug!("Sending CreateStream request."); - tx.send(EphemeralStreamRequest { - payload: Some(create_stream_request), + record_type, }) - .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)) - } - } + } +} + +#[cfg(test)] +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 = GrpcConnectionMode::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 = GrpcConnectionMode::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 8caac1bb..55b11ffc 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::GrpcConnectionMode; use types::{IngestRequest, OneshotMap, RecordLandingZone}; #[cfg(feature = "testing")] @@ -93,6 +96,14 @@ 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. + // 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, /// The stream configuration options related to recovery, fetching OAuth tokens, etc. @@ -141,9 +152,60 @@ impl ZerobusStream { table_properties: TableProperties, headers_provider: Arc, options: StreamConfigurationOptions, + ) -> ZerobusResult { + Self::new_with_kind( + channel, + table_properties, + headers_provider, + options, + StreamType::Ephemeral, + GrpcConnectionMode::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. + // 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, + 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, + GrpcConnectionMode::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: GrpcConnectionMode, ) -> 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 +238,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 +249,44 @@ 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. + // 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( + "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)); + } // 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..2431bb92 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,105 @@ 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 { + 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); + } + }; + // 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, - })) => { - 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); + 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 + { + 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)); } - }; - 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(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), - }; - - if wait_duration_ms == 0 { - info!("Server will close the stream. Triggering immediate recovery."); + 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); + + 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 +227,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; }); @@ -246,4 +242,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/sender.rs b/rust/sdk/src/stream/grpc/sender.rs index 540d1491..ad89c4fb 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,24 +10,22 @@ 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. 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, ) -> tokio::task::JoinHandle> { tokio::spawn(async move { - let physical_offset_id_generator = OffsetIdGenerator::default(); loop { let item = tokio::select! { biased; @@ -39,22 +38,12 @@ impl ZerobusStream { } } => item.clone(), }; - let offset_id = physical_offset_id_generator.next(); - let request_payload = item.payload.into_request_payload(offset_id); - - let send_result = outbound_stream - .send(EphemeralStreamRequest { - payload: Some(request_payload), - }) - .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); - 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..bf43b159 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::GrpcConnectionMode; 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,12 +52,14 @@ impl ZerobusStream { table_properties: TableProperties, headers_provider: Arc, options: StreamConfigurationOptions, + // Mutated on the persistent recovery path to flip create → resume. + mut kind: GrpcConnectionMode, 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>, @@ -84,6 +101,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 = || { @@ -92,6 +112,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 +122,7 @@ impl ZerobusStream { &table_properties, &headers_provider, record_type, + kind, options.recovery_timeout_ms, ) .await @@ -112,6 +134,7 @@ impl ZerobusStream { &table_properties, &headers_provider, record_type, + kind, ), ) .await @@ -148,8 +171,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,20 +199,57 @@ impl ZerobusStream { return Err(e); } }; + 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(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. + if let GrpcConnectionMode::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"); let _ = server_error_tx.send(None); } - // 2. Reset landing zone. + // 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. + // 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; + if !is_initial { + Self::reconcile_committed_on_resume( + &landing_zone_recovery, + &oneshot_map, + &logical_last_received_offset_id_tx, + &callback_tx, + watermark, + ) + .await; + } + } + + // 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, @@ -210,7 +270,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,9 +279,10 @@ 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(), @@ -307,6 +368,55 @@ 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 (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 + /// 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 + /// 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_observed_prefix(|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..85733393 --- /dev/null +++ b/rust/sdk/src/stream/grpc/transport.rs @@ -0,0 +1,414 @@ +//! 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** — 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. +//! +//! 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 so the sender/receiver tasks never mention a concrete RPC. + +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, RecordType, +}; +use crate::{EncodedBatch, OffsetId, OffsetIdGenerator, ZerobusError, ZerobusResult}; + +use crate::databricks::zerobus::persistent_stream_request::Payload as PersistentRequestPayload; +use crate::databricks::zerobus::persistent_stream_response::Payload as PersistentResponsePayload; +use crate::databricks::zerobus::resume_ingest_stream_request::Identifier as ResumeIdentifier; +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 GrpcConnectionMode { + /// 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. + // Constructed by the feature-gated public API introduced in the downstream PR. + #[allow(dead_code)] + Persistent { resume_stream_id: 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, + }, + 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 { + 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 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( + params.into_create_request(), + )), + }) + .await + .map_err(|_| Self::open_failed()), + OutboundSink::Persistent { + tx, + resume_stream_id, + } => { + let payload = match resume_stream_id { + None => PersistentRequestPayload::CreateStream(CreatePersistentStreamRequest { + create_stream: Some(params.into_create_request()), + }), + Some(stream_id) => { + PersistentRequestPayload::ResumeStream(ResumeIngestStreamRequest { + identifier: Some(ResumeIdentifier::StreamId(stream_id.clone())), + descriptor_proto: params.descriptor_proto, + record_type: Some(params.record_type.into()), + }) + } + }; + tx.send(PersistentStreamRequest { + payload: Some(payload), + }) + .await + .map_err(|_| Self::open_failed()) + } + } + } + + fn open_failed() -> ZerobusError { + ZerobusError::StreamClosedError(tonic::Status::internal( + "Failed to send stream-open request", + )) + } + + /// 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, + sdk_offset: OffsetId, + ) -> ZerobusResult<()> { + let wire_offset = self.next_wire_offset(sdk_offset); + match self { + OutboundSink::Ephemeral { tx, .. } => { + let payload = batch.into_request_payload(wire_offset); + tx.send(EphemeralStreamRequest { + payload: Some(payload), + }) + .await + .map_err(|_| Self::ingest_failed()) + } + OutboundSink::Persistent { tx, .. } => { + let payload = batch.into_persistent_request_payload(wire_offset); + tx.send(PersistentStreamRequest { + payload: Some(payload), + }) + .await + .map_err(|_| Self::ingest_failed()) + } + } + } + + 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 +/// 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), + 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, + })), + 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::Created { stream_id }) + } + other => Err(Self::unexpected_open(&other)), + } + } + 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::Created { stream_id }) + } + Some(PersistentResponsePayload::ResumeStreamResponse(resp)) => { + Ok(Opened::Resumed { + 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(_))) {}, + InboundStream::Persistent(s) => while matches!(s.message().await, Ok(Some(_))) {}, + } + } +} + +/// 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); + OutboundConnection::Ephemeral { + tx, + requests: ReceiverStream::new(rx), + } + } + GrpcConnectionMode::Persistent { resume_stream_id } => { + let (tx, rx) = tokio::sync::mpsc::channel(CHANNEL_BUFFER_SIZE); + OutboundConnection::Persistent { + tx, + requests: ReceiverStream::new(rx), + resume_stream_id: resume_stream_id.clone(), + } + } + } +} + +/// 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, + 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(( + OutboundSink::Ephemeral { + tx, + wire_offsets: OffsetIdGenerator::default(), + }, + InboundStream::Ephemeral(resp.into_inner()), + )) + } + 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(( + 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); + } +}