[Rust] Fix Arrow failed-enqueue finalization

rust/NEXT_CHANGELOG.md

@@ -8,6 +8,12 @@
 
 ### Bug Fixes
 
+- Arrow Flight now rolls back logical offsets and record ranges when an enqueue
+  fails with recovery disabled. `ingest_batch()` waits for terminal finalization
+  and returns the request-stream error; `flush()` and `close()` no longer wait on
+  the withdrawn offset. An already-acknowledged flush target still succeeds, while
+  `close()` preserves the terminal error and retained batches are immediately available.
+
 ### Documentation
 
 - Corrected README and rustdoc examples so their dependencies, feature flags,

rust/sdk/src/offset_generator.rs

@@ -65,11 +65,8 @@ impl OffsetIdGenerator {
 
     /// Repositions the generator so the next call to `next()` returns `next_value`.
     ///
-    /// Used by the Arrow Flight stream recovery path: when the SDK reconnects and
-    /// replays N pending batches with wire offsets `0..N-1`, the in-memory generator
-    /// must be set so that subsequent fresh batches pick up at `N` rather than the
-    /// pre-recovery monotonic counter — otherwise the server rejects the next batch
-    /// with a non-sequential-offset error.
+    /// This may move the sequence backward; callers must synchronize it with any
+    /// related state and concurrent calls to `next()` or `last()`.
     pub fn set_next(&self, next_value: OffsetId) {
         self.last_offset_id.store(next_value - 1, Ordering::SeqCst);
     }

rust/sdk/src/stream/arrow/mod.rs

@@ -89,6 +89,7 @@ struct TestHooks {
     replay_send: TestBarrierGate,
     ack_applied: TestNotifyGate,
     ack_idle: TestNotifyGate,
+    failed_enqueue: TestBarrierGate,
     close_finalize: TestBarrierGate,
 }
 
@@ -172,7 +173,7 @@ pub struct ZerobusArrowStream {
     _last_ack_rx: watch::Receiver<Option<OffsetId>>,
     /// True once the stream is terminally closed and unacknowledged batches may be retrieved.
     is_closed: Arc<AtomicBool>,
-    /// Rejects new ingests as soon as terminal finalization owns admission.
+    /// Rejects new ingests once terminal finalization or a failed enqueue owns admission.
     admission_closed: Arc<AtomicBool>,
     /// Coordinates one resumable explicit-close request with the recovery supervisor.
     close: CloseCoordinator,
@@ -408,6 +409,8 @@ impl ZerobusArrowStream {
     /// # Errors
     ///
     /// * `StreamClosedError` - If the stream is closing or closed
+    /// * The terminal request-stream error if enqueueing fails with recovery disabled;
+    ///   the call waits for terminal finalization before returning that error
     /// * `InvalidArgument` - If the batch schema doesn't match the stream schema, or the
     ///   batch has zero rows (an empty batch carries no data to send or acknowledge)
     ///
@@ -526,14 +529,32 @@ impl ZerobusArrowStream {
                     let mut pending = self.pending_batches.lock().await;
                     pending.retain(|pending_batch| pending_batch.offset_id() != offset_id);
                 }
-                let _ = timeout(
-                    Duration::from_millis(100),
-                    self.server_error_rx.clone().changed(),
-                )
-                .await;
-                return Err(Self::terminal_error_or(&self.server_error_rx, || {
-                    "Failed to send batch".to_string()
-                }));
+                // Withdraw the logical assignment before waking terminal finalization.
+                self.cumulative_records_assigned
+                    .store(start_record, Ordering::Relaxed);
+                self.offset_generator.set_next(offset_id);
+                // Claim terminal admission while close is still excluded by ingest_mutex,
+                // so the request-send failure cannot be replaced by a successful close.
+                self.admission_closed.store(true, Ordering::Release);
+                #[cfg(feature = "test-hooks")]
+                {
+                    let barrier = self.test_hooks.failed_enqueue.lock().await.take();
+                    if let Some(barrier) = barrier {
+                        barrier.reached.notify_one();
+                        barrier.proceed.notified().await;
+                    }
+                }
+                // Finalization must reacquire ingest_mutex before publishing its outcome.
+                drop(_guard);
+                self.request_send_failure.report();
+                return match self.wait_for_terminal_outcome().await {
+                    Err(error) => Err(error),
+                    // A clean outcome is unreachable after this path claims terminal
+                    // admission, but preserve any published cause defensively.
+                    Ok(()) => Err(Self::terminal_error_or(&self.server_error_rx, || {
+                        "Failed to send batch".to_string()
+                    })),
+                };
             }
         };
 
@@ -662,8 +683,7 @@ impl ZerobusArrowStream {
         })?
     }
 
-    /// Waits through the short interval where terminal finalization owns admission but
-    /// has not published `CloseState::Finalized` yet.
+    /// Waits after terminal admission is claimed but before finalization publishes its result.
     async fn wait_for_terminal_outcome(&self) -> ZerobusResult<()> {
         let mut close_rx = self.close.subscribe();
 
@@ -712,7 +732,13 @@ impl ZerobusArrowStream {
     /// ```
     #[instrument(level = "debug", skip_all, fields(table_name = %self.table_properties.table_name))]
     pub async fn flush(&self) -> ZerobusResult<()> {
-        let target_offset = match self.offset_generator.last() {
+        // Serialize the snapshot with enqueue assignment and rollback so a concurrent
+        // failed enqueue cannot leave this flush waiting on its withdrawn offset.
+        let target_offset = {
+            let _guard = self.ingest_mutex.lock().await;
+            self.offset_generator.last()
+        };
+        let target_offset = match target_offset {
             Some(offset) => offset,
             None => {
                 if self.admission_closed.load(Ordering::Acquire)
@@ -820,6 +846,10 @@ impl ZerobusArrowStream {
     /// ```
     #[instrument(level = "debug", skip_all, fields(table_name = %self.table_properties.table_name))]
     pub async fn close(&mut self) -> ZerobusResult<()> {
+        self.close_internal().await
+    }
+
+    async fn close_internal(&self) -> ZerobusResult<()> {
         info!(
             table_name = %self.table_properties.table_name,
             "Closing Arrow Flight stream"
@@ -847,8 +877,8 @@ impl ZerobusArrowStream {
                             self.close.publish(request);
                         }
                         CloseState::Open => {
-                            // Terminal finalization owns admission but publishes its result
-                            // only after the retained-batch snapshot is complete.
+                            // A failed enqueue or terminal finalization owns admission; both
+                            // complete by publishing one shared terminal result.
                             drop(guard);
                             if close_rx.changed().await.is_err() {
                                 return Err(Self::close_coordinator_stopped_error());
@@ -994,6 +1024,29 @@ impl ZerobusArrowStream {
         *self.batch_tx.lock().await = Some(closed_tx);
     }
 
+    /// Test-only: parks a recovery-disabled failed enqueue after it claims terminal
+    /// admission and before it releases `ingest_mutex` or wakes the supervisor.
+    #[cfg(feature = "test-hooks")]
+    #[doc(hidden)]
+    pub async fn arm_failed_enqueue_barrier(&self) -> (Arc<Notify>, Arc<Notify>) {
+        Self::arm_test_barrier(&self.test_hooks.failed_enqueue).await
+    }
+
+    /// Test-only: reports whether terminal admission has been claimed.
+    #[cfg(feature = "test-hooks")]
+    #[doc(hidden)]
+    pub fn admission_closed_for_test(&self) -> bool {
+        self.admission_closed.load(Ordering::Acquire)
+    }
+
+    /// Test-only: runs the close state machine through a shared reference so tests can
+    /// exercise foreign-wrapper concurrency without creating aliased Rust references.
+    #[cfg(feature = "test-hooks")]
+    #[doc(hidden)]
+    pub async fn close_concurrently_for_test(&self) -> ZerobusResult<()> {
+        self.close_internal().await
+    }
+
     /// 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")]

rust/tests/src/arrow_tests.rs

@@ -2630,6 +2630,180 @@ mod arrow_flight_tests {
             Ok(())
         }
 
+        #[tokio::test]
+        async fn test_closed_request_sender_without_recovery_finalizes_promptly(
+        ) -> Result<(), Box<dyn std::error::Error>> {
+            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::HoldResponseAfterRequestEof,
+                    ],
+                )
+                .await;
+
+            let sdk = ZerobusSdk::builder()
+                .endpoint(server_url)
+                .unity_catalog_url("https://mock-uc.com")
+                .tls_config(Arc::new(NoTlsConfig))
+                .build()?;
+            let stream = sdk
+                .stream_builder()
+                .table(TABLE_NAME)
+                .headers_provider(Arc::new(TestHeadersProvider::default()))
+                .arrow(schema.clone())
+                .flush_timeout_ms(60_000)
+                .recovery(false)
+                .build_arrow()
+                .await?;
+
+            let durable_batch =
+                create_test_record_batch(schema.clone(), vec![1], vec![Some("durable")]);
+            let durable_offset = stream.ingest_batch(durable_batch).await?;
+            stream.wait_for_offset(durable_offset).await?;
+
+            let (failed_enqueue_reached, failed_enqueue_proceed) =
+                stream.arm_failed_enqueue_barrier().await;
+            let (finalize_reached, finalize_proceed) = stream.arm_close_finalize_barrier().await;
+            stream.replace_batch_sender_with_closed_channel().await;
+            let batch = create_test_record_batch(schema, vec![2], vec![Some("withdrawn")]);
+            let mut ingest = Box::pin(stream.ingest_batch(batch));
+            tokio::select! {
+                _ = failed_enqueue_reached.notified() => {}
+                result = &mut ingest => {
+                    panic!("ingest completed before claiming terminal admission: {result:?}")
+                }
+            }
+            assert!(
+                stream.admission_closed_for_test(),
+                "failed enqueue must claim admission before close can publish success"
+            );
+            let mut close = Box::pin(stream.close_concurrently_for_test());
+            // Queue close on ingest_mutex while the failed enqueue still owns it.
+            assert!(
+                futures::poll!(close.as_mut()).is_pending(),
+                "concurrent close must wait for the failed enqueue to release ingest_mutex"
+            );
+            failed_enqueue_proceed.notify_one();
+
+            tokio::select! {
+                _ = finalize_reached.notified() => {}
+                result = &mut ingest => {
+                    panic!("ingest completed before terminal finalization: {result:?}")
+                }
+                result = &mut close => {
+                    panic!("close completed before terminal finalization: {result:?}")
+                }
+            }
+            assert!(
+                futures::poll!(ingest.as_mut()).is_pending(),
+                "ingest must wait for terminal finalization, not early error publication"
+            );
+            assert!(
+                futures::poll!(close.as_mut()).is_pending(),
+                "close must wait for the same terminal finalization"
+            );
+
+            finalize_proceed.notify_one();
+            let ingest_error = tokio::time::timeout(std::time::Duration::from_secs(1), ingest)
+                .await
+                .expect("request send failure must finish after terminal finalization")
+                .expect_err("closed request sender must reject the batch");
+            assert!(
+                ingest_error
+                    .to_string()
+                    .contains("Flight request stream closed while sending"),
+                "expected ingest to preserve the request-send failure, got: {ingest_error}"
+            );
+
+            let error = tokio::time::timeout(std::time::Duration::from_secs(1), close)
+                .await
+                .expect("concurrent close must finish after terminal finalization")
+                .expect_err("close must preserve the request-send failure");
+            assert!(
+                error
+                    .to_string()
+                    .contains("Flight request stream closed while sending"),
+                "expected the request-send failure, got: {error}"
+            );
+
+            tokio::time::timeout(std::time::Duration::from_secs(1), stream.flush())
+                .await
+                .expect("flush must not wait on the withdrawn offset")
+                .expect("flush target must roll back to the acknowledged batch");
+            assert!(stream.get_unacked_batches().await?.is_empty());
+            Ok(())
+        }
+
+        #[tokio::test]
+        async fn test_first_failed_enqueue_leaves_no_flush_target(
+        ) -> Result<(), Box<dyn std::error::Error>> {
+            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::HoldResponseAfterRequestEof],
+                )
+                .await;
+
+            let sdk = ZerobusSdk::builder()
+                .endpoint(server_url)
+                .unity_catalog_url("https://mock-uc.com")
+                .tls_config(Arc::new(NoTlsConfig))
+                .build()?;
+            let stream = sdk
+                .stream_builder()
+                .table(TABLE_NAME)
+                .headers_provider(Arc::new(TestHeadersProvider::default()))
+                .arrow(schema.clone())
+                .flush_timeout_ms(60_000)
+                .recovery(false)
+                .build_arrow()
+                .await?;
+
+            stream.replace_batch_sender_with_closed_channel().await;
+            let batch = create_test_record_batch(schema, vec![1], vec![Some("withdrawn")]);
+            let ingest_error = tokio::time::timeout(
+                std::time::Duration::from_secs(1),
+                stream.ingest_batch(batch),
+            )
+            .await
+            .expect("first failed enqueue must finalize promptly")
+            .expect_err("closed request sender must reject the first batch");
+            assert!(
+                ingest_error
+                    .to_string()
+                    .contains("Flight request stream closed while sending"),
+                "expected ingest to preserve the request-send failure, got: {ingest_error}"
+            );
+
+            let flush_error =
+                tokio::time::timeout(std::time::Duration::from_secs(1), stream.flush())
+                    .await
+                    .expect("empty-target flush must observe terminal finalization")
+                    .expect_err("flush must preserve the request-send failure");
+            assert!(
+                flush_error
+                    .to_string()
+                    .contains("Flight request stream closed while sending"),
+                "expected flush to preserve the request-send failure, got: {flush_error}"
+            );
+            assert!(stream.get_unacked_batches().await?.is_empty());
+            Ok(())
+        }
+
         #[tokio::test]
         async fn test_flush_timeout() -> Result<(), Box<dyn std::error::Error>> {
             setup_tracing();