From 6ca543b00660cd9d23e2489280addfef14048648 Mon Sep 17 00:00:00 2001 From: teodordelibasic-db Date: Wed, 12 Aug 2026 13:21:20 +0000 Subject: [PATCH 1/6] Initial commit Signed-off-by: teodordelibasic-db --- rust/NEXT_CHANGELOG.md | 5 +- rust/sdk/src/stream/arrow/acks.rs | 862 +++++++++++++++++++--- rust/sdk/src/stream/arrow/close.rs | 348 +++++++++ rust/sdk/src/stream/arrow/mod.rs | 309 +++----- rust/sdk/src/stream/arrow/options.rs | 9 +- rust/sdk/src/stream/arrow/supervisor.rs | 938 ++++++++++++------------ rust/tests/src/arrow_tests.rs | 575 ++++++++++++--- 7 files changed, 2164 insertions(+), 882 deletions(-) create mode 100644 rust/sdk/src/stream/arrow/close.rs diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index c926e11b..773d63f8 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -18,9 +18,12 @@ the full replay completes and ACK processing can resume on the replacement connection. - Arrow Flight rejects unrepresentable timeout values: stream creation returns - `InvalidArgument` when ACK or recovery deadlines exceed the platform + `InvalidArgument` when ACK, recovery, or flush deadlines exceed the platform monotonic-clock range. Server-advertised graceful-rotation periods are capped at one year. +- Arrow Flight close is cancellation-safe and half-closes the active request before + bounded response draining. Close during recovery cancels the attempt, retains the + unacknowledged suffix, and returns the error that triggered recovery. - Fixed Arrow Flight recovery sender lifetime: replacement senders are now published only after pending replay succeeds, while initial supervisor handoff and failed or cancelled replay promptly drop redundant senders instead of retaining incomplete diff --git a/rust/sdk/src/stream/arrow/acks.rs b/rust/sdk/src/stream/arrow/acks.rs index 7a1f75c5..9b78d914 100644 --- a/rust/sdk/src/stream/arrow/acks.rs +++ b/rust/sdk/src/stream/arrow/acks.rs @@ -9,7 +9,7 @@ use std::sync::Arc; use arrow_flight::error::FlightError; use arrow_flight::PutResult; -use futures::StreamExt; +use futures::{FutureExt, StreamExt}; use tokio::sync::{watch, Mutex, Notify}; use tokio::time::{sleep_until, Duration, Instant}; use tracing::{debug, error, info, warn}; @@ -17,11 +17,12 @@ use tracing::{debug, error, info, warn}; use super::batch::{ oldest_pending_ack_deadline, PendingAckDeadline, PendingBatch, PendingBatchIdentity, }; +use super::close::{CloseCoordinator, CloseRequest, CloseState}; use super::connection::{FlightResponseStream, RequestBodyControl}; use super::metadata::FlightAckMetadata; -#[cfg(feature = "test-hooks")] -use super::AckAppliedGate; use super::{ArrowStreamConfigurationOptions, BatchSender, ZerobusArrowStream}; +#[cfg(feature = "test-hooks")] +use super::{TestHooks, TestNotifyGate}; use crate::errors::ZerobusError; use crate::offset_generator::OffsetId; use crate::ZerobusResult; @@ -43,14 +44,13 @@ pub(super) struct AckProcessor { ingest_mutex: Arc>, batch_tx: BatchSender, options: ArrowStreamConfigurationOptions, + close: CloseCoordinator, #[cfg(feature = "test-hooks")] - ack_applied_gate: AckAppliedGate, - #[cfg(feature = "test-hooks")] - ack_idle_gate: super::AckIdleGate, + test_hooks: Arc, } /// State captured when rotation stops waiting for acknowledgments and begins transport -/// cleanup. The response may already have ended, but the request must still reach EOF. +/// cleanup. The response may already have ended. struct DrainState { /// Hard cutoff shared by request EOF observation and response draining. deadline: Instant, @@ -58,16 +58,41 @@ struct DrainState { response_finished: bool, /// Peer or protocol error to preserve while the remaining transport settles. terminal_error: Option, + /// Explicit close being finalized by this drain, if any. + close: Option, +} + +struct CloseDrain { + request: CloseRequest, + /// The first observed ACK, concrete error, or deadline result. + /// This remains unset only for an empty close target. + selected: Option>, } -/// Server-initiated rotation has only three phases: normal traffic, waiting for the -/// pre-signal acknowledgment snapshot, and transport drain. -enum RotationState { - Open, - WaitingForAcks { +pub(super) enum AckProcessOutcome { + Stopped, + Recovery { + error: ZerobusError, + drained: bool, + }, + Close { + request: CloseRequest, + outcome: ZerobusResult<()>, + }, +} + +enum WaitState { + Rotation { target_records: u64, deadlines: RotationDeadlines, }, + Close(CloseRequest), +} + +/// One connection lifecycle covers active traffic, ACK waiting, and transport drain. +enum ConnectionState { + Active, + Waiting(WaitState), Draining(DrainState), } @@ -84,6 +109,9 @@ enum AckEvent { PendingBatchAvailable, RequestSendFailed, AckDeadline(PendingAckDeadline), + CloseRequested(CloseRequest), + CloseFinalized, + CloseDeadline, } /// Coalesces request-sender failures until the supervisor has paused the failed @@ -119,13 +147,15 @@ struct AckProgress<'a> { last_acked_records: &'a AtomicU64, pending_batches: &'a Mutex>, last_ack_tx: &'a watch::Sender>, + close: &'a CloseCoordinator, #[cfg(feature = "test-hooks")] - ack_applied_gate: &'a AckAppliedGate, + ack_applied_gate: &'a TestNotifyGate, } impl AckProgress<'_> { /// Validates an ACK against the active connection, advances the monotonic durable /// watermark, removes fully acknowledged batches, and wakes completed offset waiters. + /// Its close-deadline time is captured after durable state is applied, before notifications. async fn apply(&self, ack: &FlightAckMetadata) -> ZerobusResult<()> { let acked_records = ack.ack_up_to_records; // `ack_up_to_records` is the durability boundary. Derive completed SDK offsets @@ -160,6 +190,7 @@ impl AckProgress<'_> { }); (effective_acked_records, max_acked_offset) }; + let applied_at = Instant::now(); debug!(target: super::LOG_TARGET, ack_up_to_offset = ack.ack_up_to_offset, @@ -177,6 +208,7 @@ impl AckProgress<'_> { if let Some(offset) = max_acked_offset { let _ = self.last_ack_tx.send(Some(offset)); + self.close.observe_ack(offset, applied_at); } Ok(()) @@ -225,8 +257,8 @@ async fn pause_and_snapshot_submitted( } impl RequestControl<'_> { - /// Atomically stops new sends, detaches queued work, then asks tonic to poll the request - /// body to EOF. [`RequestBodyControl::wait_for_eof`] observes completion separately. + /// Atomically stops new sends, detaches queued work, and shuts down the request body. + /// The caller then drains the response and final status under its bounded deadline. async fn half_close(&self) { pause_and_detach_sender(self.ingest_mutex, self.is_paused, self.batch_tx).await; self.request_body.shutdown(); @@ -248,10 +280,9 @@ impl AckProcessor { ingest_mutex: Arc::clone(&stream.ingest_mutex), batch_tx: Arc::clone(&stream.batch_tx), options: stream.options.clone(), + close: stream.close.clone(), #[cfg(feature = "test-hooks")] - ack_applied_gate: Arc::clone(&stream.ack_applied_gate), - #[cfg(feature = "test-hooks")] - ack_idle_gate: Arc::clone(&stream.ack_idle_gate), + test_hooks: Arc::clone(&stream.test_hooks), } } @@ -280,10 +311,9 @@ impl AckProcessor { server_lack_of_ack_timeout_ms: Duration::from_secs(60).as_millis() as u64, ..ArrowStreamConfigurationOptions::default() }, + close: CloseCoordinator::new(), #[cfg(feature = "test-hooks")] - ack_applied_gate: Arc::new(Mutex::new(None)), - #[cfg(feature = "test-hooks")] - ack_idle_gate: Arc::new(Mutex::new(None)), + test_hooks: Arc::new(TestHooks::default()), }; ( processor, @@ -292,7 +322,7 @@ impl AckProcessor { ) } - fn rotation_error() -> ZerobusError { + pub(super) fn rotation_error() -> ZerobusError { ZerobusError::StreamClosedError(tonic::Status::unavailable( "Server requested graceful stream rotation", )) @@ -373,38 +403,137 @@ impl AckProcessor { close_deadline.min(bounded_deadline) } - /// Half-closes the active request and drains the response under the rotation - /// deadline. Late acknowledgments are applied while tonic settles request EOF. - /// A real peer status or invalid acknowledgment outranks the synthetic retryable - /// rotation result. - async fn close_request_and_drain_response( + fn explicit_close_drain_deadline() -> Instant { + let now = Instant::now(); + now.checked_add(Duration::from_millis(ROTATION_DRAIN_TIMEOUT_MS)) + .unwrap_or(now) + } + + fn close_target_is_acknowledged(&self, request: CloseRequest) -> bool { + request.target_offset.is_some_and(|target| { + self.last_ack_tx + .borrow() + .is_some_and(|acknowledged| acknowledged >= target) + }) + } + + fn begin_close( + &self, + connection: &ConnectionState, + close_request: CloseRequest, + ) -> ConnectionState { + match connection { + ConnectionState::Active if self.close_target_is_acknowledged(close_request) => { + let outcome = if self.close.target_reached_timely() { + Ok(()) + } else { + Err(CloseCoordinator::flush_timeout_error()) + }; + ConnectionState::Draining(DrainState { + deadline: Self::explicit_close_drain_deadline(), + response_finished: false, + terminal_error: None, + close: Some(CloseDrain { + request: close_request, + selected: Some(outcome), + }), + }) + } + ConnectionState::Active if close_request.target_offset.is_none() => { + ConnectionState::Draining(DrainState { + deadline: Self::explicit_close_drain_deadline(), + response_finished: false, + terminal_error: None, + close: Some(CloseDrain { + request: close_request, + selected: None, + }), + }) + } + ConnectionState::Active => ConnectionState::Waiting(WaitState::Close(close_request)), + ConnectionState::Waiting(WaitState::Rotation { deadlines, .. }) => { + ConnectionState::Draining(DrainState { + deadline: Self::bounded_rotation_drain_deadline(deadlines.drain), + response_finished: false, + terminal_error: None, + close: Some(CloseDrain { + request: close_request, + selected: Some(Err(Self::rotation_error())), + }), + }) + } + ConnectionState::Waiting(WaitState::Close(_)) | ConnectionState::Draining(_) => { + unreachable!() + } + } + } + + fn finish_drain( + terminal_error: Option, + close: Option, + ) -> ZerobusResult { + match close { + Some(CloseDrain { + request, + selected: Some(outcome), + }) => Ok(AckProcessOutcome::Close { request, outcome }), + Some(CloseDrain { + request, + selected: None, + }) => Ok(AckProcessOutcome::Close { + request, + outcome: terminal_error.map_or(Ok(()), Err), + }), + None => Ok(AckProcessOutcome::Recovery { + error: terminal_error.unwrap_or_else(Self::rotation_error), + drained: true, + }), + } + } + + /// Half-closes the active request and drains the response under a shared deadline. + /// Late acknowledgments update retained suffixes but cannot replace a selected result. + async fn half_close_and_drain_response( + &self, response_stream: &mut FlightResponseStream, request: RequestControl<'_>, acknowledgments: AckProgress<'_>, + close_rx: &mut watch::Receiver, + mut observe_close: bool, state: DrainState, - ) -> ZerobusResult<()> { + ) -> ZerobusResult { let DrainState { deadline, mut response_finished, mut terminal_error, + mut close, } = state; request.half_close().await; - let mut request_eof = Box::pin(request.request_body.wait_for_eof()); let mut request_finished = false; loop { + // A continuously ready response must not starve close publication. + if observe_close { + if let Some(close_request) = self.close.request() { + close = Some(CloseDrain { + request: close_request, + selected: Some(Err(Self::rotation_error())), + }); + observe_close = false; + } + } if request_finished && response_finished { - return Err(terminal_error.unwrap_or_else(Self::rotation_error)); + return Self::finish_drain(terminal_error, close); } if Instant::now() >= deadline { - return Err(terminal_error.unwrap_or_else(Self::rotation_error)); + return Self::finish_drain(terminal_error, close); } tokio::select! { biased; _ = sleep_until(deadline) => { - return Err(terminal_error.unwrap_or_else(Self::rotation_error)); + return Self::finish_drain(terminal_error, close); } _ = &mut request_eof, if !request_finished => { request_finished = true; @@ -435,10 +564,43 @@ impl AckProcessor { None => response_finished = true, } } + published = self.close.wait_for_request(close_rx), if observe_close => { + match published { + Some(request) => { + close = Some(CloseDrain { + request, + selected: Some(Err(Self::rotation_error())), + }); + observe_close = false; + } + None => return Ok(AckProcessOutcome::Stopped), + } + } } } } + fn request_control<'a>(&'a self, request_body: &'a RequestBodyControl) -> RequestControl<'a> { + RequestControl { + request_body, + ingest_mutex: self.ingest_mutex.as_ref(), + is_paused: self.is_paused.as_ref(), + batch_tx: &self.batch_tx, + } + } + + fn ack_progress(&self) -> AckProgress<'_> { + AckProgress { + submitted_records: self.submitted_records.as_ref(), + last_acked_records: self.last_acked_records.as_ref(), + pending_batches: self.pending_batches.as_ref(), + last_ack_tx: &self.last_ack_tx, + close: &self.close, + #[cfg(feature = "test-hooks")] + ack_applied_gate: &self.test_hooks.ack_applied, + } + } + /// Returns the oldest submitted batch and its absolute ACK deadline while holding /// the pending lock that also synchronizes submitted-watermark publication. async fn oldest_ack_deadline( @@ -469,10 +631,30 @@ impl AckProcessor { ) } + async fn wait_for_close_deadline(request: Option) { + match request { + Some(request) => sleep_until(request.deadline).await, + None => std::future::pending().await, + } + } + + async fn wait_for_close_event(&self, close_rx: &mut watch::Receiver) -> AckEvent { + match self.close.wait_for_request(close_rx).await { + Some(request) => AckEvent::CloseRequested(request), + None => AckEvent::CloseFinalized, + } + } + /// Waits without arming an ACK timeout while no submitted batch is pending. - async fn wait_while_idle(&self, response_stream: &mut FlightResponseStream) -> AckEvent { + async fn wait_while_idle( + &self, + response_stream: &mut FlightResponseStream, + close_rx: &mut watch::Receiver, + observe_close: bool, + close_wait: Option, + ) -> AckEvent { #[cfg(feature = "test-hooks")] - if let Some(notify) = self.ack_idle_gate.lock().await.take() { + if let Some(notify) = self.test_hooks.ack_idle.lock().await.take() { notify.notify_one(); } @@ -481,6 +663,8 @@ impl AckProcessor { response = response_stream.next() => AckEvent::Response(response), _ = self.request_send_failure.notify.notified() => AckEvent::RequestSendFailed, _ = self.pending_notify.notified() => AckEvent::PendingBatchAvailable, + event = self.wait_for_close_event(close_rx), if observe_close => event, + _ = Self::wait_for_close_deadline(close_wait) => AckEvent::CloseDeadline, } } @@ -492,6 +676,9 @@ impl AckProcessor { response_stream: &mut FlightResponseStream, pending_deadline: PendingAckDeadline, expiry_tie_winner: &mut Option, + close_rx: &mut watch::Receiver, + observe_close: bool, + close_wait: Option, ) -> AckEvent { if *expiry_tie_winner == Some(pending_deadline.identity) && Instant::now() >= pending_deadline.deadline @@ -506,6 +693,8 @@ impl AckProcessor { _ = sleep_until(pending_deadline.deadline) => { AckEvent::AckDeadline(pending_deadline) } + event = self.wait_for_close_event(close_rx), if observe_close => event, + _ = Self::wait_for_close_deadline(close_wait) => AckEvent::CloseDeadline, }; *expiry_tie_winner = if matches!(event, AckEvent::Response(_)) @@ -535,103 +724,256 @@ impl AckProcessor { self.request_send_failure.clear(); } + /// Shuts down the active request and drains response data and final status for a bounded + /// interval. Late ACKs still advance the durable watermark. + pub(super) async fn close_active_connection( + &self, + response_stream: &mut FlightResponseStream, + request_body: &RequestBodyControl, + close_rx: &mut watch::Receiver, + request: CloseRequest, + selected: Option>, + ) -> ZerobusResult { + self.half_close_and_drain_response( + response_stream, + self.request_control(request_body), + self.ack_progress(), + close_rx, + false, + DrainState { + deadline: Self::explicit_close_drain_deadline(), + response_finished: false, + terminal_error: None, + close: Some(CloseDrain { request, selected }), + }, + ) + .await + } + /// Processes acknowledgments and the single server-initiated rotation path. /// /// Rotation pauses sends and snapshots submitted records, waits only for that /// connection-local target, then half-closes the request and drains late responses /// before returning a retryable result to the supervisor. + #[cfg(test)] pub(super) async fn process( &self, mut response_stream: FlightResponseStream, request_body: RequestBodyControl, ) -> ZerobusResult<()> { + let mut close_rx = self.close.subscribe(); + match self + .process_active(&mut response_stream, &request_body, &mut close_rx, false) + .await + { + Ok(AckProcessOutcome::Stopped) => Ok(()), + Ok(AckProcessOutcome::Recovery { error, .. }) => Err(error), + Ok(AckProcessOutcome::Close { .. }) => { + unreachable!("test ACK processor has no close request") + } + Err(error) => Err(error), + } + } + + /// Borrowing the connection lets the supervisor interrupt normal processing for an + /// explicit close without dropping either transport half. + pub(super) async fn process_active( + &self, + response_stream: &mut FlightResponseStream, + request_body: &RequestBodyControl, + close_rx: &mut watch::Receiver, + mut observe_close: bool, + ) -> ZerobusResult { let ack_timeout = Duration::from_millis(self.options.server_lack_of_ack_timeout_ms); - let mut rotation = RotationState::Open; + let mut connection = ConnectionState::Active; let mut expiry_tie_winner: Option = None; - let request = RequestControl { - request_body: &request_body, - ingest_mutex: self.ingest_mutex.as_ref(), - is_paused: self.is_paused.as_ref(), - batch_tx: &self.batch_tx, - }; - let acknowledgments = AckProgress { - submitted_records: self.submitted_records.as_ref(), - last_acked_records: self.last_acked_records.as_ref(), - pending_batches: self.pending_batches.as_ref(), - last_ack_tx: &self.last_ack_tx, - #[cfg(feature = "test-hooks")] - ack_applied_gate: &self.ack_applied_gate, - }; + let mut close_after_priority_response = None; + let request_control = self.request_control(request_body); + let acknowledgments = self.ack_progress(); loop { if self.is_closed.load(Ordering::Relaxed) { debug!(target: super::LOG_TARGET, "Stream closed, stopping ack processor"); - return Ok(()); + return Ok(AckProcessOutcome::Stopped); + } + + if let Some(close_request) = close_after_priority_response.take() { + connection = self.begin_close(&connection, close_request); + continue; + } + + let mut priority_response = None; + if observe_close { + if let Some(close_request) = self.close.request() { + observe_close = false; + if matches!(connection, ConnectionState::Active) + && close_request.target_offset.is_none() + { + // Give a terminal response that predates an empty close one poll. + // A pending or nonterminal response cannot postpone close again. + priority_response = response_stream.next().now_or_never(); + if priority_response.is_some() { + close_after_priority_response = Some(close_request); + } else { + connection = self.begin_close(&connection, close_request); + continue; + } + } else { + connection = self.begin_close(&connection, close_request); + continue; + } + } + } + + if let ConnectionState::Waiting(WaitState::Close(close_request)) = &connection { + if Instant::now() >= close_request.deadline { + connection = ConnectionState::Draining(DrainState { + deadline: Self::explicit_close_drain_deadline(), + response_finished: false, + terminal_error: None, + close: Some(CloseDrain { + request: *close_request, + selected: Some(Err(CloseCoordinator::flush_timeout_error())), + }), + }); + continue; + } } - if let RotationState::WaitingForAcks { + if let ConnectionState::Waiting(WaitState::Rotation { target_records, deadlines, - } = &rotation + }) = &connection { if self.last_acked_records.load(Ordering::Acquire) >= *target_records || Instant::now() >= deadlines.ack { - rotation = RotationState::Draining(DrainState { + connection = ConnectionState::Draining(DrainState { deadline: Self::bounded_rotation_drain_deadline(deadlines.drain), response_finished: false, terminal_error: None, + close: None, }); continue; } } - if matches!(rotation, RotationState::Draining(_)) { - let RotationState::Draining(state) = replace(&mut rotation, RotationState::Open) + if matches!(connection, ConnectionState::Draining(_)) { + let ConnectionState::Draining(state) = + replace(&mut connection, ConnectionState::Active) else { unreachable!() }; - return Self::close_request_and_drain_response( - &mut response_stream, - request, - acknowledgments, - state, - ) - .await; + return self + .half_close_and_drain_response( + response_stream, + request_control, + acknowledgments, + close_rx, + observe_close, + state, + ) + .await; } - let event = match &rotation { - RotationState::Open => match self.oldest_ack_deadline(ack_timeout).await? { - Some(pending_deadline) => { - self.wait_with_pending_deadline( - &mut response_stream, - pending_deadline, - &mut expiry_tie_winner, - ) - .await + let close_wait = match &connection { + ConnectionState::Waiting(WaitState::Close(request)) => Some(*request), + _ => None, + }; + let event = if let Some(response) = priority_response { + AckEvent::Response(response) + } else { + match &connection { + ConnectionState::Active => match self.oldest_ack_deadline(ack_timeout).await? { + Some(pending_deadline) => { + self.wait_with_pending_deadline( + response_stream, + pending_deadline, + &mut expiry_tie_winner, + close_rx, + observe_close, + None, + ) + .await + } + None => { + self.wait_while_idle(response_stream, close_rx, observe_close, None) + .await + } + }, + ConnectionState::Waiting(WaitState::Rotation { deadlines, .. }) => { + tokio::select! { + biased; + response = response_stream.next() => AckEvent::Response(response), + _ = self.request_send_failure.notify.notified() => { + AckEvent::RequestSendFailed + } + _ = sleep_until(deadlines.ack) => continue, + event = self.wait_for_close_event(close_rx), if observe_close => event, + } } - None => self.wait_while_idle(&mut response_stream).await, - }, - RotationState::WaitingForAcks { deadlines, .. } => { - tokio::select! { - biased; - response = response_stream.next() => AckEvent::Response(response), - _ = self.request_send_failure.notify.notified() => { - AckEvent::RequestSendFailed + ConnectionState::Waiting(WaitState::Close(_)) => { + tokio::select! { + biased; + response = response_stream.next() => AckEvent::Response(response), + _ = self.request_send_failure.notify.notified() => { + AckEvent::RequestSendFailed + } + _ = Self::wait_for_close_deadline(close_wait) => AckEvent::CloseDeadline, } - _ = sleep_until(deadlines.ack) => continue, } + ConnectionState::Draining(_) => unreachable!(), } - RotationState::Draining(_) => unreachable!(), }; match event { AckEvent::PendingBatchAvailable => continue, + AckEvent::CloseFinalized => return Ok(AckProcessOutcome::Stopped), + AckEvent::CloseRequested(close_request) => { + observe_close = false; + connection = self.begin_close(&connection, close_request); + } + AckEvent::CloseDeadline => { + let close_request = + close_wait.expect("close deadline requires a close request"); + connection = ConnectionState::Draining(DrainState { + deadline: Self::explicit_close_drain_deadline(), + response_finished: false, + terminal_error: None, + close: Some(CloseDrain { + request: close_request, + selected: Some(Err(CloseCoordinator::flush_timeout_error())), + }), + }); + } AckEvent::RequestSendFailed => { - if self.request_send_failure.take() { - return Err(Self::request_send_error()); + if !self.request_send_failure.take() { + continue; } - continue; + let error = Self::request_send_error(); + connection = match &connection { + ConnectionState::Waiting(WaitState::Rotation { deadlines, .. }) => { + ConnectionState::Draining(DrainState { + deadline: Self::bounded_rotation_drain_deadline(deadlines.drain), + response_finished: false, + terminal_error: Some(error), + close: None, + }) + } + ConnectionState::Waiting(WaitState::Close(close_request)) => { + ConnectionState::Draining(DrainState { + deadline: Self::explicit_close_drain_deadline(), + response_finished: false, + terminal_error: None, + close: Some(CloseDrain { + request: *close_request, + selected: Some(Err(error)), + }), + }) + } + ConnectionState::Active => return Err(error), + ConnectionState::Draining(_) => unreachable!(), + }; } AckEvent::AckDeadline(expected) => { if let Some(pending_count) = @@ -643,7 +985,6 @@ impl AckProcessor { ); return Err(Self::ack_timeout_error()); } - continue; } AckEvent::Response(Some(Ok(put_result))) => { let ack = match FlightAckMetadata::from_bytes(&put_result.app_metadata) { @@ -654,7 +995,7 @@ impl AckProcessor { } }; - if ack.is_close_signal() && matches!(&rotation, RotationState::Open) { + if ack.is_close_signal() && matches!(&connection, ConnectionState::Active) { let server_duration_ms = ack.close_stream_duration_ms.unwrap_or(0); let deadlines = Self::rotation_deadlines( server_duration_ms, @@ -666,10 +1007,10 @@ impl AckProcessor { &self.submitted_records, ) .await; - rotation = RotationState::WaitingForAcks { + connection = ConnectionState::Waiting(WaitState::Rotation { target_records, deadlines, - }; + }); info!(target: super::LOG_TARGET, server_duration_ms, target_records, "Server requested graceful stream rotation" @@ -677,48 +1018,112 @@ impl AckProcessor { } if ack.ack_up_to_records > 0 { - let ack_result = acknowledgments.apply(&ack).await; - if let Err(error) = ack_result { - if let RotationState::WaitingForAcks { deadlines, .. } = rotation { - rotation = RotationState::Draining(DrainState { + if let Err(error) = acknowledgments.apply(&ack).await { + connection = match &connection { + ConnectionState::Waiting(WaitState::Rotation { + deadlines, .. + }) => ConnectionState::Draining(DrainState { deadline: Self::bounded_rotation_drain_deadline( deadlines.drain, ), response_finished: false, terminal_error: Some(error), - }); - continue; - } - return Err(error); + close: None, + }), + ConnectionState::Waiting(WaitState::Close(close_request)) => { + ConnectionState::Draining(DrainState { + deadline: Self::explicit_close_drain_deadline(), + response_finished: false, + terminal_error: None, + close: Some(CloseDrain { + request: *close_request, + selected: Some(Err(error)), + }), + }) + } + ConnectionState::Active => return Err(error), + ConnectionState::Draining(_) => unreachable!(), + }; + continue; + } + } + + if let ConnectionState::Waiting(WaitState::Close(close_request)) = &connection { + if self.close_target_is_acknowledged(*close_request) { + let outcome = if self.close.target_reached_timely() { + Ok(()) + } else { + Err(CloseCoordinator::flush_timeout_error()) + }; + connection = ConnectionState::Draining(DrainState { + deadline: Self::explicit_close_drain_deadline(), + response_finished: false, + terminal_error: None, + close: Some(CloseDrain { + request: *close_request, + selected: Some(outcome), + }), + }); } } } AckEvent::Response(Some(Err(error))) => { let status: tonic::Status = error.into(); let error = ZerobusError::StreamClosedError(status); - if let RotationState::WaitingForAcks { deadlines, .. } = rotation { - rotation = RotationState::Draining(DrainState { - deadline: Self::bounded_rotation_drain_deadline(deadlines.drain), - response_finished: true, - terminal_error: Some(error), - }); - continue; - } - let _ = self.server_error_tx.send(Some(error.clone())); - return Err(error); + connection = match &connection { + ConnectionState::Waiting(WaitState::Rotation { deadlines, .. }) => { + ConnectionState::Draining(DrainState { + deadline: Self::bounded_rotation_drain_deadline(deadlines.drain), + response_finished: true, + terminal_error: Some(error), + close: None, + }) + } + ConnectionState::Waiting(WaitState::Close(close_request)) => { + ConnectionState::Draining(DrainState { + deadline: Self::explicit_close_drain_deadline(), + response_finished: true, + terminal_error: None, + close: Some(CloseDrain { + request: *close_request, + selected: Some(Err(error)), + }), + }) + } + ConnectionState::Active => { + let _ = self.server_error_tx.send(Some(error.clone())); + return Err(error); + } + ConnectionState::Draining(_) => unreachable!(), + }; } AckEvent::Response(None) => { - if let RotationState::WaitingForAcks { deadlines, .. } = rotation { - rotation = RotationState::Draining(DrainState { - deadline: Self::bounded_rotation_drain_deadline(deadlines.drain), - response_finished: true, - terminal_error: None, - }); - continue; - } - return Err(ZerobusError::StreamClosedError(tonic::Status::unknown( + let error = ZerobusError::StreamClosedError(tonic::Status::unknown( "Server closed the stream", - ))); + )); + connection = match &connection { + ConnectionState::Waiting(WaitState::Rotation { deadlines, .. }) => { + ConnectionState::Draining(DrainState { + deadline: Self::bounded_rotation_drain_deadline(deadlines.drain), + response_finished: true, + terminal_error: None, + close: None, + }) + } + ConnectionState::Waiting(WaitState::Close(close_request)) => { + ConnectionState::Draining(DrainState { + deadline: Self::explicit_close_drain_deadline(), + response_finished: true, + terminal_error: None, + close: Some(CloseDrain { + request: *close_request, + selected: Some(Err(error)), + }), + }) + } + ConnectionState::Active => return Err(error), + ConnectionState::Draining(_) => unreachable!(), + }; } } } @@ -734,14 +1139,15 @@ mod tests { use arrow_flight::error::FlightError; use arrow_flight::PutResult; use arrow_schema::{DataType, Field, Schema as ArrowSchema}; - use futures::stream::{iter, pending}; + use futures::stream::{iter, pending, repeat_with}; use futures::StreamExt as _; use tokio::sync::{watch, Mutex, Semaphore}; use tokio::time::{Duration, Instant}; - use super::super::RecordBatch; + use super::super::{CloseRequest, RecordBatch}; use super::{ - AckProcessor, FlightAckMetadata, OffsetId, PendingBatch, RequestBodyControl, ZerobusError, + AckProcessOutcome, AckProcessor, CloseDrain, ConnectionState, FlightAckMetadata, + FlightResponseStream, OffsetId, PendingBatch, RequestBodyControl, ZerobusError, MAX_SERVER_ROTATION_GRACE, ROTATION_DRAIN_TIMEOUT_MS, }; @@ -804,6 +1210,132 @@ mod tests { } } + #[tokio::test] + async fn preacked_close_preserves_latched_timeout() { + let schema = one_col_schema(); + let semaphore = Arc::new(Semaphore::new(1)); + let (processor, _request_body, _last_ack_rx) = ack_processor( + Arc::new(Mutex::new(vec![pending_batch( + &semaphore, + batch_with_rows(&schema, 1), + 0, + 0, + 1, + )])), + Arc::new(AtomicU64::new(1)), + Arc::new(AtomicU64::new(0)), + false, + ); + let request = CloseRequest { + target_offset: Some(0), + deadline: Instant::now(), + }; + processor.close.publish(request); + processor + .ack_progress() + .apply(&FlightAckMetadata { + ack_up_to_offset: 0, + ack_up_to_records: 1, + close_stream_duration_ms: None, + }) + .await + .expect("ACK application should succeed"); + + let ConnectionState::Draining(state) = + processor.begin_close(&ConnectionState::Active, request) + else { + panic!("a pre-acked close must begin draining") + }; + let Some(CloseDrain { + selected: Some(outcome), + .. + }) = state.close + else { + panic!("a pre-acked close must select an outcome") + }; + let error = outcome.expect_err("an ACK applied at the deadline must time out"); + match error { + ZerobusError::StreamClosedError(status) => { + assert_eq!(status.code(), tonic::Code::DeadlineExceeded); + } + other => panic!("expected a flush deadline error, got {other:?}"), + } + } + + #[tokio::test] + async fn published_close_is_not_starved_by_ready_malformed_responses() { + let (processor, request_body, _last_ack_rx) = ack_processor( + Arc::new(Mutex::new(Vec::new())), + Arc::new(AtomicU64::new(0)), + Arc::new(AtomicU64::new(0)), + false, + ); + let mut response_stream: FlightResponseStream = Box::pin( + repeat_with(|| { + Ok(PutResult { + app_metadata: b"not ack metadata".to_vec().into(), + }) + }) + .take(8), + ); + let mut close_rx = processor.close.subscribe(); + let request = CloseRequest { + target_offset: None, + deadline: Instant::now() + Duration::from_secs(1), + }; + processor.close.publish(request); + + let outcome = tokio::time::timeout( + Duration::from_millis(100), + processor.process_active(&mut response_stream, &request_body, &mut close_rx, true), + ) + .await + .expect("a ready response stream must not starve close") + .expect("close observation should not fail"); + + assert!(matches!( + outcome, + AckProcessOutcome::Close { + request: CloseRequest { + target_offset: None, + .. + }, + outcome: Ok(()), + } + )); + } + + #[tokio::test] + async fn ready_terminal_eof_precedes_published_empty_close() { + let (processor, request_body, _last_ack_rx) = ack_processor( + Arc::new(Mutex::new(Vec::new())), + Arc::new(AtomicU64::new(0)), + Arc::new(AtomicU64::new(0)), + false, + ); + let mut response_stream: FlightResponseStream = Box::pin(iter([])); + let mut close_rx = processor.close.subscribe(); + processor.close.publish(CloseRequest { + target_offset: None, + deadline: Instant::now() + Duration::from_secs(1), + }); + + let result = processor + .process_active(&mut response_stream, &request_body, &mut close_rx, true) + .await; + let error = match result { + Err(error) => error, + Ok(_) => panic!("terminal EOF that predates close must not become a clean close"), + }; + match error { + ZerobusError::StreamClosedError(status) => { + assert_eq!(status.code(), tonic::Code::Unknown); + assert_eq!(status.message(), "Server closed the stream"); + } + other => panic!("expected terminal stream error, got {other:?}"), + } + } + /// A buffered ACK is authoritative even when the request sender has already /// reported failure. Apply its durable watermark before asking recovery to replay. #[tokio::test] @@ -1153,6 +1685,98 @@ mod tests { .is_none()); } + #[tokio::test] + async fn close_during_rotation_ack_wait_preserves_rotation_state() { + async fn run(acked_before: bool, expired: bool, ack_during_drain: bool) { + let schema = one_col_schema(); + let semaphore = Arc::new(Semaphore::new(1)); + let pending_batches = Arc::new(Mutex::new(vec![pending_batch( + &semaphore, + batch_with_rows(&schema, 1), + 0, + 0, + 1, + )])); + let last_acked_records = Arc::new(AtomicU64::new(0)); + let (processor, request_body, _last_ack_rx) = ack_processor( + pending_batches, + Arc::new(AtomicU64::new(1)), + Arc::clone(&last_acked_records), + false, + ); + let response = |offset, records, close_stream_duration_ms| PutResult { + app_metadata: serde_json::to_vec(&FlightAckMetadata { + ack_up_to_offset: offset, + ack_up_to_records: records, + close_stream_duration_ms, + }) + .unwrap() + .into(), + }; + let (response_tx, response_rx) = futures::channel::mpsc::unbounded(); + response_tx + .unbounded_send(Ok(response(-1, 0, Some(1_000)))) + .unwrap(); + let mut response_stream: FlightResponseStream = Box::pin(response_rx); + let mut close_rx = processor.close.subscribe(); + let process = + processor.process_active(&mut response_stream, &request_body, &mut close_rx, true); + tokio::pin!(process); + + assert!(futures::poll!(process.as_mut()).is_pending()); + assert!(processor.is_paused.load(Ordering::Relaxed)); + + if acked_before { + processor.last_ack_tx.send_replace(Some(0)); + last_acked_records.store(1, Ordering::Release); + } + let request = CloseRequest { + target_offset: Some(0), + deadline: if expired { + Instant::now() + } else { + Instant::now() + Duration::from_secs(5) + }, + }; + processor.close.publish(request); + assert!(futures::poll!(process.as_mut()).is_pending()); + + if ack_during_drain { + response_tx + .unbounded_send(Ok(response(0, 1, None))) + .unwrap(); + assert!(futures::poll!(process.as_mut()).is_pending()); + } + + drop(response_tx); + let AckProcessOutcome::Close { + request: retained, + outcome, + } = process.await.expect("rotation drain outcome") + else { + panic!("close did not preserve the active rotation") + }; + assert_eq!(retained.target_offset, request.target_offset); + assert_eq!(retained.deadline, request.deadline); + let error = outcome.expect_err("rotation remains the recovery trigger"); + match error { + ZerobusError::StreamClosedError(status) => { + assert_eq!(status.code(), tonic::Code::Unavailable); + } + other => panic!("expected rotation error, got {other:?}"), + } + } + + for case in [ + (true, false, false), + (false, true, false), + (false, true, true), + (false, false, true), + ] { + run(case.0, case.1, case.2).await; + } + } + #[test] fn server_rotation_grace_is_capped() { let rotation_started_at = Instant::now(); diff --git a/rust/sdk/src/stream/arrow/close.rs b/rust/sdk/src/stream/arrow/close.rs new file mode 100644 index 00000000..e3514fc8 --- /dev/null +++ b/rust/sdk/src/stream/arrow/close.rs @@ -0,0 +1,348 @@ +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex as StdMutex}; + +use tokio::sync::{watch, Mutex}; +use tokio::time::Instant; + +use crate::errors::ZerobusError; +use crate::offset_generator::OffsetId; +use crate::ZerobusResult; + +use super::batch::PendingBatch; +use super::{BatchSender, RecordBatch, ZerobusArrowStream}; + +/// The immutable target and deadline selected by the first `close()` call. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) struct CloseRequest { + pub(super) target_offset: Option, + pub(super) deadline: Instant, +} + +/// Shared close publication and completion state. +#[derive(Clone, Debug)] +pub(super) enum CloseState { + Open, + Requested(CloseRequest), + Finalized(ZerobusResult<()>), +} + +impl CloseState { + pub(super) fn request(&self) -> Option { + match self { + Self::Requested(request) => Some(*request), + Self::Open | Self::Finalized(_) => None, + } + } +} + +#[derive(Clone)] +pub(super) struct CloseCoordinator { + state_tx: watch::Sender, + ack_state: Arc>, +} + +#[derive(Default)] +struct CloseAckState { + // Close snapshots the highest assigned offset, so the latest watermark time is sufficient. + latest: Option<(OffsetId, Instant)>, + target_reached_timely: bool, +} + +impl CloseCoordinator { + pub(super) fn new() -> Self { + let (state_tx, _state_rx) = watch::channel(CloseState::Open); + Self { + state_tx, + ack_state: Arc::new(StdMutex::new(CloseAckState::default())), + } + } + + pub(super) fn state(&self) -> CloseState { + self.state_tx.borrow().clone() + } + + pub(super) fn subscribe(&self) -> watch::Receiver { + self.state_tx.subscribe() + } + + pub(super) fn has_started(&self) -> bool { + !matches!(*self.state_tx.borrow(), CloseState::Open) + } + + pub(super) fn request(&self) -> Option { + self.state_tx.borrow().request() + } + + /// Publishes only the first request. The caller holds `ingest_mutex`. + pub(super) fn publish(&self, request: CloseRequest) { + // This lock orders publication with ACK application, covering either winner + // without relying on when the supervisor later observes the request. + let mut ack_state = self.ack_state.lock().expect("close ACK state poisoned"); + self.state_tx.send_if_modified(|state| { + if matches!(state, CloseState::Open) { + ack_state.target_reached_timely = request.target_offset.is_some_and(|target| { + ack_state.latest.is_some_and(|(offset, applied_at)| { + offset >= target && applied_at < request.deadline + }) + }); + *state = CloseState::Requested(request); + true + } else { + false + } + }); + } + + /// Records when a durable ACK was fully applied and latches timely target completion. + pub(super) fn observe_ack(&self, offset: OffsetId, applied_at: Instant) { + let mut ack_state = self.ack_state.lock().expect("close ACK state poisoned"); + match ack_state.latest { + Some((latest_offset, _)) if latest_offset > offset => {} + Some((latest_offset, latest_at)) if latest_offset == offset => { + ack_state.latest = Some((offset, latest_at.min(applied_at))); + } + _ => ack_state.latest = Some((offset, applied_at)), + } + + if let CloseState::Requested(request) = *self.state_tx.borrow() { + if request + .target_offset + .is_some_and(|target| offset >= target && applied_at < request.deadline) + { + ack_state.target_reached_timely = true; + } + } + } + + pub(super) fn target_reached_timely(&self) -> bool { + self.ack_state + .lock() + .expect("close ACK state poisoned") + .target_reached_timely + } + + /// Waits for request publication. `None` means finalization won the observation. + pub(super) async fn wait_for_request( + &self, + close_rx: &mut watch::Receiver, + ) -> Option { + loop { + match close_rx.borrow_and_update().clone() { + CloseState::Requested(request) => return Some(request), + CloseState::Finalized(_) => return None, + CloseState::Open => {} + } + if close_rx.changed().await.is_err() { + return None; + } + } + } + + /// Publishes the result once; later calls cannot replace the finalized outcome. + fn publish_finalized(&self, outcome: ZerobusResult<()>) { + self.state_tx.send_if_modified(|state| { + if matches!(state, CloseState::Finalized(_)) { + false + } else { + *state = CloseState::Finalized(outcome); + true + } + }); + } + + pub(super) fn flush_timeout_error() -> ZerobusError { + ZerobusError::StreamClosedError(tonic::Status::deadline_exceeded("Flush timed out")) + } +} + +/// Performs the one terminal state transition shared by close and failure paths. +#[derive(Clone)] +pub(super) struct CloseFinalizer { + close: CloseCoordinator, + ingest_mutex: Arc>, + batch_tx: BatchSender, + is_paused: Arc, + is_closed: Arc, + pending_batches: Arc>>, + failed_batches: Arc>>, + last_acked_records: Arc, + server_error_tx: watch::Sender>, + #[cfg(feature = "test-hooks")] + test_hooks: Arc, +} + +impl CloseFinalizer { + pub(super) fn new(stream: &ZerobusArrowStream) -> Self { + Self { + close: stream.close.clone(), + ingest_mutex: Arc::clone(&stream.ingest_mutex), + batch_tx: Arc::clone(&stream.batch_tx), + is_paused: Arc::clone(&stream.is_paused), + is_closed: Arc::clone(&stream.is_closed), + pending_batches: Arc::clone(&stream.pending_batches), + failed_batches: Arc::clone(&stream.failed_batches), + last_acked_records: Arc::clone(&stream.last_acked_records), + server_error_tx: stream.server_error_tx.clone(), + #[cfg(feature = "test-hooks")] + test_hooks: Arc::clone(&stream.test_hooks), + } + } + + pub(super) async fn finish(&self, outcome: ZerobusResult<()>) -> ZerobusResult<()> { + { + let _guard = self.ingest_mutex.lock().await; + if let CloseState::Finalized(existing) = self.close.state() { + return existing; + } + self.is_paused.store(true, Ordering::Relaxed); + *self.batch_tx.lock().await = None; + } + + self.server_error_tx + .send_replace(outcome.as_ref().err().cloned()); + + #[cfg(feature = "test-hooks")] + { + let barrier = self.test_hooks.close_finalize.lock().await.take(); + if let Some(barrier) = barrier { + barrier.reached.notify_one(); + barrier.proceed.notified().await; + } + } + + Self::finalize_closed( + &self.ingest_mutex, + &self.is_closed, + &self.pending_batches, + &self.failed_batches, + &self.last_acked_records, + ) + .await; + self.close.publish_finalized(outcome.clone()); + self.server_error_tx + .send_replace(outcome.as_ref().err().cloned()); + outcome + } + + pub(super) async fn move_pending_to_failed( + pending_batches: &Arc>>, + failed_batches: &Arc>>, + last_acked_records: &Arc, + ) { + let mut failed = failed_batches.lock().await; + let mut pending = pending_batches.lock().await; + let acked = last_acked_records.load(Ordering::Acquire); + for batch in pending.drain(..) { + if let Some(batch) = batch.unacknowledged_suffix(acked) { + failed.push(batch); + } + } + } + + pub(super) async fn finalize_closed( + ingest_mutex: &Arc>, + is_closed: &Arc, + pending_batches: &Arc>>, + failed_batches: &Arc>>, + last_acked_records: &Arc, + ) { + let _guard = ingest_mutex.lock().await; + is_closed.store(true, Ordering::Relaxed); + Self::move_pending_to_failed(pending_batches, failed_batches, last_acked_records).await; + } +} + +#[cfg(test)] +mod tests { + use super::{CloseCoordinator, CloseRequest, CloseState}; + use crate::errors::ZerobusError; + use tokio::time::{Duration, Instant}; + + fn request(target_offset: Option) -> CloseRequest { + CloseRequest { + target_offset, + deadline: Instant::now() + Duration::from_secs(1), + } + } + + #[test] + fn first_close_request_is_sticky() { + let close = CloseCoordinator::new(); + let first = request(Some(4)); + close.publish(first); + close.publish(request(Some(9))); + + assert_eq!(close.request(), Some(first)); + } + + #[test] + fn ack_before_request_is_latched_timely() { + let close = CloseCoordinator::new(); + let applied_at = Instant::now(); + close.observe_ack(4, applied_at); + close.publish(CloseRequest { + target_offset: Some(4), + deadline: applied_at + Duration::from_secs(1), + }); + + assert!(close.target_reached_timely()); + } + + #[test] + fn request_before_ack_is_latched_timely() { + let close = CloseCoordinator::new(); + let requested_at = Instant::now(); + let deadline = requested_at + Duration::from_secs(1); + close.publish(CloseRequest { + target_offset: Some(4), + deadline, + }); + close.observe_ack(4, requested_at + Duration::from_millis(500)); + + assert!(close.target_reached_timely()); + } + + #[test] + fn ack_at_deadline_is_not_latched_timely() { + let close = CloseCoordinator::new(); + let deadline = Instant::now(); + close.publish(CloseRequest { + target_offset: Some(4), + deadline, + }); + close.observe_ack(4, deadline); + + assert!(!close.target_reached_timely()); + } + + #[tokio::test] + async fn wait_observes_requested_and_finalized_states() { + let close = CloseCoordinator::new(); + let requested = request(None); + let mut requested_rx = close.subscribe(); + close.publish(requested); + assert_eq!( + close.wait_for_request(&mut requested_rx).await, + Some(requested) + ); + + let finalized = CloseCoordinator::new(); + let mut finalized_rx = finalized.subscribe(); + finalized.publish_finalized(Err(ZerobusError::InvalidStateError("terminal".to_string()))); + assert!(finalized + .wait_for_request(&mut finalized_rx) + .await + .is_none()); + } + + #[test] + fn finalized_outcome_cannot_be_replaced() { + let close = CloseCoordinator::new(); + close.publish_finalized(Ok(())); + close.publish_finalized(Err(ZerobusError::InvalidStateError( + "replacement".to_string(), + ))); + + assert!(matches!(close.state(), CloseState::Finalized(Ok(())))); + } +} diff --git a/rust/sdk/src/stream/arrow/mod.rs b/rust/sdk/src/stream/arrow/mod.rs index ad0bf57d..b511bb87 100644 --- a/rust/sdk/src/stream/arrow/mod.rs +++ b/rust/sdk/src/stream/arrow/mod.rs @@ -30,6 +30,7 @@ pub use arrow_array::RecordBatch; pub use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit}; use self::batch::{materialize_ipc, PendingBatch}; +use self::close::{CloseCoordinator, CloseFinalizer, CloseRequest, CloseState}; pub use self::options::ArrowStreamConfigurationOptions; use self::supervisor::Supervisor; use crate::errors::{should_retry_initial_connection, ZerobusError}; @@ -44,6 +45,7 @@ pub(crate) mod c_data; mod acks; mod batch; +mod close; mod connection; mod metadata; mod options; @@ -67,54 +69,27 @@ pub(super) fn configured_deadline( }) } -/// Test-only barrier used to pause `reconnect` at a precise point — the new connection -/// is established but pending ranges are not yet rebuilt — so a test can schedule a -/// concurrent ingest or `close()`. #[cfg(feature = "test-hooks")] -type ReconnectRebuildGate = Arc>>; +type TestBarrierGate = Mutex>; -/// Paired notifications for [`ReconnectRebuildGate`]: `reached` fires when reconnect -/// hits the barrier; `proceed` releases it (or a test aborts via `close()` instead). #[cfg(feature = "test-hooks")] #[derive(Clone)] -struct ReconnectRebuildBarrier { +struct TestBarrier { reached: Arc, proceed: Arc, } -/// Test-only barrier used to pause recovery after its first replay send, before the -/// remaining backlog is sent and pending ACK timestamps are refreshed. #[cfg(feature = "test-hooks")] -type ReplaySendGate = Arc>>; +type TestNotifyGate = Mutex>>; #[cfg(feature = "test-hooks")] -#[derive(Clone)] -struct ReplaySendBarrier { - reached: Arc, - proceed: Arc, -} - -/// Test-only gate: when armed, the ACK processor fires the notify right after applying a -/// non-empty ack (i.e. after storing `last_acked_records`), letting a test confirm a -/// partial ack has landed before it proceeds. -#[cfg(feature = "test-hooks")] -type AckAppliedGate = Arc>>>; - -/// Test-only gate: when armed, the ACK processor fires the notify immediately before -/// waiting without any pending-work deadline. -#[cfg(feature = "test-hooks")] -type AckIdleGate = Arc>>>; - -/// Test-only barrier that parks `close()` after the supervisor and sender are gone but -/// before pending batches are finalized, allowing cancellation-safe teardown tests. -#[cfg(feature = "test-hooks")] -type CloseFinalizeGate = Arc>>; - -#[cfg(feature = "test-hooks")] -#[derive(Clone)] -struct CloseFinalizeBarrier { - reached: Arc, - proceed: Arc, +#[derive(Default)] +struct TestHooks { + reconnect_rebuild: TestBarrierGate, + replay_send: TestBarrierGate, + ack_applied: TestNotifyGate, + ack_idle: TestNotifyGate, + close_finalize: TestBarrierGate, } /// Properties for an Arrow Flight ingestion table. @@ -197,11 +172,8 @@ pub struct ZerobusArrowStream { _last_ack_rx: watch::Receiver>, /// True once the stream is terminally closed and unacknowledged batches may be retrieved. is_closed: Arc, - /// Separates resumable teardown from final closure so retries skip flushing while - /// new ingests remain rejected. - close_teardown_started: AtomicBool, - /// Retains the first flush failure so resumed close calls return the same outcome. - close_flush_error: Mutex>, + /// Coordinates one resumable explicit-close request with the recovery supervisor. + close: CloseCoordinator, /// Handle to the supervisor task that processes acknowledgments and recovery. receiver_task: Arc>>>>, /// Accepted batches not yet fully acknowledged; retained for replay or retrieval. @@ -245,21 +217,8 @@ pub struct ZerobusArrowStream { /// Either `"zerobus-sdk-rs/"` or `"zerobus-sdk-rs/ "`. /// Re-applied to each fresh Channel built during recovery. sdk_identifier: Arc, - /// Test seam (see [`ReconnectRebuildGate`]); compiled only under `test-hooks`. - #[cfg(feature = "test-hooks")] - reconnect_rebuild_gate: ReconnectRebuildGate, - /// Test seam (see [`ReplaySendGate`]); compiled only under `test-hooks`. - #[cfg(feature = "test-hooks")] - replay_send_gate: ReplaySendGate, - /// Test seam (see [`AckAppliedGate`]); compiled only under `test-hooks`. - #[cfg(feature = "test-hooks")] - ack_applied_gate: AckAppliedGate, - /// Test seam (see [`AckIdleGate`]); compiled only under `test-hooks`. #[cfg(feature = "test-hooks")] - ack_idle_gate: AckIdleGate, - /// Test seam (see [`CloseFinalizeGate`]); compiled only under `test-hooks`. - #[cfg(feature = "test-hooks")] - close_finalize_gate: CloseFinalizeGate, + test_hooks: Arc, } impl ZerobusArrowStream { @@ -297,6 +256,11 @@ impl ZerobusArrowStream { Duration::from_millis(options.server_lack_of_ack_timeout_ms), "server_lack_of_ack_timeout_ms", )?; + configured_deadline( + validation_started_at, + Duration::from_millis(options.flush_timeout_ms), + "flush_timeout_ms", + )?; let (last_ack_tx, _last_ack_rx) = watch::channel(None); let is_closed = Arc::new(AtomicBool::new(false)); @@ -315,6 +279,7 @@ impl ZerobusArrowStream { let inflight = Arc::new(Semaphore::new(options.max_inflight_batches)); let (server_error_tx, server_error_rx) = watch::channel(None); + let close = CloseCoordinator::new(); let stream = Self { table_properties, @@ -324,8 +289,7 @@ impl ZerobusArrowStream { last_ack_tx, _last_ack_rx, is_closed, - close_teardown_started: AtomicBool::new(false), - close_flush_error: Mutex::new(None), + close, receiver_task, pending_batches, pending_notify, @@ -346,15 +310,7 @@ impl ZerobusArrowStream { is_paused, sdk_identifier, #[cfg(feature = "test-hooks")] - reconnect_rebuild_gate: Arc::new(Mutex::new(None)), - #[cfg(feature = "test-hooks")] - replay_send_gate: Arc::new(Mutex::new(None)), - #[cfg(feature = "test-hooks")] - ack_applied_gate: Arc::new(Mutex::new(None)), - #[cfg(feature = "test-hooks")] - ack_idle_gate: Arc::new(Mutex::new(None)), - #[cfg(feature = "test-hooks")] - close_finalize_gate: Arc::new(Mutex::new(None)), + test_hooks: Arc::new(TestHooks::default()), }; // Initialize the connection with retry logic. @@ -467,9 +423,7 @@ impl ZerobusArrowStream { /// ``` #[instrument(level = "debug", skip_all, fields(table_name = %self.table_properties.table_name))] pub async fn ingest_batch(&self, batch: RecordBatch) -> ZerobusResult { - if self.is_closed.load(Ordering::Relaxed) - || self.close_teardown_started.load(Ordering::Acquire) - { + if self.is_closed.load(Ordering::Relaxed) || self.close.has_started() { return Err(ZerobusError::StreamClosedError(tonic::Status::internal( "Stream is closing or closed", ))); @@ -507,9 +461,7 @@ impl ZerobusArrowStream { let _guard = self.ingest_mutex.lock().await; // May have closed while we blocked on the permit; returning drops it. - if self.is_closed.load(Ordering::Relaxed) - || self.close_teardown_started.load(Ordering::Acquire) - { + if self.is_closed.load(Ordering::Relaxed) || self.close.has_started() { return Err(ZerobusError::StreamClosedError(tonic::Status::internal( "Stream is closing or closed", ))); @@ -619,9 +571,7 @@ impl ZerobusArrowStream { /// marker after `finish()`) is allowed after that batch. #[instrument(level = "debug", skip_all, fields(table_name = %self.table_properties.table_name))] pub async fn ingest_ipc_batch(&self, ipc_bytes: Bytes) -> ZerobusResult { - if self.is_closed.load(Ordering::Relaxed) - || self.close_teardown_started.load(Ordering::Acquire) - { + if self.is_closed.load(Ordering::Relaxed) || self.close.has_started() { return Err(ZerobusError::StreamClosedError(tonic::Status::internal( "Stream is closing or closed", ))); @@ -679,9 +629,7 @@ impl ZerobusArrowStream { // state. Re-read first because the watermark can be published between the // read above and observing that state. Otherwise prefer the real terminal // error over a generic one. - if self.is_closed.load(Ordering::Relaxed) - || self.close_teardown_started.load(Ordering::Acquire) - { + if self.is_closed.load(Ordering::Relaxed) || self.close.has_started() { if let Some(ack_offset) = *offset_rx.borrow_and_update() { if ack_offset >= offset_to_wait { return Ok(()); @@ -764,9 +712,7 @@ impl ZerobusArrowStream { None => { // Nothing was ingested: report closure if closed, otherwise nothing to do. // Prefer the real terminal error over a generic closed message. - if self.is_closed.load(Ordering::Relaxed) - || self.close_teardown_started.load(Ordering::Acquire) - { + if self.is_closed.load(Ordering::Relaxed) || self.close.has_started() { if let Some(server_error) = self.server_error_rx.borrow().clone() { return Err(server_error); } @@ -829,8 +775,12 @@ impl ZerobusArrowStream { /// Flushes pending work, stops background I/O, and retains unacknowledged batches for /// retrieval. /// - /// While the stream is active, the first call attempts one flush before teardown. If - /// teardown is interrupted, a later call resumes it without flushing again. + /// The first call publishes one close request. While the active transport remains + /// usable, the supervisor continues ACK processing through the original flush deadline, + /// then owns transport cleanup and finalization. Close does not start or continue + /// recovery: a transport failure or an already-running recovery is interrupted, and + /// unacknowledged batches are retained for retrieval. Repeated calls await the same + /// request and result. An uncommitted replacement transport is dropped best-effort. /// /// # Returns /// @@ -838,15 +788,18 @@ impl ZerobusArrowStream { /// /// # Errors /// - /// Returns the initial flush error or a background terminal error. Teardown still - /// completes; use `get_unacked_batches()` to retrieve unacknowledged batches. + /// Returns a background terminal error or a timeout if the close target is not + /// acknowledged by the flush deadline. During ordinary active-connection close, a + /// timely target acknowledgment takes precedence. If close interrupts an already-active + /// server rotation or an uncommitted recovery attempt, it instead returns that attempt's + /// trigger even when the close target is durable. Teardown still completes; use + /// `get_unacked_batches()` to retrieve unacknowledged batches. /// /// # Cancellation safety /// - /// Cancelling before teardown begins does not itself close the stream, although an - /// independent terminal failure may do so. Once teardown starts, further ingests are - /// rejected; call `close()` again to resume incomplete teardown without repeating a - /// completed flush. + /// Once the close request is published, further ingests are rejected. Cancelling the + /// future does not cancel that request: call `close()` again to await the same original + /// deadline and final outcome. /// /// # Examples /// @@ -860,91 +813,40 @@ impl ZerobusArrowStream { /// ``` #[instrument(level = "debug", skip_all, fields(table_name = %self.table_properties.table_name))] pub async fn close(&mut self) -> ZerobusResult<()> { - let close_teardown_started = self.close_teardown_started.load(Ordering::Acquire); - if self.is_closed.load(Ordering::Relaxed) && !close_teardown_started { - // Already closed. If the supervisor closed it on a terminal failure, surface - // that error rather than reporting success — otherwise the common - // ingest-then-close() pattern would hide failed batches (retrievable via - // get_unacked_batches()). A clean prior close() has no stored error. - if let Some(server_error) = self.server_error_rx.borrow().clone() { - return Err(server_error); - } - if let Some(close_error) = self.close_flush_error.lock().await.clone() { - return Err(close_error); - } - return Ok(()); - } - info!( table_name = %self.table_properties.table_name, "Closing Arrow Flight stream" ); - - // Retain a completed flush result before publishing teardown so retries after - // teardown starts skip another flush and return the same outcome. - let flush_result = if close_teardown_started { - match self.close_flush_error.lock().await.clone() { - Some(error) => Err(error), - None => Ok(()), - } - } else { - let result = self.flush().await; - *self.close_flush_error.lock().await = result.as_ref().err().cloned(); - self.close_teardown_started.store(true, Ordering::Release); - result - }; - if let Err(e) = &flush_result { - warn!( - "Flush failed during close: {}. Draining pending batches to the failed set.", - e - ); - } - - // Reap the supervisor (abort + await) BEFORE clearing the sender, so an in-flight - // reconnect can't reinstall batch_tx after we clear it, and no ACK processing / - // reconnect mutates pending_batches or last_acked_records while we drain. Join in - // place and only clear receiver_task once the join completes, so a close() - // cancelled during the await doesn't drop the handle — a retry re-joins it. - { - let mut task = self.receiver_task.lock().await; - if let Some(handle) = task.as_mut() { - handle.abort(); - let _ = handle.await; - } - *task = None; - } - - // Detach the sender now that nothing can reinstall it. - { - let mut tx = self.batch_tx.lock().await; - *tx = None; - } - - // Test seam: cancel close after teardown became irreversible but before finalization. - #[cfg(feature = "test-hooks")] - { - let barrier = self.close_finalize_gate.lock().await.take(); - if let Some(barrier) = barrier { - barrier.reached.notify_one(); - barrier.proceed.notified().await; + let mut close_rx = self.close.subscribe(); + + loop { + let state = { close_rx.borrow_and_update().clone() }; + match state { + CloseState::Open => { + // This mutex makes the target snapshot and request publication atomic + // with ingest admission and replacement-sender publication. + let _guard = self.ingest_mutex.lock().await; + let deadline = configured_deadline( + Instant::now(), + Duration::from_millis(self.options.flush_timeout_ms), + "flush_timeout_ms", + )?; + let request = CloseRequest { + target_offset: self.offset_generator.last(), + deadline, + }; + self.close.publish(request); + } + CloseState::Requested(_) => { + if close_rx.changed().await.is_err() { + return Err(ZerobusError::StreamClosedError(tonic::Status::internal( + "Close coordinator stopped unexpectedly", + ))); + } + } + CloseState::Finalized(result) => return result, } } - - // Finalize under ingest_mutex so the pending drain is serialized with - // ingest_batch. Keep close_teardown_started set while finalization is in flight, - // then clear it immediately afterward; cancellation before completion remains - // resumable even if closure was already published. - Supervisor::finalize_closed( - &self.ingest_mutex, - &self.is_closed, - &self.pending_batches, - &self.failed_batches, - &self.last_acked_records, - ) - .await; - self.close_teardown_started.store(false, Ordering::Release); - - flush_result } /// Returns the un-acknowledged batches after the stream has been closed, for manual @@ -958,8 +860,8 @@ impl ZerobusArrowStream { /// /// # Errors /// - /// * `InvalidStateError` - If closure has not been finalized, including after - /// interrupted teardown; call `close()` again first. + /// * `InvalidStateError` - If closure has not been finalized; call `close()` first, + /// or call it again to await a previously requested close. /// /// # Examples /// @@ -995,7 +897,7 @@ impl ZerobusArrowStream { // failed set, then return the consolidated snapshot. move_pending_to_failed locks // failed first, so this serializes with a concurrent terminal drain and repeated // calls are idempotent (pending is already empty on the second call). - Supervisor::move_pending_to_failed( + CloseFinalizer::move_pending_to_failed( &self.pending_batches, &self.failed_batches, &self.last_acked_records, @@ -1004,42 +906,45 @@ impl ZerobusArrowStream { Ok(self.failed_batches.lock().await.clone()) } - /// Returns true once terminal finalization publishes closure. Interrupted teardown - /// remains false until finalization begins; cancellation during finalization may leave - /// this true while `close_teardown_started` marks teardown as resumable. + /// Returns true once supervisor-owned terminal finalization publishes closure. pub fn is_closed(&self) -> bool { self.is_closed.load(Ordering::Relaxed) } - /// Test-only: arms the reconnect rebuild barrier. The next `reconnect` pauses after - /// establishing the connection but before rebuilding pending ranges/watermark, - /// firing the returned `reached` notify, then waits on `proceed`. A test either - /// releases `proceed` to let recovery finish, or drives a concurrent `close()` - /// (which reaps the paused supervisor) without releasing it. #[cfg(feature = "test-hooks")] - #[doc(hidden)] - pub async fn arm_reconnect_rebuild_barrier(&self) -> (Arc, Arc) { + async fn arm_test_barrier(gate: &TestBarrierGate) -> (Arc, Arc) { let reached = Arc::new(Notify::new()); let proceed = Arc::new(Notify::new()); - *self.reconnect_rebuild_gate.lock().await = Some(ReconnectRebuildBarrier { + *gate.lock().await = Some(TestBarrier { reached: Arc::clone(&reached), proceed: Arc::clone(&proceed), }); (reached, proceed) } - /// Test-only: pauses the next recovery after its first replay send and before the - /// remaining backlog is sent or its pending ACK timestamps are refreshed. + #[cfg(feature = "test-hooks")] + async fn arm_test_notify(gate: &TestNotifyGate) -> Arc { + let notify = Arc::new(Notify::new()); + *gate.lock().await = Some(Arc::clone(¬ify)); + notify + } + + /// Test-only: arms the reconnect rebuild barrier. The next `reconnect` pauses after + /// establishing the connection but before rebuilding pending ranges/watermark, + /// firing the returned `reached` notify, then waits on `proceed`. Cancellation drops + /// the uncommitted replacement transport best-effort. + #[cfg(feature = "test-hooks")] + #[doc(hidden)] + pub async fn arm_reconnect_rebuild_barrier(&self) -> (Arc, Arc) { + Self::arm_test_barrier(&self.test_hooks.reconnect_rebuild).await + } + + /// Test-only: pauses the next recovery after its first replay handoff and before the + /// remaining backlog or sender publication can be committed. #[cfg(feature = "test-hooks")] #[doc(hidden)] pub async fn arm_replay_send_barrier(&self) -> (Arc, Arc) { - let reached = Arc::new(Notify::new()); - let proceed = Arc::new(Notify::new()); - *self.replay_send_gate.lock().await = Some(ReplaySendBarrier { - reached: Arc::clone(&reached), - proceed: Arc::clone(&proceed), - }); - (reached, proceed) + Self::arm_test_barrier(&self.test_hooks.replay_send).await } /// Test-only: arms a notify that fires each time the ACK processor applies a non-empty @@ -1048,9 +953,7 @@ impl ZerobusArrowStream { #[cfg(feature = "test-hooks")] #[doc(hidden)] pub async fn arm_ack_applied_notify(&self) -> Arc { - let notify = Arc::new(Notify::new()); - *self.ack_applied_gate.lock().await = Some(Arc::clone(¬ify)); - notify + Self::arm_test_notify(&self.test_hooks.ack_applied).await } /// Test-only: arms a one-shot notification for the next time the ACK processor @@ -1058,9 +961,7 @@ impl ZerobusArrowStream { #[cfg(feature = "test-hooks")] #[doc(hidden)] pub async fn arm_ack_idle_notify(&self) -> Arc { - let notify = Arc::new(Notify::new()); - *self.ack_idle_gate.lock().await = Some(Arc::clone(¬ify)); - notify + Self::arm_test_notify(&self.test_hooks.ack_idle).await } /// Test-only: replaces the active batch sender with a sender whose receiver is @@ -1074,18 +975,12 @@ impl ZerobusArrowStream { *self.batch_tx.lock().await = Some(closed_tx); } - /// Test-only: parks the next `close()` after supervisor/sender teardown but before - /// finalization. Dropping the close future at that point simulates cancellation. + /// Test-only: parks close finalization after choosing the local outcome and before + /// moving pending batches into the final failed-batch snapshot. #[cfg(feature = "test-hooks")] #[doc(hidden)] pub async fn arm_close_finalize_barrier(&self) -> (Arc, Arc) { - let reached = Arc::new(Notify::new()); - let proceed = Arc::new(Notify::new()); - *self.close_finalize_gate.lock().await = Some(CloseFinalizeBarrier { - reached: Arc::clone(&reached), - proceed: Arc::clone(&proceed), - }); - (reached, proceed) + Self::arm_test_barrier(&self.test_hooks.close_finalize).await } /// Returns the table name for this stream. diff --git a/rust/sdk/src/stream/arrow/options.rs b/rust/sdk/src/stream/arrow/options.rs index 49189e43..26cf44d9 100644 --- a/rust/sdk/src/stream/arrow/options.rs +++ b/rust/sdk/src/stream/arrow/options.rs @@ -85,6 +85,8 @@ pub struct ArrowStreamConfigurationOptions { /// Timeout in milliseconds for flush operations. /// /// If a `flush()` call cannot complete within this time, it will return a timeout error. + /// Values whose absolute deadline cannot be represented by the platform's + /// monotonic clock are rejected when the stream is built. /// /// Default: 300,000 (5 minutes) pub flush_timeout_ms: u64, @@ -135,9 +137,10 @@ pub struct ArrowStreamConfigurationOptions { /// cleanup and terminates without reconnecting. Batches accepted while paused remain /// available through `get_unacked_batches()`. /// - /// The clean half-close guarantee applies to the active connection during normal - /// server rotation. Explicit close during recovery remains best-effort and may abort - /// an incomplete replacement request. + /// The clean half-close guarantee applies to the active connection. Explicit close + /// during an already-active rotation or recovery retains that attempt's trigger even if + /// the explicit close target is already acknowledged. Any uncommitted replacement request + /// is dropped best-effort. /// /// Default: `None` (use the available server grace period) pub stream_paused_max_wait_time_ms: Option, diff --git a/rust/sdk/src/stream/arrow/supervisor.rs b/rust/sdk/src/stream/arrow/supervisor.rs index cd7ecff6..b6c8771c 100644 --- a/rust/sdk/src/stream/arrow/supervisor.rs +++ b/rust/sdk/src/stream/arrow/supervisor.rs @@ -9,11 +9,12 @@ use std::sync::Arc; use arrow_flight::error::FlightError; use tokio::sync::{mpsc, watch, Mutex}; use tokio::task::{spawn, JoinHandle}; -use tokio::time::{sleep, timeout_at, Duration, Instant}; +use tokio::time::{sleep, sleep_until, timeout_at, Duration, Instant}; use tracing::{debug, error, info, warn}; -use super::acks::{pause_and_detach_sender, AckProcessor}; +use super::acks::{pause_and_detach_sender, AckProcessOutcome, AckProcessor}; use super::batch::{rebuild_pending_for_replay, refresh_pending_ack_deadlines, PendingBatch}; +use super::close::{CloseCoordinator, CloseFinalizer, CloseState}; use super::connection::{FlightConnection, FlightResponseStream, RequestBodyControl}; use super::{ configured_deadline, ArrowStreamConfigurationOptions, ArrowTableProperties, BatchSender, @@ -35,8 +36,9 @@ pub(super) struct Supervisor { ack_processor: AckProcessor, batch_tx: BatchSender, is_closed: Arc, + close: CloseCoordinator, + close_finalizer: CloseFinalizer, pending_batches: Arc>>, - failed_batches: Arc>>, recovery_attempts: Arc, server_error_tx: watch::Sender>, cumulative_records_assigned: Arc, @@ -46,9 +48,7 @@ pub(super) struct Supervisor { ingest_mutex: Arc>, sdk_identifier: Arc, #[cfg(feature = "test-hooks")] - reconnect_rebuild_gate: super::ReconnectRebuildGate, - #[cfg(feature = "test-hooks")] - replay_send_gate: super::ReplaySendGate, + test_hooks: Arc, } impl Supervisor { @@ -63,8 +63,9 @@ impl Supervisor { ack_processor: AckProcessor::new(stream), batch_tx: Arc::clone(&stream.batch_tx), is_closed: Arc::clone(&stream.is_closed), + close: stream.close.clone(), + close_finalizer: CloseFinalizer::new(stream), pending_batches: Arc::clone(&stream.pending_batches), - failed_batches: Arc::clone(&stream.failed_batches), recovery_attempts: Arc::clone(&stream.recovery_attempts), server_error_tx: stream.server_error_tx.clone(), cumulative_records_assigned: Arc::clone(&stream.cumulative_records_assigned), @@ -74,9 +75,7 @@ impl Supervisor { ingest_mutex: Arc::clone(&stream.ingest_mutex), sdk_identifier: Arc::clone(&stream.sdk_identifier), #[cfg(feature = "test-hooks")] - reconnect_rebuild_gate: Arc::clone(&stream.reconnect_rebuild_gate), - #[cfg(feature = "test-hooks")] - replay_send_gate: Arc::clone(&stream.replay_send_gate), + test_hooks: Arc::clone(&stream.test_hooks), } } @@ -88,6 +87,66 @@ impl Supervisor { spawn(self.run(response_stream, request_body)) } + fn spawn_headers_invalidation(&self, deadline: Instant) -> JoinHandle { + let headers_provider = Arc::clone(&self.headers_provider); + let timeout_ms = self.options.recovery_timeout_ms; + spawn(async move { + if timeout_at(deadline, headers_provider.invalidate()) + .await + .is_ok() + { + true + } else { + warn!(target: super::LOG_TARGET, + timeout_ms, + "Headers provider invalidation timed out" + ); + false + } + }) + } + + fn spawn_detached_headers_invalidation(&self) { + match configured_deadline( + Instant::now(), + Duration::from_millis(self.options.recovery_timeout_ms), + "recovery_timeout_ms", + ) { + Ok(deadline) => { + drop(self.spawn_headers_invalidation(deadline)); + } + Err(error) => { + warn!(target: super::LOG_TARGET, + error = %error, + "Skipping headers provider invalidation because its deadline is unrepresentable" + ); + } + } + } + + fn spawn_detached_auth_invalidation(&self, error: &ZerobusError) { + if error.is_auth_rejection() { + self.spawn_detached_headers_invalidation(); + } + } + + async fn finish(&self, outcome: ZerobusResult<()>) -> ZerobusResult<()> { + self.close_finalizer.finish(outcome).await + } + + fn finalized_result(&self) -> ZerobusResult<()> { + Self::result_from_close_state(self.close.state()) + } + + fn result_from_close_state(state: CloseState) -> ZerobusResult<()> { + match state { + CloseState::Finalized(result) => result, + CloseState::Open | CloseState::Requested(_) => Err(ZerobusError::InvalidStateError( + "Supervisor exited before close finalization".to_string(), + )), + } + } + async fn run( self, initial_response_stream: FlightResponseStream, @@ -95,82 +154,100 @@ impl Supervisor { ) -> ZerobusResult<()> { let mut response_stream = Some(initial_response_stream); let mut request_body = Some(initial_request_body); - // Carries a failed reconnect's real error into the next iteration's handling - // instead of round-tripping a synthetic error through a dummy stream. let mut pending_error: Option = None; - // True when `pending_error` is a reconnect auth rejection: the cached token was - // invalidated and we want to retry (mint a fresh one) even though auth errors - // classify as non-retryable — while still surfacing the original error if - // retries are ultimately exhausted. let mut reconnect_auth_retry = false; + let mut close_rx = self.close.subscribe(); loop { if self.is_closed.load(Ordering::Relaxed) { debug!(target: super::LOG_TARGET, "Supervisor: Stream closed, exiting"); - return Ok(()); + return self.finalized_result(); } - // Run ACK processing until it returns — unless a prior reconnect attempt - // failed, in which case carry that real error into the handling below - // (preserving its message and retry classification). - let result = if let Some(e) = pending_error.take() { - Err(e) + let mut active_was_drained = false; + let result = if let Some(error) = pending_error.take() { + Err(error) } else { - self.ack_processor - .process( - response_stream - .take() - .expect("response_stream present when no pending reconnect error"), - request_body - .take() - .expect("request_body present when no pending reconnect error"), - ) + let active_response = response_stream + .as_mut() + .expect("response stream present outside recovery"); + let active_request = request_body + .as_ref() + .expect("request body present outside recovery"); + match self + .ack_processor + .process_active(active_response, active_request, &mut close_rx, true) .await + { + Ok(AckProcessOutcome::Stopped) => return self.finalized_result(), + Ok(AckProcessOutcome::Recovery { error, drained }) => { + active_was_drained = drained; + Err(error) + } + Ok(AckProcessOutcome::Close { request, outcome }) => { + debug_assert_eq!(self.close.request(), Some(request)); + if let Err(error) = &outcome { + self.spawn_detached_auth_invalidation(error); + } + return self.finish(outcome).await; + } + Err(error) => Err(error), + } }; - // Check if stream was closed during processing. - if self.is_closed.load(Ordering::Relaxed) { - debug!(target: super::LOG_TARGET, "Supervisor: Stream closed after process_acks, exiting"); - return result; + if let Err(error) = &result { + if !reconnect_auth_retry { + self.spawn_detached_auth_invalidation(error); + } } - // Handle the result. - match result { - Ok(()) => { - // Stream ended gracefully. - debug!(target: super::LOG_TARGET, "Supervisor: process_acks completed successfully"); - return Ok(()); + if let Some(close_request) = self.close.request() { + if !active_was_drained { + if let (Some(active_response), Some(active_request)) = + (response_stream.as_mut(), request_body.as_ref()) + { + let selected = result.clone(); + match self + .ack_processor + .close_active_connection( + active_response, + active_request, + &mut close_rx, + close_request, + Some(selected), + ) + .await + { + Ok(AckProcessOutcome::Close { request, outcome }) => { + debug_assert_eq!(self.close.request(), Some(request)); + return self.finish(outcome).await; + } + Ok(AckProcessOutcome::Stopped) => { + return self.finalized_result(); + } + Ok(AckProcessOutcome::Recovery { error, .. }) | Err(error) => { + return self.finish(Err(error)).await; + } + } + } } + return self.finish(result).await; + } + + match result { + Ok(()) => return self.finish(Ok(())).await, Err(ref error) if (error.is_retryable() || reconnect_auth_retry) && self.options.recovery => { - // Retriable error (or a reconnect auth rejection we've chosen to - // retry with re-minted credentials) - attempt recovery. reconnect_auth_retry = false; let attempts = self.recovery_attempts.fetch_add(1, Ordering::Relaxed); if attempts >= self.options.recovery_retries { error!(target: super::LOG_TARGET, - attempts = attempts, + attempts, max_retries = self.options.recovery_retries, "Supervisor: Max recovery retries exceeded" ); - // Publish the terminal error before finalization (so a waiter - // checking is_closed right after it already sees the real error; - // reconnect-failure errors carried via pending_error are never - // pre-published by ACK processing) and again after (to wake - // already-parked waiters). finalize_closed also drains pending - // under ingest_mutex so a concurrent ingest can't be omitted. - let _ = self.server_error_tx.send(Some(error.clone())); - Self::finalize_closed( - &self.ingest_mutex, - &self.is_closed, - &self.pending_batches, - &self.failed_batches, - &self.last_acked_records, - ) - .await; - let _ = self.server_error_tx.send(Some(error.clone())); - return result; + return self.finish(Err(error.clone())).await; } info!(target: super::LOG_TARGET, @@ -180,152 +257,114 @@ impl Supervisor { "Supervisor: Attempting recovery after retriable error" ); - // Atomically pause ingest and detach the sender under - // ingest_mutex, so an in-flight ingest_batch either completes - // before the pause or observes is_paused and buffers — it never - // sees is_paused=false with a detached sender. Successful replay - // lifts the gate; failed attempts remain paused for retry/finalization. pause_and_detach_sender(&self.ingest_mutex, &self.is_paused, &self.batch_tx) .await; + response_stream = None; + request_body = None; self.ack_processor.clear_request_send_failure(); + // Close that cancels this attempt reports its trigger. Once an attempt + // failure is accepted below, that failure becomes the next trigger. + let recovery_error = error.clone(); + + let close_during_backoff = tokio::select! { + biased; + request = self.close.wait_for_request(&mut close_rx) => request.is_some(), + _ = sleep(Duration::from_millis(self.options.recovery_backoff_ms)) => false, + }; + if close_during_backoff { + return self.finish(Err(recovery_error)).await; + } + if matches!(self.close.state(), CloseState::Finalized(_)) { + return self.finalized_result(); + } - sleep(Duration::from_millis(self.options.recovery_backoff_ms)).await; - - let _ = self.server_error_tx.send(None); - - // Share one absolute timeout budget across reconnect and - // auth-rejection invalidation. - let recovery_timeout = Duration::from_millis(self.options.recovery_timeout_ms); - let recovery_started = Instant::now(); + self.server_error_tx.send_replace(None); let recovery_deadline = match configured_deadline( - recovery_started, - recovery_timeout, + Instant::now(), + Duration::from_millis(self.options.recovery_timeout_ms), "recovery_timeout_ms", ) { Ok(deadline) => deadline, - Err(error) => { - pending_error = Some(error); + Err(deadline_error) => { + if self.close.has_started() { + return self.finish(Err(recovery_error)).await; + } + pending_error = Some(deadline_error); continue; } }; - let reconnect_result = timeout_at(recovery_deadline, self.reconnect()).await; + + // A ready attempt wins a simultaneous close. If sender commit completed, + // the replacement is active and receives the normal graceful-close path. + let reconnect_result = tokio::select! { + biased; + result = self.reconnect() => Some(result), + request = self.close.wait_for_request(&mut close_rx) => { + if request.is_some() { + return self.finish(Err(recovery_error)).await; + } + return self.finalized_result(); + } + _ = sleep_until(recovery_deadline) => None, + }; match reconnect_result { - Ok(Ok((new_response_stream, new_request_body))) => { + Some(Ok(Some(connection))) => { info!(target: super::LOG_TARGET, "Supervisor: Recovery successful, resuming"); self.recovery_attempts.store(0, Ordering::Relaxed); - // is_paused was already cleared inside reconnect(). + let (new_response_stream, new_request_body) = + connection.into_supervisor_io(); response_stream = Some(new_response_stream); request_body = Some(new_request_body); } - Ok(Err(e)) => { - warn!(target: super::LOG_TARGET, "Supervisor: Reconnection failed: {}", e); - // Ask the provider to invalidate cached authentication - // state after an auth rejection, then retry even though - // such errors are otherwise non-retryable. Preserve this - // reconnect error if refresh or later recovery cannot proceed. - if e.is_auth_rejection() { - match timeout_at( - recovery_deadline, - self.headers_provider.invalidate(), - ) - .await - { - Ok(()) => reconnect_auth_retry = true, - Err(_) => { - warn!(target: super::LOG_TARGET, - timeout_ms = self.options.recovery_timeout_ms, - "Recovery deadline reached while invalidating \ - the headers provider; terminating recovery" - ); - // A custom provider must not stall recovery - // indefinitely. Close with the original auth - // rejection; publish before and after - // finalization for waiter race-freedom. - let _ = self.server_error_tx.send(Some(e.clone())); - Self::finalize_closed( - &self.ingest_mutex, - &self.is_closed, - &self.pending_batches, - &self.failed_batches, - &self.last_acked_records, - ) - .await; - let _ = self.server_error_tx.send(Some(e.clone())); - return Err(e); - } - } - } - pending_error = Some(e); + Some(Ok(None)) => { + return self.finish(Err(recovery_error)).await; } - Err(_timeout) => { + None => { warn!(target: super::LOG_TARGET, "Supervisor: Reconnection timed out"); pending_error = Some(ZerobusError::ConnectionTimeout(format!( "Reconnection timed out after {}ms", self.options.recovery_timeout_ms ))); } - } - } - Err(error) => { - error!(target: super::LOG_TARGET, "Supervisor: Non-retriable error, closing stream: {}", error); - // Publish the terminal error before finalization (so a waiter - // checking is_closed right after it already sees the real error; - // reconnect-failure errors carried via pending_error are never - // pre-published by ACK processing) and again after (to wake - // already-parked waiters). finalize_closed drains pending under - // ingest_mutex so a concurrent ingest can't be omitted. - let _ = self.server_error_tx.send(Some(error.clone())); - Self::finalize_closed( - &self.ingest_mutex, - &self.is_closed, - &self.pending_batches, - &self.failed_batches, - &self.last_acked_records, - ) - .await; - let _ = self.server_error_tx.send(Some(error.clone())); - // Ask the provider to invalidate cached authentication state after - // a terminal rejection. The stream is already finalized and waiters - // have the real error; bound the callback so the supervisor cannot - // remain alive indefinitely. - if error.is_auth_rejection() { - match configured_deadline( - Instant::now(), - Duration::from_millis(self.options.recovery_timeout_ms), - "recovery_timeout_ms", - ) { - Ok(deadline) => { - if timeout_at(deadline, self.headers_provider.invalidate()) - .await - .is_err() - { - warn!(target: super::LOG_TARGET, - timeout_ms = self.options.recovery_timeout_ms, - "Terminal headers provider invalidation timed out" - ); + Some(Err(reconnect_error)) => { + if reconnect_error.is_auth_rejection() { + let mut invalidation = + self.spawn_headers_invalidation(recovery_deadline); + let invalidated = tokio::select! { + biased; + request = self.close.wait_for_request(&mut close_rx) => { + if request.is_some() { + return self.finish(Err(recovery_error)).await; + } + return self.finalized_result(); + } + result = &mut invalidation => result, + }; + match invalidated { + Ok(true) => reconnect_auth_retry = true, + Ok(false) | Err(_) => { + return self.finish(Err(reconnect_error)).await; + } } + } else if self.close.has_started() { + return self.finish(Err(recovery_error)).await; } - Err(deadline_error) => { - warn!(target: super::LOG_TARGET, - error = %deadline_error, - "Skipping terminal headers provider invalidation because its deadline is unrepresentable" - ); - } + pending_error = Some(reconnect_error); } } - return Err(error); + } + Err(error) => { + error!(target: super::LOG_TARGET, "Supervisor: Non-retriable error, closing stream: {}", error); + return self.finish(Err(error)).await; } } } } - /// Reconnects to the server and replays pending batches. - /// - /// On successful replay, holds `ingest_mutex` until `is_paused` is cleared so - /// subsequently admitted ingests send normally. Error paths remain paused for - /// supervisor retry or finalization. - async fn reconnect(&self) -> ZerobusResult<(FlightResponseStream, RequestBodyControl)> { + /// Completes setup, READY, replay, and sender publication as one cancellable attempt. + /// Cancellation before publication drops an established replacement best-effort. + async fn reconnect(&self) -> ZerobusResult> { let connection = ZerobusArrowStream::reconnect_transport( &self.endpoint, &self.tls_config, @@ -336,122 +375,164 @@ impl Supervisor { &self.sdk_identifier, ) .await?; - let (response_stream, tx, request_body) = connection.into_parts(); - - // Counters are reset atomically with the range rebuild inside - // replay_pending_batches, so a concurrent ingest can't fetch_add a reset counter, - // fabricate a low range, and have replay drop it as fully-acked. + let tx = connection.sender(); let acked_before_disconnect = self.last_acked_records.load(Ordering::Acquire); - // Test seam: pause after the connection is established but before ingest_mutex is - // held and ranges/watermark are rebuilt, so a test can schedule a paused ingest - // that wins ingest_mutex first (reset/rebase race) or drive a concurrent close(). + if self.replay_and_commit(&tx, acked_before_disconnect).await? { + Ok(Some(connection)) + } else { + Ok(None) + } + } + + async fn replay_and_commit( + &self, + tx: &mpsc::Sender>, + acked_before_disconnect: u64, + ) -> ZerobusResult { #[cfg(feature = "test-hooks")] { - let barrier = self.reconnect_rebuild_gate.lock().await.take(); + let barrier = self.test_hooks.reconnect_rebuild.lock().await.take(); if let Some(barrier) = barrier { barrier.reached.notify_one(); barrier.proceed.notified().await; } } - // Hold ingest_mutex across the replay so no concurrent ingest interleaves. - let _ingest_guard = self.ingest_mutex.lock().await; - let replay_result = Self::replay_pending_batches( - &tx, - &self.pending_batches, - &self.cumulative_records_assigned, + let replay_batches = { + let _ingest_guard = self.ingest_mutex.lock().await; + if self.close.has_started() { + return Ok(false); + } + Self::prepare_pending_replay( + &self.pending_batches, + &self.cumulative_records_assigned, + &self.submitted_records, + &self.last_acked_records, + acked_before_disconnect, + ) + .await + }; + + if !Self::send_replay_batches( + tx, + replay_batches, &self.submitted_records, - &self.last_acked_records, - acked_before_disconnect, + &self.ingest_mutex, + &self.close, #[cfg(feature = "test-hooks")] - Some(&self.replay_send_gate), + Some(&self.test_hooks.replay_send), ) - .await; - - // Commit the replacement sender only after replay succeeds. While ingest_mutex - // remains held, publish the sender before clearing the pause gate so normal ingest - // cannot observe an unpaused stream without its active sender. - Self::commit_reconnect_after_replay(replay_result, tx, &self.batch_tx, &self.is_paused) - .await?; + .await? + { + return Ok(false); + } - // ACK processing cannot resume until reconnect returns. Refresh only after replay - // is fully committed so connection setup, backlog sends, and sender publication do - // not consume any batch's new ACK budget. - let mut pending = self.pending_batches.lock().await; - refresh_pending_ack_deadlines(&mut pending, Instant::now()); - drop(pending); + loop { + let buffered = { + let ingest_guard = self.ingest_mutex.lock().await; + if self.close.has_started() { + return Ok(false); + } + let submitted = self.submitted_records.load(Ordering::Acquire); + let buffered = self + .pending_batches + .lock() + .await + .iter() + .find_map(|batch| batch.unacknowledged_suffix(submitted)); + if buffered.is_none() { + return Ok(Self::commit_reconnect( + tx.clone(), + &self.pending_batches, + &self.batch_tx, + &self.is_paused, + &self.close, + &ingest_guard, + ) + .await); + } + buffered + }; - Ok((response_stream, request_body)) + if !Self::send_replay_batch( + tx, + buffered.expect("buffered batch was selected"), + &self.submitted_records, + &self.ingest_mutex, + &self.close, + "Failed to replay buffered batch during recovery", + ) + .await? + { + return Ok(false); + } + } } - /// Commits the replacement sender after replay. The caller holds `ingest_mutex`. - async fn commit_reconnect_after_replay( - replay_result: ZerobusResult<()>, + async fn commit_reconnect( tx: mpsc::Sender>, + pending_batches: &Arc>>, batch_tx: &BatchSender, is_paused: &AtomicBool, - ) -> ZerobusResult<()> { - replay_result?; - { - let mut tx_guard = batch_tx.lock().await; - *tx_guard = Some(tx.clone()); + close: &CloseCoordinator, + _ingest_guard: &tokio::sync::MutexGuard<'_, ()>, + ) -> bool { + if close.has_started() { + return false; } + let mut pending = pending_batches.lock().await; + let mut sender = batch_tx.lock().await; + refresh_pending_ack_deadlines(&mut pending, Instant::now()); + *sender = Some(tx); is_paused.store(false, Ordering::Relaxed); - Ok(()) + true } - /// Rebuilds `pending_batches` for replay after a reconnect and replays them over - /// `tx`: partially-acked batches (vs `acked_before_disconnect`) are sliced to their - /// un-acked suffix, fully-acked ones dropped. - /// - /// The rebuilt pending set and the counter reset are installed together under the - /// `pending_batches` lock, before any send: a replay-send failure keeps pending (and - /// permits) intact, and no concurrent ingest can observe reset counters against stale - /// ranges. The caller refreshes pending ACK timestamps together only after replay is - /// committed, immediately before ACK processing can resume. Caller holds `ingest_mutex`. - async fn replay_pending_batches( - tx: &mpsc::Sender>, + async fn prepare_pending_replay( pending_batches: &Arc>>, cumulative_records_assigned: &Arc, submitted_records: &Arc, last_acked_records: &Arc, acked_before_disconnect: u64, - #[cfg(feature = "test-hooks")] replay_send_gate: Option<&super::ReplaySendGate>, - ) -> ZerobusResult<()> { - let replay_batches: Vec = { - let mut pending = pending_batches.lock().await; - - if !pending.is_empty() { - info!(target: super::LOG_TARGET, - batch_count = pending.len(), - acked_records = acked_before_disconnect, - "Replaying pending batches after recovery" - ); - } - let (replay, new_cumulative) = - rebuild_pending_for_replay(&mut pending, acked_before_disconnect); - - // Reset counters together with the range install, before any send. - cumulative_records_assigned.store(new_cumulative, Ordering::Relaxed); - submitted_records.store(0, Ordering::Release); - last_acked_records.store(0, Ordering::Release); - - replay - }; + ) -> Vec { + let mut pending = pending_batches.lock().await; + if !pending.is_empty() { + info!(target: super::LOG_TARGET, + batch_count = pending.len(), + acked_records = acked_before_disconnect, + "Replaying pending batches after recovery" + ); + } + let (replay, new_cumulative) = + rebuild_pending_for_replay(&mut pending, acked_before_disconnect); + cumulative_records_assigned.store(new_cumulative, Ordering::Relaxed); + submitted_records.store(0, Ordering::Release); + last_acked_records.store(0, Ordering::Release); + replay + } - // Send only after the pending_batches lock is released (ingest_mutex is still - // held by the caller); pending stays intact on failure. The replacement response - // stream is not polled until replay returns, so publishing after each successful - // handoff cannot race a valid acknowledgement on this connection. + async fn send_replay_batches( + tx: &mpsc::Sender>, + replay_batches: Vec, + submitted_records: &Arc, + ingest_mutex: &Arc>, + close: &CloseCoordinator, + #[cfg(feature = "test-hooks")] replay_send_gate: Option<&super::TestBarrierGate>, + ) -> ZerobusResult { for batch in replay_batches { - let record_count = batch.num_rows() as u64; - if tx.send(Ok(batch)).await.is_err() { - return Err(ZerobusError::StreamClosedError(tonic::Status::internal( - "Failed to replay batch during recovery", - ))); + if !Self::send_replay_batch( + tx, + batch, + submitted_records, + ingest_mutex, + close, + "Failed to replay batch during recovery", + ) + .await? + { + return Ok(false); } - submitted_records.fetch_add(record_count, Ordering::Release); #[cfg(feature = "test-hooks")] { @@ -465,51 +546,31 @@ impl Supervisor { } } } - - Ok(()) - } - - /// Moves each pending batch's unacknowledged suffix to the failed list, dropping - /// fully acknowledged batches. - pub(super) async fn move_pending_to_failed( - pending_batches: &Arc>>, - failed_batches: &Arc>>, - last_acked_records: &Arc, - ) { - // Lock failed first and hold it across the pending drain so this serializes with - // get_unacked_batches (which uses the same order): whichever runs first drains - // pending; the other then sees an empty pending and the same failed snapshot. - // Lock order is always failed -> pending; no path takes them in the reverse. - let mut failed = failed_batches.lock().await; - let mut pending = pending_batches.lock().await; - let acked = last_acked_records.load(Ordering::Acquire); - for pb in pending.drain(..) { - // Slice off any durably-acked prefix so a manual retry via - // get_unacked_batches doesn't re-send already-persisted records. - if let Some(batch) = pb.unacknowledged_suffix(acked) { - failed.push(batch); - } - } + Ok(true) } - /// Publishes stream closure and drains pending -> failed atomically with respect to - /// `ingest_batch`. Holding `ingest_mutex` across the `is_closed` store and the drain - /// means an ingest either finishes its append before this runs (and is drained here) - /// or observes `is_closed` after the mutex is released (and refuses to append), so a - /// retrieval snapshot can never omit an accepted batch that a later call reveals. - pub(super) async fn finalize_closed( + async fn send_replay_batch( + tx: &mpsc::Sender>, + batch: RecordBatch, + submitted_records: &Arc, ingest_mutex: &Arc>, - is_closed: &Arc, - pending_batches: &Arc>>, - failed_batches: &Arc>>, - last_acked_records: &Arc, - ) { - let _guard = ingest_mutex.lock().await; - is_closed.store(true, Ordering::Relaxed); - Self::move_pending_to_failed(pending_batches, failed_batches, last_acked_records).await; + close: &CloseCoordinator, + failure_message: &'static str, + ) -> ZerobusResult { + let permit = tx.reserve().await.map_err(|_| { + ZerobusError::StreamClosedError(tonic::Status::internal(failure_message)) + })?; + // Capacity waits stay outside the ingest lock. The close check and handoff share + // that lock with publication, ordering each replay batch wholly before or after it. + let _ingest_guard = ingest_mutex.lock().await; + if close.has_started() { + return Ok(false); + } + submitted_records.fetch_add(batch.num_rows() as u64, Ordering::Release); + permit.send(Ok(batch)); + Ok(true) } } - #[cfg(test)] mod tests { use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; @@ -517,16 +578,13 @@ mod tests { use arrow_array::Int32Array; use arrow_flight::error::FlightError; - use arrow_flight::PutResult; use arrow_schema::{DataType, Field, Schema as ArrowSchema}; - use futures::stream::iter; use tokio::sync::{mpsc, Mutex, Semaphore}; use tokio::time::{timeout, Duration, Instant}; - use super::super::metadata::FlightAckMetadata; + use super::super::close::{CloseCoordinator, CloseFinalizer, CloseRequest, CloseState}; use super::{ - pause_and_detach_sender, AckProcessor, BatchSender, PendingBatch, RecordBatch, Supervisor, - ZerobusError, + pause_and_detach_sender, BatchSender, PendingBatch, RecordBatch, Supervisor, ZerobusError, }; use crate::offset_generator::OffsetId; @@ -559,90 +617,103 @@ mod tests { ) } - #[tokio::test] - async fn failed_replay_leaves_sender_detached_and_closes_request_channel() { - let (tx, mut request_rx) = mpsc::channel::>(1); - let batch_tx: BatchSender = Arc::new(Mutex::new(None)); - let is_paused = AtomicBool::new(true); - let replay_result = Err(ZerobusError::StreamClosedError(tonic::Status::internal( - "failed replay", - ))); - - let result = - Supervisor::commit_reconnect_after_replay(replay_result, tx, &batch_tx, &is_paused) - .await; + #[test] + fn non_finalized_supervisor_exit_is_an_invariant_error() { + let request = CloseRequest { + target_offset: Some(0), + deadline: Instant::now() + Duration::from_secs(1), + }; - assert!(result.is_err()); - assert!(batch_tx.lock().await.is_none()); - assert!(is_paused.load(Ordering::Relaxed)); - assert!( - request_rx.recv().await.is_none(), - "failed replay must drop the only replacement sender" - ); + for state in [CloseState::Open, CloseState::Requested(request)] { + assert!(matches!( + Supervisor::result_from_close_state(state), + Err(ZerobusError::InvalidStateError(_)) + )); + } + assert!(Supervisor::result_from_close_state(CloseState::Finalized(Ok(()))).is_ok()); } #[tokio::test] - async fn regressive_ack_replays_only_unacknowledged_suffix() { + async fn close_publication_precedes_queued_replay_handoff() { let schema = one_col_schema(); - let sem = Arc::new(Semaphore::new(1)); - let pending_batches = Arc::new(Mutex::new(vec![pending_batch( - &sem, - batch_with_rows(&schema, 10), - 0, - 0, - 10, - )])); - let response_stream = iter([5, 0].map(|acked_records| { - Ok(PutResult { - app_metadata: serde_json::to_vec(&FlightAckMetadata { - ack_up_to_offset: 0, - ack_up_to_records: acked_records, - close_stream_duration_ms: None, - }) - .unwrap() - .into(), - }) - })); - let cumulative_records_assigned = Arc::new(AtomicU64::new(10)); - let submitted_records = Arc::new(AtomicU64::new(10)); - let last_acked_records = Arc::new(AtomicU64::new(0)); - let (processor, request_body, _last_ack_rx) = AckProcessor::for_test( - Arc::clone(&pending_batches), - Arc::clone(&submitted_records), - Arc::clone(&last_acked_records), - false, - ); + let ingest_mutex = Arc::new(Mutex::new(())); + let close = CloseCoordinator::new(); + let submitted = Arc::new(AtomicU64::new(0)); + let (tx, mut rx) = mpsc::channel::>(1); + let request = CloseRequest { + target_offset: Some(0), + deadline: Instant::now() + Duration::from_secs(30), + }; - let _stream_closed = processor - .process(Box::pin(response_stream), request_body) - .await; + let guard = ingest_mutex.lock().await; + let publish = async { + let _guard = ingest_mutex.lock().await; + close.publish(request); + }; + tokio::pin!(publish); + assert!(futures::poll!(publish.as_mut()).is_pending()); - let acked_before_disconnect = last_acked_records.load(Ordering::Acquire); - assert_eq!(acked_before_disconnect, 5); - let (tx, mut rx) = mpsc::channel::>(1); - Supervisor::replay_pending_batches( + let send = Supervisor::send_replay_batch( &tx, - &pending_batches, - &cumulative_records_assigned, - &submitted_records, - &last_acked_records, + batch_with_rows(&schema, 1), + &submitted, + &ingest_mutex, + &close, + "replay failed", + ); + tokio::pin!(send); + assert!(futures::poll!(send.as_mut()).is_pending()); + assert_eq!(tx.capacity(), 0, "the replay send must reserve capacity"); + + drop(guard); + timeout(Duration::from_secs(1), publish) + .await + .expect("close publication should acquire the mutex first"); + assert!(close.has_started()); + assert!( + !timeout(Duration::from_secs(1), send) + .await + .expect("replay send should resume after close publication") + .expect("replay send should not fail"), + "a published close must reject the queued replay handoff" + ); + assert_eq!(submitted.load(Ordering::Acquire), 0); + assert!(matches!( + rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + } + + async fn replay_pending_batches( + tx: &mpsc::Sender>, + pending_batches: &Arc>>, + cumulative_records_assigned: &Arc, + submitted_records: &Arc, + last_acked_records: &Arc, + acked_before_disconnect: u64, + ) -> crate::ZerobusResult<()> { + let ingest_mutex = Arc::new(Mutex::new(())); + let close = CloseCoordinator::new(); + let replay_batches = Supervisor::prepare_pending_replay( + pending_batches, + cumulative_records_assigned, + submitted_records, + last_acked_records, acked_before_disconnect, + ) + .await; + let sent = Supervisor::send_replay_batches( + tx, + replay_batches, + submitted_records, + &ingest_mutex, + &close, #[cfg(feature = "test-hooks")] None, ) - .await - .expect("replay should succeed"); - - let pending = pending_batches.lock().await; - assert_eq!(pending.len(), 1); - assert_eq!(pending[0].record_count(), 5); - assert_eq!(pending[0].record_range(), (0, 5)); - drop(pending); - assert_eq!(cumulative_records_assigned.load(Ordering::Relaxed), 5); - assert_eq!(submitted_records.load(Ordering::Acquire), 5); - assert_eq!(last_acked_records.load(Ordering::Acquire), 0); - assert_eq!(rx.try_recv().unwrap().unwrap().num_rows(), 5); - assert!(rx.try_recv().is_err()); + .await?; + assert!(sent, "close was not published in the unit helper"); + Ok(()) } /// A replay-send failure must not drop pending batches, their permits, or desync @@ -675,17 +746,8 @@ mod tests { let (tx, rx) = mpsc::channel::>(4); drop(rx); - let res = Supervisor::replay_pending_batches( - &tx, - &pending, - &cumulative, - &submitted, - &last_acked, - 0, - #[cfg(feature = "test-hooks")] - None, - ) - .await; + let res = + replay_pending_batches(&tx, &pending, &cumulative, &submitted, &last_acked, 0).await; assert!(res.is_err(), "replay must surface the send failure"); let guard = pending.lock().await; @@ -726,59 +788,6 @@ mod tests { ); } - /// With an open receiver, both batches remain pending, replay in order, and reset the - /// connection-relative counters. - #[tokio::test] - async fn replay_success_reinstalls_and_sends_all() { - let schema = one_col_schema(); - let sem = Arc::new(Semaphore::new(4)); - let original_enqueued_at = Instant::now() - Duration::from_secs(1); - let mut pending_batches = vec![ - pending_batch(&sem, batch_with_rows(&schema, 3), 0, 0, 3), - pending_batch(&sem, batch_with_rows(&schema, 2), 1, 3, 5), - ]; - for batch in &mut pending_batches { - batch.refresh_enqueued_at(original_enqueued_at); - } - let pending = Arc::new(Mutex::new(pending_batches)); - let cumulative = Arc::new(AtomicU64::new(0)); - let submitted = Arc::new(AtomicU64::new(0)); - let last_acked = Arc::new(AtomicU64::new(9)); - - let (tx, mut rx) = mpsc::channel::>(4); - - let res = Supervisor::replay_pending_batches( - &tx, - &pending, - &cumulative, - &submitted, - &last_acked, - 0, - #[cfg(feature = "test-hooks")] - None, - ) - .await; - assert!(res.is_ok()); - - let pending_guard = pending.lock().await; - assert_eq!(pending_guard.len(), 2); - assert!( - pending_guard - .iter() - .all(|batch| batch.enqueued_at() == original_enqueued_at), - "the replay send phase must not start pending ACK deadlines" - ); - drop(pending_guard); - assert_eq!(cumulative.load(Ordering::Relaxed), 5); - assert_eq!(submitted.load(Ordering::Acquire), 5); - assert_eq!(last_acked.load(Ordering::Relaxed), 0); - - let first = rx.try_recv().expect("first replay batch"); - assert_eq!(first.unwrap().num_rows(), 3); - let second = rx.try_recv().expect("second replay batch"); - assert_eq!(second.unwrap().num_rows(), 2); - } - /// A fully-acked batch is dropped during replay (permit released), and a partially /// acked batch is sliced to its un-acked suffix. #[tokio::test] @@ -797,22 +806,14 @@ mod tests { let last_acked = Arc::new(AtomicU64::new(4)); let (tx, mut rx) = mpsc::channel::>(4); - let res = Supervisor::replay_pending_batches( - &tx, - &pending, - &cumulative, - &submitted, - &last_acked, - 4, - #[cfg(feature = "test-hooks")] - None, - ) - .await; + let res = + replay_pending_batches(&tx, &pending, &cumulative, &submitted, &last_acked, 4).await; assert!(res.is_ok()); // Only the partially-acked batch remains, rebuilt from cumulative 0. let guard = pending.lock().await; assert_eq!(guard.len(), 1); + assert_eq!(guard[0].record_count(), 2); assert_eq!(guard[0].record_range(), (0, 2)); drop(guard); assert_eq!(cumulative.load(Ordering::Relaxed), 2); @@ -884,8 +885,13 @@ mod tests { // about to append): hold ingest_mutex. let guard = ingest_mutex.lock().await; - let fut = - Supervisor::finalize_closed(&ingest_mutex, &is_closed, &pending, &failed, &last_acked); + let fut = CloseFinalizer::finalize_closed( + &ingest_mutex, + &is_closed, + &pending, + &failed, + &last_acked, + ); tokio::pin!(fut); // Finalization must block while the ingest holds ingest_mutex, and must not @@ -920,4 +926,34 @@ mod tests { ); assert!(pending.lock().await.is_empty()); } + + #[tokio::test] + async fn published_close_prevents_sender_commit() { + let close = CloseCoordinator::new(); + let request = CloseRequest { + target_offset: Some(7), + deadline: Instant::now() + Duration::from_secs(30), + }; + close.publish(request); + + let (tx, _rx) = mpsc::channel(1); + let pending = Arc::new(Mutex::new(Vec::new())); + let batch_tx: BatchSender = Arc::new(Mutex::new(None)); + let is_paused = AtomicBool::new(true); + let ingest_mutex = Mutex::new(()); + let ingest_guard = ingest_mutex.lock().await; + assert!( + !Supervisor::commit_reconnect( + tx, + &pending, + &batch_tx, + &is_paused, + &close, + &ingest_guard, + ) + .await + ); + assert!(batch_tx.lock().await.is_none()); + assert!(is_paused.load(Ordering::Relaxed)); + } } diff --git a/rust/tests/src/arrow_tests.rs b/rust/tests/src/arrow_tests.rs index eda68c56..221aec6a 100644 --- a/rust/tests/src/arrow_tests.rs +++ b/rust/tests/src/arrow_tests.rs @@ -2,6 +2,7 @@ mod mock_arrow_flight; mod utils; mod arrow_flight_tests { + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use arrow_array::{Array, RecordBatch, StructArray}; @@ -9,7 +10,11 @@ mod arrow_flight_tests { use databricks_zerobus_ingest_sdk::internal::arrow_c_data::{ import_c_data_record_batch, FFI_ArrowArray, FFI_ArrowSchema, }; - use databricks_zerobus_ingest_sdk::{NoTlsConfig, ZerobusError, ZerobusSdk}; + use databricks_zerobus_ingest_sdk::{ + ConnectorFactory, NoTlsConfig, ProxyConnector, ZerobusError, ZerobusSdk, + }; + use tokio::io::AsyncReadExt; + use tokio::net::TcpListener; use tracing::info; use crate::mock_arrow_flight::{start_mock_flight_server, MockFlightResponse}; @@ -22,6 +27,23 @@ mod arrow_flight_tests { const TABLE_NAME: &str = "test_catalog.test_schema.test_table"; + struct PausedTimeAutoAdvanceGuard(tokio::task::JoinHandle<()>); + + impl PausedTimeAutoAdvanceGuard { + fn start() -> Self { + Self(tokio::spawn(std::future::poll_fn(|cx| { + cx.waker().wake_by_ref(); + std::task::Poll::<()>::Pending + }))) + } + } + + impl Drop for PausedTimeAutoAdvanceGuard { + fn drop(&mut self) { + self.0.abort(); + } + } + /// Extracts the `id` (Int64) column values from a test batch, for asserting that a /// recovered/sliced batch contains the expected rows (not just the right row count). fn batch_ids(batch: &arrow_array::RecordBatch) -> Vec { @@ -146,7 +168,7 @@ mod arrow_flight_tests { } #[tokio::test(start_paused = true)] - async fn test_unrepresentable_recovery_timeout_is_rejected( + async fn test_unrepresentable_arrow_timeouts_are_rejected( ) -> Result<(), Box> { setup_tracing(); @@ -160,56 +182,108 @@ mod arrow_flight_tests { advance_tokio_time_near_instant_limit().await; - let stream_result = sdk - .stream_builder() - .table(TABLE_NAME) - .headers_provider(Arc::new(TestHeadersProvider::default())) - .arrow(schema) - .recovery_timeout_ms(u64::MAX) - .build_arrow() - .await; + for option_name in [ + "recovery_timeout_ms", + "server_lack_of_ack_timeout_ms", + "flush_timeout_ms", + ] { + let builder = sdk + .stream_builder() + .table(TABLE_NAME) + .headers_provider(Arc::new(TestHeadersProvider::default())) + .arrow(Arc::clone(&schema)); + let stream_result = match option_name { + "recovery_timeout_ms" => { + builder.recovery_timeout_ms(u64::MAX).build_arrow().await + } + "server_lack_of_ack_timeout_ms" => { + builder + .server_lack_of_ack_timeout_ms(u64::MAX) + .build_arrow() + .await + } + "flush_timeout_ms" => builder.flush_timeout_ms(u64::MAX).build_arrow().await, + _ => unreachable!(), + }; - match stream_result { - Err(ZerobusError::InvalidArgument(message)) => { - assert!(message.contains("recovery_timeout_ms")); + match stream_result { + Err(ZerobusError::InvalidArgument(message)) => { + assert!(message.contains(option_name)); + } + Err(error) => panic!("unexpected stream creation error: {error}"), + Ok(_) => panic!("unrepresentable {option_name} was accepted"), } - Err(error) => panic!("unexpected stream creation error: {error}"), - Ok(_) => panic!("unrepresentable recovery timeout was accepted"), } Ok(()) } - #[tokio::test(start_paused = true)] - async fn test_unrepresentable_ack_timeout_is_rejected( + #[tokio::test] + async fn test_unrepresentable_runtime_close_deadline_does_not_publish_close( ) -> Result<(), Box> { + const FLUSH_TIMEOUT_MS: u64 = 32 * 365 * 24 * 60 * 60 * 1_000; + setup_tracing(); - let (_mock_server, server_url) = start_mock_flight_server().await?; + let (mock_server, server_url) = start_mock_flight_server().await?; let schema = create_test_arrow_schema(); + mock_server + .inject_responses( + TABLE_NAME, + vec![MockFlightResponse::Error { + status: Status::invalid_argument("permanent cleanup error"), + delay_ms: 0, + }], + ) + .await; + let sdk = ZerobusSdk::builder() .endpoint(server_url) .unity_catalog_url("https://mock-uc.com") .tls_config(Arc::new(NoTlsConfig)) .build()?; - - advance_tokio_time_near_instant_limit().await; - - let stream_result = sdk + let mut stream = sdk .stream_builder() .table(TABLE_NAME) .headers_provider(Arc::new(TestHeadersProvider::default())) - .arrow(schema) - .server_lack_of_ack_timeout_ms(u64::MAX) + .arrow(schema.clone()) + .recovery(false) + .flush_timeout_ms(FLUSH_TIMEOUT_MS) .build_arrow() - .await; + .await?; + + // Real tonic setup must finish before Tokio time is paused. + tokio::time::pause(); + advance_tokio_time_near_instant_limit().await; - match stream_result { + match stream.close().await { Err(ZerobusError::InvalidArgument(message)) => { - assert!(message.contains("server_lack_of_ack_timeout_ms")); + assert!(message.contains("flush_timeout_ms")); } - Err(error) => panic!("unexpected stream creation error: {error}"), - Ok(_) => panic!("unrepresentable ACK timeout was accepted"), + Err(error) => panic!("unexpected close error: {error}"), + Ok(()) => panic!("unrepresentable close deadline was accepted"), + } + assert!(!stream.is_closed()); + + // A successful ingest proves the failed deadline calculation did not publish close. + let batch = create_test_record_batch(schema, vec![1], vec![Some("still usable")]); + stream.ingest_batch(batch).await?; + run_with_paused_time_watchdog(async { + while !stream.is_closed() { + tokio::task::yield_now().await; + } + }) + .await; + + let error = run_with_paused_time_watchdog(stream.close()) + .await + .expect_err("repeated close must observe the stored terminal error"); + match error { + ZerobusError::StreamClosedError(status) => { + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert_eq!(status.message(), "permanent cleanup error"); + } + other => panic!("expected the stored peer error, got {other:?}"), } Ok(()) @@ -1174,14 +1248,14 @@ mod arrow_flight_tests { assert!(!stream.is_closed()); stream.close().await?; assert!(stream.is_closed()); + assert_eq!(mock_server.get_request_half_close_count(), 1); + assert_eq!(mock_server.get_request_reset_count(), 0); Ok(()) } - /// Cancelling close after supervisor/sender teardown must leave the stream in a - /// resumable Closing state: new ingests are rejected, retrieval remains disabled, - /// and resumed/repeated close calls return the original flush error without waiting - /// for flush_timeout again. + /// Cancelling the caller after close publication leaves supervisor-owned teardown + /// running. A repeated call waits for the same deadline and result. #[tokio::test] async fn test_cancelled_close_rejects_ingest_and_resumes_teardown( ) -> Result<(), Box> { @@ -1220,7 +1294,7 @@ mod arrow_flight_tests { create_test_record_batch(schema.clone(), vec![1], vec![Some("pending")]); stream.ingest_batch(pending_batch).await?; - let (reached, _proceed) = stream.arm_close_finalize_barrier().await; + let (reached, proceed) = stream.arm_close_finalize_barrier().await; let mut close_future = Box::pin(stream.close()); tokio::time::timeout(std::time::Duration::from_secs(5), async { tokio::select! { @@ -1248,6 +1322,7 @@ mod arrow_flight_tests { "unacked retrieval is allowed only after Closed" ); + proceed.notify_one(); let resumed = tokio::time::timeout(std::time::Duration::from_secs(1), stream.close()) .await .expect("resumed close must skip flush and finish promptly") @@ -2455,22 +2530,19 @@ mod arrow_flight_tests { } #[tokio::test] - async fn test_close_propagates_flush_error() -> Result<(), Box> { + async fn test_late_drain_ack_does_not_replace_close_timeout( + ) -> Result<(), Box> { setup_tracing(); - info!("Starting test_close_propagates_flush_error"); + info!("Starting test_late_drain_ack_does_not_replace_close_timeout"); let (mock_server, server_url) = start_mock_flight_server().await?; let schema = create_test_arrow_schema(); - // Ack is delayed well past the flush timeout, so the flush performed by - // close() times out while the stream is still open (no server error, so the - // idempotent close guard does not short-circuit). mock_server .inject_responses( TABLE_NAME, - vec![MockFlightResponse::BatchAck { + vec![MockFlightResponse::BatchAckAfterRequestEof { ack_up_to_offset: 0, - delay_ms: 5000, ack_up_to_records: 1, }], ) @@ -2494,16 +2566,25 @@ mod arrow_flight_tests { let batch = create_test_record_batch(schema, vec![1], vec![Some("unacked")]); let _offset = stream.ingest_batch(batch).await?; - let close_result = stream.close().await; - assert!( - close_result.is_err(), - "close() must propagate the error when its flush fails" - ); + let error = stream + .close() + .await + .expect_err("the post-deadline ACK must not replace the selected timeout"); + match error { + ZerobusError::StreamClosedError(status) => { + assert_eq!(status.code(), tonic::Code::DeadlineExceeded); + assert_eq!(status.message(), "Flush timed out"); + } + other => panic!("expected a flush timeout, got {other:?}"), + } - // The stream is still torn down, and the unacked batch is recoverable. assert!(stream.is_closed()); - let unacked = stream.get_unacked_batches().await?; - assert_eq!(unacked.len(), 1, "Should have 1 unacked batch"); + assert!( + stream.get_unacked_batches().await?.is_empty(), + "the late ACK still advances the retained suffix" + ); + assert_eq!(mock_server.get_request_half_close_count(), 1); + assert_eq!(mock_server.get_request_reset_count(), 0); Ok(()) } @@ -3040,7 +3121,7 @@ mod arrow_flight_tests { /// The supervisor must surface a non-retryable configuration error if a timeout /// that was representable at construction can no longer form a runtime deadline. - #[tokio::test(start_paused = true)] + #[tokio::test] async fn test_unrepresentable_runtime_recovery_deadline_is_rejected( ) -> Result<(), Box> { const RECOVERY_TIMEOUT_MS: u64 = 32 * 365 * 24 * 60 * 60 * 1_000; @@ -3076,6 +3157,8 @@ mod arrow_flight_tests { .build_arrow() .await?; + // Real tonic setup must finish before Tokio time is paused. + tokio::time::pause(); advance_tokio_time_near_instant_limit().await; let batch = create_test_record_batch(schema, vec![1], vec![Some("replay")]); @@ -3099,6 +3182,165 @@ mod arrow_flight_tests { Ok(()) } + #[tokio::test] + async fn test_close_during_reconnect_setup_preserves_recovery_trigger( + ) -> Result<(), Box> { + setup_tracing(); + + let (mock_server, server_url) = start_mock_flight_server().await?; + let schema = create_test_arrow_schema(); + mock_server + .inject_responses( + TABLE_NAME, + vec![ + MockFlightResponse::Error { + status: tonic::Status::unavailable("active transport failed"), + delay_ms: 0, + }, + MockFlightResponse::FailSetupAfter { + status: tonic::Status::unavailable("replacement setup failed"), + delay_ms: 60_000, + }, + ], + ) + .await; + let setup_armed = mock_server.delayed_setup_armed(); + + let sdk = ZerobusSdk::builder() + .endpoint(server_url) + .unity_catalog_url("https://mock-uc.com") + .tls_config(Arc::new(NoTlsConfig)) + .build()?; + let mut stream = sdk + .stream_builder() + .table(TABLE_NAME) + .headers_provider(Arc::new(TestHeadersProvider::default())) + .arrow(schema.clone()) + .recovery(true) + .recovery_backoff_ms(0) + .recovery_timeout_ms(120_000) + .recovery_retries(1) + .flush_timeout_ms(60_000) + .build_arrow() + .await?; + + let batch = create_test_record_batch(schema, vec![1], vec![Some("pending")]); + stream.ingest_batch(batch).await?; + tokio::time::timeout(std::time::Duration::from_secs(5), setup_armed.notified()) + .await + .expect("reconnect must enter replacement setup"); + + let error = tokio::time::timeout(std::time::Duration::from_secs(1), stream.close()) + .await + .expect("close must cancel replacement setup") + .expect_err("close during recovery must preserve the recovery trigger"); + assert!( + error.to_string().contains("active transport failed"), + "expected the active recovery trigger, got: {error}" + ); + assert_eq!(stream.get_unacked_batches().await?.len(), 1); + Ok(()) + } + + #[tokio::test] + async fn test_close_during_reconnect_transport_handshake_preserves_trigger_and_suffix( + ) -> Result<(), Box> { + setup_tracing(); + + let (mock_server, server_url) = start_mock_flight_server().await?; + let schema = create_test_arrow_schema(); + mock_server + .inject_responses( + TABLE_NAME, + vec![MockFlightResponse::Error { + status: tonic::Status::unavailable("active transport failed"), + delay_ms: 0, + }], + ) + .await; + + let listener = TcpListener::bind("127.0.0.1:0").await?; + let proxy_url = format!("http://{}", listener.local_addr()?); + let connect_received = Arc::new(tokio::sync::Notify::new()); + let proxy_task = { + let connect_received = Arc::clone(&connect_received); + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("proxy must accept"); + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + while !request.ends_with(b"\r\n\r\n") { + let read = socket + .read(&mut buffer) + .await + .expect("read CONNECT request"); + assert!(read > 0, "proxy client closed before sending CONNECT"); + request.extend_from_slice(&buffer[..read]); + assert!(request.len() <= 8 * 1024, "CONNECT request is too large"); + } + assert!(request.starts_with(b"CONNECT ")); + connect_received.notify_one(); + std::future::pending::<()>().await; + }) + }; + + let factory_calls = Arc::new(AtomicUsize::new(0)); + let calls = Arc::clone(&factory_calls); + let connector_factory: ConnectorFactory = Arc::new(move |_| { + if calls.fetch_add(1, Ordering::SeqCst) == 0 { + None + } else { + Some(ProxyConnector::new(&proxy_url).expect("valid local proxy URL")) + } + }); + let sdk = ZerobusSdk::builder() + .endpoint(server_url) + .unity_catalog_url("https://mock-uc.com") + .tls_config(Arc::new(NoTlsConfig)) + .connector_factory(connector_factory) + .build()?; + let mut stream = sdk + .stream_builder() + .table(TABLE_NAME) + .headers_provider(Arc::new(TestHeadersProvider::default())) + .arrow(schema.clone()) + .recovery(true) + .recovery_backoff_ms(0) + .recovery_timeout_ms(120_000) + .recovery_retries(1) + .flush_timeout_ms(60_000) + .build_arrow() + .await?; + + let batch = create_test_record_batch(schema, vec![7, 8], vec![Some("a"), Some("b")]); + stream.ingest_batch(batch).await?; + tokio::time::timeout( + std::time::Duration::from_secs(5), + connect_received.notified(), + ) + .await + .expect("reconnect must wait for the proxy CONNECT response"); + + let error = tokio::time::timeout(std::time::Duration::from_secs(1), stream.close()) + .await + .expect("close must cancel the replacement transport handshake") + .expect_err("close during recovery must preserve the recovery trigger"); + assert!( + error.to_string().contains("active transport failed"), + "expected the active recovery trigger, got: {error}" + ); + let unacked: Vec> = stream + .get_unacked_batches() + .await? + .iter() + .map(batch_ids) + .collect(); + assert_eq!(unacked, vec![vec![7, 8]]); + + proxy_task.abort(); + let _ = proxy_task.await; + Ok(()) + } + /// Time spent replaying a recovered backlog must not consume the batches' ACK /// budget before the replacement response processor can observe acknowledgments. #[tokio::test(start_paused = true)] @@ -3109,8 +3351,10 @@ mod arrow_flight_tests { const ACK_DELAY_MS: u64 = 50; setup_tracing(); - let (mock_server, server_url) = start_mock_flight_server().await?; + // Real tonic setup under paused time needs a ready task to prevent Tokio from + // auto-advancing connection deadlines while kernel I/O is still pending. + let _auto_advance_guard = PausedTimeAutoAdvanceGuard::start(); let schema = create_test_arrow_schema(); mock_server .inject_responses( @@ -3178,11 +3422,8 @@ mod arrow_flight_tests { Ok(()) } - /// close() during the reconnect rebuild window must slice with the pre-reconnect - /// watermark. A barrier parks reconnect after the new connection is established but - /// before pending ranges/watermark are rebuilt; close() then reaps the parked - /// supervisor and drains with the pre-reconnect watermark/ranges, so A is sliced to - /// its un-acked suffix (ids [2, 3]) and B is retained whole. + /// Close after replacement setup and READY preserves the active recovery trigger + /// and slices pending work with the pre-reconnect watermark. #[tokio::test] async fn test_close_during_reconnect_rebuild_window_slices_correctly( ) -> Result<(), Box> { @@ -3225,8 +3466,7 @@ mod arrow_flight_tests { .recovery(true) .recovery_backoff_ms(0) .recovery_retries(5) - // Short flush timeout so close()'s flush returns quickly while parked. - .flush_timeout_ms(200) + .flush_timeout_ms(60_000) .build_arrow() .await?; @@ -3258,8 +3498,14 @@ mod arrow_flight_tests { .await .expect("reconnect should reach the rebuild barrier"); - // close() reaps the parked supervisor and drains with the pre-rebuild state. - let _ = stream.close().await; + let error = tokio::time::timeout(std::time::Duration::from_secs(1), stream.close()) + .await + .expect("close must interrupt recovery after READY") + .expect_err("close during recovery must preserve the recovery trigger"); + assert!( + error.to_string().contains("Connection lost"), + "expected the active recovery trigger, got: {error}" + ); let unacked = stream.get_unacked_batches().await?; assert_eq!(unacked.len(), 2, "expected sliced A suffix + full B"); @@ -3273,6 +3519,82 @@ mod arrow_flight_tests { Ok(()) } + #[tokio::test] + async fn test_close_after_partial_replay_preserves_trigger_and_suffix( + ) -> Result<(), Box> { + setup_tracing(); + + let (mock_server, server_url) = start_mock_flight_server().await?; + let schema = create_test_arrow_schema(); + mock_server + .inject_responses( + TABLE_NAME, + vec![ + MockFlightResponse::BatchAck { + ack_up_to_offset: 0, + delay_ms: 0, + ack_up_to_records: 1, + }, + MockFlightResponse::Error { + status: tonic::Status::unavailable("active transport failed"), + delay_ms: 0, + }, + ], + ) + .await; + + let sdk = ZerobusSdk::builder() + .endpoint(server_url) + .unity_catalog_url("https://mock-uc.com") + .tls_config(Arc::new(NoTlsConfig)) + .build()?; + let mut stream = sdk + .stream_builder() + .table(TABLE_NAME) + .headers_provider(Arc::new(TestHeadersProvider::default())) + .arrow(schema.clone()) + .recovery(true) + .recovery_backoff_ms(0) + .recovery_retries(1) + .flush_timeout_ms(60_000) + .build_arrow() + .await?; + + let ack_applied = stream.arm_ack_applied_notify().await; + let (replay_reached, _replay_proceed) = stream.arm_replay_send_barrier().await; + let first = create_test_record_batch( + schema.clone(), + vec![1, 2, 3], + vec![Some("a"), Some("b"), Some("c")], + ); + stream.ingest_batch(first).await?; + tokio::time::timeout(std::time::Duration::from_secs(5), ack_applied.notified()) + .await + .expect("partial ACK must be applied before recovery"); + let second = create_test_record_batch(schema, vec![4], vec![Some("d")]); + stream.ingest_batch(second).await?; + tokio::time::timeout(std::time::Duration::from_secs(5), replay_reached.notified()) + .await + .expect("recovery must hand off the first replay batch"); + + let error = tokio::time::timeout(std::time::Duration::from_secs(1), stream.close()) + .await + .expect("close must interrupt partial replay") + .expect_err("close during recovery must preserve the recovery trigger"); + assert!( + error.to_string().contains("active transport failed"), + "expected the active recovery trigger, got: {error}" + ); + let unacked: Vec> = stream + .get_unacked_batches() + .await? + .iter() + .map(batch_ids) + .collect(); + assert_eq!(unacked, vec![vec![2, 3], vec![4]]); + Ok(()) + } + /// Regression test for the mock contract: `ack_up_to_records` must be /// connection-relative. A retriable error forces a second DoPut connection /// for the same table; the auto-ack on that connection must count only the @@ -3459,76 +3781,93 @@ mod arrow_flight_tests { Ok(()) } - /// close() while a reconnect is parked at the rebuild barrier must tear the stream - /// down and move pending batches to the failed set, without panicking or hanging - /// beyond the flush timeout. #[tokio::test] - async fn test_close_during_reconnect_window_moves_pending_to_failed( + async fn test_close_during_recovery_backoff_preserves_trigger_and_suffix( ) -> Result<(), Box> { setup_tracing(); - info!("Starting test_close_during_reconnect_window_moves_pending_to_failed"); let (mock_server, server_url) = start_mock_flight_server().await?; let schema = create_test_arrow_schema(); - - // A retriable error triggers the reconnect we park at the rebuild barrier. mock_server .inject_responses( TABLE_NAME, - vec![MockFlightResponse::Error { - status: tonic::Status::unavailable("Connection lost"), - delay_ms: 0, - }], + vec![ + MockFlightResponse::BatchAck { + ack_up_to_offset: 0, + delay_ms: 0, + ack_up_to_records: 1, + }, + MockFlightResponse::Error { + status: tonic::Status::unavailable("active transport failed"), + delay_ms: 0, + }, + ], ) .await; let sdk = ZerobusSdk::builder() - .endpoint(server_url.clone()) + .endpoint(server_url) .unity_catalog_url("https://mock-uc.com") .tls_config(Arc::new(NoTlsConfig)) .build()?; - let mut stream = sdk .stream_builder() .table(TABLE_NAME) .headers_provider(Arc::new(TestHeadersProvider::default())) .arrow(schema.clone()) .recovery(true) - .recovery_backoff_ms(0) - .recovery_retries(5) - .flush_timeout_ms(200) + .recovery_backoff_ms(60_000) + .recovery_retries(1) + .flush_timeout_ms(120_000) .build_arrow() .await?; - // Park the reconnect; do not release it — close() reaps it instead. - let (reached, _proceed) = stream.arm_reconnect_rebuild_barrier().await; + // Real tonic setup must finish before Tokio time is paused. + tokio::time::pause(); + let ack_applied = stream.arm_ack_applied_notify().await; + let first = create_test_record_batch( + schema.clone(), + vec![1, 2, 3], + vec![Some("a"), Some("b"), Some("c")], + ); + stream.ingest_batch(first).await?; + run_with_paused_time_watchdog(ack_applied.notified()).await; - let batch = create_test_record_batch(schema.clone(), vec![1], vec![Some("a")]); - stream.ingest_batch(batch).await?; + let second = create_test_record_batch(schema, vec![4], vec![Some("d")]); + stream.ingest_batch(second).await?; - tokio::time::timeout(std::time::Duration::from_secs(5), reached.notified()) - .await - .expect("reconnect should reach the rebuild barrier"); + // Keep paused time from advancing through the long backoff while the mock + // delivers the active error and the supervisor enters recovery. + let watchdog_started = std::time::Instant::now(); + while mock_server.get_batch_count().await < 2 { + assert!( + watchdog_started.elapsed() < std::time::Duration::from_secs(5), + "active connection did not receive both batches" + ); + tokio::task::yield_now().await; + } + for _ in 0..32 { + tokio::task::yield_now().await; + } - // close() must return (its flush times out at 200ms) without hanging or - // panicking, reaping the parked supervisor along the way. - let close_result = - tokio::time::timeout(std::time::Duration::from_secs(5), stream.close()) - .await - .expect("close() must not hang while a reconnect is parked"); + let close_error = run_with_paused_time_watchdog(stream.close()) + .await + .expect_err("close during backoff must preserve the recovery trigger"); assert!( - close_result.is_err(), - "close() should surface the flush timeout" + close_error.to_string().contains("active transport failed"), + "expected the active recovery trigger, got: {close_error}" ); - - // The un-acked batch was moved to the failed set and is retrievable. - let unacked = stream.get_unacked_batches().await?; + let unacked: Vec> = stream + .get_unacked_batches() + .await? + .iter() + .map(batch_ids) + .collect(); assert_eq!( - unacked.len(), - 1, - "pending batch must be moved to the failed set" + unacked, + vec![vec![2, 3], vec![4]], + "close during backoff must retain only the exact unacknowledged suffix" ); - Ok(()) } @@ -3569,7 +3908,7 @@ mod arrow_flight_tests { .tls_config(Arc::new(NoTlsConfig)) .build()?; - let stream = sdk + let mut stream = sdk .stream_builder() .table(TABLE_NAME) .headers_provider(Arc::new(TestHeadersProvider::default())) @@ -3592,6 +3931,20 @@ mod arrow_flight_tests { let result = stream.wait_for_offset(offset2).await; assert!(result.is_ok(), "Expected recovery to succeed: {:?}", result); + let half_closes_before = mock_server.get_request_half_close_count(); + let resets_before = mock_server.get_request_reset_count(); + stream.close().await?; + assert_eq!( + mock_server.get_request_half_close_count(), + half_closes_before + 1, + "close after sender commit must half-close the replacement exactly once" + ); + assert_eq!( + mock_server.get_request_reset_count(), + resets_before, + "close after sender commit must not reset the replacement" + ); + Ok(()) } @@ -5393,6 +5746,19 @@ mod arrow_flight_tests { create_test_record_batch(schema, vec![1, 2, 3], vec![Some("a"), Some("b"), Some("c")]) } + fn assert_rotation_error(error: ZerobusError) { + match error { + ZerobusError::StreamClosedError(status) => { + assert_eq!(status.code(), tonic::Code::Unavailable); + assert_eq!( + status.message(), + "Server requested graceful stream rotation" + ); + } + other => panic!("expected rotation error, got {other:?}"), + } + } + #[tokio::test] async fn test_server_rotation_half_closes_before_reconnect( ) -> Result<(), Box> { @@ -5810,7 +6176,11 @@ mod arrow_flight_tests { assert_eq!(mock_server.get_request_half_close_count(), 1); assert_eq!(mock_server.get_request_reset_count(), 0); - stream.close().await?; + // The ACK is observable before rotation recovery commits. Cleanup may close + // the replacement or interrupt the still-active rotation. + if let Err(error) = stream.close().await { + assert_rotation_error(error); + } Ok(()) } @@ -5931,7 +6301,10 @@ mod arrow_flight_tests { ); assert_eq!(mock_server.get_request_half_close_count(), 1); assert_eq!(mock_server.get_request_reset_count(), 0); - stream.close().await?; + // The post-EOF ACK is observable before rotation recovery commits. + if let Err(error) = stream.close().await { + assert_rotation_error(error); + } Ok(()) } From 1dca65f57c7689b20367cb99a73c60c89f96ccec Mon Sep 17 00:00:00 2001 From: teodordelibasic-db Date: Thu, 13 Aug 2026 12:32:28 +0000 Subject: [PATCH 2/6] [Rust] Finalize close on abort Supervisor panic or abort left close() waiting on CloseState forever. A detached reaper now finalizes the coordinator, and admission closes before the unacked snapshot is published. Signed-off-by: teodordelibasic-db --- rust/NEXT_CHANGELOG.md | 8 +- rust/sdk/src/stream/arrow/close.rs | 5 + rust/sdk/src/stream/arrow/mod.rs | 108 +++++++-- rust/sdk/src/stream/arrow/supervisor.rs | 89 ++++++-- rust/tests/src/arrow_tests.rs | 289 +++++++++++++++++++++++- 5 files changed, 446 insertions(+), 53 deletions(-) diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index 773d63f8..34c28900 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -22,8 +22,12 @@ monotonic-clock range. Server-advertised graceful-rotation periods are capped at one year. - Arrow Flight close is cancellation-safe and half-closes the active request before - bounded response draining. Close during recovery cancels the attempt, retains the - unacknowledged suffix, and returns the error that triggered recovery. + bounded response draining. ACK success is decided when the durable watermark is + applied relative to the original flush deadline. Close during recovery cancels the + attempt, retains the unacknowledged suffix, and returns the error that triggered the + current attempt. Close during an existing recovery or server-requested rotation keeps + that trigger even if every record is durable, so an error can coexist with an empty + unacknowledged-batch set. - Fixed Arrow Flight recovery sender lifetime: replacement senders are now published only after pending replay succeeds, while initial supervisor handoff and failed or cancelled replay promptly drop redundant senders instead of retaining incomplete diff --git a/rust/sdk/src/stream/arrow/close.rs b/rust/sdk/src/stream/arrow/close.rs index e3514fc8..62876651 100644 --- a/rust/sdk/src/stream/arrow/close.rs +++ b/rust/sdk/src/stream/arrow/close.rs @@ -162,6 +162,7 @@ pub(super) struct CloseFinalizer { ingest_mutex: Arc>, batch_tx: BatchSender, is_paused: Arc, + admission_closed: Arc, is_closed: Arc, pending_batches: Arc>>, failed_batches: Arc>>, @@ -178,6 +179,7 @@ impl CloseFinalizer { ingest_mutex: Arc::clone(&stream.ingest_mutex), batch_tx: Arc::clone(&stream.batch_tx), is_paused: Arc::clone(&stream.is_paused), + admission_closed: Arc::clone(&stream.admission_closed), is_closed: Arc::clone(&stream.is_closed), pending_batches: Arc::clone(&stream.pending_batches), failed_batches: Arc::clone(&stream.failed_batches), @@ -194,6 +196,9 @@ impl CloseFinalizer { if let CloseState::Finalized(existing) = self.close.state() { return existing; } + // `is_closed` remains false until the retained-batch snapshot is complete. + // Release: empty flush loads this flag without ingest_mutex. + self.admission_closed.store(true, Ordering::Release); self.is_paused.store(true, Ordering::Relaxed); *self.batch_tx.lock().await = None; } diff --git a/rust/sdk/src/stream/arrow/mod.rs b/rust/sdk/src/stream/arrow/mod.rs index b511bb87..99e11f8e 100644 --- a/rust/sdk/src/stream/arrow/mod.rs +++ b/rust/sdk/src/stream/arrow/mod.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use arrow_flight::error::FlightError; use bytes::Bytes; use tokio::sync::{mpsc, watch, Mutex, Notify, Semaphore}; -use tokio::task::JoinHandle; +use tokio::task::AbortHandle; use tokio::time::{timeout, Duration, Instant}; use tokio_retry::strategy::FixedInterval; use tokio_retry::RetryIf; @@ -172,10 +172,12 @@ pub struct ZerobusArrowStream { _last_ack_rx: watch::Receiver>, /// True once the stream is terminally closed and unacknowledged batches may be retrieved. is_closed: Arc, + /// Rejects new ingests as soon as terminal finalization owns admission. + admission_closed: Arc, /// Coordinates one resumable explicit-close request with the recovery supervisor. close: CloseCoordinator, - /// Handle to the supervisor task that processes acknowledgments and recovery. - receiver_task: Arc>>>>, + /// Abort handle for the supervisor worker; its detached reaper remains independent. + supervisor_abort: Arc>>, /// Accepted batches not yet fully acknowledged; retained for replay or retrieval. pending_batches: Arc>>, /// Wakes the ACK processor when a batch is submitted after an idle period. @@ -264,13 +266,14 @@ impl ZerobusArrowStream { let (last_ack_tx, _last_ack_rx) = watch::channel(None); let is_closed = Arc::new(AtomicBool::new(false)); + let admission_closed = Arc::new(AtomicBool::new(false)); let pending_batches = Arc::new(Mutex::new(Vec::new())); let pending_notify = Arc::new(Notify::new()); let request_send_failure = Arc::new(acks::RequestSendFailure::default()); let failed_batches = Arc::new(Mutex::new(Vec::new())); let recovery_attempts = Arc::new(AtomicU32::new(0)); let batch_tx = Arc::new(Mutex::new(None)); - let receiver_task = Arc::new(Mutex::new(None)); + let supervisor_abort = Arc::new(Mutex::new(None)); let cumulative_records_assigned = Arc::new(AtomicU64::new(0)); let submitted_records = Arc::new(AtomicU64::new(0)); let last_acked_records = Arc::new(AtomicU64::new(0)); @@ -289,8 +292,9 @@ impl ZerobusArrowStream { last_ack_tx, _last_ack_rx, is_closed, + admission_closed, close, - receiver_task, + supervisor_abort, pending_batches, pending_notify, request_send_failure, @@ -371,8 +375,8 @@ impl ZerobusArrowStream { let task = Supervisor::new(&stream).spawn(connection); { - let mut receiver_task = stream.receiver_task.lock().await; - *receiver_task = Some(task); + let mut supervisor_abort = stream.supervisor_abort.lock().await; + *supervisor_abort = Some(task); } info!( @@ -423,7 +427,10 @@ impl ZerobusArrowStream { /// ``` #[instrument(level = "debug", skip_all, fields(table_name = %self.table_properties.table_name))] pub async fn ingest_batch(&self, batch: RecordBatch) -> ZerobusResult { - if self.is_closed.load(Ordering::Relaxed) || self.close.has_started() { + if self.admission_closed.load(Ordering::Acquire) + || self.is_closed.load(Ordering::Relaxed) + || self.close.has_started() + { return Err(ZerobusError::StreamClosedError(tonic::Status::internal( "Stream is closing or closed", ))); @@ -461,7 +468,10 @@ impl ZerobusArrowStream { let _guard = self.ingest_mutex.lock().await; // May have closed while we blocked on the permit; returning drops it. - if self.is_closed.load(Ordering::Relaxed) || self.close.has_started() { + if self.admission_closed.load(Ordering::Acquire) + || self.is_closed.load(Ordering::Relaxed) + || self.close.has_started() + { return Err(ZerobusError::StreamClosedError(tonic::Status::internal( "Stream is closing or closed", ))); @@ -571,7 +581,10 @@ impl ZerobusArrowStream { /// marker after `finish()`) is allowed after that batch. #[instrument(level = "debug", skip_all, fields(table_name = %self.table_properties.table_name))] pub async fn ingest_ipc_batch(&self, ipc_bytes: Bytes) -> ZerobusResult { - if self.is_closed.load(Ordering::Relaxed) || self.close.has_started() { + if self.admission_closed.load(Ordering::Acquire) + || self.is_closed.load(Ordering::Relaxed) + || self.close.has_started() + { return Err(ZerobusError::StreamClosedError(tonic::Status::internal( "Stream is closing or closed", ))); @@ -673,6 +686,24 @@ impl ZerobusArrowStream { })? } + /// Waits through the short interval where terminal finalization owns admission but + /// has not published `CloseState::Finalized` yet. + async fn wait_for_terminal_outcome(&self) -> ZerobusResult<()> { + let mut close_rx = self.close.subscribe(); + + loop { + if let CloseState::Finalized(result) = close_rx.borrow_and_update().clone() { + return result; + } + + if close_rx.changed().await.is_err() { + return Err(ZerobusError::StreamClosedError(tonic::Status::internal( + "Close coordinator stopped unexpectedly", + ))); + } + } + } + /// Flushes all currently pending batches and waits for their acknowledgments. /// /// Snapshots the highest assigned offset when it begins and waits through that offset. @@ -710,6 +741,11 @@ impl ZerobusArrowStream { let target_offset = match self.offset_generator.last() { Some(offset) => offset, None => { + if self.admission_closed.load(Ordering::Acquire) + && matches!(self.close.state(), CloseState::Open) + { + return self.wait_for_terminal_outcome().await; + } // Nothing was ingested: report closure if closed, otherwise nothing to do. // Prefer the real terminal error over a generic closed message. if self.is_closed.load(Ordering::Relaxed) || self.close.has_started() { @@ -825,17 +861,35 @@ impl ZerobusArrowStream { CloseState::Open => { // This mutex makes the target snapshot and request publication atomic // with ingest admission and replacement-sender publication. - let _guard = self.ingest_mutex.lock().await; - let deadline = configured_deadline( - Instant::now(), - Duration::from_millis(self.options.flush_timeout_ms), - "flush_timeout_ms", - )?; - let request = CloseRequest { - target_offset: self.offset_generator.last(), - deadline, - }; - self.close.publish(request); + let guard = self.ingest_mutex.lock().await; + match self.close.state() { + CloseState::Open if !self.admission_closed.load(Ordering::Acquire) => { + let deadline = configured_deadline( + Instant::now(), + Duration::from_millis(self.options.flush_timeout_ms), + "flush_timeout_ms", + )?; + let request = CloseRequest { + target_offset: self.offset_generator.last(), + deadline, + }; + self.close.publish(request); + } + CloseState::Open => { + // Terminal finalization owns admission but publishes its result + // only after the retained-batch snapshot is complete. + drop(guard); + if close_rx.changed().await.is_err() { + return Err(ZerobusError::StreamClosedError( + tonic::Status::internal( + "Close coordinator stopped unexpectedly", + ), + )); + } + } + CloseState::Requested(_) => {} + CloseState::Finalized(result) => return result, + } } CloseState::Requested(_) => { if close_rx.changed().await.is_err() { @@ -983,6 +1037,15 @@ impl ZerobusArrowStream { Self::arm_test_barrier(&self.test_hooks.close_finalize).await } + /// Test-only: aborts the supervisor worker while leaving its finalizer reaper running. + #[cfg(feature = "test-hooks")] + #[doc(hidden)] + pub async fn abort_supervisor_for_test(&self) { + if let Some(handle) = self.supervisor_abort.lock().await.as_ref() { + handle.abort(); + } + } + /// Returns the table name for this stream. pub fn table_name(&self) -> &str { &self.table_properties.table_name @@ -1005,10 +1068,11 @@ impl ZerobusArrowStream { impl Drop for ZerobusArrowStream { fn drop(&mut self) { + self.admission_closed.store(true, Ordering::Release); self.is_closed.store(true, Ordering::Relaxed); // Best-effort abort the supervisor. Drop does not preserve pending batches for // retrieval; call close() or let recovery reach terminal finalization first. - if let Ok(mut guard) = self.receiver_task.try_lock() { + if let Ok(mut guard) = self.supervisor_abort.try_lock() { if let Some(handle) = guard.take() { handle.abort(); } diff --git a/rust/sdk/src/stream/arrow/supervisor.rs b/rust/sdk/src/stream/arrow/supervisor.rs index b6c8771c..72024194 100644 --- a/rust/sdk/src/stream/arrow/supervisor.rs +++ b/rust/sdk/src/stream/arrow/supervisor.rs @@ -8,7 +8,7 @@ use std::sync::Arc; use arrow_flight::error::FlightError; use tokio::sync::{mpsc, watch, Mutex}; -use tokio::task::{spawn, JoinHandle}; +use tokio::task::{spawn, AbortHandle, JoinError, JoinHandle}; use tokio::time::{sleep, sleep_until, timeout_at, Duration, Instant}; use tracing::{debug, error, info, warn}; @@ -79,12 +79,38 @@ impl Supervisor { } } - pub(super) fn spawn( - self, - initial_connection: FlightConnection, - ) -> JoinHandle> { + pub(super) fn spawn(self, initial_connection: FlightConnection) -> AbortHandle { let (response_stream, request_body) = initial_connection.into_supervisor_io(); - spawn(self.run(response_stream, request_body)) + let close = self.close.clone(); + let finalizer = self.close_finalizer.clone(); + let worker = spawn(self.run(response_stream, request_body)); + let abort_handle = worker.abort_handle(); + // The detached reaper owns the JoinHandle so cancelling a close caller cannot + // lose observation of an abnormal supervisor exit. + spawn(async move { + let joined = worker.await; + if matches!(close.state(), CloseState::Finalized(_)) { + return; + } + let outcome = Self::unfinalized_exit_outcome(joined); + let _ = finalizer.finish(outcome).await; + }); + abort_handle + } + + fn unfinalized_exit_outcome(joined: Result, JoinError>) -> ZerobusResult<()> { + match joined { + Ok(Err(error)) => Err(error), + Ok(Ok(())) => Err(ZerobusError::InvalidStateError( + "Supervisor exited successfully before close finalization".to_string(), + )), + Err(error) if error.is_cancelled() => Err(ZerobusError::InvalidStateError( + "Supervisor task was cancelled before close finalization".to_string(), + )), + Err(_) => Err(ZerobusError::InvalidStateError( + "Supervisor task panicked before close finalization".to_string(), + )), + } } fn spawn_headers_invalidation(&self, deadline: Instant) -> JoinHandle { @@ -165,8 +191,8 @@ impl Supervisor { } let mut active_was_drained = false; - let result = if let Some(error) = pending_error.take() { - Err(error) + let active_error = if let Some(error) = pending_error.take() { + error } else { let active_response = response_stream .as_mut() @@ -182,7 +208,7 @@ impl Supervisor { Ok(AckProcessOutcome::Stopped) => return self.finalized_result(), Ok(AckProcessOutcome::Recovery { error, drained }) => { active_was_drained = drained; - Err(error) + error } Ok(AckProcessOutcome::Close { request, outcome }) => { debug_assert_eq!(self.close.request(), Some(request)); @@ -191,14 +217,12 @@ impl Supervisor { } return self.finish(outcome).await; } - Err(error) => Err(error), + Err(error) => error, } }; - if let Err(error) = &result { - if !reconnect_auth_retry { - self.spawn_detached_auth_invalidation(error); - } + if !reconnect_auth_retry { + self.spawn_detached_auth_invalidation(&active_error); } if let Some(close_request) = self.close.request() { @@ -206,7 +230,7 @@ impl Supervisor { if let (Some(active_response), Some(active_request)) = (response_stream.as_mut(), request_body.as_ref()) { - let selected = result.clone(); + let selected = Err(active_error.clone()); match self .ack_processor .close_active_connection( @@ -231,12 +255,11 @@ impl Supervisor { } } } - return self.finish(result).await; + return self.finish(Err(active_error)).await; } - match result { - Ok(()) => return self.finish(Ok(())).await, - Err(ref error) + match active_error { + error if (error.is_retryable() || reconnect_auth_retry) && self.options.recovery => { reconnect_auth_retry = false; @@ -354,7 +377,7 @@ impl Supervisor { } } } - Err(error) => { + error => { error!(target: super::LOG_TARGET, "Supervisor: Non-retriable error, closing stream: {}", error); return self.finish(Err(error)).await; } @@ -580,6 +603,7 @@ mod tests { use arrow_flight::error::FlightError; use arrow_schema::{DataType, Field, Schema as ArrowSchema}; use tokio::sync::{mpsc, Mutex, Semaphore}; + use tokio::task::JoinHandle; use tokio::time::{timeout, Duration, Instant}; use super::super::close::{CloseCoordinator, CloseFinalizer, CloseRequest, CloseState}; @@ -633,6 +657,31 @@ mod tests { assert!(Supervisor::result_from_close_state(CloseState::Finalized(Ok(()))).is_ok()); } + #[test] + fn reaper_preserves_worker_error_and_rejects_unfinalized_success() { + let returned = ZerobusError::ConnectionTimeout("worker error".to_string()); + assert!(matches!( + Supervisor::unfinalized_exit_outcome(Ok(Err(returned))), + Err(ZerobusError::ConnectionTimeout(message)) if message == "worker error" + )); + assert!(matches!( + Supervisor::unfinalized_exit_outcome(Ok(Ok(()))), + Err(ZerobusError::InvalidStateError(_)) + )); + } + + #[tokio::test] + async fn panicked_supervisor_exit_is_an_invariant_error() { + let worker: JoinHandle> = + tokio::spawn(async { panic!("supervisor test panic") }); + let joined = worker.await; + + assert!(matches!( + Supervisor::unfinalized_exit_outcome(joined), + Err(ZerobusError::InvalidStateError(_)) + )); + } + #[tokio::test] async fn close_publication_precedes_queued_replay_handoff() { let schema = one_col_schema(); diff --git a/rust/tests/src/arrow_tests.rs b/rust/tests/src/arrow_tests.rs index 221aec6a..6abb3e8a 100644 --- a/rust/tests/src/arrow_tests.rs +++ b/rust/tests/src/arrow_tests.rs @@ -265,19 +265,20 @@ mod arrow_flight_tests { } assert!(!stream.is_closed()); - // A successful ingest proves the failed deadline calculation did not publish close. + let (reached, proceed) = stream.arm_close_finalize_barrier().await; let batch = create_test_record_batch(schema, vec![1], vec![Some("still usable")]); stream.ingest_batch(batch).await?; - run_with_paused_time_watchdog(async { - while !stream.is_closed() { - tokio::task::yield_now().await; - } - }) - .await; + run_with_paused_time_watchdog(reached.notified()).await; - let error = run_with_paused_time_watchdog(stream.close()) + let mut close_future = Box::pin(stream.close()); + assert!( + futures::poll!(close_future.as_mut()).is_pending(), + "an unrepresentable deadline must not mask in-progress terminal finalization" + ); + proceed.notify_one(); + let error = run_with_paused_time_watchdog(close_future) .await - .expect_err("repeated close must observe the stored terminal error"); + .expect_err("close must observe the stored terminal error"); match error { ZerobusError::StreamClosedError(status) => { assert_eq!(status.code(), tonic::Code::InvalidArgument); @@ -1347,6 +1348,73 @@ mod arrow_flight_tests { Ok(()) } + + #[tokio::test] + async fn test_supervisor_abort_after_close_request_finalizes_exact_suffix( + ) -> Result<(), Box> { + setup_tracing(); + + let (mock_server, server_url) = start_mock_flight_server().await?; + let schema = create_test_arrow_schema(); + mock_server + .inject_responses( + TABLE_NAME, + vec![MockFlightResponse::BatchAck { + ack_up_to_offset: 0, + delay_ms: 0, + ack_up_to_records: 1, + }], + ) + .await; + + let sdk = ZerobusSdk::builder() + .endpoint(server_url) + .unity_catalog_url("https://mock-uc.com") + .tls_config(Arc::new(NoTlsConfig)) + .build()?; + let mut stream = sdk + .stream_builder() + .table(TABLE_NAME) + .headers_provider(Arc::new(TestHeadersProvider::default())) + .arrow(schema.clone()) + .flush_timeout_ms(60_000) + .build_arrow() + .await?; + + let ack_applied = stream.arm_ack_applied_notify().await; + let batch = create_test_record_batch( + schema, + vec![1, 2, 3], + vec![Some("a"), Some("b"), Some("c")], + ); + stream.ingest_batch(batch).await?; + tokio::time::timeout(std::time::Duration::from_secs(5), ack_applied.notified()) + .await + .expect("partial ACK must be applied before close"); + + let mut close_future = Box::pin(stream.close()); + assert!( + futures::poll!(close_future.as_mut()).is_pending(), + "close must publish its request before waiting" + ); + drop(close_future); + stream.abort_supervisor_for_test().await; + + let error = tokio::time::timeout(std::time::Duration::from_secs(1), stream.close()) + .await + .expect("the supervisor reaper must finalize an aborted worker") + .expect_err("aborted supervisor must produce an invariant error"); + assert!(matches!(error, ZerobusError::InvalidStateError(_))); + let unacked: Vec> = stream + .get_unacked_batches() + .await? + .iter() + .map(batch_ids) + .collect(); + assert_eq!(unacked, vec![vec![2, 3]]); + + Ok(()) + } } mod error_handling_tests { @@ -1396,6 +1464,118 @@ mod arrow_flight_tests { Ok(()) } + #[tokio::test] + async fn test_terminal_finalization_rejects_new_ingest_before_snapshot_publish( + ) -> Result<(), Box> { + setup_tracing(); + + let (mock_server, server_url) = start_mock_flight_server().await?; + let schema = create_test_arrow_schema(); + mock_server + .inject_responses( + TABLE_NAME, + vec![MockFlightResponse::Error { + status: Status::invalid_argument("terminal peer error"), + delay_ms: 0, + }], + ) + .await; + + let sdk = ZerobusSdk::builder() + .endpoint(server_url) + .unity_catalog_url("https://mock-uc.com") + .tls_config(Arc::new(NoTlsConfig)) + .build()?; + let mut stream = sdk + .stream_builder() + .table(TABLE_NAME) + .headers_provider(Arc::new(TestHeadersProvider::default())) + .arrow(schema.clone()) + .recovery(false) + .build_arrow() + .await?; + + let (reached, proceed) = stream.arm_close_finalize_barrier().await; + let first = create_test_record_batch( + schema.clone(), + vec![1], + vec![Some("accepted before failure")], + ); + stream.ingest_batch(first).await?; + tokio::time::timeout(std::time::Duration::from_secs(5), reached.notified()) + .await + .expect("terminal finalization must reach the snapshot barrier"); + assert!(!stream.is_closed(), "the snapshot is not yet published"); + + let late = create_test_record_batch(schema, vec![2], vec![Some("too late")]); + assert!( + stream.ingest_batch(late).await.is_err(), + "terminal finalization must close admission before publishing is_closed" + ); + + proceed.notify_one(); + let error = stream + .close() + .await + .expect_err("close must return the terminal peer error"); + match error { + ZerobusError::StreamClosedError(status) => { + assert_eq!(status.code(), tonic::Code::InvalidArgument); + } + other => panic!("expected terminal stream error, got {other:?}"), + } + assert_eq!(stream.get_unacked_batches().await?.len(), 1); + + Ok(()) + } + + #[tokio::test] + async fn test_empty_flush_waits_for_terminal_outcome_during_finalization( + ) -> Result<(), Box> { + setup_tracing(); + + let (_mock_server, server_url) = start_mock_flight_server().await?; + let sdk = ZerobusSdk::builder() + .endpoint(server_url) + .unity_catalog_url("https://mock-uc.com") + .tls_config(Arc::new(NoTlsConfig)) + .build()?; + let mut stream = sdk + .stream_builder() + .table(TABLE_NAME) + .headers_provider(Arc::new(TestHeadersProvider::default())) + .arrow(create_test_arrow_schema()) + .build_arrow() + .await?; + + let (reached, proceed) = stream.arm_close_finalize_barrier().await; + stream.abort_supervisor_for_test().await; + tokio::time::timeout(std::time::Duration::from_secs(5), reached.notified()) + .await + .expect("abnormal-exit finalization must reach the barrier"); + + let mut flush = Box::pin(stream.flush()); + assert!( + futures::poll!(flush.as_mut()).is_pending(), + "empty flush must wait while terminal outcome publication is pending" + ); + proceed.notify_one(); + let error = tokio::time::timeout(std::time::Duration::from_secs(1), flush) + .await + .expect("empty flush must complete after outcome publication") + .expect_err("empty flush must return the supervisor exit error"); + assert!(matches!(error, ZerobusError::InvalidStateError(_))); + + let close_error = + tokio::time::timeout(std::time::Duration::from_secs(1), stream.close()) + .await + .expect("terminal finalization must complete") + .expect_err("close must return the same supervisor exit error"); + assert!(matches!(close_error, ZerobusError::InvalidStateError(_))); + + Ok(()) + } + /// A terminal server error must reach a blocked wait_for_offset()/flush() as the /// real error, not a generic close/timeout. process_acks publishes the error while /// is_closed is still false (the waiter keeps waiting), so the supervisor @@ -3871,6 +4051,97 @@ mod arrow_flight_tests { Ok(()) } + #[tokio::test] + async fn test_close_during_second_recovery_attempt_preserves_latest_trigger( + ) -> Result<(), Box> { + setup_tracing(); + + let (mock_server, server_url) = start_mock_flight_server().await?; + let schema = create_test_arrow_schema(); + mock_server + .inject_responses( + TABLE_NAME, + vec![ + MockFlightResponse::BatchAck { + ack_up_to_offset: 0, + delay_ms: 0, + ack_up_to_records: 1, + }, + MockFlightResponse::Error { + status: tonic::Status::unavailable("active transport failed"), + delay_ms: 0, + }, + MockFlightResponse::FailSetup { + status: tonic::Status::unavailable("first reconnect failed"), + }, + MockFlightResponse::FailSetupAfter { + status: tonic::Status::unavailable("second reconnect failed"), + delay_ms: 60_000, + }, + ], + ) + .await; + + let delayed_setup_armed = mock_server.delayed_setup_armed(); + let sdk = ZerobusSdk::builder() + .endpoint(server_url) + .unity_catalog_url("https://mock-uc.com") + .tls_config(Arc::new(NoTlsConfig)) + .build()?; + let mut stream = sdk + .stream_builder() + .table(TABLE_NAME) + .headers_provider(Arc::new(TestHeadersProvider::default())) + .arrow(schema.clone()) + .recovery(true) + .recovery_backoff_ms(0) + .recovery_timeout_ms(120_000) + .recovery_retries(3) + .flush_timeout_ms(60_000) + .build_arrow() + .await?; + + let ack_applied = stream.arm_ack_applied_notify().await; + let first = create_test_record_batch( + schema.clone(), + vec![1, 2, 3], + vec![Some("a"), Some("b"), Some("c")], + ); + stream.ingest_batch(first).await?; + tokio::time::timeout(std::time::Duration::from_secs(5), ack_applied.notified()) + .await + .expect("partial ACK must be applied before recovery"); + let second = create_test_record_batch(schema, vec![4], vec![Some("d")]); + stream.ingest_batch(second).await?; + tokio::time::timeout( + std::time::Duration::from_secs(5), + delayed_setup_armed.notified(), + ) + .await + .expect("the second recovery attempt must enter setup"); + + let error = tokio::time::timeout(std::time::Duration::from_secs(1), stream.close()) + .await + .expect("close must interrupt the second recovery attempt") + .expect_err("close during recovery must preserve its current trigger"); + match error { + ZerobusError::CreateStreamError(status) => { + assert_eq!(status.code(), tonic::Code::Unavailable); + assert_eq!(status.message(), "first reconnect failed"); + } + other => panic!("expected the first attempt's reconnect error, got: {other:?}"), + } + let unacked: Vec> = stream + .get_unacked_batches() + .await? + .iter() + .map(batch_ids) + .collect(); + assert_eq!(unacked, vec![vec![2, 3], vec![4]]); + + Ok(()) + } + #[tokio::test] async fn test_supervisor_recovery_after_retriable_error( ) -> Result<(), Box> { From fd602116a22160c3bd6b895f60840f305b72f3fd Mon Sep 17 00:00:00 2001 From: teodordelibasic-db Date: Thu, 13 Aug 2026 12:32:39 +0000 Subject: [PATCH 3/6] [Rust] Unblock send-failure recovery Continuously ready no-progress responses could starve recovery after a request send failure. One response-first tie is allowed; a buffered terminal status or EOF still wins over the local Unavailable error. Signed-off-by: teodordelibasic-db --- rust/sdk/src/stream/arrow/acks.rs | 109 +++++++++++++++++++++++++++++- 1 file changed, 108 insertions(+), 1 deletion(-) diff --git a/rust/sdk/src/stream/arrow/acks.rs b/rust/sdk/src/stream/arrow/acks.rs index 9b78d914..7466cece 100644 --- a/rust/sdk/src/stream/arrow/acks.rs +++ b/rust/sdk/src/stream/arrow/acks.rs @@ -133,6 +133,10 @@ impl RequestSendFailure { self.pending.swap(false, Ordering::AcqRel) } + fn is_pending(&self) -> bool { + self.pending.load(Ordering::Acquire) + } + fn clear(&self) { self.pending.store(false, Ordering::Release); } @@ -788,6 +792,7 @@ impl AckProcessor { let mut connection = ConnectionState::Active; let mut expiry_tie_winner: Option = None; let mut close_after_priority_response = None; + let mut response_deferred_send_failure = false; let request_control = self.request_control(request_body); let acknowledgments = self.ack_progress(); @@ -880,9 +885,21 @@ impl AckProcessor { ConnectionState::Waiting(WaitState::Close(request)) => Some(*request), _ => None, }; - let event = if let Some(response) = priority_response { + let event = if response_deferred_send_failure && self.request_send_failure.is_pending() + { + response_deferred_send_failure = false; + // The one-tie already consumed a ready item. Prefer a buffered + // terminal status or EOF over local send failure, but do not poll + // another non-progress Ok. + match response_stream.next().now_or_never() { + Some(Some(Err(error))) => AckEvent::Response(Some(Err(error))), + Some(None) => AckEvent::Response(None), + _ => AckEvent::RequestSendFailed, + } + } else if let Some(response) = priority_response { AckEvent::Response(response) } else { + response_deferred_send_failure = false; match &connection { ConnectionState::Active => match self.oldest_ack_deadline(ack_timeout).await? { Some(pending_deadline) => { @@ -926,6 +943,13 @@ impl AckProcessor { } }; + // A ready response wins one tie so its ACK or terminal status is observed. + // A no-progress winner cannot postpone send failure past one buffered + // terminal status or EOF. + if matches!(event, AckEvent::Response(_)) && self.request_send_failure.is_pending() { + response_deferred_send_failure = true; + } + match event { AckEvent::PendingBatchAvailable => continue, AckEvent::CloseFinalized => return Ok(AckProcessOutcome::Stopped), @@ -1413,6 +1437,89 @@ mod tests { } } + /// A terminal peer status buffered behind one no-progress ACK must not be + /// rewritten as the local send-failure error. + #[tokio::test] + async fn terminal_status_behind_no_progress_ack_wins_reported_send_failure() { + let no_progress = PutResult { + app_metadata: serde_json::to_vec(&FlightAckMetadata { + ack_up_to_offset: -1, + ack_up_to_records: 0, + close_stream_duration_ms: None, + }) + .unwrap() + .into(), + }; + let response_stream = iter([ + Ok(no_progress), + Err(tonic::Status::permission_denied("permanent server rejection").into()), + ]); + let (processor, request_body, _last_ack_rx) = ack_processor( + Arc::new(Mutex::new(Vec::new())), + Arc::new(AtomicU64::new(0)), + Arc::new(AtomicU64::new(0)), + false, + ); + processor.request_send_failure.report(); + + let error = processor + .process(Box::pin(response_stream), request_body) + .await + .expect_err("the buffered server rejection must be returned"); + + assert!(!error.is_retryable()); + match error { + ZerobusError::StreamClosedError(status) => { + assert_eq!(status.code(), tonic::Code::PermissionDenied); + assert_eq!(status.message(), "permanent server rejection"); + } + other => panic!("expected a stream-closed error, got {other:?}"), + } + } + + /// At most one ready nonterminal response may defer a reported request-send failure. + #[tokio::test] + async fn continuously_ready_nonprogress_responses_do_not_starve_send_failure() { + let valid_no_progress = PutResult { + app_metadata: serde_json::to_vec(&FlightAckMetadata { + ack_up_to_offset: -1, + ack_up_to_records: 0, + close_stream_duration_ms: None, + }) + .unwrap() + .into(), + }; + let malformed = PutResult { + app_metadata: b"not ack metadata".to_vec().into(), + }; + + for response in [valid_no_progress, malformed] { + let response_stream = repeat_with(move || Ok(response.clone())); + let (processor, request_body, _last_ack_rx) = ack_processor( + Arc::new(Mutex::new(Vec::new())), + Arc::new(AtomicU64::new(0)), + Arc::new(AtomicU64::new(0)), + false, + ); + processor.request_send_failure.report(); + + let error = tokio::time::timeout( + Duration::from_millis(100), + processor.process(Box::pin(response_stream), request_body), + ) + .await + .expect("ready responses must not starve request-send failure") + .expect_err("the reported request-send failure must trigger recovery"); + + match error { + ZerobusError::StreamClosedError(status) => { + assert_eq!(status.code(), tonic::Code::Unavailable); + } + other => panic!("expected a stream-closed error, got {other:?}"), + } + } + } + /// An acknowledgement beyond the connection-local submitted-record count is a protocol /// violation and must not make unsent records appear durable. #[tokio::test] From 0b807918c333a8362b7a4a527077c7f57cac20f5 Mon Sep 17 00:00:00 2001 From: teodordelibasic-db Date: Thu, 13 Aug 2026 14:16:58 +0000 Subject: [PATCH 4/6] [Rust] Preserve ACKs after send-failure A second now_or_never poll consumed a ready PutResult without applying it. Force send-failure on the next loop without polling another response. Signed-off-by: teodordelibasic-db --- rust/sdk/src/stream/arrow/acks.rs | 24 +++++++----------------- 1 file changed, 7 insertions(+), 17 deletions(-) diff --git a/rust/sdk/src/stream/arrow/acks.rs b/rust/sdk/src/stream/arrow/acks.rs index 7466cece..b0806a99 100644 --- a/rust/sdk/src/stream/arrow/acks.rs +++ b/rust/sdk/src/stream/arrow/acks.rs @@ -888,14 +888,7 @@ impl AckProcessor { let event = if response_deferred_send_failure && self.request_send_failure.is_pending() { response_deferred_send_failure = false; - // The one-tie already consumed a ready item. Prefer a buffered - // terminal status or EOF over local send failure, but do not poll - // another non-progress Ok. - match response_stream.next().now_or_never() { - Some(Some(Err(error))) => AckEvent::Response(Some(Err(error))), - Some(None) => AckEvent::Response(None), - _ => AckEvent::RequestSendFailed, - } + AckEvent::RequestSendFailed } else if let Some(response) = priority_response { AckEvent::Response(response) } else { @@ -944,8 +937,7 @@ impl AckProcessor { }; // A ready response wins one tie so its ACK or terminal status is observed. - // A no-progress winner cannot postpone send failure past one buffered - // terminal status or EOF. + // The next loop forces send failure without polling another response. if matches!(event, AckEvent::Response(_)) && self.request_send_failure.is_pending() { response_deferred_send_failure = true; } @@ -1437,10 +1429,9 @@ mod tests { } } - /// A terminal peer status buffered behind one no-progress ACK must not be - /// rewritten as the local send-failure error. + /// After the one response-first tie, a later terminal status is not polled. #[tokio::test] - async fn terminal_status_behind_no_progress_ack_wins_reported_send_failure() { + async fn later_terminal_status_does_not_replace_reported_send_failure() { let no_progress = PutResult { app_metadata: serde_json::to_vec(&FlightAckMetadata { ack_up_to_offset: -1, @@ -1465,13 +1456,12 @@ mod tests { let error = processor .process(Box::pin(response_stream), request_body) .await - .expect_err("the buffered server rejection must be returned"); + .expect_err("the reported request-send failure must trigger recovery"); - assert!(!error.is_retryable()); + assert!(error.is_retryable()); match error { ZerobusError::StreamClosedError(status) => { - assert_eq!(status.code(), tonic::Code::PermissionDenied); - assert_eq!(status.message(), "permanent server rejection"); + assert_eq!(status.code(), tonic::Code::Unavailable); } other => panic!("expected a stream-closed error, got {other:?}"), } From afe07a60ed3ff79fff4106ce45f90e7db8cdae8a Mon Sep 17 00:00:00 2001 From: teodordelibasic-db Date: Thu, 13 Aug 2026 14:53:10 +0000 Subject: [PATCH 5/6] [Rust] Drop redundant send-failure test Signed-off-by: teodordelibasic-db --- rust/sdk/src/stream/arrow/acks.rs | 38 ------------------------------- 1 file changed, 38 deletions(-) diff --git a/rust/sdk/src/stream/arrow/acks.rs b/rust/sdk/src/stream/arrow/acks.rs index b0806a99..760fcc57 100644 --- a/rust/sdk/src/stream/arrow/acks.rs +++ b/rust/sdk/src/stream/arrow/acks.rs @@ -1429,44 +1429,6 @@ mod tests { } } - /// After the one response-first tie, a later terminal status is not polled. - #[tokio::test] - async fn later_terminal_status_does_not_replace_reported_send_failure() { - let no_progress = PutResult { - app_metadata: serde_json::to_vec(&FlightAckMetadata { - ack_up_to_offset: -1, - ack_up_to_records: 0, - close_stream_duration_ms: None, - }) - .unwrap() - .into(), - }; - let response_stream = iter([ - Ok(no_progress), - Err(tonic::Status::permission_denied("permanent server rejection").into()), - ]); - let (processor, request_body, _last_ack_rx) = ack_processor( - Arc::new(Mutex::new(Vec::new())), - Arc::new(AtomicU64::new(0)), - Arc::new(AtomicU64::new(0)), - false, - ); - processor.request_send_failure.report(); - - let error = processor - .process(Box::pin(response_stream), request_body) - .await - .expect_err("the reported request-send failure must trigger recovery"); - - assert!(error.is_retryable()); - match error { - ZerobusError::StreamClosedError(status) => { - assert_eq!(status.code(), tonic::Code::Unavailable); - } - other => panic!("expected a stream-closed error, got {other:?}"), - } - } - /// At most one ready nonterminal response may defer a reported request-send failure. #[tokio::test] async fn continuously_ready_nonprogress_responses_do_not_starve_send_failure() { From d1b99c856df523a90f9dc6ec5db3800c6d07a670 Mon Sep 17 00:00:00 2001 From: teodordelibasic-db Date: Fri, 14 Aug 2026 07:59:00 +0000 Subject: [PATCH 6/6] [Rust] Final pass Signed-off-by: teodordelibasic-db --- rust/NEXT_CHANGELOG.md | 3 +- rust/sdk/src/stream/arrow/acks.rs | 118 ++++++++++++++++++++++++++++++ 2 files changed, 120 insertions(+), 1 deletion(-) diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index 34c28900..6338da1f 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -27,7 +27,8 @@ attempt, retains the unacknowledged suffix, and returns the error that triggered the current attempt. Close during an existing recovery or server-requested rotation keeps that trigger even if every record is durable, so an error can coexist with an empty - unacknowledged-batch set. + unacknowledged-batch set. After a request-send failure, one ready response may still + be applied; later stream items are not discarded in order to start recovery. - Fixed Arrow Flight recovery sender lifetime: replacement senders are now published only after pending replay succeeds, while initial supervisor handoff and failed or cancelled replay promptly drop redundant senders instead of retaining incomplete diff --git a/rust/sdk/src/stream/arrow/acks.rs b/rust/sdk/src/stream/arrow/acks.rs index 760fcc57..5b6df48c 100644 --- a/rust/sdk/src/stream/arrow/acks.rs +++ b/rust/sdk/src/stream/arrow/acks.rs @@ -813,9 +813,12 @@ impl AckProcessor { observe_close = false; if matches!(connection, ConnectionState::Active) && close_request.target_offset.is_none() + && !(response_deferred_send_failure + && self.request_send_failure.is_pending()) { // Give a terminal response that predates an empty close one poll. // A pending or nonterminal response cannot postpone close again. + // Skip if the send-failure one-tie already consumed a response. priority_response = response_stream.next().now_or_never(); if priority_response.is_some() { close_after_priority_response = Some(close_request); @@ -1150,6 +1153,7 @@ impl AckProcessor { mod tests { use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; + use std::task::Poll; use arrow_array::Int32Array; use arrow_flight::error::FlightError; @@ -1429,6 +1433,120 @@ mod tests { } } + /// After the one response-first tie, a later terminal status is not polled. + #[tokio::test] + async fn later_terminal_status_does_not_replace_reported_send_failure() { + let no_progress = PutResult { + app_metadata: serde_json::to_vec(&FlightAckMetadata { + ack_up_to_offset: -1, + ack_up_to_records: 0, + close_stream_duration_ms: None, + }) + .unwrap() + .into(), + }; + let response_stream = iter([ + Ok(no_progress), + Err(tonic::Status::permission_denied("permanent server rejection").into()), + ]); + let (processor, request_body, _last_ack_rx) = ack_processor( + Arc::new(Mutex::new(Vec::new())), + Arc::new(AtomicU64::new(0)), + Arc::new(AtomicU64::new(0)), + false, + ); + processor.request_send_failure.report(); + + let error = processor + .process(Box::pin(response_stream), request_body) + .await + .expect_err("the reported request-send failure must trigger recovery"); + + assert!(error.is_retryable()); + match error { + ZerobusError::StreamClosedError(status) => { + assert_eq!(status.code(), tonic::Code::Unavailable); + } + other => panic!("expected a stream-closed error, got {other:?}"), + } + } + + /// Empty-close's one-shot peek must not consume a later item after send-failure + /// has already used its response-first tie. + #[tokio::test] + async fn deferred_send_failure_skips_empty_close_peek() { + let no_progress = PutResult { + app_metadata: serde_json::to_vec(&FlightAckMetadata { + ack_up_to_offset: -1, + ack_up_to_records: 0, + close_stream_duration_ms: None, + }) + .unwrap() + .into(), + }; + let (processor, request_body, _last_ack_rx) = ack_processor( + Arc::new(Mutex::new(Vec::new())), + Arc::new(AtomicU64::new(0)), + Arc::new(AtomicU64::new(0)), + false, + ); + processor.request_send_failure.report(); + + let taken = Arc::new(AtomicU64::new(0)); + let taken_for_stream = Arc::clone(&taken); + let close = processor.close.clone(); + let request = CloseRequest { + target_offset: None, + deadline: Instant::now() + Duration::from_secs(1), + }; + let mut items = vec![ + Ok(no_progress), + Err(tonic::Status::permission_denied("permanent server rejection").into()), + ] + .into_iter(); + let mut response_stream: FlightResponseStream = + Box::pin(futures::stream::poll_fn(move |_cx| { + let n = taken_for_stream.fetch_add(1, Ordering::SeqCst) + 1; + if n == 1 { + close.publish(request); + } + Poll::Ready(items.next()) + })); + let mut close_rx = processor.close.subscribe(); + + let result = processor + .process_active(&mut response_stream, &request_body, &mut close_rx, true) + .await; + match result { + Ok(AckProcessOutcome::Close { + request: + CloseRequest { + target_offset: None, + .. + }, + outcome: Err(error), + }) => { + assert!(!error.is_retryable()); + match error { + ZerobusError::StreamClosedError(status) => { + assert_eq!(status.code(), tonic::Code::PermissionDenied); + } + other => panic!("expected drained peer status, got {other:?}"), + } + } + Err(ZerobusError::StreamClosedError(status)) + if status.code() == tonic::Code::Unavailable => + { + panic!("empty-close peek discarded the buffered peer status"); + } + _ => panic!("expected empty close to surface the buffered peer status"), + } + assert!( + taken.load(Ordering::SeqCst) >= 2, + "drain must observe the buffered peer status instead of dropping it" + ); + } + /// At most one ready nonterminal response may defer a reported request-send failure. #[tokio::test] async fn continuously_ready_nonprogress_responses_do_not_starve_send_failure() {