diff --git a/arrow-array/src/ffi_stream.rs b/arrow-array/src/ffi_stream.rs index 9a09c3753dd4..a119d51b2ef9 100644 --- a/arrow-array/src/ffi_stream.rs +++ b/arrow-array/src/ffi_stream.rs @@ -97,17 +97,24 @@ const ENOSYS: i32 = 38; #[derive(Debug)] #[allow(non_camel_case_types)] pub struct FFI_ArrowArrayStream { + // Fields are private so safe code can't install a bogus callback that import + // or [`Drop`] would invoke. Write the release fields with the unsafe setters. /// C function to get schema from the stream - pub get_schema: - Option c_int>, + get_schema: Option c_int>, /// C function to get next array from the stream - pub get_next: Option c_int>, + get_next: Option c_int>, /// C function to get the error from last operation on the stream - pub get_last_error: Option *const c_char>, + get_last_error: Option *const c_char>, /// C function to release the stream - pub release: Option, - /// Private data used by the stream - pub private_data: *mut c_void, + /// + /// Private so safe code can't install a callback that [`Drop`] would invoke. + /// Use [`FFI_ArrowArrayStream::release`] and [`FFI_ArrowArrayStream::set_release`]. + release: Option, + /// Private data used by the stream, owned by the release callback. + /// + /// Private for the same reason as `release`. Use + /// [`FFI_ArrowArrayStream::private_data`] and [`FFI_ArrowArrayStream::set_private_data`]. + private_data: *mut c_void, } unsafe impl Send for FFI_ArrowArrayStream {} @@ -213,6 +220,48 @@ impl FFI_ArrowArrayStream { private_data: std::ptr::null_mut(), } } + + /// Returns the producer-provided release callback, if any. + /// + /// Lets a consumer wrap release: save this callback, install its own with + /// [`FFI_ArrowArrayStream::set_release`], and chain back to it on drop. See + /// . + pub fn release(&self) -> Option { + self.release + } + + /// Returns the opaque producer-provided private data pointer. + /// + /// See [`FFI_ArrowArrayStream::release`] for the intended use. + pub fn private_data(&self) -> *mut c_void { + self.private_data + } + + /// Replaces the release callback, returning the previous one. + /// + /// # Safety + /// + /// [`Drop`] calls this callback with a pointer to `self`. The new callback + /// must correctly release this stream (usually by chaining to the returned + /// one) and must match the [`FFI_ArrowArrayStream::private_data`] it reads. + /// A wrong callback is undefined behavior on drop. + pub unsafe fn set_release( + &mut self, + release: Option, + ) -> Option { + std::mem::replace(&mut self.release, release) + } + + /// Replaces the private data pointer, returning the previous one. + /// + /// # Safety + /// + /// The old pointer is returned without being freed; the caller owns it from + /// here. The new pointer must match what the current + /// [`FFI_ArrowArrayStream::release`] callback expects. + pub unsafe fn set_private_data(&mut self, private_data: *mut c_void) -> *mut c_void { + std::mem::replace(&mut self.private_data, private_data) + } } struct ExportedArrayStream { @@ -566,4 +615,48 @@ mod tests { Ok(()) } + + // A consumer wraps the release callback with its own, then chains back to + // the original on drop. This is the same wrap-release pattern the + // release/private_data accessors exist for (#9771). + static STREAM_WRAPPER_RAN: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + + struct StreamWrapperData { + original_release: Option, + original_private_data: *mut c_void, + } + + unsafe extern "C" fn wrapping_release(stream: *mut FFI_ArrowArrayStream) { + use std::sync::atomic::Ordering; + let stream = unsafe { &mut *stream }; + let data = unsafe { Box::from_raw(stream.private_data() as *mut StreamWrapperData) }; + STREAM_WRAPPER_RAN.store(true, Ordering::SeqCst); + unsafe { stream.set_release(data.original_release) }; + unsafe { stream.set_private_data(data.original_private_data) }; + if let Some(release) = stream.release() { + unsafe { release(stream) }; + } + } + + #[test] + fn test_wrap_release_callback() { + use std::sync::atomic::Ordering; + + let batch_reader = TestRecordBatchReader::new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])), + Box::new(std::iter::empty()), + ); + let mut stream = FFI_ArrowArrayStream::new(batch_reader); + + let data = Box::new(StreamWrapperData { + original_release: stream.release(), + original_private_data: stream.private_data(), + }); + unsafe { stream.set_release(Some(wrapping_release)) }; + unsafe { stream.set_private_data(Box::into_raw(data) as *mut c_void) }; + + drop(stream); // runs wrapping_release, which chains to the original + assert!(STREAM_WRAPPER_RAN.load(Ordering::SeqCst)); + } } diff --git a/arrow-data/src/ffi.rs b/arrow-data/src/ffi.rs index 80ccee6c8a53..e6afaf31f8f0 100644 --- a/arrow-data/src/ffi.rs +++ b/arrow-data/src/ffi.rs @@ -37,31 +37,41 @@ use std::ffi::c_void; #[repr(C)] #[derive(Debug)] pub struct FFI_ArrowArray { + // Fields are private so safe code can't set them to values that break import + // or the release callback (e.g. a bad `buffers` pointer gets dereferenced on + // import). Read them with the getters; write the callback fields with the + // unsafe setters. /// Logical length of the array - pub length: i64, + length: i64, /// Number of null items in the array - pub null_count: i64, + null_count: i64, /// logical offset inside the array - pub offset: i64, + offset: i64, /// Number of physical buffers backing this array - pub n_buffers: i64, + n_buffers: i64, /// Number of children this array has - pub n_children: i64, + n_children: i64, /// C array of pointers to the start of each physical buffer backing this array - pub buffers: *mut *const c_void, + buffers: *mut *const c_void, /// C array of pointers to each child array of this array - pub children: *mut *mut FFI_ArrowArray, + children: *mut *mut FFI_ArrowArray, /// Pointer to the underlying array of dictionary values - pub dictionary: *mut FFI_ArrowArray, - /// Pointer to a producer-provided release callback - pub release: Option, + dictionary: *mut FFI_ArrowArray, + /// Producer-provided release callback. + /// + /// Private so safe code can't install a callback that [`Drop`] would invoke. + /// Use [`FFI_ArrowArray::release`] and [`FFI_ArrowArray::set_release`]. + release: Option, /// Opaque pointer to producer-provided private data /// When exported, this MUST contain everything that is owned by this array. /// For example, any buffer pointed to in `buffers` must be here, as well /// as the `buffers` pointer itself. /// In other words, everything in [FFI_ArrowArray] must be owned by /// `private_data` and can assume that they do not outlive `private_data`. - pub private_data: *mut c_void, + /// + /// Private for the same reason as `release`. Use + /// [`FFI_ArrowArray::private_data`] and [`FFI_ArrowArray::set_private_data`]. + private_data: *mut c_void, } impl Drop for FFI_ArrowArray { @@ -251,6 +261,48 @@ impl FFI_ArrowArray { } } + /// Returns the producer-provided release callback, if any. + /// + /// Lets a consumer wrap release: save this callback, install its own with + /// [`FFI_ArrowArray::set_release`], and chain back to it on drop. See + /// . + pub fn release(&self) -> Option { + self.release + } + + /// Returns the opaque producer-provided private data pointer. + /// + /// See [`FFI_ArrowArray::release`] for the intended use. + pub fn private_data(&self) -> *mut c_void { + self.private_data + } + + /// Replaces the release callback, returning the previous one. + /// + /// # Safety + /// + /// [`Drop`] calls this callback with a pointer to `self`. The new callback + /// must correctly release this array (usually by chaining to the returned + /// one) and must match the [`FFI_ArrowArray::private_data`] it reads. A + /// wrong callback is undefined behavior on drop. + pub unsafe fn set_release( + &mut self, + release: Option, + ) -> Option { + std::mem::replace(&mut self.release, release) + } + + /// Replaces the private data pointer, returning the previous one. + /// + /// # Safety + /// + /// The old pointer is returned without being freed; the caller owns it from + /// here. The new pointer must match what the current + /// [`FFI_ArrowArray::release`] callback expects. + pub unsafe fn set_private_data(&mut self, private_data: *mut c_void) -> *mut c_void { + std::mem::replace(&mut self.private_data, private_data) + } + /// the length of the array #[inline] pub fn len(&self) -> usize { diff --git a/arrow-schema/src/ffi.rs b/arrow-schema/src/ffi.rs index f2cfaea0263f..958946ff54fb 100644 --- a/arrow-schema/src/ffi.rs +++ b/arrow-schema/src/ffi.rs @@ -75,25 +75,35 @@ bitflags! { #[derive(Debug)] #[allow(non_camel_case_types)] pub struct FFI_ArrowSchema { + // Fields are private so safe code can't set them to values that break the + // release callback (e.g. a `format`/`name` that isn't a `CString`, or a bad + // pointer). Read them with the getters; write the callback fields with the + // unsafe setters. /// Null-terminated, UTF8-encoded string describing the data type - pub format: *const c_char, + format: *const c_char, /// Null-terminated, UTF8-encoded string of the field or array name - pub name: *const c_char, + name: *const c_char, /// Binary string describing the type’s metadata - pub metadata: *const c_char, + metadata: *const c_char, /// A bitfield of flags enriching the type description /// Refer to [Arrow Flags](https://arrow.apache.org/docs/format/CDataInterface.html#c.ArrowSchema.flags) - pub flags: i64, + flags: i64, /// The number of children this type has - pub n_children: i64, + n_children: i64, /// C array of pointers to each child type of this type - pub children: *mut *mut FFI_ArrowSchema, + children: *mut *mut FFI_ArrowSchema, /// Pointer to the type of dictionary values - pub dictionary: *mut FFI_ArrowSchema, - /// Pointer to a producer-provided release callback - pub release: Option, - /// Opaque pointer to producer-provided private data - pub private_data: *mut c_void, + dictionary: *mut FFI_ArrowSchema, + /// Producer-provided release callback. + /// + /// Private so safe code can't install a callback that [`Drop`] would invoke. + /// Use [`FFI_ArrowSchema::release`] and [`FFI_ArrowSchema::set_release`]. + release: Option, + /// Opaque producer-provided private data, owned by the release callback. + /// + /// Private for the same reason as `release`. Use + /// [`FFI_ArrowSchema::private_data`] and [`FFI_ArrowSchema::set_private_data`]. + private_data: *mut c_void, } struct SchemaPrivateData { @@ -276,6 +286,48 @@ impl FFI_ArrowSchema { } } + /// Returns the producer-provided release callback, if any. + /// + /// Lets a consumer wrap release: save this callback, install its own with + /// [`FFI_ArrowSchema::set_release`], and chain back to it on drop. See + /// . + pub fn release(&self) -> Option { + self.release + } + + /// Returns the opaque producer-provided private data pointer. + /// + /// See [`FFI_ArrowSchema::release`] for the intended use. + pub fn private_data(&self) -> *mut c_void { + self.private_data + } + + /// Replaces the release callback, returning the previous one. + /// + /// # Safety + /// + /// [`Drop`] calls this callback with a pointer to `self`. The new callback + /// must correctly release this schema (usually by chaining to the returned + /// one) and must match the [`FFI_ArrowSchema::private_data`] it reads. A + /// wrong callback is undefined behavior on drop. + pub unsafe fn set_release( + &mut self, + release: Option, + ) -> Option { + std::mem::replace(&mut self.release, release) + } + + /// Replaces the private data pointer, returning the previous one. + /// + /// # Safety + /// + /// The old pointer is returned without being freed; the caller owns it from + /// here. The new pointer must match what the current + /// [`FFI_ArrowSchema::release`] callback expects. + pub unsafe fn set_private_data(&mut self, private_data: *mut c_void) -> *mut c_void { + std::mem::replace(&mut self.private_data, private_data) + } + /// Returns the format of this schema. pub fn format(&self) -> &str { assert!(!self.format.is_null()); @@ -861,6 +913,7 @@ impl TryFrom for FFI_ArrowSchema { mod tests { use super::*; use crate::Fields; + use std::sync::atomic::{AtomicBool, Ordering}; fn round_trip_type(dtype: DataType) { let c_schema = FFI_ArrowSchema::try_from(&dtype).unwrap(); @@ -1029,4 +1082,41 @@ mod tests { let field = Field::try_from(&c_schema).unwrap(); assert_eq!(field.name(), ""); } + + // A consumer wraps the release callback with its own, then chains back to + // the original on drop. This is the cross-thread use case from #9771 and is + // what the release/private_data accessors exist for. + static WRAPPER_RAN: AtomicBool = AtomicBool::new(false); + + struct WrapperData { + original_release: Option, + original_private_data: *mut c_void, + } + + unsafe extern "C" fn wrapping_release(schema: *mut FFI_ArrowSchema) { + let schema = unsafe { &mut *schema }; + let data = unsafe { Box::from_raw(schema.private_data() as *mut WrapperData) }; + WRAPPER_RAN.store(true, Ordering::SeqCst); + // restore the originals, then let the original callback free everything + unsafe { schema.set_release(data.original_release) }; + unsafe { schema.set_private_data(data.original_private_data) }; + if let Some(release) = schema.release() { + unsafe { release(schema) }; + } + } + + #[test] + fn test_wrap_release_callback() { + let mut schema = FFI_ArrowSchema::try_from(&DataType::Int32).unwrap(); + + let data = Box::new(WrapperData { + original_release: schema.release(), + original_private_data: schema.private_data(), + }); + unsafe { schema.set_release(Some(wrapping_release)) }; + unsafe { schema.set_private_data(Box::into_raw(data) as *mut c_void) }; + + drop(schema); // runs wrapping_release, which chains to the original + assert!(WRAPPER_RAN.load(Ordering::SeqCst)); + } }