From d5e950f211b9109bf9f489734e4efb9485997dca Mon Sep 17 00:00:00 2001 From: nikolaobradovic-db Date: Tue, 18 Aug 2026 08:04:25 +0000 Subject: [PATCH 1/2] [FFI] Complete Arrow shutdown before free returns --- CLAUDE.md | 3 +- cpp/NEXT_CHANGELOG.md | 5 + cpp/README.md | 5 + cpp/include/zerobus/arrow_stream.hpp | 6 +- go/CLAUDE.md | 3 +- go/NEXT_CHANGELOG.md | 8 + go/README.md | 7 +- go/arrow_stream.go | 26 +- rust/ffi/Cargo.toml | 2 +- rust/ffi/NEXT_CHANGELOG.md | 2 + rust/ffi/src/arrow.rs | 110 +++++- rust/ffi/zerobus.h | 24 +- rust/sdk/src/internal.rs | 7 + rust/sdk/src/stream/arrow/c_data.rs | 5 +- rust/sdk/src/stream/arrow/connection.rs | 440 ++++++++++++++++++++--- rust/sdk/src/stream/arrow/mod.rs | 159 ++++++-- rust/sdk/src/stream/arrow/supervisor.rs | 82 ++++- rust/tests/src/arrow_c_data_ffi_tests.rs | 398 +++++++++++++++++++- 18 files changed, 1170 insertions(+), 122 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 9b7ff843..4af03b38 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -68,7 +68,8 @@ Because every non-Rust SDK crosses a foreign-function boundary, keep these in mi ### Thread safety - Rust core is async (tokio). The FFI layer manages its own tokio runtime. -- Go: safe for concurrent goroutines. +- Go: concurrent ingestion is internally synchronized. Serialize `Close` with + every other operation on the same stream. - Python: async streams are safe; sync streams are single-threaded per instance. - TypeScript: async-safe via Node.js event loop. - Java: **not thread-safe** — external synchronization required for concurrent access. diff --git a/cpp/NEXT_CHANGELOG.md b/cpp/NEXT_CHANGELOG.md index 19dd253d..d4094102 100644 --- a/cpp/NEXT_CHANGELOG.md +++ b/cpp/NEXT_CHANGELOG.md @@ -16,6 +16,11 @@ `create_arrow_stream` still take a `std::shared_ptr` — and you no longer need to keep your own reference alive past `create_stream`. +- Fixed `ArrowStream` destruction so it waits while native background + shutdown finishes. Destruction must not race another operation on the stream. + Internal native shutdown infrastructure failures terminate the process rather + than returning after incomplete native destruction. + ### Documentation - Corrected custom-header examples and clarified that acknowledgment callbacks diff --git a/cpp/README.md b/cpp/README.md index 8488fe46..bcdfad51 100644 --- a/cpp/README.md +++ b/cpp/README.md @@ -460,6 +460,11 @@ access externally (the same contract as the Rust core). A single `Sdk` may create many streams. See [`CLAUDE.md`](https://github.com/databricks/zerobus-sdk/blob/main/cpp/CLAUDE.md) for the full memory-ownership and threading contract. +Destroying an `ArrowStream` can block while native background shutdown +finishes. Destruction must not race another operation on the same stream. An +internal native shutdown infrastructure failure terminates the process rather +than returning after incomplete native destruction. + ## HTTP proxy support Like the other SDKs, the C++ SDK honors the standard proxy environment variables diff --git a/cpp/include/zerobus/arrow_stream.hpp b/cpp/include/zerobus/arrow_stream.hpp index e30a0292..06141a57 100644 --- a/cpp/include/zerobus/arrow_stream.hpp +++ b/cpp/include/zerobus/arrow_stream.hpp @@ -24,9 +24,13 @@ class Sdk; /// errors (the destructor swallows them) and flushes synchronously, which can /// block up to `flush_timeout_ms` (default 5 minutes) if the server is /// unresponsive. Letting the object fall out of scope drags that blocking close -/// into the destructor. +/// into the destructor. Native destruction also waits for background shutdown +/// to finish and does not return on a shutdown timeout. An internal native +/// shutdown infrastructure failure terminates the process rather than returning +/// after incomplete native destruction. /// /// Thread safety: not safe for concurrent use from multiple threads. +/// Destruction must not race another operation on the same stream. /// /// Only the IPC-bytes ingest path is wrapped; /// `zerobus_arrow_stream_ingest_batch_via_record_batch` is a functional diff --git a/go/CLAUDE.md b/go/CLAUDE.md index c3902d20..cc89b09e 100644 --- a/go/CLAUDE.md +++ b/go/CLAUDE.md @@ -110,7 +110,8 @@ Public API is everything exported (capitalized) in the `go/` package: ## Thread safety -Go SDK is safe for concurrent use from multiple goroutines. Internal synchronization handles concurrent `Ingest` calls. +Concurrent `Ingest` calls are internally synchronized. Serialize `Close` with +every other operation on the same stream. ## Changelog and documentation diff --git a/go/NEXT_CHANGELOG.md b/go/NEXT_CHANGELOG.md index 3c0fc07e..ddcc9d74 100644 --- a/go/NEXT_CHANGELOG.md +++ b/go/NEXT_CHANGELOG.md @@ -8,6 +8,10 @@ ### Bug Fixes +- Arrow stream finalizers now hand potentially blocking native shutdown to a + separate goroutine, so one stalled cleanup cannot block Go's global finalizer + goroutine. Explicit `Close` remains synchronous. + ### Documentation - Corrected consumer installation so tagged releases do not require Rust or @@ -19,6 +23,10 @@ - Batch examples name the offset returned by `IngestRecordsOffset` `batchOffset`. Example and test modules keep the CI toolchain Go versions; the documented SDK minimum remains Go 1.21+. +- Documented that closing an Arrow stream can block while native background + shutdown finishes and must not race another operation on the same stream. + Internal native shutdown infrastructure failures terminate the process rather + than returning after incomplete native destruction. ### Internal Changes diff --git a/go/README.md b/go/README.md index 3d8e2493..973075f6 100644 --- a/go/README.md +++ b/go/README.md @@ -1442,7 +1442,12 @@ Waits for all pending batches to be acknowledged. func (s *ZerobusArrowStream) Close() error ``` -Flushes and closes the stream. +Flushes and closes the stream. Native destruction can block while background +shutdown finishes. Serialize `Close` with every other operation on the same +stream. The GC finalizer schedules this cleanup asynchronously; call `Close` +explicitly when completion and errors matter. +An internal native shutdown infrastructure failure terminates the process +rather than returning after incomplete native destruction. ```go func (s *ZerobusArrowStream) GetUnackedBatches() ([][]byte, error) diff --git a/go/arrow_stream.go b/go/arrow_stream.go index 76336a9e..68de4f9f 100644 --- a/go/arrow_stream.go +++ b/go/arrow_stream.go @@ -105,6 +105,17 @@ type ZerobusArrowStream struct { ptr unsafe.Pointer } +func setArrowStreamFinalizer(stream *ZerobusArrowStream) { + runtime.SetFinalizer(stream, func(st *ZerobusArrowStream) { + // Native destruction may wait indefinitely to guarantee that no foreign + // release callback runs after it returns. Never block Go's sole finalizer + // goroutine on that wait. + go func() { + _ = st.Close() + }() + }) +} + // CreateArrowStream creates an Arrow Flight stream authenticated with OAuth client credentials. // // schemaIpcBytes must be Arrow IPC stream bytes containing only the schema (no data batches). @@ -129,9 +140,7 @@ func (s *ZerobusSdk) CreateArrowStream( } stream := &ZerobusArrowStream{ptr: ptr} - runtime.SetFinalizer(stream, func(st *ZerobusArrowStream) { - st.Close() //nolint:errcheck - }) + setArrowStreamFinalizer(stream) return stream, nil } @@ -156,9 +165,7 @@ func (s *ZerobusSdk) CreateArrowStreamWithHeadersProvider( } stream := &ZerobusArrowStream{ptr: ptr} - runtime.SetFinalizer(stream, func(st *ZerobusArrowStream) { - st.Close() //nolint:errcheck - }) + setArrowStreamFinalizer(stream) return stream, nil } @@ -189,7 +196,12 @@ func (st *ZerobusArrowStream) Flush() error { } // Close gracefully shuts down the stream after flushing all pending batches. -// The stream must not be used after Close returns. Close is idempotent. +// Native destruction can block while background shutdown finishes. Close must +// be serialized with IngestBatch, Flush, and every other operation on this +// stream. An internal native shutdown infrastructure failure terminates the +// process rather than returning after incomplete native destruction. +// The stream must not be used after Close returns. +// Close is idempotent. func (st *ZerobusArrowStream) Close() error { if st.ptr == nil { return nil diff --git a/rust/ffi/Cargo.toml b/rust/ffi/Cargo.toml index ea8b45c6..78887689 100644 --- a/rust/ffi/Cargo.toml +++ b/rust/ffi/Cargo.toml @@ -18,7 +18,7 @@ databricks-zerobus-ingest-sdk = { path = "../sdk", version = "2.0.1", features = arrow-ipc.workspace = true # FFI helpers -tokio = { workspace = true, features = ["rt", "rt-multi-thread"] } +tokio = { workspace = true, features = ["rt", "rt-multi-thread", "time"] } once_cell.workspace = true prost.workspace = true prost-types.workspace = true diff --git a/rust/ffi/NEXT_CHANGELOG.md b/rust/ffi/NEXT_CHANGELOG.md index 696dd1b7..18384762 100644 --- a/rust/ffi/NEXT_CHANGELOG.md +++ b/rust/ffi/NEXT_CHANGELOG.md @@ -19,6 +19,8 @@ ### Behavior Changes +- `zerobus_arrow_stream_free` now blocks until Arrow background shutdown completes, every Flight request body reaches EOF or is dropped, and all retained Arrow C Data owners are released. Previously, a request body could retain an owner and run its release callback after free returned on an unacknowledged/failure path, risking callback-after-free use of wrapper state. When the calling restrictions below are respected, no Arrow C Data release callback for that stream can run after free returns. The function logs a warning every 30 seconds while shutdown remains incomplete; it does not return on a timeout. Free must not race another operation on the same stream handle. Freeing the same stream reentrantly from one of its SDK callbacks is unsupported because shutdown would wait for that callback; freeing a different stream from a callback remains supported. An internal native shutdown panic, a required helper-thread spawn failure, or a helper-thread panic terminates the process rather than returning without the release-callback guarantee. + ### Breaking Changes ### Deprecations diff --git a/rust/ffi/src/arrow.rs b/rust/ffi/src/arrow.rs index ee1d4cad..154c93b9 100644 --- a/rust/ffi/src/arrow.rs +++ b/rust/ffi/src/arrow.rs @@ -4,6 +4,7 @@ use crate::arrow_c_data::{CArrowArray, CArrowSchema}; use crate::common::*; use arrow_ipc::{reader::StreamReader, writer::StreamWriter, CompressionType}; use bytes::Bytes; +use databricks_zerobus_ingest_sdk::internal::abort_arrow_stream_and_wait; use databricks_zerobus_ingest_sdk::internal::arrow_c_data::{ import_c_data_record_batch, FFI_ArrowArray, FFI_ArrowSchema, }; @@ -11,14 +12,23 @@ use databricks_zerobus_ingest_sdk::{ HeadersProvider, RecordBatch, StreamBuilder, ZerobusArrowStream, ZerobusError, ZerobusResult, }; use std::os::raw::c_char; +use std::panic::{catch_unwind, AssertUnwindSafe}; use std::ptr; -use std::sync::Arc; +use std::sync::{Arc, Mutex as StdMutex}; +use std::thread; +use std::time::Duration; +use tracing::{error, warn}; // ============================================================================ // Arrow Flight FFI // ============================================================================ +const ARROW_STREAM_SHUTDOWN_WARNING_AFTER: Duration = Duration::from_secs(30); + /// Opaque handle for an Arrow Flight stream. +/// +/// FFI pointers are `Box` addresses cast to this type. +/// Creation, validation, test-hook casts, and `Box::from_raw` must stay coupled. #[repr(C)] pub struct CArrowStream { _private: [u8; 0], @@ -291,13 +301,102 @@ pub extern "C" fn zerobus_sdk_create_arrow_stream_with_headers_provider( }) } +fn abort_and_drop_arrow_stream(stream: Box) { + let mut stream = Some(stream); + let shutdown = catch_unwind(AssertUnwindSafe(|| { + let stream_ref = stream + .as_ref() + .expect("Arrow stream must remain owned during shutdown"); + RUNTIME.block_on(async { + let shutdown = abort_arrow_stream_and_wait(stream_ref); + tokio::pin!(shutdown); + while tokio::time::timeout(ARROW_STREAM_SHUTDOWN_WARNING_AFTER, &mut shutdown) + .await + .is_err() + { + warn!( + threshold_seconds = ARROW_STREAM_SHUTDOWN_WARNING_AFTER.as_secs(), + "Arrow stream background shutdown is still running; continuing to wait" + ); + } + }); + drop( + stream + .take() + .expect("Arrow stream must remain owned until shutdown completes"), + ); + })); + + if shutdown.is_err() { + error!("Arrow stream shutdown or destruction panicked; aborting process"); + std::process::abort(); + } +} + +fn abort_and_drop_arrow_stream_on_thread(stream: Box) { + let stream_slot = Arc::new(StdMutex::new(Some(stream))); + let worker_slot = Arc::clone(&stream_slot); + let worker = thread::Builder::new() + .name("zerobus-arrow-free".to_string()) + .spawn(move || { + let mut slot = worker_slot + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let stream = slot + .take() + .expect("Arrow stream shutdown has exactly one owner"); + drop(slot); + abort_and_drop_arrow_stream(stream); + }); + + let worker = match worker { + Ok(worker) => worker, + Err(error) => { + // Leaking only the stream cannot preserve the callback guarantee: an ACK + // can already have moved the last SDK-owned batch into the request body, + // which may release it after this function returns. + error!(%error, "Unable to start Arrow stream shutdown thread; aborting process"); + std::process::abort(); + } + }; + + if worker.join().is_err() { + error!("Arrow stream shutdown thread panicked; aborting process"); + std::process::abort(); + } +} + /// Frees an Arrow Flight stream instance. +/// +/// This call blocks until background shutdown completes, every Flight request body reaches EOF +/// or is dropped, and all retained Arrow C Data owners are released. If shutdown takes longer +/// than 30 seconds, it logs a warning every 30 seconds while continuing to wait. It does not +/// return on a timeout. +/// When the calling restrictions below are respected, no Arrow C Data release callback for this +/// stream can run after this function returns. +/// An internal shutdown panic, a required helper-thread spawn failure, or a helper-thread panic +/// terminates the process rather than returning without that guarantee. +/// +/// Do not race this function with another operation on the same stream handle. Freeing this same +/// stream reentrantly from one of its SDK callbacks is unsupported because shutdown would wait +/// for the callback that is making the call. Freeing a different stream from a callback remains +/// supported. #[no_mangle] pub extern "C" fn zerobus_arrow_stream_free(stream: *mut CArrowStream) { ffi_guard(ptr::null_mut(), (), move || { if !stream.is_null() { unsafe { - let _ = Box::from_raw(stream as *mut ZerobusArrowStream); + let stream = Box::from_raw(stream as *mut ZerobusArrowStream); + match tokio::runtime::Handle::try_current() { + Ok(handle) + if handle.runtime_flavor() + == tokio::runtime::RuntimeFlavor::MultiThread => + { + tokio::task::block_in_place(|| abort_and_drop_arrow_stream(stream)); + } + Ok(_) => abort_and_drop_arrow_stream_on_thread(stream), + Err(_) => abort_and_drop_arrow_stream(stream), + } } } }) @@ -365,9 +464,10 @@ pub extern "C" fn zerobus_arrow_stream_ingest_batch( /// `ArrowArray` / `ArrowSchema` structure satisfying the Arrow C Data /// Interface. All referenced children, dictionaries, buffers, `private_data`, /// and release callbacks must remain valid for the lifetime required by the -/// producer contract. After ownership transfer, the SDK may invoke release -/// asynchronously on an internal runtime thread. Release callbacks must -/// therefore be thread-safe and must not unwind or throw across the C ABI. +/// producer contract. After ownership transfer, release callbacks may run on +/// any thread that drops the final owner, including SDK runtime/transport +/// threads or the thread calling `zerobus_arrow_stream_free`. They must be +/// thread-safe and must not unwind or throw across the C ABI. /// /// Malformed, dangling, or malicious structures are caller undefined behavior /// and cannot be safely validated by this function. diff --git a/rust/ffi/zerobus.h b/rust/ffi/zerobus.h index 26865f57..a8e363f9 100644 --- a/rust/ffi/zerobus.h +++ b/rust/ffi/zerobus.h @@ -37,6 +37,9 @@ typedef struct CHeaders { /** * Opaque handle for an Arrow Flight stream. + * + * FFI pointers are `Box` addresses cast to this type. + * Creation, validation, test-hook casts, and `Box::from_raw` must stay coupled. */ typedef struct CArrowStream { uint8_t _private[0]; @@ -302,6 +305,20 @@ struct CArrowStream *zerobus_sdk_create_arrow_stream_with_headers_provider(struc /** * Frees an Arrow Flight stream instance. + * + * This call blocks until background shutdown completes, every Flight request body reaches EOF + * or is dropped, and all retained Arrow C Data owners are released. If shutdown takes longer + * than 30 seconds, it logs a warning every 30 seconds while continuing to wait. It does not + * return on a timeout. + * When the calling restrictions below are respected, no Arrow C Data release callback for this + * stream can run after this function returns. + * An internal shutdown panic, a required helper-thread spawn failure, or a helper-thread panic + * terminates the process rather than returning without that guarantee. + * + * Do not race this function with another operation on the same stream handle. Freeing this same + * stream reentrantly from one of its SDK callbacks is unsupported because shutdown would wait + * for the callback that is making the call. Freeing a different stream from a callback remains + * supported. */ void zerobus_arrow_stream_free(struct CArrowStream *stream); @@ -329,9 +346,10 @@ int64_t zerobus_arrow_stream_ingest_batch(struct CArrowStream *stream, * `ArrowArray` / `ArrowSchema` structure satisfying the Arrow C Data * Interface. All referenced children, dictionaries, buffers, `private_data`, * and release callbacks must remain valid for the lifetime required by the - * producer contract. After ownership transfer, the SDK may invoke release - * asynchronously on an internal runtime thread. Release callbacks must - * therefore be thread-safe and must not unwind or throw across the C ABI. + * producer contract. After ownership transfer, release callbacks may run on + * any thread that drops the final owner, including SDK runtime/transport + * threads or the thread calling `zerobus_arrow_stream_free`. They must be + * thread-safe and must not unwind or throw across the C ABI. * * Malformed, dangling, or malicious structures are caller undefined behavior * and cannot be safely validated by this function. diff --git a/rust/sdk/src/internal.rs b/rust/sdk/src/internal.rs index a93937f8..5c923c90 100644 --- a/rust/sdk/src/internal.rs +++ b/rust/sdk/src/internal.rs @@ -10,3 +10,10 @@ pub mod arrow_c_data { import_c_data_record_batch, FFI_ArrowArray, FFI_ArrowSchema, }; } + +/// Stops Arrow stream background work without flushing and waits for shutdown to complete. +#[cfg(feature = "internal-arrow-c-data")] +#[doc(hidden)] +pub async fn abort_arrow_stream_and_wait(stream: &crate::ZerobusArrowStream) { + stream.abort_and_wait().await; +} diff --git a/rust/sdk/src/stream/arrow/c_data.rs b/rust/sdk/src/stream/arrow/c_data.rs index d0f919c9..08621eaa 100644 --- a/rust/sdk/src/stream/arrow/c_data.rs +++ b/rust/sdk/src/stream/arrow/c_data.rs @@ -15,8 +15,9 @@ use crate::{ZerobusError, ZerobusResult}; /// Both values are consumed on success and error. /// All retained buffers, children, dictionaries, `private_data`, and release /// callbacks must support asynchronous cross-thread retention until released. -/// Release callbacks may run on an arbitrary SDK runtime thread; they must be -/// thread-safe and must not unwind. +/// Release callbacks may run on any thread that drops the final owner, +/// including SDK runtime/transport threads or a caller performing destructive +/// stream teardown; they must be thread-safe and must not unwind. pub unsafe fn import_c_data_record_batch( array: FFI_ArrowArray, schema: FFI_ArrowSchema, diff --git a/rust/sdk/src/stream/arrow/connection.rs b/rust/sdk/src/stream/arrow/connection.rs index ef671673..4e4632a1 100644 --- a/rust/sdk/src/stream/arrow/connection.rs +++ b/rust/sdk/src/stream/arrow/connection.rs @@ -3,7 +3,7 @@ //! A `FlightConnection` owns both halves of one DoPut exchange. Request shutdown //! is observable so rotation can half-close without dropping the HTTP/2 stream. -use std::future::{pending, Future}; +use std::future::Future; use std::pin::Pin; use std::sync::atomic::{AtomicI64, Ordering}; use std::sync::Arc; @@ -13,7 +13,7 @@ use arrow_flight::encode::FlightDataEncoderBuilder; use arrow_flight::error::FlightError; use arrow_flight::{FlightClient, FlightData, PutResult}; use futures::{stream::poll_fn, Stream, StreamExt}; -use tokio::sync::{mpsc, watch}; +use tokio::sync::{mpsc, watch, Mutex}; use tokio::time::{timeout, timeout_at, Duration, Instant}; use tokio_util::sync::CancellationToken; use tonic::metadata::MetadataValue; @@ -23,8 +23,8 @@ use tracing::{error, info, warn}; use super::batch::make_ipc_write_options; use super::metadata::{FlightAckMetadata, FlightBatchMetadata}; use super::{ - configured_deadline, ArrowStreamConfigurationOptions, ArrowTableProperties, RecordBatch, - ZerobusArrowStream, + configured_deadline, ArrowStreamConfigurationOptions, ArrowTableProperties, + FlightConnectionParameters, RecordBatch, ZerobusArrowStream, }; use crate::errors::ZerobusError; use crate::headers_provider::HeadersProvider; @@ -38,6 +38,7 @@ type FlightRequestStream = Pin, @@ -57,6 +58,10 @@ impl RequestBodyControl { self.shutdown.cancel(); } + pub(super) fn is_finished(&self) -> bool { + *self.eof_rx.borrow() || self.eof_rx.has_changed().is_err() + } + pub(super) async fn wait_for_eof(&self) { let mut eof_rx = self.eof_rx.clone(); loop { @@ -64,9 +69,144 @@ impl RequestBodyControl { return; } if eof_rx.changed().await.is_err() { - pending::<()>().await; + // RequestBodyStreamState::finish/Drop drops encoded channel state and + // publishes EOF before dropping this sender, so closure also proves + // every request-body owner was released. + return; + } + } + } +} + +#[derive(Clone, Default)] +pub(super) struct RequestBodyRegistry { + controls: Arc>>, +} + +impl RequestBodyRegistry { + pub(super) async fn register(&self, control: &RequestBodyControl) { + let mut controls = self.controls.lock().await; + controls.retain(|registered| !registered.is_finished()); + controls.push(control.clone()); + } + + #[cfg(any(feature = "internal-arrow-c-data", test))] + pub(super) async fn shutdown_all(&self) { + let mut controls = self.controls.lock().await; + controls.retain(|control| !control.is_finished()); + for control in controls.iter() { + control.shutdown(); + } + } + + pub(super) fn try_shutdown_all(&self) { + if let Ok(mut controls) = self.controls.try_lock() { + controls.retain(|control| !control.is_finished()); + for control in controls.iter() { + control.shutdown(); + } + } + } + + #[cfg(any(feature = "internal-arrow-c-data", test))] + pub(super) async fn wait_for_all_eof(&self) { + loop { + let controls = { + let mut registered = self.controls.lock().await; + registered.retain(|control| !control.is_finished()); + registered.clone() + }; + if controls.is_empty() { + return; + } + for control in controls { + control.wait_for_eof().await; + } + } + } + + #[cfg(test)] + async fn registered_count(&self) -> usize { + self.controls.lock().await.len() + } +} + +#[cfg(feature = "test-hooks")] +struct RequestBodyTestState { + hooks: Arc, + shutdown_barrier: Option + Send>>>, + before_batch_poll_barrier: Option + Send>>>, +} + +#[cfg(feature = "test-hooks")] +impl RequestBodyTestState { + fn new(hooks: Arc) -> Self { + Self { + hooks, + shutdown_barrier: None, + before_batch_poll_barrier: None, + } + } + + fn poll_shutdown_barrier(&mut self, cx: &mut std::task::Context<'_>) -> Poll<()> { + let barrier = self.shutdown_barrier.get_or_insert_with(|| { + let hooks = Arc::clone(&self.hooks); + Box::pin(async move { + let barrier = hooks.request_body_shutdown.lock().await.take(); + if let Some(barrier) = barrier { + barrier.reached.notify_one(); + barrier.proceed.notified().await; + } + }) + }); + barrier.as_mut().poll(cx) + } + + fn poll_before_batch_barrier(&mut self, cx: &mut std::task::Context<'_>) -> Poll<()> { + if self.before_batch_poll_barrier.is_none() { + let barrier = self + .hooks + .request_body_before_batch_poll + .try_lock() + .ok() + .and_then(|mut gate| gate.take()); + if let Some(barrier) = barrier { + self.before_batch_poll_barrier = Some(Box::pin(async move { + barrier.reached.notify_one(); + barrier.proceed.notified().await; + })); } } + if let Some(barrier) = self.before_batch_poll_barrier.as_mut() { + if barrier.as_mut().poll(cx).is_pending() { + return Poll::Pending; + } + self.before_batch_poll_barrier = None; + } + Poll::Ready(()) + } +} + +struct RequestBodyStreamState { + encoded: Option, + cancelled: Pin + Send>>, + eof_tx: Option>, + #[cfg(feature = "test-hooks")] + test: RequestBodyTestState, +} + +impl RequestBodyStreamState { + fn finish(&mut self) { + drop(self.encoded.take()); + if let Some(eof_tx) = self.eof_tx.take() { + eof_tx.send_replace(true); + } + } +} + +impl Drop for RequestBodyStreamState { + fn drop(&mut self) { + self.finish(); } } @@ -104,34 +244,36 @@ impl ZerobusArrowStream { /// Attempts to establish a Flight connection. /// Returns the complete connection on success. pub(super) async fn try_connect( - endpoint: &str, - tls_config: &Arc, - connector_factory: Option<&ConnectorFactory>, - table_properties: &ArrowTableProperties, - options: &ArrowStreamConfigurationOptions, - headers_provider: &Arc, - sdk_identifier: &str, + parameters: &FlightConnectionParameters<'_>, ) -> ZerobusResult { // Share one deadline across connection setup and auth-rejection invalidation. // This preserves the original auth error if a custom provider stalls instead of // reclassifying the attempt as a retryable setup timeout. - let attempt_timeout = Duration::from_millis(options.recovery_timeout_ms); + let attempt_timeout = Duration::from_millis(parameters.options.recovery_timeout_ms); let attempt_started = Instant::now(); let attempt_deadline = configured_deadline(attempt_started, attempt_timeout, "recovery_timeout_ms")?; let result = timeout_at(attempt_deadline, async { let client = Self::create_flight_client( - endpoint, - tls_config, - connector_factory, - table_properties, - options, - headers_provider, - sdk_identifier, + parameters.endpoint, + parameters.tls_config, + parameters.connector_factory, + parameters.table_properties, + parameters.options, + parameters.headers_provider, + parameters.sdk_identifier, ) .await?; - Self::start_stream_connection(client, table_properties, options).await + Self::start_stream_connection( + client, + parameters.table_properties, + parameters.options, + parameters.request_bodies, + #[cfg(feature = "test-hooks")] + Arc::clone(parameters.test_hooks), + ) + .await }) .await .map_err(|_| { @@ -147,12 +289,12 @@ impl ZerobusArrowStream { // not be able to turn a known auth rejection into repeated generic // timeout retries by stalling here. if error.is_auth_rejection() - && timeout_at(attempt_deadline, headers_provider.invalidate()) + && timeout_at(attempt_deadline, parameters.headers_provider.invalidate()) .await .is_err() { warn!(target: super::LOG_TARGET, - timeout_ms = options.recovery_timeout_ms, + timeout_ms = parameters.options.recovery_timeout_ms, "Initial headers provider invalidation timed out; preserving auth rejection" ); } @@ -238,13 +380,14 @@ impl ZerobusArrowStream { batch_rx: mpsc::Receiver>, table_properties: &ArrowTableProperties, options: &ArrowStreamConfigurationOptions, + #[cfg(feature = "test-hooks")] test_hooks: Arc, ) -> ZerobusResult<(FlightRequestStream, RequestBodyControl)> { let ipc_write_options = make_ipc_write_options(options.ipc_compression)?; let schema = Arc::clone(&table_properties.schema); let batch_stream = tokio_stream::wrappers::ReceiverStream::new(batch_rx); let offset_counter = Arc::new(AtomicI64::new(0)); let offset_counter_clone = Arc::clone(&offset_counter); - let mut encoded: FlightRequestStream = Box::pin( + let encoded: FlightRequestStream = Box::pin( FlightDataEncoderBuilder::new() .with_schema(schema) .with_options(ipc_write_options) @@ -265,16 +408,37 @@ impl ZerobusArrowStream { ); let shutdown = CancellationToken::new(); - let mut cancelled = Box::pin(shutdown.clone().cancelled_owned()); let (eof_tx, eof_rx) = watch::channel(false); + let mut state = RequestBodyStreamState { + encoded: Some(encoded), + cancelled: Box::pin(shutdown.clone().cancelled_owned()), + eof_tx: Some(eof_tx), + #[cfg(feature = "test-hooks")] + test: RequestBodyTestState::new(test_hooks), + }; let controlled = poll_fn(move |cx| { - if cancelled.as_mut().poll(cx).is_ready() { - eof_tx.send_replace(true); + if state.encoded.is_none() { + return Poll::Ready(None); + } + if state.cancelled.as_mut().poll(cx).is_ready() { + #[cfg(feature = "test-hooks")] + if state.test.poll_shutdown_barrier(cx).is_pending() { + return Poll::Pending; + } + state.finish(); return Poll::Ready(None); } + #[cfg(feature = "test-hooks")] + if state.test.poll_before_batch_barrier(cx).is_pending() { + return Poll::Pending; + } + let encoded = state + .encoded + .as_mut() + .expect("request body completion checked before polling"); match encoded.as_mut().poll_next(cx) { Poll::Ready(None) => { - eof_tx.send_replace(true); + state.finish(); Poll::Ready(None) } result => result, @@ -298,13 +462,23 @@ impl ZerobusArrowStream { mut client: FlightClient, table_properties: &ArrowTableProperties, options: &ArrowStreamConfigurationOptions, + request_bodies: &RequestBodyRegistry, + #[cfg(feature = "test-hooks")] test_hooks: Arc, ) -> ZerobusResult { // Create channel for sending RecordBatches. let (batch_tx, batch_rx) = mpsc::channel::>(options.max_inflight_batches); - let (flight_data_stream, request_body) = - Self::make_request_stream(batch_rx, table_properties, options)?; + let (flight_data_stream, request_body) = Self::make_request_stream( + batch_rx, + table_properties, + options, + #[cfg(feature = "test-hooks")] + test_hooks, + )?; + // Register before tonic takes ownership so every do_put/setup exit remains + // observable by destructive free and cannot outlive the C Data owner guarantee. + request_bodies.register(&request_body).await; // Start the DoPut stream. let mut response_stream = client @@ -382,29 +556,32 @@ impl ZerobusArrowStream { /// Establishes a replacement DoPut transport and validates its ready signal. pub(super) async fn reconnect_transport( - endpoint: &str, - tls_config: &Arc, - connector_factory: Option<&ConnectorFactory>, - table_properties: &ArrowTableProperties, - options: &ArrowStreamConfigurationOptions, - headers_provider: &Arc, - sdk_identifier: &str, + parameters: &FlightConnectionParameters<'_>, ) -> ZerobusResult { let client = Self::create_flight_client( - endpoint, - tls_config, - connector_factory, - table_properties, - options, - headers_provider, - sdk_identifier, + parameters.endpoint, + parameters.tls_config, + parameters.connector_factory, + parameters.table_properties, + parameters.options, + parameters.headers_provider, + parameters.sdk_identifier, ) .await?; - let (batch_tx, batch_rx) = - mpsc::channel::>(options.max_inflight_batches); - let (flight_data_stream, request_body) = - Self::make_request_stream(batch_rx, table_properties, options)?; + let (batch_tx, batch_rx) = mpsc::channel::>( + parameters.options.max_inflight_batches, + ); + let (flight_data_stream, request_body) = Self::make_request_stream( + batch_rx, + parameters.table_properties, + parameters.options, + #[cfg(feature = "test-hooks")] + Arc::clone(parameters.test_hooks), + )?; + // Register before tonic takes ownership so replay can never target an + // untracked request body, including when do_put or READY setup fails. + parameters.request_bodies.register(&request_body).await; let mut flight_client = client; let mut response_stream = flight_client @@ -414,7 +591,7 @@ impl ZerobusArrowStream { // flatten it to `Unknown` and break auth/retry classification. .map_err(|e| ZerobusError::CreateStreamError(e.into()))?; - let setup_timeout = Duration::from_millis(options.connection_timeout_ms); + let setup_timeout = Duration::from_millis(parameters.options.connection_timeout_ms); match timeout(setup_timeout, response_stream.next()).await { Ok(Some(Ok(put_result))) => { match FlightAckMetadata::from_bytes(&put_result.app_metadata) { @@ -453,11 +630,11 @@ impl ZerobusArrowStream { Err(_timeout) => { error!(target: super::LOG_TARGET, "Timed out waiting for server reconnect confirmation ({}ms)", - options.connection_timeout_ms + parameters.options.connection_timeout_ms ); return Err(ZerobusError::ConnectionTimeout(format!( "Timed out waiting for server reconnect confirmation ({}ms)", - options.connection_timeout_ms + parameters.options.connection_timeout_ms ))); } } @@ -477,14 +654,15 @@ mod tests { use arrow_flight::error::FlightError; use async_trait::async_trait; + use futures::StreamExt; use tokio::sync::mpsc; use tokio::time::{timeout, Duration}; use super::super::{ArrowSchema, RecordBatch}; use super::{ - ArrowStreamConfigurationOptions, ArrowTableProperties, FlightConnection, - FlightResponseStream, HeadersProvider, RequestBodyControl, TlsConfig, ZerobusArrowStream, - ZerobusResult, + ArrowStreamConfigurationOptions, ArrowTableProperties, FlightClient, FlightConnection, + FlightResponseStream, HeadersProvider, RequestBodyControl, RequestBodyRegistry, TlsConfig, + ZerobusArrowStream, ZerobusResult, }; struct PassthroughTlsConfig; @@ -530,6 +708,160 @@ mod tests { assert!(received.is_none()); } + #[tokio::test] + async fn request_body_control_clone_stops_body_and_reports_owner_drop() { + let table_properties = ArrowTableProperties { + table_name: "catalog.schema.table".to_string(), + schema: Arc::new(ArrowSchema::empty()), + }; + let (_batch_tx, batch_rx) = mpsc::channel(1); + let (mut request_body, control) = ZerobusArrowStream::make_request_stream( + batch_rx, + &table_properties, + &ArrowStreamConfigurationOptions::default(), + #[cfg(feature = "test-hooks")] + Arc::new(super::super::TestHooks::default()), + ) + .unwrap(); + let cloned = control.clone(); + + cloned.shutdown(); + assert!(request_body.next().await.is_none()); + timeout(Duration::from_secs(1), control.wait_for_eof()) + .await + .expect("request body completion must reach every control clone"); + } + + #[tokio::test] + async fn request_body_remains_finished_after_natural_eof() { + let table_properties = ArrowTableProperties { + table_name: "catalog.schema.table".to_string(), + schema: Arc::new(ArrowSchema::empty()), + }; + let (batch_tx, batch_rx) = mpsc::channel(1); + let (mut request_body, control) = ZerobusArrowStream::make_request_stream( + batch_rx, + &table_properties, + &ArrowStreamConfigurationOptions::default(), + #[cfg(feature = "test-hooks")] + Arc::new(super::super::TestHooks::default()), + ) + .unwrap(); + drop(batch_tx); + + while request_body.next().await.is_some() {} + assert!(request_body.next().await.is_none()); + assert!(control.is_finished()); + } + + #[cfg(feature = "test-hooks")] + #[tokio::test] + async fn request_body_ignores_shutdown_hooks_after_natural_eof() { + let table_properties = ArrowTableProperties { + table_name: "catalog.schema.table".to_string(), + schema: Arc::new(ArrowSchema::empty()), + }; + let (batch_tx, batch_rx) = mpsc::channel(1); + let test_hooks = Arc::new(super::super::TestHooks::default()); + let (mut request_body, control) = ZerobusArrowStream::make_request_stream( + batch_rx, + &table_properties, + &ArrowStreamConfigurationOptions::default(), + Arc::clone(&test_hooks), + ) + .unwrap(); + drop(batch_tx); + while request_body.next().await.is_some() {} + + let reached = Arc::new(tokio::sync::Notify::new()); + let proceed = Arc::new(tokio::sync::Notify::new()); + *test_hooks.request_body_shutdown.lock().await = Some(super::super::TestBarrier { + reached, + proceed: Arc::clone(&proceed), + }); + control.shutdown(); + + for _ in 0..2 { + let next = timeout(Duration::from_millis(100), request_body.next()).await; + proceed.notify_one(); + assert!(next + .expect("completed request body must not enter shutdown hooks") + .is_none()); + } + } + + #[test] + fn closed_eof_watch_is_finished() { + let (eof_tx, eof_rx) = tokio::sync::watch::channel(false); + drop(eof_tx); + let control = RequestBodyControl { + shutdown: tokio_util::sync::CancellationToken::new(), + eof_rx, + }; + + assert!(control.is_finished()); + } + + #[tokio::test] + async fn request_body_registry_prunes_dropped_bodies() { + let table_properties = ArrowTableProperties { + table_name: "catalog.schema.table".to_string(), + schema: Arc::new(ArrowSchema::empty()), + }; + let (_batch_tx, batch_rx) = mpsc::channel(1); + let (request_body, control) = ZerobusArrowStream::make_request_stream( + batch_rx, + &table_properties, + &ArrowStreamConfigurationOptions::default(), + #[cfg(feature = "test-hooks")] + Arc::new(super::super::TestHooks::default()), + ) + .unwrap(); + let registry = RequestBodyRegistry::default(); + registry.register(&control).await; + assert_eq!(registry.registered_count().await, 1); + + drop(request_body); + registry.wait_for_all_eof().await; + assert_eq!(registry.registered_count().await, 0); + } + + #[tokio::test] + async fn request_body_is_registered_before_do_put_completes() { + let table_properties = ArrowTableProperties { + table_name: "catalog.schema.table".to_string(), + schema: Arc::new(ArrowSchema::empty()), + }; + let registry = RequestBodyRegistry::default(); + let client = FlightClient::new( + tonic::transport::Channel::from_static("http://127.0.0.1:1").connect_lazy(), + ); + let options = ArrowStreamConfigurationOptions::default(); + + { + let attempt = ZerobusArrowStream::start_stream_connection( + client, + &table_properties, + &options, + ®istry, + #[cfg(feature = "test-hooks")] + Arc::new(super::super::TestHooks::default()), + ); + tokio::pin!(attempt); + let _ = futures::poll!(attempt.as_mut()); + assert_eq!( + registry.registered_count().await, + 1, + "request body must be registered before do_put can complete" + ); + } + + timeout(Duration::from_secs(1), registry.wait_for_all_eof()) + .await + .expect("dropping the setup attempt must release its request body"); + assert_eq!(registry.registered_count().await, 0); + } + #[tokio::test] async fn authorization_metadata_is_sensitive() { let table_properties = ArrowTableProperties { diff --git a/rust/sdk/src/stream/arrow/mod.rs b/rust/sdk/src/stream/arrow/mod.rs index c922a11d..fe12d9ec 100644 --- a/rust/sdk/src/stream/arrow/mod.rs +++ b/rust/sdk/src/stream/arrow/mod.rs @@ -20,7 +20,6 @@ use std::sync::Arc; use arrow_flight::error::FlightError; use bytes::Bytes; use tokio::sync::{mpsc, watch, Mutex, Notify, Semaphore}; -use tokio::task::AbortHandle; use tokio::time::{timeout, Duration, Instant}; use tokio_retry::strategy::FixedInterval; use tokio_retry::RetryIf; @@ -31,8 +30,9 @@ pub use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit}; use self::batch::{materialize_ipc, PendingBatch}; use self::close::{CloseCoordinator, CloseFinalizer, CloseRequest, CloseState}; +use self::connection::RequestBodyRegistry; pub use self::options::ArrowStreamConfigurationOptions; -use self::supervisor::Supervisor; +use self::supervisor::{Supervisor, SupervisorTaskHandle}; use crate::errors::{should_retry_initial_connection, ZerobusError}; use crate::headers_provider::HeadersProvider; use crate::offset_generator::{OffsetId, OffsetIdGenerator}; @@ -55,6 +55,19 @@ const LOG_TARGET: &str = module_path!(); type BatchSender = Arc>>>>; +struct FlightConnectionParameters<'a> { + endpoint: &'a str, + tls_config: &'a Arc, + connector_factory: Option<&'a ConnectorFactory>, + table_properties: &'a ArrowTableProperties, + options: &'a ArrowStreamConfigurationOptions, + headers_provider: &'a Arc, + sdk_identifier: &'a str, + request_bodies: &'a RequestBodyRegistry, + #[cfg(feature = "test-hooks")] + test_hooks: &'a Arc, +} + /// Converts a configured relative timeout into an absolute monotonic-clock deadline. pub(super) fn configured_deadline( started_at: Instant, @@ -91,6 +104,10 @@ struct TestHooks { ack_idle: TestNotifyGate, failed_enqueue: TestBarrierGate, close_finalize: TestBarrierGate, + request_body_before_batch_poll: TestBarrierGate, + request_body_shutdown: TestBarrierGate, + retained_batches_cleared: TestNotifyGate, + free_shutdown_complete: TestBarrierGate, } /// Properties for an Arrow Flight ingestion table. @@ -177,8 +194,10 @@ pub struct ZerobusArrowStream { admission_closed: Arc, /// Coordinates one resumable explicit-close request with the recovery supervisor. close: CloseCoordinator, - /// Abort handle for the supervisor worker; its detached reaper remains independent. - supervisor_abort: Arc>>, + /// Supervisor worker and reaper ownership, consumed once by terminal shutdown. + supervisor_task: Arc>>, + /// Tracks every tonic-owned Flight request body until its queued owners are dropped. + request_bodies: RequestBodyRegistry, /// 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. @@ -274,7 +293,8 @@ impl ZerobusArrowStream { 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 supervisor_abort = Arc::new(Mutex::new(None)); + let supervisor_task = Arc::new(Mutex::new(None)); + let request_bodies = RequestBodyRegistry::default(); 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)); @@ -295,7 +315,8 @@ impl ZerobusArrowStream { is_closed, admission_closed, close, - supervisor_abort, + supervisor_task, + request_bodies, pending_batches, pending_notify, request_send_failure, @@ -325,6 +346,7 @@ impl ZerobusArrowStream { let table_properties = stream.table_properties.clone(); let options = stream.options.clone(); let headers_provider = Arc::clone(&stream.headers_provider); + let request_bodies = stream.request_bodies.clone(); let strategy = FixedInterval::from_millis(options.recovery_backoff_ms) .take(options.recovery_retries as usize); @@ -336,18 +358,24 @@ impl ZerobusArrowStream { let options = options.clone(); let headers_provider = Arc::clone(&headers_provider); let sdk_identifier = Arc::clone(&stream.sdk_identifier); + let request_bodies = request_bodies.clone(); + #[cfg(feature = "test-hooks")] + let test_hooks = Arc::clone(&stream.test_hooks); async move { - Self::try_connect( - &endpoint, - &tls_config, - connector_factory.as_ref(), - &table_properties, - &options, - &headers_provider, - &sdk_identifier, - ) - .await + let parameters = FlightConnectionParameters { + endpoint: &endpoint, + tls_config: &tls_config, + connector_factory: connector_factory.as_ref(), + table_properties: &table_properties, + options: &options, + headers_provider: &headers_provider, + sdk_identifier: &sdk_identifier, + request_bodies: &request_bodies, + #[cfg(feature = "test-hooks")] + test_hooks: &test_hooks, + }; + Self::try_connect(¶meters).await } }; // Keep auth errors globally non-retryable, but let initial setup refresh one @@ -376,8 +404,8 @@ impl ZerobusArrowStream { let task = Supervisor::new(&stream).spawn(connection); { - let mut supervisor_abort = stream.supervisor_abort.lock().await; - *supervisor_abort = Some(task); + let mut supervisor_task = stream.supervisor_task.lock().await; + *supervisor_task = Some(task); } info!( @@ -955,6 +983,43 @@ impl ZerobusArrowStream { Ok(self.failed_batches.lock().await.clone()) } + #[cfg(feature = "internal-arrow-c-data")] + pub(crate) async fn abort_and_wait(&self) { + self.admission_closed.store(true, Ordering::Release); + self.request_bodies.shutdown_all().await; + + let task = self.supervisor_task.lock().await.take(); + if let Some(task) = task { + task.abort_and_wait().await; + } + + // A reconnect can register after the first snapshot but not after worker termination. + self.request_bodies.shutdown_all().await; + + let (pending_batches, failed_batches) = { + let mut failed = self.failed_batches.lock().await; + let mut pending = self.pending_batches.lock().await; + (std::mem::take(&mut *pending), std::mem::take(&mut *failed)) + }; + drop((pending_batches, failed_batches)); + + #[cfg(feature = "test-hooks")] + if let Some(notify) = self.test_hooks.retained_batches_cleared.lock().await.take() { + notify.notify_one(); + } + + self.request_bodies.wait_for_all_eof().await; + + #[cfg(feature = "test-hooks")] + { + let barrier = self.test_hooks.free_shutdown_complete.lock().await.take(); + if let Some(barrier) = barrier { + barrier.reached.notify_one(); + barrier.proceed.notified().await; + } + } + } + /// Returns true once supervisor-owned terminal finalization publishes closure. pub fn is_closed(&self) -> bool { self.is_closed.load(Ordering::Relaxed) @@ -1055,12 +1120,58 @@ impl ZerobusArrowStream { Self::arm_test_barrier(&self.test_hooks.close_finalize).await } + /// Test-only: parks the active request body after it observes forced shutdown but before + /// it reports EOF or drops transport-owned batches. + #[cfg(feature = "test-hooks")] + #[doc(hidden)] + pub async fn arm_request_body_shutdown_barrier(&self) -> (Arc, Arc) { + Self::arm_test_barrier(&self.test_hooks.request_body_shutdown).await + } + + /// Test-only: parks the active request body before it polls a newly queued batch. + #[cfg(feature = "test-hooks")] + #[doc(hidden)] + pub async fn arm_request_body_before_batch_poll_barrier(&self) -> (Arc, Arc) { + Self::arm_test_barrier(&self.test_hooks.request_body_before_batch_poll).await + } + + /// Test-only: notifies after destructive free clears SDK-retained batch collections. + #[cfg(feature = "test-hooks")] + #[doc(hidden)] + pub async fn arm_retained_batches_cleared_notify(&self) -> Arc { + Self::arm_test_notify(&self.test_hooks.retained_batches_cleared).await + } + + /// Test-only: retains a batch in the failed collection for destructive-free tests. + #[cfg(feature = "test-hooks")] + #[doc(hidden)] + pub async fn retain_failed_batch_for_test(&self, batch: RecordBatch) { + self.failed_batches.lock().await.push(batch); + } + + /// Test-only: reports whether destructive-free batch collection locks are available. + #[cfg(feature = "test-hooks")] + #[doc(hidden)] + pub fn retained_batch_locks_available_for_test(&self) -> bool { + let Ok(_failed) = self.failed_batches.try_lock() else { + return false; + }; + self.pending_batches.try_lock().is_ok() + } + + /// Test-only: parks destructive free after all request bodies and retained batches finish. + #[cfg(feature = "test-hooks")] + #[doc(hidden)] + pub async fn arm_free_shutdown_complete_barrier(&self) -> (Arc, Arc) { + Self::arm_test_barrier(&self.test_hooks.free_shutdown_complete).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(); + if let Some(task) = self.supervisor_task.lock().await.as_ref() { + task.abort(); } } @@ -1114,11 +1225,13 @@ impl Drop for ZerobusArrowStream { fn drop(&mut self) { self.admission_closed.store(true, Ordering::Release); self.is_closed.store(true, Ordering::Relaxed); + self.request_bodies.try_shutdown_all(); // 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.supervisor_abort.try_lock() { - if let Some(handle) = guard.take() { - handle.abort(); + if let Ok(mut guard) = self.supervisor_task.try_lock() { + if let Some(task) = guard.take() { + task.abort(); + // Dropping the reaper JoinHandle intentionally detaches ordinary Drop cleanup. } } } diff --git a/rust/sdk/src/stream/arrow/supervisor.rs b/rust/sdk/src/stream/arrow/supervisor.rs index 72024194..7cbb7de4 100644 --- a/rust/sdk/src/stream/arrow/supervisor.rs +++ b/rust/sdk/src/stream/arrow/supervisor.rs @@ -15,10 +15,12 @@ use tracing::{debug, error, info, warn}; 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::connection::{ + FlightConnection, FlightResponseStream, RequestBodyControl, RequestBodyRegistry, +}; use super::{ configured_deadline, ArrowStreamConfigurationOptions, ArrowTableProperties, BatchSender, - RecordBatch, ZerobusArrowStream, + FlightConnectionParameters, RecordBatch, ZerobusArrowStream, }; use crate::errors::ZerobusError; use crate::headers_provider::HeadersProvider; @@ -38,6 +40,7 @@ pub(super) struct Supervisor { is_closed: Arc, close: CloseCoordinator, close_finalizer: CloseFinalizer, + request_bodies: RequestBodyRegistry, pending_batches: Arc>>, recovery_attempts: Arc, server_error_tx: watch::Sender>, @@ -51,6 +54,27 @@ pub(super) struct Supervisor { test_hooks: Arc, } +pub(super) struct SupervisorTaskHandle { + worker: AbortHandle, + // Dropping this handle intentionally detaches the reaper during ordinary Drop. + #[cfg_attr(not(feature = "internal-arrow-c-data"), allow(dead_code))] + reaper: JoinHandle<()>, +} + +impl SupervisorTaskHandle { + pub(super) fn abort(&self) { + self.worker.abort(); + } + + #[cfg(feature = "internal-arrow-c-data")] + pub(super) async fn abort_and_wait(self) { + self.worker.abort(); + self.reaper + .await + .expect("Arrow supervisor reaper failed during shutdown"); + } +} + impl Supervisor { pub(super) fn new(stream: &ZerobusArrowStream) -> Self { Self { @@ -65,6 +89,7 @@ impl Supervisor { is_closed: Arc::clone(&stream.is_closed), close: stream.close.clone(), close_finalizer: CloseFinalizer::new(stream), + request_bodies: stream.request_bodies.clone(), pending_batches: Arc::clone(&stream.pending_batches), recovery_attempts: Arc::clone(&stream.recovery_attempts), server_error_tx: stream.server_error_tx.clone(), @@ -79,15 +104,15 @@ impl Supervisor { } } - pub(super) fn spawn(self, initial_connection: FlightConnection) -> AbortHandle { + pub(super) fn spawn(self, initial_connection: FlightConnection) -> SupervisorTaskHandle { let (response_stream, request_body) = initial_connection.into_supervisor_io(); 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(); + let worker_abort = 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 reaper = spawn(async move { let joined = worker.await; if matches!(close.state(), CloseState::Finalized(_)) { return; @@ -95,7 +120,10 @@ impl Supervisor { let outcome = Self::unfinalized_exit_outcome(joined); let _ = finalizer.finish(outcome).await; }); - abort_handle + SupervisorTaskHandle { + worker: worker_abort, + reaper, + } } fn unfinalized_exit_outcome(joined: Result, JoinError>) -> ZerobusResult<()> { @@ -388,16 +416,19 @@ impl Supervisor { /// 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, - self.connector_factory.as_ref(), - &self.table_properties, - &self.options, - &self.headers_provider, - &self.sdk_identifier, - ) - .await?; + let parameters = FlightConnectionParameters { + endpoint: &self.endpoint, + tls_config: &self.tls_config, + connector_factory: self.connector_factory.as_ref(), + table_properties: &self.table_properties, + options: &self.options, + headers_provider: &self.headers_provider, + sdk_identifier: &self.sdk_identifier, + request_bodies: &self.request_bodies, + #[cfg(feature = "test-hooks")] + test_hooks: &self.test_hooks, + }; + let connection = ZerobusArrowStream::reconnect_transport(¶meters).await?; let tx = connection.sender(); let acked_before_disconnect = self.last_acked_records.load(Ordering::Acquire); @@ -607,6 +638,8 @@ mod tests { use tokio::time::{timeout, Duration, Instant}; use super::super::close::{CloseCoordinator, CloseFinalizer, CloseRequest, CloseState}; + #[cfg(feature = "internal-arrow-c-data")] + use super::SupervisorTaskHandle; use super::{ pause_and_detach_sender, BatchSender, PendingBatch, RecordBatch, Supervisor, ZerobusError, }; @@ -670,6 +703,23 @@ mod tests { )); } + #[cfg(feature = "internal-arrow-c-data")] + #[tokio::test] + #[should_panic(expected = "Arrow supervisor reaper failed during shutdown")] + async fn abort_and_wait_propagates_reaper_failure() { + let worker = tokio::spawn(std::future::pending::<()>()); + let worker_abort = worker.abort_handle(); + drop(worker); + let reaper = tokio::spawn(async { panic!("reaper test panic") }); + + SupervisorTaskHandle { + worker: worker_abort, + reaper, + } + .abort_and_wait() + .await; + } + #[tokio::test] async fn panicked_supervisor_exit_is_an_invariant_error() { let worker: JoinHandle> = diff --git a/rust/tests/src/arrow_c_data_ffi_tests.rs b/rust/tests/src/arrow_c_data_ffi_tests.rs index 2f55d0b0..c689a5cd 100644 --- a/rust/tests/src/arrow_c_data_ffi_tests.rs +++ b/rust/tests/src/arrow_c_data_ffi_tests.rs @@ -16,16 +16,19 @@ const TABLE_NAME: &str = "test_catalog.test_schema.test_table"; mod ffi_c_data_lifetime_tests { use std::ffi::{c_void, CStr, CString}; use std::ptr; - use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::sync::mpsc; use std::time::Duration; + use databricks_zerobus_ingest_sdk::internal::arrow_c_data::import_c_data_record_batch; + use databricks_zerobus_ingest_sdk::ZerobusArrowStream; use tonic::Status; use zerobus_ffi::{ - zerobus_arrow_get_default_config, zerobus_arrow_stream_flush, zerobus_arrow_stream_free, - zerobus_arrow_stream_ingest_c_data, zerobus_free_error_message, zerobus_sdk_builder_build, - zerobus_sdk_builder_disable_tls, zerobus_sdk_builder_endpoint, zerobus_sdk_builder_new, - zerobus_sdk_create_arrow_stream_with_headers_provider, zerobus_sdk_free, CArrowArray, - CArrowSchema, CHeaders, CResult, + zerobus_arrow_get_default_config, zerobus_arrow_stream_close, zerobus_arrow_stream_flush, + zerobus_arrow_stream_free, zerobus_arrow_stream_ingest_c_data, zerobus_free_error_message, + zerobus_sdk_builder_build, zerobus_sdk_builder_disable_tls, zerobus_sdk_builder_endpoint, + zerobus_sdk_builder_new, zerobus_sdk_create_arrow_stream_with_headers_provider, + zerobus_sdk_free, CArrowArray, CArrowSchema, CHeaders, CResult, }; use super::{ @@ -68,6 +71,98 @@ mod ffi_c_data_lifetime_tests { } } + struct TimingArrayRelease { + inner_release: Option, + inner_private_data: *mut c_void, + release_count: Arc, + } + + unsafe extern "C" fn timing_array_release(array: *mut FFI_ArrowArray) { + let array = unsafe { &mut *array }; + let state = unsafe { Box::from_raw(array.private_data.cast::()) }; + array.release = state.inner_release; + array.private_data = state.inner_private_data; + state.release_count.fetch_add(1, Ordering::SeqCst); + if let Some(release) = state.inner_release { + unsafe { release(array) }; + } + } + + struct LockCheckingArrayRelease { + inner_release: Option, + inner_private_data: *mut c_void, + stream: usize, + locks_available: Arc, + release_count: Arc, + } + + unsafe extern "C" fn lock_checking_array_release(array: *mut FFI_ArrowArray) { + let array = unsafe { &mut *array }; + let state = unsafe { Box::from_raw(array.private_data.cast::()) }; + // This intentionally mirrors the opaque-handle cast in arrow.rs. If that handle + // gains an outer wrapper, this test hook and the production cast must change together. + let stream = unsafe { &*(state.stream as *const ZerobusArrowStream) }; + state.locks_available.store( + stream.retained_batch_locks_available_for_test(), + Ordering::SeqCst, + ); + state.release_count.fetch_add(1, Ordering::SeqCst); + array.release = state.inner_release; + array.private_data = state.inner_private_data; + if let Some(release) = state.inner_release { + unsafe { release(array) }; + } + } + + fn exported_lock_checking_batch( + batch: RecordBatch, + stream: usize, + locks_available: Arc, + release_count: Arc, + ) -> (FFI_ArrowArray, FFI_ArrowSchema) { + let schema = batch.schema(); + let struct_array = StructArray::from(batch); + let mut array = FFI_ArrowArray::new(&struct_array.to_data()); + let array_state = Box::new(LockCheckingArrayRelease { + inner_release: array.release, + inner_private_data: array.private_data, + stream, + locks_available, + release_count, + }); + array.release = Some(lock_checking_array_release); + array.private_data = Box::into_raw(array_state).cast(); + let schema = FFI_ArrowSchema::try_from(schema.as_ref()).unwrap(); + (array, schema) + } + + fn exported_timing_batch( + batch: RecordBatch, + release_count: Arc, + schema_releases: Arc, + ) -> (FFI_ArrowArray, FFI_ArrowSchema) { + let schema = batch.schema(); + let struct_array = StructArray::from(batch); + let mut array = FFI_ArrowArray::new(&struct_array.to_data()); + let array_state = Box::new(TimingArrayRelease { + inner_release: array.release, + inner_private_data: array.private_data, + release_count, + }); + array.release = Some(timing_array_release); + array.private_data = Box::into_raw(array_state).cast(); + + let mut schema = FFI_ArrowSchema::try_from(schema.as_ref()).unwrap(); + let schema_state = Box::new(CountingSchemaRelease { + inner_release: schema.release, + inner_private_data: schema.private_data, + releases: schema_releases, + }); + schema.release = Some(counting_schema_release); + schema.private_data = Box::into_raw(schema_state).cast(); + (array, schema) + } + fn exported_counting_batch( batch: RecordBatch, array_releases: Arc, @@ -217,6 +312,24 @@ mod ffi_c_data_lifetime_tests { .unwrap() } + async fn close_ffi_stream(stream: usize) -> Result<(), String> { + tokio::task::spawn_blocking(move || { + let mut result = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + if zerobus_arrow_stream_close(stream as *mut _, &mut result) { + Ok(()) + } else { + Err(take_result_error(&mut result) + .unwrap_or_else(|| "Arrow stream close failed without a message".to_string())) + } + }) + .await + .unwrap() + } + async fn free_ffi_handles(sdk: usize, stream: usize) { tokio::task::spawn_blocking(move || { zerobus_arrow_stream_free(stream as *mut _); @@ -236,6 +349,38 @@ mod ffi_c_data_lifetime_tests { .expect("release callback did not run before timeout"); } + struct TestBarrierGuard(Option>); + + impl TestBarrierGuard { + fn release(&mut self) { + if let Some(proceed) = self.0.take() { + proceed.notify_one(); + } + } + } + + impl Drop for TestBarrierGuard { + fn drop(&mut self) { + self.release(); + } + } + + async fn wait_for_free_completion( + completion: &mut mpsc::Receiver<()>, + ) -> Result<(), &'static str> { + loop { + match completion.try_recv() { + Ok(()) => return Ok(()), + Err(mpsc::TryRecvError::Empty) => { + tokio::time::sleep(Duration::from_millis(1)).await + } + Err(mpsc::TryRecvError::Disconnected) => { + return Err("stream free thread exited without reporting completion"); + } + } + } + } + #[tokio::test] async fn ffi_c_data_owner_releases_once_after_ack_and_flush( ) -> Result<(), Box> { @@ -317,8 +462,247 @@ mod ffi_c_data_lifetime_tests { ); free_ffi_handles(sdk, stream).await; - wait_for_release_count(array_releases.as_ref(), 1).await; + assert_eq!( + array_releases.load(Ordering::SeqCst), + 1, + "failed unacknowledged owner must be released during stream free, not later" + ); assert_eq!(schema_releases.load(Ordering::SeqCst), 1); Ok(()) } + + #[tokio::test] + async fn ffi_c_data_owner_remains_released_after_close_then_free( + ) -> Result<(), Box> { + 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, stream) = create_ffi_stream(server_url, Arc::clone(&schema)).await?; + let array_releases = Arc::new(AtomicUsize::new(0)); + let schema_releases = Arc::new(AtomicUsize::new(0)); + let batch = create_test_record_batch(schema, vec![1], vec![Some("closed")]); + let (array, schema) = exported_counting_batch( + batch, + Arc::clone(&array_releases), + Arc::clone(&schema_releases), + ); + + assert_eq!(ingest_ffi_batch(stream, array, schema).await?, 0); + close_ffi_stream(stream).await?; + assert_eq!(array_releases.load(Ordering::SeqCst), 1); + assert_eq!(schema_releases.load(Ordering::SeqCst), 1); + + free_ffi_handles(sdk, stream).await; + assert_eq!(array_releases.load(Ordering::SeqCst), 1); + assert_eq!(schema_releases.load(Ordering::SeqCst), 1); + Ok(()) + } + + #[tokio::test] + async fn ffi_c_data_release_callbacks_run_without_retained_batch_locks( + ) -> Result<(), Box> { + let (_mock_server, server_url) = start_mock_flight_server().await?; + let schema = create_test_arrow_schema(); + let (sdk, stream) = create_ffi_stream(server_url, Arc::clone(&schema)).await?; + let locks_available = Arc::new(AtomicBool::new(false)); + let release_count = Arc::new(AtomicUsize::new(0)); + let batch = create_test_record_batch(schema, vec![1], vec![Some("retained")]); + let (array, schema_ffi) = exported_lock_checking_batch( + batch, + stream, + Arc::clone(&locks_available), + Arc::clone(&release_count), + ); + let batch = unsafe { import_c_data_record_batch(array, schema_ffi) }?; + + // This intentionally mirrors the opaque-handle cast in arrow.rs. If that handle + // gains an outer wrapper, this test hook and the production cast must change together. + let stream_ref = unsafe { &*(stream as *const ZerobusArrowStream) }; + stream_ref.retain_failed_batch_for_test(batch).await; + assert_eq!(release_count.load(Ordering::SeqCst), 0); + + free_ffi_handles(sdk, stream).await; + + assert_eq!(release_count.load(Ordering::SeqCst), 1); + assert!( + locks_available.load(Ordering::SeqCst), + "release callback must not run while retained-batch locks are held" + ); + Ok(()) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn ffi_c_data_free_from_multithread_runtime_releases_before_return( + ) -> Result<(), Box> { + let (_mock_server, server_url) = start_mock_flight_server().await?; + let schema = create_test_arrow_schema(); + let (sdk, stream) = create_ffi_stream(server_url, Arc::clone(&schema)).await?; + let array_releases = Arc::new(AtomicUsize::new(0)); + let schema_releases = Arc::new(AtomicUsize::new(0)); + let batch = create_test_record_batch(schema, vec![1], vec![Some("multi-thread free")]); + let (array, schema_ffi) = exported_counting_batch( + batch, + Arc::clone(&array_releases), + Arc::clone(&schema_releases), + ); + let batch = unsafe { import_c_data_record_batch(array, schema_ffi) }?; + // This mirrors the opaque-handle cast in arrow.rs and must change with it. + let stream_ref = unsafe { &*(stream as *const ZerobusArrowStream) }; + stream_ref.retain_failed_batch_for_test(batch).await; + + zerobus_arrow_stream_free(stream as *mut _); + + assert_eq!(array_releases.load(Ordering::SeqCst), 1); + assert_eq!(schema_releases.load(Ordering::SeqCst), 1); + tokio::task::spawn_blocking(move || zerobus_sdk_free(sdk as *mut _)) + .await + .expect("sdk free task panicked"); + Ok(()) + } + + #[tokio::test] + async fn ffi_c_data_request_body_owner_releases_before_free_returns( + ) -> Result<(), Box> { + 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: 5_000, + ack_up_to_records: 3, + }], + ) + .await; + let (sdk, stream) = create_ffi_stream(server_url, Arc::clone(&schema)).await?; + let release_count = Arc::new(AtomicUsize::new(0)); + let schema_releases = Arc::new(AtomicUsize::new(0)); + let batch = + create_test_record_batch(schema, vec![1, 2, 3], vec![Some("a"), Some("b"), Some("c")]); + let (array, schema_ffi) = exported_timing_batch( + batch, + Arc::clone(&release_count), + Arc::clone(&schema_releases), + ); + + let (before_batch_poll_reached, before_batch_poll_proceed) = { + // This mirrors the opaque-handle cast in arrow.rs and must change with it. + let stream_ref = unsafe { &*(stream as *const ZerobusArrowStream) }; + stream_ref + .arm_request_body_before_batch_poll_barrier() + .await + }; + let mut before_batch_poll_guard = TestBarrierGuard(Some(before_batch_poll_proceed)); + assert_eq!(ingest_ffi_batch(stream, array, schema_ffi).await?, 0); + tokio::time::timeout(Duration::from_secs(5), before_batch_poll_reached.notified()) + .await + .expect("request body must park before consuming the queued batch"); + assert_eq!(schema_releases.load(Ordering::SeqCst), 1); + assert_eq!( + release_count.load(Ordering::SeqCst), + 0, + "array owner must remain alive while the batch is pending" + ); + + let ( + request_shutdown_reached, + request_shutdown_proceed, + retained_batches_cleared, + free_shutdown_complete_reached, + free_shutdown_complete_proceed, + ) = { + // This mirrors the opaque-handle cast in arrow.rs and must change with it. + let stream_ref = unsafe { &*(stream as *const ZerobusArrowStream) }; + let (request_shutdown_reached, request_shutdown_proceed) = + stream_ref.arm_request_body_shutdown_barrier().await; + let retained_batches_cleared = stream_ref.arm_retained_batches_cleared_notify().await; + let (free_shutdown_complete_reached, free_shutdown_complete_proceed) = + stream_ref.arm_free_shutdown_complete_barrier().await; + ( + request_shutdown_reached, + request_shutdown_proceed, + retained_batches_cleared, + free_shutdown_complete_reached, + free_shutdown_complete_proceed, + ) + }; + let mut request_shutdown_guard = TestBarrierGuard(Some(request_shutdown_proceed)); + let mut free_shutdown_complete_guard = + TestBarrierGuard(Some(free_shutdown_complete_proceed)); + + let (free_completed_tx, mut free_completed_rx) = mpsc::channel(); + drop(std::thread::spawn(move || { + zerobus_arrow_stream_free(stream as *mut _); + let _ = free_completed_tx.send(()); + })); + + tokio::time::timeout(Duration::from_secs(5), async { + tokio::select! { + _ = request_shutdown_reached.notified() => Ok(()), + _ = free_shutdown_complete_reached.notified() => Err( + "free reached its pre-return boundary without shutting down the request body" + ), + result = wait_for_free_completion(&mut free_completed_rx) => { + result.and(Err( + "zerobus_arrow_stream_free returned before request-body shutdown" + )) + }, + } + }) + .await + .expect("request body must observe forced shutdown")?; + + tokio::time::timeout(Duration::from_secs(5), retained_batches_cleared.notified()) + .await + .expect("destructive free must clear retained batches before waiting for request EOF"); + assert_eq!( + release_count.load(Ordering::SeqCst), + 0, + "the blocked request body must retain the owner after SDK collections are cleared" + ); + + before_batch_poll_guard.release(); + request_shutdown_guard.release(); + tokio::time::timeout( + Duration::from_secs(5), + free_shutdown_complete_reached.notified(), + ) + .await + .expect("free must reach its completed pre-return boundary"); + + assert_eq!( + release_count.load(Ordering::SeqCst), + 1, + "all request-body and SDK owners must release before the pre-return boundary" + ); + assert!(matches!( + free_completed_rx.try_recv(), + Err(mpsc::TryRecvError::Empty) + )); + + free_shutdown_complete_guard.release(); + tokio::time::timeout( + Duration::from_secs(5), + wait_for_free_completion(&mut free_completed_rx), + ) + .await + .expect("zerobus_arrow_stream_free did not complete after shutdown proceeded")?; + + tokio::task::spawn_blocking(move || { + zerobus_sdk_free(sdk as *mut _); + }) + .await + .expect("sdk free task panicked"); + Ok(()) + } } From df539ca398dc4762f1d8c5ead2d111f457b4b6d2 Mon Sep 17 00:00:00 2001 From: nikolaobradovic-db Date: Tue, 18 Aug 2026 09:35:27 +0000 Subject: [PATCH 2/2] Preserve old behaviour for IPC only streams. --- CLAUDE.md | 3 +- cpp/NEXT_CHANGELOG.md | 5 - cpp/README.md | 5 - cpp/include/zerobus/arrow_stream.hpp | 6 +- go/CLAUDE.md | 3 +- go/NEXT_CHANGELOG.md | 8 -- go/README.md | 7 +- go/arrow_stream.go | 26 +--- rust/ffi/NEXT_CHANGELOG.md | 2 +- rust/ffi/src/arrow.rs | 38 ++++-- rust/ffi/zerobus.h | 30 ++-- rust/sdk/src/internal.rs | 14 ++ rust/sdk/src/stream/arrow/mod.rs | 20 ++- rust/tests/src/arrow_c_data_ffi_tests.rs | 166 ++++++++++++++++++++++- 14 files changed, 248 insertions(+), 85 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 4af03b38..9b7ff843 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -68,8 +68,7 @@ Because every non-Rust SDK crosses a foreign-function boundary, keep these in mi ### Thread safety - Rust core is async (tokio). The FFI layer manages its own tokio runtime. -- Go: concurrent ingestion is internally synchronized. Serialize `Close` with - every other operation on the same stream. +- Go: safe for concurrent goroutines. - Python: async streams are safe; sync streams are single-threaded per instance. - TypeScript: async-safe via Node.js event loop. - Java: **not thread-safe** — external synchronization required for concurrent access. diff --git a/cpp/NEXT_CHANGELOG.md b/cpp/NEXT_CHANGELOG.md index d4094102..19dd253d 100644 --- a/cpp/NEXT_CHANGELOG.md +++ b/cpp/NEXT_CHANGELOG.md @@ -16,11 +16,6 @@ `create_arrow_stream` still take a `std::shared_ptr` — and you no longer need to keep your own reference alive past `create_stream`. -- Fixed `ArrowStream` destruction so it waits while native background - shutdown finishes. Destruction must not race another operation on the stream. - Internal native shutdown infrastructure failures terminate the process rather - than returning after incomplete native destruction. - ### Documentation - Corrected custom-header examples and clarified that acknowledgment callbacks diff --git a/cpp/README.md b/cpp/README.md index bcdfad51..8488fe46 100644 --- a/cpp/README.md +++ b/cpp/README.md @@ -460,11 +460,6 @@ access externally (the same contract as the Rust core). A single `Sdk` may create many streams. See [`CLAUDE.md`](https://github.com/databricks/zerobus-sdk/blob/main/cpp/CLAUDE.md) for the full memory-ownership and threading contract. -Destroying an `ArrowStream` can block while native background shutdown -finishes. Destruction must not race another operation on the same stream. An -internal native shutdown infrastructure failure terminates the process rather -than returning after incomplete native destruction. - ## HTTP proxy support Like the other SDKs, the C++ SDK honors the standard proxy environment variables diff --git a/cpp/include/zerobus/arrow_stream.hpp b/cpp/include/zerobus/arrow_stream.hpp index 06141a57..e30a0292 100644 --- a/cpp/include/zerobus/arrow_stream.hpp +++ b/cpp/include/zerobus/arrow_stream.hpp @@ -24,13 +24,9 @@ class Sdk; /// errors (the destructor swallows them) and flushes synchronously, which can /// block up to `flush_timeout_ms` (default 5 minutes) if the server is /// unresponsive. Letting the object fall out of scope drags that blocking close -/// into the destructor. Native destruction also waits for background shutdown -/// to finish and does not return on a shutdown timeout. An internal native -/// shutdown infrastructure failure terminates the process rather than returning -/// after incomplete native destruction. +/// into the destructor. /// /// Thread safety: not safe for concurrent use from multiple threads. -/// Destruction must not race another operation on the same stream. /// /// Only the IPC-bytes ingest path is wrapped; /// `zerobus_arrow_stream_ingest_batch_via_record_batch` is a functional diff --git a/go/CLAUDE.md b/go/CLAUDE.md index cc89b09e..c3902d20 100644 --- a/go/CLAUDE.md +++ b/go/CLAUDE.md @@ -110,8 +110,7 @@ Public API is everything exported (capitalized) in the `go/` package: ## Thread safety -Concurrent `Ingest` calls are internally synchronized. Serialize `Close` with -every other operation on the same stream. +Go SDK is safe for concurrent use from multiple goroutines. Internal synchronization handles concurrent `Ingest` calls. ## Changelog and documentation diff --git a/go/NEXT_CHANGELOG.md b/go/NEXT_CHANGELOG.md index ddcc9d74..3c0fc07e 100644 --- a/go/NEXT_CHANGELOG.md +++ b/go/NEXT_CHANGELOG.md @@ -8,10 +8,6 @@ ### Bug Fixes -- Arrow stream finalizers now hand potentially blocking native shutdown to a - separate goroutine, so one stalled cleanup cannot block Go's global finalizer - goroutine. Explicit `Close` remains synchronous. - ### Documentation - Corrected consumer installation so tagged releases do not require Rust or @@ -23,10 +19,6 @@ - Batch examples name the offset returned by `IngestRecordsOffset` `batchOffset`. Example and test modules keep the CI toolchain Go versions; the documented SDK minimum remains Go 1.21+. -- Documented that closing an Arrow stream can block while native background - shutdown finishes and must not race another operation on the same stream. - Internal native shutdown infrastructure failures terminate the process rather - than returning after incomplete native destruction. ### Internal Changes diff --git a/go/README.md b/go/README.md index 973075f6..3d8e2493 100644 --- a/go/README.md +++ b/go/README.md @@ -1442,12 +1442,7 @@ Waits for all pending batches to be acknowledged. func (s *ZerobusArrowStream) Close() error ``` -Flushes and closes the stream. Native destruction can block while background -shutdown finishes. Serialize `Close` with every other operation on the same -stream. The GC finalizer schedules this cleanup asynchronously; call `Close` -explicitly when completion and errors matter. -An internal native shutdown infrastructure failure terminates the process -rather than returning after incomplete native destruction. +Flushes and closes the stream. ```go func (s *ZerobusArrowStream) GetUnackedBatches() ([][]byte, error) diff --git a/go/arrow_stream.go b/go/arrow_stream.go index 68de4f9f..76336a9e 100644 --- a/go/arrow_stream.go +++ b/go/arrow_stream.go @@ -105,17 +105,6 @@ type ZerobusArrowStream struct { ptr unsafe.Pointer } -func setArrowStreamFinalizer(stream *ZerobusArrowStream) { - runtime.SetFinalizer(stream, func(st *ZerobusArrowStream) { - // Native destruction may wait indefinitely to guarantee that no foreign - // release callback runs after it returns. Never block Go's sole finalizer - // goroutine on that wait. - go func() { - _ = st.Close() - }() - }) -} - // CreateArrowStream creates an Arrow Flight stream authenticated with OAuth client credentials. // // schemaIpcBytes must be Arrow IPC stream bytes containing only the schema (no data batches). @@ -140,7 +129,9 @@ func (s *ZerobusSdk) CreateArrowStream( } stream := &ZerobusArrowStream{ptr: ptr} - setArrowStreamFinalizer(stream) + runtime.SetFinalizer(stream, func(st *ZerobusArrowStream) { + st.Close() //nolint:errcheck + }) return stream, nil } @@ -165,7 +156,9 @@ func (s *ZerobusSdk) CreateArrowStreamWithHeadersProvider( } stream := &ZerobusArrowStream{ptr: ptr} - setArrowStreamFinalizer(stream) + runtime.SetFinalizer(stream, func(st *ZerobusArrowStream) { + st.Close() //nolint:errcheck + }) return stream, nil } @@ -196,12 +189,7 @@ func (st *ZerobusArrowStream) Flush() error { } // Close gracefully shuts down the stream after flushing all pending batches. -// Native destruction can block while background shutdown finishes. Close must -// be serialized with IngestBatch, Flush, and every other operation on this -// stream. An internal native shutdown infrastructure failure terminates the -// process rather than returning after incomplete native destruction. -// The stream must not be used after Close returns. -// Close is idempotent. +// The stream must not be used after Close returns. Close is idempotent. func (st *ZerobusArrowStream) Close() error { if st.ptr == nil { return nil diff --git a/rust/ffi/NEXT_CHANGELOG.md b/rust/ffi/NEXT_CHANGELOG.md index 18384762..67801e42 100644 --- a/rust/ffi/NEXT_CHANGELOG.md +++ b/rust/ffi/NEXT_CHANGELOG.md @@ -19,7 +19,7 @@ ### Behavior Changes -- `zerobus_arrow_stream_free` now blocks until Arrow background shutdown completes, every Flight request body reaches EOF or is dropped, and all retained Arrow C Data owners are released. Previously, a request body could retain an owner and run its release callback after free returned on an unacknowledged/failure path, risking callback-after-free use of wrapper state. When the calling restrictions below are respected, no Arrow C Data release callback for that stream can run after free returns. The function logs a warning every 30 seconds while shutdown remains incomplete; it does not return on a timeout. Free must not race another operation on the same stream handle. Freeing the same stream reentrantly from one of its SDK callbacks is unsupported because shutdown would wait for that callback; freeing a different stream from a callback remains supported. An internal native shutdown panic, a required helper-thread spawn failure, or a helper-thread panic terminates the process rather than returning without the release-callback guarantee. +- `zerobus_arrow_stream_free` now selects destruction behavior based on how the stream was used. IPC-only streams preserve best-effort, nonblocking destruction. Once a stream accepts an Arrow C Data batch, free blocks until Arrow background shutdown completes, every Flight request body reaches EOF or is dropped, and all retained foreign owners are released. Previously, a request body could retain an owner and run its release callback after free returned on an unacknowledged/failure path, risking callback-after-free use of producer state. When the calling restrictions below are respected, no Arrow C Data release callback for that stream can run after free returns. The function logs a warning every 30 seconds while required C Data shutdown remains incomplete; it does not return on a timeout. Free must not race another operation on the same stream handle. After C Data import, freeing the same stream reentrantly from one of its SDK callbacks is unsupported because complete shutdown would wait for that callback. IPC-only concurrent or reentrant free remains invalid because the opaque handle has single ownership; freeing a different stream from a callback remains supported. During required C Data shutdown, an internal native shutdown panic, a required helper-thread spawn failure, or a helper-thread panic terminates the process rather than returning without the release-callback guarantee. ### Breaking Changes diff --git a/rust/ffi/src/arrow.rs b/rust/ffi/src/arrow.rs index 154c93b9..513f4456 100644 --- a/rust/ffi/src/arrow.rs +++ b/rust/ffi/src/arrow.rs @@ -4,10 +4,13 @@ use crate::arrow_c_data::{CArrowArray, CArrowSchema}; use crate::common::*; use arrow_ipc::{reader::StreamReader, writer::StreamWriter, CompressionType}; use bytes::Bytes; -use databricks_zerobus_ingest_sdk::internal::abort_arrow_stream_and_wait; use databricks_zerobus_ingest_sdk::internal::arrow_c_data::{ import_c_data_record_batch, FFI_ArrowArray, FFI_ArrowSchema, }; +use databricks_zerobus_ingest_sdk::internal::{ + abort_arrow_stream_and_wait, arrow_stream_has_ingested_c_data, + mark_arrow_stream_c_data_ingested, +}; use databricks_zerobus_ingest_sdk::{ HeadersProvider, RecordBatch, StreamBuilder, ZerobusArrowStream, ZerobusError, ZerobusResult, }; @@ -368,25 +371,31 @@ fn abort_and_drop_arrow_stream_on_thread(stream: Box) { /// Frees an Arrow Flight stream instance. /// -/// This call blocks until background shutdown completes, every Flight request body reaches EOF -/// or is dropped, and all retained Arrow C Data owners are released. If shutdown takes longer -/// than 30 seconds, it logs a warning every 30 seconds while continuing to wait. It does not -/// return on a timeout. -/// When the calling restrictions below are respected, no Arrow C Data release callback for this -/// stream can run after this function returns. -/// An internal shutdown panic, a required helper-thread spawn failure, or a helper-thread panic -/// terminates the process rather than returning without that guarantee. +/// IPC-only streams preserve best-effort, nonblocking destruction. Once a stream accepts an +/// Arrow C Data batch, this call instead blocks until background shutdown completes, every Flight +/// request body reaches EOF or is dropped, and all retained foreign owners are released. If that +/// shutdown takes longer than 30 seconds, it logs a warning every 30 seconds while continuing to +/// wait; it does not return on a timeout. +/// When the calling restrictions below are respected, no Arrow C Data release callback for that +/// stream can run after this function returns. During required C Data shutdown, an internal +/// shutdown panic, a required helper-thread spawn failure, or a helper-thread panic terminates the +/// process rather than returning without that guarantee. /// -/// Do not race this function with another operation on the same stream handle. Freeing this same -/// stream reentrantly from one of its SDK callbacks is unsupported because shutdown would wait -/// for the callback that is making the call. Freeing a different stream from a callback remains -/// supported. +/// Do not race this function with another operation on the same stream handle. After C Data import, +/// freeing this same stream reentrantly from one of its SDK callbacks is unsupported because +/// complete shutdown would wait for the callback making the call. IPC-only concurrent or reentrant +/// free remains invalid because the opaque handle has single ownership. Freeing a different stream +/// from a callback remains supported. #[no_mangle] pub extern "C" fn zerobus_arrow_stream_free(stream: *mut CArrowStream) { ffi_guard(ptr::null_mut(), (), move || { if !stream.is_null() { unsafe { let stream = Box::from_raw(stream as *mut ZerobusArrowStream); + if !arrow_stream_has_ingested_c_data(&stream) { + drop(stream); + return; + } match tokio::runtime::Handle::try_current() { Ok(handle) if handle.runtime_flavor() @@ -459,6 +468,8 @@ pub extern "C" fn zerobus_arrow_stream_ingest_batch( /// every success or error path. Their release callbacks are cleared before /// validation, and the imported buffers may remain owned by the stream until /// acknowledgment, recovery finalization, or stream destruction. +/// Once valid C Data is imported, `zerobus_arrow_stream_free` uses complete +/// shutdown for that stream; later IPC ingestion does not revert this mode. /// /// Every non-null pointer must address a valid, properly aligned canonical /// `ArrowArray` / `ArrowSchema` structure satisfying the Arrow C Data @@ -517,6 +528,7 @@ pub extern "C" fn zerobus_arrow_stream_ingest_c_data( return -1; } }; + mark_arrow_stream_c_data_ingested(stream_ref); match RUNTIME.block_on(stream_ref.ingest_batch(batch)) { Ok(offset) => { diff --git a/rust/ffi/zerobus.h b/rust/ffi/zerobus.h index a8e363f9..73092baa 100644 --- a/rust/ffi/zerobus.h +++ b/rust/ffi/zerobus.h @@ -306,19 +306,21 @@ struct CArrowStream *zerobus_sdk_create_arrow_stream_with_headers_provider(struc /** * Frees an Arrow Flight stream instance. * - * This call blocks until background shutdown completes, every Flight request body reaches EOF - * or is dropped, and all retained Arrow C Data owners are released. If shutdown takes longer - * than 30 seconds, it logs a warning every 30 seconds while continuing to wait. It does not - * return on a timeout. - * When the calling restrictions below are respected, no Arrow C Data release callback for this - * stream can run after this function returns. - * An internal shutdown panic, a required helper-thread spawn failure, or a helper-thread panic - * terminates the process rather than returning without that guarantee. - * - * Do not race this function with another operation on the same stream handle. Freeing this same - * stream reentrantly from one of its SDK callbacks is unsupported because shutdown would wait - * for the callback that is making the call. Freeing a different stream from a callback remains - * supported. + * IPC-only streams preserve best-effort, nonblocking destruction. Once a stream accepts an + * Arrow C Data batch, this call instead blocks until background shutdown completes, every Flight + * request body reaches EOF or is dropped, and all retained foreign owners are released. If that + * shutdown takes longer than 30 seconds, it logs a warning every 30 seconds while continuing to + * wait; it does not return on a timeout. + * When the calling restrictions below are respected, no Arrow C Data release callback for that + * stream can run after this function returns. During required C Data shutdown, an internal + * shutdown panic, a required helper-thread spawn failure, or a helper-thread panic terminates the + * process rather than returning without that guarantee. + * + * Do not race this function with another operation on the same stream handle. After C Data import, + * freeing this same stream reentrantly from one of its SDK callbacks is unsupported because + * complete shutdown would wait for the callback making the call. IPC-only concurrent or reentrant + * free remains invalid because the opaque handle has single ownership. Freeing a different stream + * from a callback remains supported. */ void zerobus_arrow_stream_free(struct CArrowStream *stream); @@ -341,6 +343,8 @@ int64_t zerobus_arrow_stream_ingest_batch(struct CArrowStream *stream, * every success or error path. Their release callbacks are cleared before * validation, and the imported buffers may remain owned by the stream until * acknowledgment, recovery finalization, or stream destruction. + * Once valid C Data is imported, `zerobus_arrow_stream_free` uses complete + * shutdown for that stream; later IPC ingestion does not revert this mode. * * Every non-null pointer must address a valid, properly aligned canonical * `ArrowArray` / `ArrowSchema` structure satisfying the Arrow C Data diff --git a/rust/sdk/src/internal.rs b/rust/sdk/src/internal.rs index 5c923c90..7cf2710e 100644 --- a/rust/sdk/src/internal.rs +++ b/rust/sdk/src/internal.rs @@ -17,3 +17,17 @@ pub mod arrow_c_data { pub async fn abort_arrow_stream_and_wait(stream: &crate::ZerobusArrowStream) { stream.abort_and_wait().await; } + +/// Marks that an Arrow stream may retain foreign Arrow C Data owners. +#[cfg(feature = "internal-arrow-c-data")] +#[doc(hidden)] +pub fn mark_arrow_stream_c_data_ingested(stream: &crate::ZerobusArrowStream) { + stream.mark_c_data_ingested(); +} + +/// Returns whether Arrow stream destruction must wait for foreign C Data owners. +#[cfg(feature = "internal-arrow-c-data")] +#[doc(hidden)] +pub fn arrow_stream_has_ingested_c_data(stream: &crate::ZerobusArrowStream) -> bool { + stream.has_ingested_c_data() +} diff --git a/rust/sdk/src/stream/arrow/mod.rs b/rust/sdk/src/stream/arrow/mod.rs index fe12d9ec..8dc7040c 100644 --- a/rust/sdk/src/stream/arrow/mod.rs +++ b/rust/sdk/src/stream/arrow/mod.rs @@ -196,6 +196,9 @@ pub struct ZerobusArrowStream { close: CloseCoordinator, /// Supervisor worker and reaper ownership, consumed once by terminal shutdown. supervisor_task: Arc>>, + /// Once true, FFI destruction must wait until all foreign C Data owners are released. + #[cfg(feature = "internal-arrow-c-data")] + has_ingested_c_data: AtomicBool, /// Tracks every tonic-owned Flight request body until its queued owners are dropped. request_bodies: RequestBodyRegistry, /// Accepted batches not yet fully acknowledged; retained for replay or retrieval. @@ -316,6 +319,8 @@ impl ZerobusArrowStream { admission_closed, close, supervisor_task, + #[cfg(feature = "internal-arrow-c-data")] + has_ingested_c_data: AtomicBool::new(false), request_bodies, pending_batches, pending_notify, @@ -983,6 +988,16 @@ impl ZerobusArrowStream { Ok(self.failed_batches.lock().await.clone()) } + #[cfg(feature = "internal-arrow-c-data")] + pub(crate) fn mark_c_data_ingested(&self) { + self.has_ingested_c_data.store(true, Ordering::Release); + } + + #[cfg(feature = "internal-arrow-c-data")] + pub(crate) fn has_ingested_c_data(&self) -> bool { + self.has_ingested_c_data.load(Ordering::Acquire) + } + #[cfg(feature = "internal-arrow-c-data")] pub(crate) async fn abort_and_wait(&self) { self.admission_closed.store(true, Ordering::Release); @@ -1142,10 +1157,11 @@ impl ZerobusArrowStream { Self::arm_test_notify(&self.test_hooks.retained_batches_cleared).await } - /// Test-only: retains a batch in the failed collection for destructive-free tests. - #[cfg(feature = "test-hooks")] + /// Test-only: marks and retains a foreign-owned batch for destructive-free tests. + #[cfg(all(feature = "test-hooks", feature = "internal-arrow-c-data"))] #[doc(hidden)] pub async fn retain_failed_batch_for_test(&self, batch: RecordBatch) { + self.mark_c_data_ingested(); self.failed_batches.lock().await.push(batch); } diff --git a/rust/tests/src/arrow_c_data_ffi_tests.rs b/rust/tests/src/arrow_c_data_ffi_tests.rs index c689a5cd..a0649200 100644 --- a/rust/tests/src/arrow_c_data_ffi_tests.rs +++ b/rust/tests/src/arrow_c_data_ffi_tests.rs @@ -21,14 +21,16 @@ mod ffi_c_data_lifetime_tests { use std::time::Duration; use databricks_zerobus_ingest_sdk::internal::arrow_c_data::import_c_data_record_batch; + use databricks_zerobus_ingest_sdk::internal::arrow_stream_has_ingested_c_data; use databricks_zerobus_ingest_sdk::ZerobusArrowStream; use tonic::Status; use zerobus_ffi::{ zerobus_arrow_get_default_config, zerobus_arrow_stream_close, zerobus_arrow_stream_flush, - zerobus_arrow_stream_free, zerobus_arrow_stream_ingest_c_data, zerobus_free_error_message, - zerobus_sdk_builder_build, zerobus_sdk_builder_disable_tls, zerobus_sdk_builder_endpoint, - zerobus_sdk_builder_new, zerobus_sdk_create_arrow_stream_with_headers_provider, - zerobus_sdk_free, CArrowArray, CArrowSchema, CHeaders, CResult, + zerobus_arrow_stream_free, zerobus_arrow_stream_ingest_batch, + zerobus_arrow_stream_ingest_c_data, zerobus_free_error_message, zerobus_sdk_builder_build, + zerobus_sdk_builder_disable_tls, zerobus_sdk_builder_endpoint, zerobus_sdk_builder_new, + zerobus_sdk_create_arrow_stream_with_headers_provider, zerobus_sdk_free, CArrowArray, + CArrowSchema, CHeaders, CResult, }; use super::{ @@ -197,6 +199,15 @@ mod ffi_c_data_lifetime_tests { bytes } + fn batch_ipc_bytes(batch: &RecordBatch) -> Vec { + let mut bytes = Vec::new(); + let mut writer = + arrow_ipc::writer::StreamWriter::try_new(&mut bytes, batch.schema().as_ref()).unwrap(); + writer.write(batch).unwrap(); + writer.finish().unwrap(); + bytes + } + extern "C" fn empty_headers(_user_data: *mut c_void) -> CHeaders { CHeaders { headers: ptr::null_mut(), @@ -293,6 +304,30 @@ mod ffi_c_data_lifetime_tests { .unwrap() } + async fn ingest_ffi_ipc_batch(stream: usize, ipc: Vec) -> Result { + tokio::task::spawn_blocking(move || { + let mut result = CResult { + success: true, + error_message: ptr::null_mut(), + is_retryable: false, + }; + let offset = zerobus_arrow_stream_ingest_batch( + stream as *mut _, + ipc.as_ptr(), + ipc.len(), + &mut result, + ); + if result.success { + Ok(offset) + } else { + Err(take_result_error(&mut result) + .unwrap_or_else(|| "IPC ingest failed without a message".to_string())) + } + }) + .await + .unwrap() + } + async fn flush_ffi_stream(stream: usize) -> Result<(), String> { tokio::task::spawn_blocking(move || { let mut result = CResult { @@ -381,6 +416,126 @@ mod ffi_c_data_lifetime_tests { } } + async fn assert_mixed_mode_free_waits( + c_data_first: bool, + ) -> Result<(), Box> { + let (_mock_server, server_url) = start_mock_flight_server().await?; + let schema = create_test_arrow_schema(); + let ipc_batch = create_test_record_batch(Arc::clone(&schema), vec![1], vec![Some("IPC")]); + let c_data_batch = + create_test_record_batch(Arc::clone(&schema), vec![2], vec![Some("C Data")]); + let ipc = batch_ipc_bytes(&ipc_batch); + let array_releases = Arc::new(AtomicUsize::new(0)); + let schema_releases = Arc::new(AtomicUsize::new(0)); + let (array, schema_ffi) = exported_counting_batch( + c_data_batch, + Arc::clone(&array_releases), + Arc::clone(&schema_releases), + ); + let (sdk, stream) = create_ffi_stream(server_url, schema).await?; + + if c_data_first { + assert_eq!(ingest_ffi_batch(stream, array, schema_ffi).await?, 0); + assert_eq!(ingest_ffi_ipc_batch(stream, ipc).await?, 1); + } else { + assert_eq!(ingest_ffi_ipc_batch(stream, ipc).await?, 0); + assert_eq!(ingest_ffi_batch(stream, array, schema_ffi).await?, 1); + } + + // This mirrors the opaque-handle cast in arrow.rs and must change with it. + let stream_ref = unsafe { &*(stream as *const ZerobusArrowStream) }; + assert!(arrow_stream_has_ingested_c_data(stream_ref)); + let (shutdown_reached, shutdown_proceed) = + stream_ref.arm_free_shutdown_complete_barrier().await; + let mut shutdown_guard = TestBarrierGuard(Some(shutdown_proceed)); + + let (free_completed_tx, mut free_completed_rx) = mpsc::channel(); + drop(std::thread::spawn(move || { + zerobus_arrow_stream_free(stream as *mut _); + let _ = free_completed_tx.send(()); + })); + + tokio::time::timeout(Duration::from_secs(5), shutdown_reached.notified()) + .await + .expect("mixed-mode free must use complete C Data shutdown"); + assert!(matches!( + free_completed_rx.try_recv(), + Err(mpsc::TryRecvError::Empty) + )); + shutdown_guard.release(); + tokio::time::timeout( + Duration::from_secs(5), + wait_for_free_completion(&mut free_completed_rx), + ) + .await + .expect("mixed-mode free did not complete after shutdown proceeded")?; + + assert_eq!(array_releases.load(Ordering::SeqCst), 1); + assert_eq!(schema_releases.load(Ordering::SeqCst), 1); + tokio::task::spawn_blocking(move || zerobus_sdk_free(sdk as *mut _)) + .await + .expect("sdk free task panicked"); + Ok(()) + } + + #[tokio::test] + async fn ffi_ipc_only_free_does_not_wait_for_complete_shutdown( + ) -> Result<(), Box> { + let (_mock_server, server_url) = start_mock_flight_server().await?; + let schema = create_test_arrow_schema(); + let batch = create_test_record_batch(Arc::clone(&schema), vec![1], vec![Some("IPC")]); + let (sdk, stream) = create_ffi_stream(server_url, schema).await?; + assert_eq!( + ingest_ffi_ipc_batch(stream, batch_ipc_bytes(&batch)).await?, + 0 + ); + // This mirrors the opaque-handle cast in arrow.rs and must change with it. + let stream_ref = unsafe { &*(stream as *const ZerobusArrowStream) }; + assert!(!arrow_stream_has_ingested_c_data(stream_ref)); + let (_shutdown_reached, shutdown_proceed) = + stream_ref.arm_free_shutdown_complete_barrier().await; + let mut shutdown_guard = TestBarrierGuard(Some(shutdown_proceed)); + + let (free_completed_tx, mut free_completed_rx) = mpsc::channel(); + drop(std::thread::spawn(move || { + zerobus_arrow_stream_free(stream as *mut _); + let _ = free_completed_tx.send(()); + })); + + let completion = tokio::time::timeout( + Duration::from_secs(1), + wait_for_free_completion(&mut free_completed_rx), + ) + .await; + shutdown_guard.release(); + if completion.is_err() { + tokio::time::timeout( + Duration::from_secs(5), + wait_for_free_completion(&mut free_completed_rx), + ) + .await + .expect("free did not complete after releasing the test barrier")?; + } + + tokio::task::spawn_blocking(move || zerobus_sdk_free(sdk as *mut _)) + .await + .expect("sdk free task panicked"); + completion.expect("IPC-only free must preserve best-effort nonblocking destruction")?; + Ok(()) + } + + #[tokio::test] + async fn ffi_ipc_then_c_data_free_waits_for_complete_shutdown( + ) -> Result<(), Box> { + assert_mixed_mode_free_waits(false).await + } + + #[tokio::test] + async fn ffi_c_data_then_ipc_free_waits_for_complete_shutdown( + ) -> Result<(), Box> { + assert_mixed_mode_free_waits(true).await + } + #[tokio::test] async fn ffi_c_data_owner_releases_once_after_ack_and_flush( ) -> Result<(), Box> { @@ -604,6 +759,9 @@ mod ffi_c_data_lifetime_tests { }; let mut before_batch_poll_guard = TestBarrierGuard(Some(before_batch_poll_proceed)); assert_eq!(ingest_ffi_batch(stream, array, schema_ffi).await?, 0); + // This mirrors the opaque-handle cast in arrow.rs and must change with it. + let stream_ref = unsafe { &*(stream as *const ZerobusArrowStream) }; + assert!(arrow_stream_has_ingested_c_data(stream_ref)); tokio::time::timeout(Duration::from_secs(5), before_batch_poll_reached.notified()) .await .expect("request body must park before consuming the queued batch");