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..f8c2b9d6 100644 --- a/rust/ffi/NEXT_CHANGELOG.md +++ b/rust/ffi/NEXT_CHANGELOG.md @@ -19,6 +19,8 @@ ### Behavior Changes +- `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. Callers must not block the only thread, event loop, or runtime lock needed by a release callback: offload free, release required runtime locks, and continue servicing callback dependencies until it completes. 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 ### Deprecations diff --git a/rust/ffi/src/arrow.rs b/rust/ffi/src/arrow.rs index ee1d4cad..577a69e3 100644 --- a/rust/ffi/src/arrow.rs +++ b/rust/ffi/src/arrow.rs @@ -7,18 +7,31 @@ use bytes::Bytes; 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, }; 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 +304,114 @@ 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. +/// +/// 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 call this function from the only thread, event loop, or runtime lock that a release +/// callback needs in order to complete. For example, a caller must not hold the Python GIL if a +/// release callback must reacquire it, or block a single-threaded event loop/runtime used by that +/// callback. Offload this synchronous function to a blocking or OS thread, release any required +/// runtime locks, and continue servicing the event loop until the call completes. +/// +/// 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 _ = Box::from_raw(stream as *mut ZerobusArrowStream); + 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() + == 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), + } } } }) @@ -360,14 +474,17 @@ 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 /// 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. @@ -417,6 +534,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 26865f57..cae4dc92 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,28 @@ struct CArrowStream *zerobus_sdk_create_arrow_stream_with_headers_provider(struc /** * Frees an Arrow Flight stream instance. + * + * 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 call this function from the only thread, event loop, or runtime lock that a release + * callback needs in order to complete. For example, a caller must not hold the Python GIL if a + * release callback must reacquire it, or block a single-threaded event loop/runtime used by that + * callback. Offload this synchronous function to a blocking or OS thread, release any required + * runtime locks, and continue servicing the event loop until the call completes. + * + * 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); @@ -324,14 +349,17 @@ 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 * 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..7cf2710e 100644 --- a/rust/sdk/src/internal.rs +++ b/rust/sdk/src/internal.rs @@ -10,3 +10,24 @@ 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; +} + +/// 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/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..8dc7040c 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,13 @@ 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>>, + /// 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. pending_batches: Arc>>, /// Wakes the ACK processor when a batch is submitted after an idle period. @@ -274,7 +296,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 +318,10 @@ impl ZerobusArrowStream { is_closed, admission_closed, close, - supervisor_abort, + supervisor_task, + #[cfg(feature = "internal-arrow-c-data")] + has_ingested_c_data: AtomicBool::new(false), + request_bodies, pending_batches, pending_notify, request_send_failure, @@ -325,6 +351,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 +363,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 +409,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 +988,53 @@ 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); + 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 +1135,59 @@ 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: 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); + } + + /// 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 +1241,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..c84853af 100644 --- a/rust/tests/src/arrow_c_data_ffi_tests.rs +++ b/rust/tests/src/arrow_c_data_ffi_tests.rs @@ -16,12 +16,17 @@ 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::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_flush, zerobus_arrow_stream_free, + zerobus_arrow_get_default_config, zerobus_arrow_stream_close, zerobus_arrow_stream_flush, + 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, @@ -68,6 +73,148 @@ 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) + } + + struct RuntimeWaitingArrayRelease { + inner_release: Option, + inner_private_data: *mut c_void, + runtime: tokio::runtime::Handle, + runtime_progressed: Arc, + release_count: Arc, + } + + unsafe extern "C" fn runtime_waiting_array_release(array: *mut FFI_ArrowArray) { + let array = unsafe { &mut *array }; + let state = + unsafe { Box::from_raw(array.private_data.cast::()) }; + let (completed_tx, completed_rx) = mpsc::sync_channel(1); + state.runtime.spawn(async move { + let _ = completed_tx.send(()); + }); + state.runtime_progressed.store( + completed_rx.recv_timeout(Duration::from_secs(2)).is_ok(), + 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_runtime_waiting_batch( + batch: RecordBatch, + runtime: tokio::runtime::Handle, + runtime_progressed: 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(RuntimeWaitingArrayRelease { + inner_release: array.release, + inner_private_data: array.private_data, + runtime, + runtime_progressed, + release_count, + }); + array.release = Some(runtime_waiting_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, @@ -102,6 +249,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(), @@ -198,6 +354,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 { @@ -217,6 +397,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 +434,158 @@ 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"); + } + } + } + } + + 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> { @@ -317,8 +667,292 @@ 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(flavor = "current_thread")] + async fn ffi_c_data_free_offloaded_from_current_thread_runtime_allows_callback_progress( + ) -> 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 runtime_progressed = Arc::new(AtomicBool::new(false)); + let release_count = Arc::new(AtomicUsize::new(0)); + let batch = + create_test_record_batch(schema, vec![1], vec![Some("current-thread callback")]); + let (array, schema_ffi) = exported_runtime_waiting_batch( + batch, + tokio::runtime::Handle::current(), + Arc::clone(&runtime_progressed), + Arc::clone(&release_count), + ); + 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; + + tokio::time::timeout( + Duration::from_secs(5), + tokio::task::spawn_blocking(move || { + zerobus_arrow_stream_free(stream as *mut _); + }), + ) + .await + .expect("offloaded free timed out") + .expect("offloaded free task panicked"); + + assert_eq!(release_count.load(Ordering::SeqCst), 1); + assert!( + runtime_progressed.load(Ordering::SeqCst), + "offloaded free must let the current-thread runtime drive callback work" + ); + 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); + // 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"); + 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(()) } }