[FFI] Complete Arrow shutdown before free returns

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

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. 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

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<ZerobusArrowStream>` 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,108 @@ pub extern "C" fn zerobus_sdk_create_arrow_stream_with_headers_provider(
     })
 }
 
+fn abort_and_drop_arrow_stream(stream: Box<ZerobusArrowStream>) {
+    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<ZerobusArrowStream>) {
+    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 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 +468,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 +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) => {

rust/ffi/zerobus.h

@@ -37,6 +37,9 @@ typedef struct CHeaders {
 
 /**
  * Opaque handle for an Arrow Flight stream.
+ *
+ * FFI pointers are `Box<ZerobusArrowStream>` 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,22 @@ 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 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 +343,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.

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()
+}

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,