From 77ae1395e0844e4091e1db634aa34fc55ee94619 Mon Sep 17 00:00:00 2001 From: Martin Prammer Date: Mon, 21 Sep 2026 09:50:05 -0400 Subject: [PATCH 1/2] Expose callback reads and row-position projections in C FFI Co-authored-by: Claude Signed-off-by: Martin Prammer --- vortex-ffi/cinclude/vortex.h | 88 +++++++++++++++ vortex-ffi/src/data_source.rs | 207 ++++++++++++++++++++++++++++++++++ vortex-ffi/src/expression.rs | 125 ++++++++++++++++++++ vortex-ffi/src/lib.rs | 1 + vortex-ffi/src/read_at.rs | 204 +++++++++++++++++++++++++++++++++ vortex-ffi/src/scan.rs | 110 ++++++++++++++++++ 6 files changed, 735 insertions(+) create mode 100644 vortex-ffi/src/read_at.rs diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index fa214e8efbe..043cb390f3a 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -543,6 +543,45 @@ typedef struct { size_t paths_len; } vx_data_source_options; +/** + * A random-access byte source implemented by the caller. + * + * "read_at" must tolerate concurrent calls from arbitrary threads. The struct + * is read only during the call it is passed to, but "ctx" and the callbacks + * must stay valid until "release" runs. + */ +typedef struct { + /** + * Opaque caller state passed to every callback. May be NULL. + */ + void *ctx; + /** + * Total length of the source in bytes. Must be exact: the footer is read + * relative to the end, so a wrong length surfaces as a corrupt file. + */ + uint64_t len; + /** + * Maximum number of concurrent "read_at" calls. 0 selects a default. The + * cap is per-source, so opening many files multiplies it. + */ + size_t concurrency; + /** + * Optional name, typically the URI, used for cache keys and error messages; + * it should be stable and unique. Copied. Zero-length means anonymous. + */ + vx_view name; + /** + * Required. Must write all "length" bytes at "offset" into "dst" and return + * 0, or return non-zero; success without filling "dst" leaks uninitialized + * memory into the scan. + */ + int32_t (*read_at)(void *ctx, uint64_t offset, uint8_t *dst, size_t length); + /** + * Optional. Called once, after Vortex has dropped the source. + */ + void (*release)(void *ctx); +} vx_readat; + /** * Used for estimating number of partitions in a data source or number of rows * in a partition. @@ -889,6 +928,22 @@ vx_data_source_new(const vx_session *session, const vx_data_source_options *opti const vx_data_source * vx_data_source_new_buffer(const vx_session *session, const void *buffer, size_t buffer_len, vx_error **err); +/** + * Create a data source that reads through caller-supplied callbacks instead of + * Vortex's own I/O. + * + * Unlike vx_data_source_new_buffer, this keeps I/O pruning: only the segments a + * scan needs are fetched, rather than the whole file up front. + * + * "reader" is read during this call only; its callbacks and context must stay + * valid until "release" runs. A rejected descriptor leaves ownership with the + * caller and never calls "release"; once accepted, "release" always runs. + * + * On error, returns NULL and sets "err". + */ +const vx_data_source * +vx_data_source_new_readat(const vx_session *session, const vx_readat *reader, vx_error **err); + /** * Increase reference count on vx_data_source */ @@ -1115,6 +1170,39 @@ vx_expression *vx_expression_clone(const vx_expression *ptr); */ vx_expression *vx_expression_literal(const vx_scalar *scalar, vx_error **err); +/** + * Create an expression yielding each row's position within the file scanned. + * + * Recovers a row's original position after a filter dropped the rows around it, + * as Iceberg positional deletes need. Only valid inside a scan; else it errors. + */ +vx_expression *vx_expression_row_idx(void); + +/** + * Create a struct-valued expression from named child expressions. + * + * Where vx_expression_select trims a struct to some of its fields, pack builds + * one from arbitrary expressions - how fields inside a nested struct get pruned. + * + * "names" and "expressions" must both point to arrays of "len" entries, paired + * by position. "nullable" sets the resulting struct's nullability. Names are + * copied. + * + * Returns NULL if len == 0, if either array is NULL, if any entry of + * "expressions" is NULL, or if a name is not valid UTF-8. + * + * Example: + * + * vx_expression* root = vx_expression_root(); + * vx_expression* addr = vx_expression_get_item(vx_view_from_cstr("addr"), root); + * vx_expression* city = vx_expression_get_item(vx_view_from_cstr("city"), addr); + * vx_view names[] = {vx_view_from_cstr("city")}; + * const vx_expression* parts[] = {city}; + * vx_expression* packed = vx_expression_pack(names, parts, 1, false); + */ +vx_expression * +vx_expression_pack(const vx_view *names, const vx_expression *const *expressions, size_t len, bool nullable); + /** * Create an expression that selects (includes) specific fields from a child * expression. Child expression must have a DTYPE_STRUCT dtype. Errors in diff --git a/vortex-ffi/src/data_source.rs b/vortex-ffi/src/data_source.rs index a881b3ac085..6339e7fe20e 100644 --- a/vortex-ffi/src/data_source.rs +++ b/vortex-ffi/src/data_source.rs @@ -33,6 +33,8 @@ use crate::box_wrapper; use crate::dtype::vx_dtype; use crate::error::try_or; use crate::error::vx_error; +use crate::read_at::read_at_from_ffi; +use crate::read_at::vx_readat; use crate::scan::vx_estimate; use crate::scan::vx_estimate_type; use crate::session::vx_session; @@ -192,6 +194,46 @@ pub unsafe extern "C-unwind" fn vx_data_source_new_buffer( }) } +/// Create a data source that reads through caller-supplied callbacks instead of +/// Vortex's own I/O. +/// +/// Unlike vx_data_source_new_buffer, this keeps I/O pruning: only the segments a +/// scan needs are fetched, rather than the whole file up front. +/// +/// "reader" is read during this call only; its callbacks and context must stay +/// valid until "release" runs. A rejected descriptor leaves ownership with the +/// caller and never calls "release"; once accepted, "release" always runs. +/// +/// On error, returns NULL and sets "err". +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_data_source_new_readat( + session: *const vx_session, + reader: *const vx_readat, + err: *mut *mut vx_error, +) -> *const vx_data_source { + try_or(err, ptr::null(), || { + vortex_ensure!(!session.is_null()); + + let session = vx_session::as_ref(session); + let source = unsafe { read_at_from_ffi(reader) }?; + + let (file, len) = RUNTIME.block_on(async { + let len = source.size().await?; + let file = session.open_options().open(source).await?; + VortexResult::Ok((file, len)) + })?; + + let ds = MultiLayoutDataSource::new_with_first( + file.layout_reader()?, + Vec::new(), + vec![Some(len)], + session, + ); + + Ok(vx_data_source::new(ds)) + }) +} + /// Increase reference count on vx_data_source #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_data_source_clone( @@ -237,17 +279,22 @@ mod tests { use std::ffi::c_void; use std::fs::read; use std::ptr; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; use crate::data_source::vx_data_source_dtype; use crate::data_source::vx_data_source_free; use crate::data_source::vx_data_source_get_row_count; use crate::data_source::vx_data_source_new; use crate::data_source::vx_data_source_new_buffer; + use crate::data_source::vx_data_source_new_readat; use crate::data_source::vx_data_source_options; use crate::dtype::vx_dtype; use crate::dtype::vx_dtype_free; + use crate::read_at::vx_readat; use crate::scan::vx_estimate; use crate::scan::vx_estimate_type; + use crate::session::vx_session; use crate::session::vx_session_free; use crate::session::vx_session_new; use crate::string::vx_view; @@ -429,4 +476,164 @@ mod tests { vx_session_free(session); } } + + /// Backing store for the callback reader: the whole file in memory, plus a + /// count of `release` calls so the test can prove it runs exactly once. + struct ReadAtCtx { + data: Vec, + releases: AtomicUsize, + } + + unsafe extern "C" fn read_at_cb( + ctx: *mut c_void, + offset: u64, + dst: *mut u8, + length: usize, + ) -> i32 { + let ctx = unsafe { &*ctx.cast::() }; + let Ok(start) = usize::try_from(offset) else { + return 1; + }; + let Some(end) = start.checked_add(length) else { + return 1; + }; + if end > ctx.data.len() { + return 1; + } + unsafe { ptr::copy_nonoverlapping(ctx.data[start..end].as_ptr(), dst, length) }; + 0 + } + + unsafe extern "C" fn release_cb(ctx: *mut c_void) { + let ctx = unsafe { &*ctx.cast::() }; + ctx.releases.fetch_add(1, Ordering::SeqCst); + } + + /// A reader whose callback always fails, to check the error surfaces + /// instead of handing uninitialized bytes to the scan. + unsafe extern "C" fn failing_read_at_cb( + _ctx: *mut c_void, + _offset: u64, + _dst: *mut u8, + _length: usize, + ) -> i32 { + -7 + } + + fn readat_ctx(session: *const vx_session) -> (Box, u64, tempfile::NamedTempFile) { + let (sample, _) = unsafe { write_sample(session) }; + let data = read(sample.path()).unwrap(); + let len = data.len() as u64; + ( + Box::new(ReadAtCtx { + data, + releases: AtomicUsize::new(0), + }), + len, + sample, + ) + } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_create_readat() { + unsafe { + let session = vx_session_new(); + let (ctx, len, _sample) = readat_ctx(session); + + let reader = vx_readat { + ctx: (&raw const *ctx).cast::().cast_mut(), + len, + concurrency: 0, + name: vx_view::from_str("test://sample.vortex"), + read_at: Some(read_at_cb), + release: Some(release_cb), + }; + + let mut error = ptr::null_mut(); + let ds = vx_data_source_new_readat(session, &raw const reader, &raw mut error); + assert_no_error(error); + assert!(!ds.is_null()); + + let ffi_dtype = vx_data_source_dtype(ds); + let mut row_count = vx_estimate::default(); + vx_data_source_get_row_count(ds, &raw mut row_count); + assert_eq!(row_count.r#type, vx_estimate_type::VX_ESTIMATE_EXACT); + assert_eq!(row_count.estimate, SAMPLE_ROWS as u64); + + assert_eq!(ctx.releases.load(Ordering::SeqCst), 0); + + vx_dtype_free(ffi_dtype); + vx_data_source_free(ds); + vx_session_free(session); + + assert_eq!(ctx.releases.load(Ordering::SeqCst), 1); + } + } + + /// Descriptor validation fails before ownership transfers, so the caller + /// keeps the context and `release` must not run. + #[test] + #[cfg_attr(miri, ignore)] + fn test_create_readat_invalid() { + unsafe { + let session = vx_session_new(); + let (ctx, len, _sample) = readat_ctx(session); + let mut error = ptr::null_mut(); + + let ds = vx_data_source_new_readat(ptr::null(), ptr::null(), &raw mut error); + assert_error(error); + assert!(ds.is_null()); + + let mut error = ptr::null_mut(); + let ds = vx_data_source_new_readat(session, ptr::null(), &raw mut error); + assert_error(error); + assert!(ds.is_null()); + + // read_at is required. + let mut error = ptr::null_mut(); + let reader = vx_readat { + ctx: (&raw const *ctx).cast::().cast_mut(), + len, + concurrency: 0, + name: vx_view::from_str(""), + read_at: None, + release: Some(release_cb), + }; + let ds = vx_data_source_new_readat(session, &raw const reader, &raw mut error); + assert_error(error); + assert!(ds.is_null()); + + assert_eq!(ctx.releases.load(Ordering::SeqCst), 0); + vx_session_free(session); + } + } + + /// A failing callback must surface as an error. Ownership has already + /// transferred by then, so `release` still runs. + #[test] + #[cfg_attr(miri, ignore)] + fn test_create_readat_callback_failure() { + unsafe { + let session = vx_session_new(); + let (ctx, len, _sample) = readat_ctx(session); + + let reader = vx_readat { + ctx: (&raw const *ctx).cast::().cast_mut(), + len, + concurrency: 0, + name: vx_view::from_str(""), + read_at: Some(failing_read_at_cb), + release: Some(release_cb), + }; + + let mut error = ptr::null_mut(); + let ds = vx_data_source_new_readat(session, &raw const reader, &raw mut error); + assert_error(error); + assert!(ds.is_null()); + assert_eq!(ctx.releases.load(Ordering::SeqCst), 1); + + vx_session_free(session); + } + } } diff --git a/vortex-ffi/src/expression.rs b/vortex-ffi/src/expression.rs index 1c3eaad3b72..47b796c18ac 100644 --- a/vortex-ffi/src/expression.rs +++ b/vortex-ffi/src/expression.rs @@ -16,8 +16,10 @@ use vortex::expr::list_contains; use vortex::expr::lit; use vortex::expr::not; use vortex::expr::or_collect; +use vortex::expr::pack; use vortex::expr::root; use vortex::expr::select; +use vortex::layout::layouts::row_idx::row_idx; use vortex::scalar_fn::ScalarFnVTableExt; use vortex::scalar_fn::fns::binary::Binary; use vortex::scalar_fn::fns::operators::Operator; @@ -109,6 +111,66 @@ pub unsafe extern "C-unwind" fn vx_expression_literal( }) } +/// Create an expression yielding each row's position within the file scanned. +/// +/// Recovers a row's original position after a filter dropped the rows around it, +/// as Iceberg positional deletes need. Only valid inside a scan; else it errors. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_expression_row_idx() -> *mut vx_expression { + vx_expression::new(row_idx()) +} + +/// Create a struct-valued expression from named child expressions. +/// +/// Where vx_expression_select trims a struct to some of its fields, pack builds +/// one from arbitrary expressions - how fields inside a nested struct get pruned. +/// +/// "names" and "expressions" must both point to arrays of "len" entries, paired +/// by position. "nullable" sets the resulting struct's nullability. Names are +/// copied. +/// +/// Returns NULL if len == 0, if either array is NULL, if any entry of +/// "expressions" is NULL, or if a name is not valid UTF-8. +/// +/// Example: +/// +/// vx_expression* root = vx_expression_root(); +/// vx_expression* addr = vx_expression_get_item(vx_view_from_cstr("addr"), root); +/// vx_expression* city = vx_expression_get_item(vx_view_from_cstr("city"), addr); +/// vx_view names[] = {vx_view_from_cstr("city")}; +/// const vx_expression* parts[] = {city}; +/// vx_expression* packed = vx_expression_pack(names, parts, 1, false); +#[unsafe(no_mangle)] +pub unsafe extern "C" fn vx_expression_pack( + names: *const vx_view, + expressions: *const *const vx_expression, + len: usize, + nullable: bool, +) -> *mut vx_expression { + if len == 0 || names.is_null() || expressions.is_null() { + return ptr::null_mut(); + } + + let names = match unsafe { to_field_names(names, len) } { + Ok(names) => names, + Err(_) => return ptr::null_mut(), + }; + + let exprs = unsafe { slice::from_raw_parts(expressions, len) }; + // `as_ref` panics on NULL and this is `extern "C"`, so an unchecked entry + // would abort the process instead of reporting an error. + if exprs.iter().any(|expr| expr.is_null()) { + return ptr::null_mut(); + } + + let elements = names + .into_iter() + .zip(exprs.iter().map(|e| vx_expression::as_ref(*e).clone())) + .collect::>(); + + vx_expression::new(pack(elements, nullable.into())) +} + /// Create an expression that selects (includes) specific fields from a child /// expression. Child expression must have a DTYPE_STRUCT dtype. Errors in /// vx_array_apply if the child expression doesn't have a specified field. @@ -319,6 +381,7 @@ mod tests { use vortex::array::arrays::StructArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::bool::BoolArrayExt; + use vortex::array::assert_arrays_eq; use vortex::array::validity::Validity; use vortex::buffer::Buffer; use vortex::buffer::buffer; @@ -338,6 +401,7 @@ mod tests { use crate::expression::vx_expression_list_contains; use crate::expression::vx_expression_literal; use crate::expression::vx_expression_or; + use crate::expression::vx_expression_pack; use crate::expression::vx_expression_root; use crate::expression::vx_expression_select; use crate::scalar::vx_scalar_free; @@ -619,4 +683,65 @@ mod tests { vx_expression_free(root); } } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_pack() { + let mut ctx = array_session().create_execution_ctx(); + let (array, names_array, ages_array) = struct_array(); + unsafe { + let root = vx_expression_root(); + let name = vx_expression_get_item(vx_view::from_str("name"), root); + let age = vx_expression_get_item(vx_view::from_str("age"), root); + + let names = [vx_view::from_str("who"), vx_view::from_str("how_old")]; + let parts = [name.cast_const(), age.cast_const()]; + let packed = vx_expression_pack(names.as_ptr(), parts.as_ptr(), 2, false); + assert!(!packed.is_null()); + + let array = vx_array::new(array.into_array()); + let mut error = ptr::null_mut(); + let applied = vx_array_apply(array, packed, &raw mut error); + assert!(error.is_null()); + assert!(!applied.is_null()); + { + let expected = StructArray::try_new( + ["who", "how_old"].into(), + vec![names_array.into_array(), ages_array.into_array()], + 3, + Validity::NonNullable, + ) + .unwrap(); + assert_arrays_eq!(vx_array::as_ref(applied), expected, &mut ctx); + } + + vx_array_free(applied); + vx_array_free(array); + vx_expression_free(packed); + vx_expression_free(age); + vx_expression_free(name); + vx_expression_free(root); + } + } + + /// A NULL entry must be reported as NULL, not abort. `vx_expression_pack` + /// is `extern "C"`, so the panic from dereferencing one cannot unwind out. + #[test] + #[cfg_attr(miri, ignore)] + fn test_pack_rejects_invalid_arguments() { + unsafe { + let root = vx_expression_root(); + let names = [vx_view::from_str("a"), vx_view::from_str("b")]; + let parts = [root.cast_const(), root.cast_const()]; + + assert!(vx_expression_pack(names.as_ptr(), parts.as_ptr(), 0, false).is_null()); + assert!(vx_expression_pack(ptr::null(), parts.as_ptr(), 2, false).is_null()); + assert!(vx_expression_pack(names.as_ptr(), ptr::null(), 2, false).is_null()); + + let with_null = [root.cast_const(), ptr::null()]; + assert!(vx_expression_pack(names.as_ptr(), with_null.as_ptr(), 2, false).is_null()); + + vx_expression_free(root); + } + } } diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index 1ef423bb383..8f009c21893 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -15,6 +15,7 @@ mod error; mod expression; mod log; mod ptype; +mod read_at; mod scalar; mod scan; mod session; diff --git a/vortex-ffi/src/read_at.rs b/vortex-ffi/src/read_at.rs new file mode 100644 index 00000000000..e952f57e5d9 --- /dev/null +++ b/vortex-ffi/src/read_at.rs @@ -0,0 +1,204 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A [`VortexReadAt`] backed by caller-supplied C callbacks, so a host with its +//! own configured I/O stack can serve the bytes instead of duplicating that +//! configuration in Vortex. C form of the Java `dev.vortex.io.NativeReadable`, +//! see . + +use std::ffi::c_void; +use std::sync::Arc; + +use futures::FutureExt; +use futures::future::BoxFuture; +use vortex::array::buffer::BufferHandle; +use vortex::buffer::Alignment; +use vortex::buffer::ByteBufferMut; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; +use vortex::io::CoalesceConfig; +use vortex::io::VortexReadAt; +use vortex::io::runtime::BlockingRuntime; +use vortex::io::runtime::Handle; + +use crate::string::vx_view; + +/// Concurrency used when the caller passes 0. Matches the object-store default, +/// since a host that delegates its I/O is usually talking to remote storage. +const DEFAULT_CONCURRENCY: usize = 192; + +/// A random-access byte source implemented by the caller. +/// +/// "read_at" must tolerate concurrent calls from arbitrary threads. The struct +/// is read only during the call it is passed to, but "ctx" and the callbacks +/// must stay valid until "release" runs. +#[repr(C)] +pub struct vx_readat { + /// Opaque caller state passed to every callback. May be NULL. + pub ctx: *mut c_void, + /// Total length of the source in bytes. Must be exact: the footer is read + /// relative to the end, so a wrong length surfaces as a corrupt file. + pub len: u64, + /// Maximum number of concurrent "read_at" calls. 0 selects a default. The + /// cap is per-source, so opening many files multiplies it. + pub concurrency: usize, + /// Optional name, typically the URI, used for cache keys and error messages; + /// it should be stable and unique. Copied. Zero-length means anonymous. + pub name: vx_view, + /// Required. Must write all "length" bytes at "offset" into "dst" and return + /// 0, or return non-zero; success without filling "dst" leaks uninitialized + /// memory into the scan. + pub read_at: Option< + unsafe extern "C" fn(ctx: *mut c_void, offset: u64, dst: *mut u8, length: usize) -> i32, + >, + /// Optional. Called once, after Vortex has dropped the source. + pub release: Option, +} + +/// The callbacks and context, owned for as long as Vortex holds the source. +struct CReadAtInner { + ctx: *mut c_void, + len: u64, + concurrency: usize, + name: Option>, + read_at: unsafe extern "C" fn(*mut c_void, u64, *mut u8, usize) -> i32, + release: Option, +} + +// SAFETY: `vx_readat` requires `read_at` to be callable from any thread and `ctx` +// to stay valid until `release`, which runs only once the last `Arc` drops. +unsafe impl Send for CReadAtInner {} +unsafe impl Sync for CReadAtInner {} + +impl Drop for CReadAtInner { + fn drop(&mut self) { + if let Some(release) = self.release { + // SAFETY: last owner, so every read has finished - reads hold this `Arc`. + unsafe { release(self.ctx) }; + } + } +} + +/// A [`VortexReadAt`] that forwards positional reads to C callbacks. +/// +/// Reads run on the blocking pool: the callback blocks, and the FFI runtime is +/// current-thread, so calling it from a future would stall every other task. +struct CReadAt { + inner: Arc, + handle: Handle, +} + +impl VortexReadAt for CReadAt { + fn uri(&self) -> Option<&Arc> { + self.inner.name.as_ref() + } + + fn coalesce_config(&self) -> Option { + // Host I/O behind the callback is usually remote: prefer fewer, larger reads. + Some(CoalesceConfig::object_storage()) + } + + fn concurrency(&self) -> usize { + self.inner.concurrency + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + let len = self.inner.len; + + async move { Ok(len) }.boxed() + } + + fn read_at( + &self, + offset: u64, + length: usize, + alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + let inner = Arc::clone(&self.inner); + let handle = self.handle.clone(); + + async move { + handle + .spawn_blocking(move || { + let end = offset + .checked_add(length as u64) + .ok_or_else(|| vortex_err!("read {offset}+{length} overflows u64"))?; + if end > inner.len { + vortex_bail!( + "read {offset}..{end} out of bounds for source of length {}", + inner.len + ); + } + + let mut buffer = ByteBufferMut::with_capacity_aligned(length, alignment); + + if length > 0 { + // SAFETY: spare capacity covers `length` bytes; the pointer is + // not retained past the call. + let rc = unsafe { + (inner.read_at)( + inner.ctx, + offset, + buffer.spare_capacity_mut().as_mut_ptr().cast(), + length, + ) + }; + if rc != 0 { + vortex_bail!( + "read_at callback failed with code {rc} for {offset}..{end}" + ); + } + } + + // SAFETY: the callback contract requires all `length` bytes written + // whenever it reports success. + unsafe { buffer.set_len(length) }; + + Ok(BufferHandle::new_host(buffer.freeze())) + }) + .await + } + .boxed() + } +} + +/// Build an owned [`VortexReadAt`] from a caller-supplied `vx_readat`. +/// +/// On error the caller keeps `ctx`: `release` is wired up only once the source +/// exists, so a rejected descriptor never double-frees. +pub(crate) unsafe fn read_at_from_ffi( + reader: *const vx_readat, +) -> VortexResult> { + vortex_ensure!(!reader.is_null(), "null vx_readat"); + + let reader = unsafe { &*reader }; + let read_at = reader + .read_at + .ok_or_else(|| vortex_err!("vx_readat.read_at is required"))?; + + let name = if reader.name.ptr.is_null() || reader.name.len == 0 { + None + } else { + Some(Arc::from(unsafe { reader.name.as_str() }?)) + }; + + let concurrency = if reader.concurrency == 0 { + DEFAULT_CONCURRENCY + } else { + reader.concurrency + }; + + Ok(Arc::new(CReadAt { + inner: Arc::new(CReadAtInner { + ctx: reader.ctx, + len: reader.len, + concurrency, + name, + read_at, + release: reader.release, + }), + handle: crate::RUNTIME.handle(), + })) +} diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index e5e15aa7f05..20829b7732d 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -452,7 +452,9 @@ mod tests { use vortex::session::VortexSession; use vortex_array::VortexSessionExecute; use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::struct_::StructArrayExt; + use vortex_array::arrays::struct_::StructArraySlotsExt; use vortex_array::assert_arrays_eq; use crate::array::vx_array; @@ -465,7 +467,9 @@ mod tests { use crate::expression::vx_expression_free; use crate::expression::vx_expression_get_item; use crate::expression::vx_expression_literal; + use crate::expression::vx_expression_pack; use crate::expression::vx_expression_root; + use crate::expression::vx_expression_row_idx; use crate::scalar::vx_scalar_free; use crate::scalar::vx_scalar_new_u64; use crate::scan::vx_data_source_scan; @@ -760,4 +764,110 @@ mod tests { vx_session_free(session); } } + + /// `row_idx` is only meaningful inside a scan: executed directly it errors, + /// and the scan substitutes it for the row's position in the file. + #[test] + #[cfg_attr(miri, ignore)] + fn test_project_row_idx() { + let mut ctx = array_session().create_execution_ctx(); + unsafe { + let root = vx_expression_root(); + let idx = vx_expression_row_idx(); + let age = vx_expression_get_item(vx_view::from_str("age"), root); + + let names = [vx_view::from_str("idx"), vx_view::from_str("age")]; + let parts = [idx.cast_const(), age.cast_const()]; + let projection = vx_expression_pack(names.as_ptr(), parts.as_ptr(), 2, false); + assert!(!projection.is_null()); + + let opts = vx_scan_options { + projection, + ..Default::default() + }; + let (array, _) = scan(&raw const opts); + { + let array = vx_array::as_ref(array) + .clone() + .execute::(&mut ctx) + .unwrap(); + let idx = array + .fields() + .get(0) + .unwrap() + .clone() + .execute::(&mut ctx) + .unwrap(); + assert_eq!( + idx.to_buffer::().as_slice(), + (0..SAMPLE_ROWS as u64).collect::>().as_slice() + ); + } + vx_array_free(array); + + vx_expression_free(projection); + vx_expression_free(age); + vx_expression_free(idx); + vx_expression_free(root); + } + } + + /// The point of projecting `row_idx`: a filter drops the rows around a + /// match, and the surviving rows still carry their original position. + #[test] + #[cfg_attr(miri, ignore)] + fn test_row_idx_survives_filter() { + let mut ctx = array_session().create_execution_ctx(); + unsafe { + let root = vx_expression_root(); + let idx = vx_expression_row_idx(); + + let names = [vx_view::from_str("idx")]; + let parts = [idx.cast_const()]; + let projection = vx_expression_pack(names.as_ptr(), parts.as_ptr(), 1, false); + + let age_expr = vx_expression_get_item(vx_view::from_str("age"), root); + let value = vx_scalar_new_u64(100, false); + let mut error = ptr::null_mut(); + let lit_100 = vx_expression_literal(value, &raw mut error); + assert_no_error(error); + vx_scalar_free(value); + let filter = + vx_expression_binary(vx_binary_operator::VX_OPERATOR_GTE, age_expr, lit_100); + + let opts = vx_scan_options { + projection, + filter, + ..Default::default() + }; + let (array, _) = scan(&raw const opts); + { + let array = vx_array::as_ref(array) + .clone() + .execute::(&mut ctx) + .unwrap(); + let idx = array + .fields() + .get(0) + .unwrap() + .clone() + .execute::(&mut ctx) + .unwrap(); + // age == row position in write_sample, so the filter keeps + // exactly rows 100..SAMPLE_ROWS, at their original positions. + assert_eq!( + idx.to_buffer::().as_slice(), + (100..SAMPLE_ROWS as u64).collect::>().as_slice() + ); + } + vx_array_free(array); + + vx_expression_free(filter); + vx_expression_free(lit_100); + vx_expression_free(age_expr); + vx_expression_free(projection); + vx_expression_free(idx); + vx_expression_free(root); + } + } } From be7f8d7ab6e68af64883daa09eee2571a9cfd605 Mon Sep 17 00:00:00 2001 From: Martin Prammer Date: Mon, 21 Sep 2026 15:38:19 -0400 Subject: [PATCH 2/2] Detect short reads, bound callback concurrency, and release at free Addresses review feedback on the C FFI callback reader. The read callback now returns the number of bytes written rather than a status code, so a host that reports success without filling the buffer is rejected instead of handing uninitialized memory to the scan. The Java bindings already guard this via `remaining()`; the previous C signature made it undetectable. A source was owned by the segment source's spawned read driver, which the FFI's current-thread runtime never polled again once a call returned, so `release` ran at an unpredictable later point or not at all. Hosts freeing their context after `vx_data_source_free` could therefore hit a use-after-free. `CurrentThreadRuntime::drain` runs queued tasks to a stop, and the FFI drains after dropping a data source and after a failed open, so `release` now runs before the call that drops the source returns. Callback concurrency is additionally capped process-wide rather than only per source, matching how the Java bindings bound upcalls across a filesystem, and a name view of NULL with a non-zero length is rejected rather than silently treated as anonymous. Co-authored-by: Claude Signed-off-by: Martin Prammer --- Cargo.lock | 1 + vortex-ffi/Cargo.toml | 1 + vortex-ffi/cinclude/vortex.h | 17 +-- vortex-ffi/src/data_source.rs | 229 ++++++++++++++++++++++++------- vortex-ffi/src/lib.rs | 11 +- vortex-ffi/src/read_at.rs | 58 +++++--- vortex-io/src/runtime/current.rs | 8 ++ 7 files changed, 245 insertions(+), 80 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ca9cd626f1b..f31f5a1107a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11165,6 +11165,7 @@ dependencies = [ "arrow-array 59.3.0", "arrow-schema 59.3.0", "async-fs", + "async-lock", "bytes", "cbindgen", "futures", diff --git a/vortex-ffi/Cargo.toml b/vortex-ffi/Cargo.toml index 75d644571d1..677b5cff0c7 100644 --- a/vortex-ffi/Cargo.toml +++ b/vortex-ffi/Cargo.toml @@ -20,6 +20,7 @@ categories = { workspace = true } all-features = true [dependencies] +async-lock = { workspace = true } arrow-array = { workspace = true } arrow-schema = { workspace = true } async-fs = { workspace = true } diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index 043cb390f3a..c95abc685a9 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -561,8 +561,8 @@ typedef struct { */ uint64_t len; /** - * Maximum number of concurrent "read_at" calls. 0 selects a default. The - * cap is per-source, so opening many files multiplies it. + * Maximum number of concurrent "read_at" calls for this source. 0 selects a + * default. A process-wide ceiling applies across all sources as well. */ size_t concurrency; /** @@ -571,13 +571,13 @@ typedef struct { */ vx_view name; /** - * Required. Must write all "length" bytes at "offset" into "dst" and return - * 0, or return non-zero; success without filling "dst" leaks uninitialized - * memory into the scan. + * Required. Writes "length" bytes at "offset" into "dst" and returns the + * count written; a short count or a negative value fails the read. */ - int32_t (*read_at)(void *ctx, uint64_t offset, uint8_t *dst, size_t length); + int64_t (*read_at)(void *ctx, uint64_t offset, uint8_t *dst, size_t length); /** - * Optional. Called once, after Vortex has dropped the source. + * Optional. Called once, before the call that drops the source returns - + * on that thread, or on a worker thread if any are configured. */ void (*release)(void *ctx); } vx_readat; @@ -937,7 +937,8 @@ vx_data_source_new_buffer(const vx_session *session, const void *buffer, size_t * * "reader" is read during this call only; its callbacks and context must stay * valid until "release" runs. A rejected descriptor leaves ownership with the - * caller and never calls "release"; once accepted, "release" always runs. + * caller and never calls "release"; once accepted, "release" runs before this + * call returns if it fails, and otherwise before vx_data_source_free returns. * * On error, returns NULL and sets "err". */ diff --git a/vortex-ffi/src/data_source.rs b/vortex-ffi/src/data_source.rs index 6339e7fe20e..e06d63da3a3 100644 --- a/vortex-ffi/src/data_source.rs +++ b/vortex-ffi/src/data_source.rs @@ -50,7 +50,8 @@ box_wrapper!( /// /// Copying a vx_data_source via vx_data_source_clone is a cheap operation. MultiLayoutDataSource, - vx_data_source + vx_data_source, + drain_on_free ); /// Options for creating a data source. @@ -202,7 +203,8 @@ pub unsafe extern "C-unwind" fn vx_data_source_new_buffer( /// /// "reader" is read during this call only; its callbacks and context must stay /// valid until "release" runs. A rejected descriptor leaves ownership with the -/// caller and never calls "release"; once accepted, "release" always runs. +/// caller and never calls "release"; once accepted, "release" runs before this +/// call returns if it fails, and otherwise before vx_data_source_free returns. /// /// On error, returns NULL and sets "err". #[unsafe(no_mangle)] @@ -211,7 +213,7 @@ pub unsafe extern "C-unwind" fn vx_data_source_new_readat( reader: *const vx_readat, err: *mut *mut vx_error, ) -> *const vx_data_source { - try_or(err, ptr::null(), || { + let ds = try_or(err, ptr::null(), || { vortex_ensure!(!session.is_null()); let session = vx_session::as_ref(session); @@ -231,7 +233,14 @@ pub unsafe extern "C-unwind" fn vx_data_source_new_readat( ); Ok(vx_data_source::new(ds)) - }) + }); + + // A failure past descriptor validation leaves the reader owned by a spawned + // task with no data source to free, so release it here instead. + if ds.is_null() { + RUNTIME.drain(); + } + ds } /// Increase reference count on vx_data_source @@ -282,6 +291,13 @@ mod tests { use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; + use vortex::array::array_session; + use vortex::array::arrays::StructArray; + use vortex::array::assert_arrays_eq; + use vortex_array::VortexSessionExecute; + + use crate::array::vx_array; + use crate::array::vx_array_free; use crate::data_source::vx_data_source_dtype; use crate::data_source::vx_data_source_free; use crate::data_source::vx_data_source_get_row_count; @@ -292,8 +308,13 @@ mod tests { use crate::dtype::vx_dtype; use crate::dtype::vx_dtype_free; use crate::read_at::vx_readat; + use crate::scan::vx_data_source_scan; use crate::scan::vx_estimate; use crate::scan::vx_estimate_type; + use crate::scan::vx_partition_free; + use crate::scan::vx_partition_next; + use crate::scan::vx_scan_free; + use crate::scan::vx_scan_next_partition; use crate::session::vx_session; use crate::session::vx_session_free; use crate::session::vx_session_new; @@ -489,19 +510,19 @@ mod tests { offset: u64, dst: *mut u8, length: usize, - ) -> i32 { + ) -> i64 { let ctx = unsafe { &*ctx.cast::() }; let Ok(start) = usize::try_from(offset) else { - return 1; + return -1; }; let Some(end) = start.checked_add(length) else { - return 1; + return -1; }; if end > ctx.data.len() { - return 1; + return -1; } unsafe { ptr::copy_nonoverlapping(ctx.data[start..end].as_ptr(), dst, length) }; - 0 + length as i64 } unsafe extern "C" fn release_cb(ctx: *mut c_void) { @@ -516,22 +537,65 @@ mod tests { _offset: u64, _dst: *mut u8, _length: usize, - ) -> i32 { + ) -> i64 { -7 } - fn readat_ctx(session: *const vx_session) -> (Box, u64, tempfile::NamedTempFile) { - let (sample, _) = unsafe { write_sample(session) }; - let data = read(sample.path()).unwrap(); - let len = data.len() as u64; - ( - Box::new(ReadAtCtx { - data, - releases: AtomicUsize::new(0), - }), - len, - sample, - ) + /// Reports success while leaving the buffer untouched - what a naive + /// `pread(2)` wrapper does at EOF or on EINTR. + unsafe extern "C" fn short_read_at_cb( + _ctx: *mut c_void, + _offset: u64, + _dst: *mut u8, + length: usize, + ) -> i64 { + (length / 2) as i64 + } + + /// A written sample file, the callback state that serves it, and the array + /// it should read back as. + struct Sample { + ctx: Box, + len: u64, + array: StructArray, + _file: tempfile::NamedTempFile, + } + + impl Sample { + fn new(session: *const vx_session) -> Self { + let (file, array) = unsafe { write_sample(session) }; + let data = read(file.path()).unwrap(); + let len = data.len() as u64; + Self { + ctx: Box::new(ReadAtCtx { + data, + releases: AtomicUsize::new(0), + }), + len, + array, + _file: file, + } + } + + /// An anonymous descriptor over this sample. Tests needing a name or a + /// missing callback adjust the returned struct. + fn reader( + &self, + read_at: unsafe extern "C" fn(*mut c_void, u64, *mut u8, usize) -> i64, + ) -> vx_readat { + vx_readat { + ctx: (&raw const *self.ctx).cast::().cast_mut(), + len: self.len, + concurrency: 0, + name: vx_view::from_str(""), + read_at: Some(read_at), + release: Some(release_cb), + } + } + + fn releases(&self) -> usize { + self.ctx.releases.load(Ordering::SeqCst) + } } #[test] @@ -539,16 +603,9 @@ mod tests { fn test_create_readat() { unsafe { let session = vx_session_new(); - let (ctx, len, _sample) = readat_ctx(session); - - let reader = vx_readat { - ctx: (&raw const *ctx).cast::().cast_mut(), - len, - concurrency: 0, - name: vx_view::from_str("test://sample.vortex"), - read_at: Some(read_at_cb), - release: Some(release_cb), - }; + let sample = Sample::new(session); + let mut reader = sample.reader(read_at_cb); + reader.name = vx_view::from_str("test://sample.vortex"); let mut error = ptr::null_mut(); let ds = vx_data_source_new_readat(session, &raw const reader, &raw mut error); @@ -561,13 +618,13 @@ mod tests { assert_eq!(row_count.r#type, vx_estimate_type::VX_ESTIMATE_EXACT); assert_eq!(row_count.estimate, SAMPLE_ROWS as u64); - assert_eq!(ctx.releases.load(Ordering::SeqCst), 0); + assert_eq!(sample.releases(), 0); vx_dtype_free(ffi_dtype); vx_data_source_free(ds); vx_session_free(session); - assert_eq!(ctx.releases.load(Ordering::SeqCst), 1); + assert_eq!(sample.releases(), 1); } } @@ -578,7 +635,7 @@ mod tests { fn test_create_readat_invalid() { unsafe { let session = vx_session_new(); - let (ctx, len, _sample) = readat_ctx(session); + let sample = Sample::new(session); let mut error = ptr::null_mut(); let ds = vx_data_source_new_readat(ptr::null(), ptr::null(), &raw mut error); @@ -592,19 +649,13 @@ mod tests { // read_at is required. let mut error = ptr::null_mut(); - let reader = vx_readat { - ctx: (&raw const *ctx).cast::().cast_mut(), - len, - concurrency: 0, - name: vx_view::from_str(""), - read_at: None, - release: Some(release_cb), - }; + let mut reader = sample.reader(read_at_cb); + reader.read_at = None; let ds = vx_data_source_new_readat(session, &raw const reader, &raw mut error); assert_error(error); assert!(ds.is_null()); - assert_eq!(ctx.releases.load(Ordering::SeqCst), 0); + assert_eq!(sample.releases(), 0); vx_session_free(session); } } @@ -616,23 +667,95 @@ mod tests { fn test_create_readat_callback_failure() { unsafe { let session = vx_session_new(); - let (ctx, len, _sample) = readat_ctx(session); + let sample = Sample::new(session); + let reader = sample.reader(failing_read_at_cb); - let reader = vx_readat { - ctx: (&raw const *ctx).cast::().cast_mut(), - len, - concurrency: 0, - name: vx_view::from_str(""), - read_at: Some(failing_read_at_cb), - release: Some(release_cb), + let mut error = ptr::null_mut(); + let ds = vx_data_source_new_readat(session, &raw const reader, &raw mut error); + assert_error(error); + assert!(ds.is_null()); + assert_eq!(sample.releases(), 1); + + vx_session_free(session); + } + } + + /// A short read must fail rather than reach `set_len`, which would expose + /// the unwritten tail of the buffer to the scan. + #[test] + #[cfg_attr(miri, ignore)] + fn test_create_readat_short_read() { + unsafe { + let session = vx_session_new(); + let sample = Sample::new(session); + let reader = sample.reader(short_read_at_cb); + + let mut error = ptr::null_mut(); + let ds = vx_data_source_new_readat(session, &raw const reader, &raw mut error); + assert_error(error); + assert!(ds.is_null()); + + // The descriptor was accepted before the read failed, so Vortex owns + // the context and must have released it before returning. + assert_eq!(sample.releases(), 1); + + vx_session_free(session); + } + } + + /// A name view of NULL with a non-zero length is a caller bug, not an + /// anonymous source. + #[test] + #[cfg_attr(miri, ignore)] + fn test_create_readat_invalid_name() { + unsafe { + let session = vx_session_new(); + let sample = Sample::new(session); + let mut reader = sample.reader(read_at_cb); + reader.name = vx_view { + ptr: ptr::null(), + len: 5, }; let mut error = ptr::null_mut(); let ds = vx_data_source_new_readat(session, &raw const reader, &raw mut error); assert_error(error); assert!(ds.is_null()); - assert_eq!(ctx.releases.load(Ordering::SeqCst), 1); + assert_eq!(sample.releases(), 0); + + vx_session_free(session); + } + } + + /// Scan data through the callbacks, not just the footer: this is the path + /// that exercises data segments, coalescing and concurrent reads. + #[test] + #[cfg_attr(miri, ignore)] + fn test_scan_readat() { + let mut ctx_exec = array_session().create_execution_ctx(); + unsafe { + let session = vx_session_new(); + let sample = Sample::new(session); + let reader = sample.reader(read_at_cb); + let mut error = ptr::null_mut(); + let ds = vx_data_source_new_readat(session, &raw const reader, &raw mut error); + assert_no_error(error); + + let scan = vx_data_source_scan(ds, ptr::null(), ptr::null_mut(), &raw mut error); + assert_no_error(error); + let partition = vx_scan_next_partition(scan, &raw mut error); + assert_no_error(error); + let array = vx_partition_next(partition, &raw mut error); + assert_no_error(error); + assert!(!array.is_null()); + + assert_arrays_eq!(vx_array::as_ref(array), sample.array, &mut ctx_exec); + + vx_array_free(array); + vx_partition_free(partition); + vx_scan_free(scan); + vx_data_source_free(ds); vx_session_free(session); } } diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index 8f009c21893..b82b144f745 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -70,6 +70,14 @@ static POOL: LazyLock = LazyLock::new(|| RUNTIME.new_po #[macro_export] macro_rules! box_wrapper { ($(#[$meta:meta])* $T:ty, $ffi_ident:ident) => { + $crate::box_wrapper!(@impl $(#[$meta])* $T, $ffi_ident, {}); + }; + // Drain the runtime after freeing, for types that own a spawned task. See + // `CurrentThreadRuntime::drain`. + ($(#[$meta:meta])* $T:ty, $ffi_ident:ident, drain_on_free) => { + $crate::box_wrapper!(@impl $(#[$meta])* $T, $ffi_ident, { $crate::RUNTIME.drain(); }); + }; + (@impl $(#[$meta:meta])* $T:ty, $ffi_ident:ident, $after_free:block) => { paste::paste! { $(#[$meta])* pub struct $ffi_ident($T); @@ -121,7 +129,8 @@ macro_rules! box_wrapper { #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn [<$ffi_ident _free>](ptr: *const $ffi_ident) { if !ptr.is_null() { - std::mem::drop(unsafe { Box::from_raw(ptr.cast::<$T>().cast_mut()) }) + std::mem::drop(unsafe { Box::from_raw(ptr.cast::<$T>().cast_mut()) }); + $after_free } } } diff --git a/vortex-ffi/src/read_at.rs b/vortex-ffi/src/read_at.rs index e952f57e5d9..66856ea9b6f 100644 --- a/vortex-ffi/src/read_at.rs +++ b/vortex-ffi/src/read_at.rs @@ -8,7 +8,9 @@ use std::ffi::c_void; use std::sync::Arc; +use std::sync::LazyLock; +use async_lock::Semaphore; use futures::FutureExt; use futures::future::BoxFuture; use vortex::array::buffer::BufferHandle; @@ -29,6 +31,11 @@ use crate::string::vx_view; /// since a host that delegates its I/O is usually talking to remote storage. const DEFAULT_CONCURRENCY: usize = 192; +/// Ceiling on callbacks in flight across every source in the process, since +/// per-source limits otherwise multiply. The Java bindings cap upcalls likewise. +static READ_LIMITER: LazyLock> = + LazyLock::new(|| Arc::new(Semaphore::new(DEFAULT_CONCURRENCY))); + /// A random-access byte source implemented by the caller. /// /// "read_at" must tolerate concurrent calls from arbitrary threads. The struct @@ -41,19 +48,19 @@ pub struct vx_readat { /// Total length of the source in bytes. Must be exact: the footer is read /// relative to the end, so a wrong length surfaces as a corrupt file. pub len: u64, - /// Maximum number of concurrent "read_at" calls. 0 selects a default. The - /// cap is per-source, so opening many files multiplies it. + /// Maximum number of concurrent "read_at" calls for this source. 0 selects a + /// default. A process-wide ceiling applies across all sources as well. pub concurrency: usize, /// Optional name, typically the URI, used for cache keys and error messages; /// it should be stable and unique. Copied. Zero-length means anonymous. pub name: vx_view, - /// Required. Must write all "length" bytes at "offset" into "dst" and return - /// 0, or return non-zero; success without filling "dst" leaks uninitialized - /// memory into the scan. + /// Required. Writes "length" bytes at "offset" into "dst" and returns the + /// count written; a short count or a negative value fails the read. pub read_at: Option< - unsafe extern "C" fn(ctx: *mut c_void, offset: u64, dst: *mut u8, length: usize) -> i32, + unsafe extern "C" fn(ctx: *mut c_void, offset: u64, dst: *mut u8, length: usize) -> i64, >, - /// Optional. Called once, after Vortex has dropped the source. + /// Optional. Called once, before the call that drops the source returns - + /// on that thread, or on a worker thread if any are configured. pub release: Option, } @@ -63,7 +70,7 @@ struct CReadAtInner { len: u64, concurrency: usize, name: Option>, - read_at: unsafe extern "C" fn(*mut c_void, u64, *mut u8, usize) -> i32, + read_at: unsafe extern "C" fn(*mut c_void, u64, *mut u8, usize) -> i64, release: Option, } @@ -120,8 +127,18 @@ impl VortexReadAt for CReadAt { let handle = self.handle.clone(); async move { + // Take a permit before occupying a blocking thread. The lock-free path + // barges ahead of waiters: the limiter must not become the bottleneck. + let permit = match READ_LIMITER.try_acquire_arc() { + Some(permit) => permit, + None => READ_LIMITER.acquire_arc().await, + }; + handle .spawn_blocking(move || { + // Cancelling the read drops the task handle but cannot interrupt a + // callback already running, so the permit stays with the work. + let _permit = permit; let end = offset .checked_add(length as u64) .ok_or_else(|| vortex_err!("read {offset}+{length} overflows u64"))?; @@ -137,7 +154,7 @@ impl VortexReadAt for CReadAt { if length > 0 { // SAFETY: spare capacity covers `length` bytes; the pointer is // not retained past the call. - let rc = unsafe { + let written = unsafe { (inner.read_at)( inner.ctx, offset, @@ -145,15 +162,21 @@ impl VortexReadAt for CReadAt { length, ) }; - if rc != 0 { + if written < 0 { vortex_bail!( - "read_at callback failed with code {rc} for {offset}..{end}" + "read_at callback failed with code {written} for {offset}..{end}" + ); + } + // A host returning short without saying so would leave the tail of + // `dst` uninitialized for the scan to read. + if written as u64 != length as u64 { + vortex_bail!( + "read_at callback wrote {written} of {length} bytes for {offset}..{end}" ); } } - // SAFETY: the callback contract requires all `length` bytes written - // whenever it reports success. + // SAFETY: the callback reported writing all `length` bytes, checked above. unsafe { buffer.set_len(length) }; Ok(BufferHandle::new_host(buffer.freeze())) @@ -178,11 +201,10 @@ pub(crate) unsafe fn read_at_from_ffi( .read_at .ok_or_else(|| vortex_err!("vx_readat.read_at is required"))?; - let name = if reader.name.ptr.is_null() || reader.name.len == 0 { - None - } else { - Some(Arc::from(unsafe { reader.name.as_str() }?)) - }; + // `as_str` rejects a null pointer with a non-zero length, which a manual + // null check here would silently accept as anonymous. + let name = unsafe { reader.name.as_str() }?; + let name = (!name.is_empty()).then(|| Arc::from(name)); let concurrency = if reader.concurrency == 0 { DEFAULT_CONCURRENCY diff --git a/vortex-io/src/runtime/current.rs b/vortex-io/src/runtime/current.rs index bb1c369af77..ea292ed8ab9 100644 --- a/vortex-io/src/runtime/current.rs +++ b/vortex-io/src/runtime/current.rs @@ -52,6 +52,14 @@ impl CurrentThreadRuntime { CurrentThreadWorkerPool::new(Arc::clone(&self.executor)) } + /// Run already-queued tasks until none is immediately runnable. + /// + /// Spawned tasks only advance inside `block_on`, so one holding the last + /// reference to a resource keeps it alive until the runtime is driven again. + pub fn drain(&self) { + while self.executor.async_executor().try_tick() {} + } + /// Returns an iterator wrapper around a stream, blocking the current thread for each item. /// /// ## Multi-threaded Usage