From 44e3a5ddaadc0d08c2413cb4c75b167977de53ca Mon Sep 17 00:00:00 2001 From: kid Date: Fri, 18 Sep 2026 12:04:13 +0800 Subject: [PATCH 1/2] feat: add index build progress callbacks to segment builders Bridge lance core's IndexBuildProgress trait to a C callback on LanceIndexSegmentBuilder (E3 of the distributed-build track): - lance_index_segment_builder_set_progress_callback with START/PROGRESS/ COMPLETE events; total/unit meaningful on START, completed on PROGRESS. - Thread-safe/reentrant, non-blocking, no-reentrancy contract; callback and context must outlive the builder (conservative: spawned worker tasks and error paths may still deliver events). - Advisory only: the callback cannot abort the build. - C++ fluent IndexSegmentBuilder::progress_callback wrapper. - Direct async-trait dependency for the trait impl (already transitive via lance-index). Refs #55. --- Cargo.lock | 1 + Cargo.toml | 4 + include/lance/lance.h | 76 +++++ include/lance/lance.hpp | 19 ++ src/index_segment.rs | 177 ++++++++++- tests/c_api_test.rs | 628 +++++++++++++++++++++++++++++++++++++ tests/cpp/test_c_api.c | 99 ++++++ tests/cpp/test_cpp_api.cpp | 61 ++++ 8 files changed, 1064 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 3135cef..6cf44e4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3788,6 +3788,7 @@ dependencies = [ "arrow", "arrow-array", "arrow-schema", + "async-trait", "chrono", "datafusion", "futures", diff --git a/Cargo.toml b/Cargo.toml index 330bfd4..489a609 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,10 @@ rust-version = "1.91.0" crate-type = ["cdylib", "staticlib", "rlib"] [dependencies] +# Direct for the `#[async_trait::async_trait]` impl of lance's +# `lance_index::progress::IndexBuildProgress`; already in the graph +# transitively via lance-index. +async-trait = "0.1" lance = { git = "https://github.com/lance-format/lance.git", rev = "356acb0d333c96e970f6f84b97314fc5bc4193f7", features = ["substrait"] } lance-core = { git = "https://github.com/lance-format/lance.git", rev = "356acb0d333c96e970f6f84b97314fc5bc4193f7" } lance-file = { git = "https://github.com/lance-format/lance.git", rev = "356acb0d333c96e970f6f84b97314fc5bc4193f7" } diff --git a/include/lance/lance.h b/include/lance/lance.h index e9a8a84..136910a 100644 --- a/include/lance/lance.h +++ b/include/lance/lance.h @@ -1816,6 +1816,82 @@ int32_t lance_index_segment_builder_execute_uncommitted( size_t* out_len ); +/** + * Event codes for LanceIndexBuildProgressCallback, passed as the `event` + * argument. Exactly one stage is active at a time: a stage's + * LANCE_INDEX_BUILD_PROGRESS_STAGE_COMPLETE is always delivered before the + * next stage's LANCE_INDEX_BUILD_PROGRESS_STAGE_START. + */ +typedef enum { + LANCE_INDEX_BUILD_PROGRESS_STAGE_START = 0, + LANCE_INDEX_BUILD_PROGRESS_STAGE_PROGRESS = 1, + LANCE_INDEX_BUILD_PROGRESS_STAGE_COMPLETE = 2, +} LanceIndexBuildProgressEvent; + +/** + * Receives index build progress events while + * lance_index_segment_builder_execute_uncommitted runs. + * + * `stage` is non-NULL, NUL-terminated, and borrowed: it is valid only for the + * duration of this call. `unit` is non-NULL and NUL-terminated, but is the + * empty string ("") for LANCE_INDEX_BUILD_PROGRESS_STAGE_PROGRESS and + * LANCE_INDEX_BUILD_PROGRESS_STAGE_COMPLETE. The parameter mapping is: + * + * - LANCE_INDEX_BUILD_PROGRESS_STAGE_START: `stage` is the stage name, + * `total` is the number of work units (0 = unknown), `unit` describes what + * is being counted (e.g. "partitions", "batches", "rows"; "" = unknown), + * and `completed` is 0. + * - LANCE_INDEX_BUILD_PROGRESS_STAGE_PROGRESS: `total` is 0, `unit` is "", + * and `completed` is the number of units completed so far. + * - LANCE_INDEX_BUILD_PROGRESS_STAGE_COMPLETE: `total` is 0, `unit` is "", + * and `completed` is 0. + * + * The callback is invoked from lance-c's internal tokio runtime worker + * threads. Certain stages report progress concurrently from parallel worker + * tasks, so the callback MUST be thread-safe and reentrant. It must be + * non-blocking and must not call back into any `lance_*` function (no + * reentrancy). + * + * The callback is invoked without a panic guard: it must return normally, + * because unwinding or throwing across this boundary can abort the host + * process. The callback cannot abort the build; progress reporting is + * advisory and diagnostic and cannot affect the build outcome. + */ +typedef void (*LanceIndexBuildProgressCallback)( + void* callback_ctx, + int32_t event, + const char* stage, + uint64_t total, + const char* unit, + uint64_t completed +); + +/** + * Register the index-build progress callback for a segment builder. + * + * Must be called before the builder is executed; the builder is single-use, + * so calling it after lance_index_segment_builder_execute_uncommitted has + * been called (even if that call failed) returns -1. `callback` must not be + * NULL. `callback_ctx` may be NULL and is passed through to the callback + * opaquely. Setting a callback replaces any previously set callback. + * + * `callback` and `callback_ctx` must remain valid and safe to invoke until the + * builder is released with lance_index_segment_builder_free. Invocations only + * occur while lance_index_segment_builder_execute_uncommitted is executing. + * Lance core may deliver events from spawned worker tasks, and error paths can + * detach them before they finish; keep callback and callback_ctx valid until + * the builder is released rather than retiring them when execute returns. See + * LanceIndexBuildProgressCallback for the full threading and reentrancy + * contract. + * + * @return 0 on success, -1 on error. + */ +int32_t lance_index_segment_builder_set_progress_callback( + LanceIndexSegmentBuilder* builder, + LanceIndexBuildProgressCallback callback, + void* callback_ctx +); + /** Free metadata bytes returned by an uncommitted segment build. NULL-safe. */ void lance_free_bytes(uint8_t* bytes); diff --git a/include/lance/lance.hpp b/include/lance/lance.hpp index fcb1197..752b4e1 100644 --- a/include/lance/lance.hpp +++ b/include/lance/lance.hpp @@ -1280,6 +1280,25 @@ class IndexSegmentBuilder { return std::vector(guard.bytes, guard.bytes + len); } + /// Register a non-null index-build progress callback. Must be called + /// before execute_uncommitted; the builder is single-use. The callback is + /// invoked from internal worker threads and may be called concurrently + /// from parallel worker tasks, so it must be thread-safe, non-blocking, + /// and must not re-enter any lance_* function. It must return normally; + /// unwinding or throwing across this boundary can abort the host process. + /// The callback and a non-null context must remain valid until the + /// builder is released: core may deliver events from spawned worker + /// tasks, and error paths can detach them before they finish, so do not + /// retire them when execute_uncommitted returns. Progress reporting is + /// advisory and cannot affect the build outcome. + IndexSegmentBuilder& progress_callback(LanceIndexBuildProgressCallback callback, + void* callback_ctx) { + if (lance_index_segment_builder_set_progress_callback(handle_.get(), callback, + callback_ctx) != 0) + check_error(); + return *this; + } + LanceIndexSegmentBuilder* c_handle() { return handle_.get(); } }; diff --git a/src/index_segment.rs b/src/index_segment.rs index 26db56c..e2acb36 100644 --- a/src/index_segment.rs +++ b/src/index_segment.rs @@ -4,7 +4,7 @@ //! Distributed index segment build and metadata C API. use std::collections::HashSet; -use std::ffi::{CStr, CString, c_char}; +use std::ffi::{CStr, CString, c_char, c_void}; use std::ptr; use std::slice; use std::sync::Arc; @@ -16,6 +16,7 @@ use chrono::{DateTime, Utc}; use lance::Dataset; use lance::index::DatasetIndexExt; use lance_core::{Error, Result}; +use lance_index::progress::IndexBuildProgress; use lance_index::scalar::ScalarIndexParams; use lance_table::format::{IndexMetadata, pb}; use prost::Message; @@ -117,6 +118,7 @@ pub struct LanceIndexSegmentBuilder { kind: SegmentKind, fragment_ids: Option>, index_uuid: Option, + progress: Option, executed: bool, } @@ -132,6 +134,120 @@ enum SegmentKind { }, } +/// Event code delivered to `LanceIndexBuildProgressCallback`; mirrors +/// `LanceIndexBuildProgressEvent` in lance.h. +#[repr(i32)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum IndexBuildProgressEvent { + /// A build stage began. + StageStart = 0, + /// Work units completed within the active stage. + StageProgress = 1, + /// The active stage finished. + StageComplete = 2, +} + +/// Index build progress callback bridged to a C function pointer. +/// +/// The public C type is an `Option` of the raw function pointer so NULL can be +/// rejected at the setter; once stored here the callback is known non-NULL. +pub type LanceIndexBuildProgressCallback = Option< + unsafe extern "C" fn( + callback_ctx: *mut c_void, + event: i32, + stage: *const c_char, + total: u64, + unit: *const c_char, + completed: u64, + ), +>; + +#[derive(Debug, Clone)] +struct SendIndexBuildProgressCallback { + callback: unsafe extern "C" fn( + callback_ctx: *mut c_void, + event: i32, + stage: *const c_char, + total: u64, + unit: *const c_char, + completed: u64, + ), + ctx: *mut c_void, +} + +// SAFETY: The C API requires the callback and its context to remain valid and +// safe to invoke until the segment builder is freed, and the builder is +// single-use, so no invocation can outlive the builder. Certain build stages +// report progress concurrently from parallel worker tasks, so the callback +// may be invoked concurrently; the C contract requires it to be thread-safe. +unsafe impl Send for SendIndexBuildProgressCallback {} +unsafe impl Sync for SendIndexBuildProgressCallback {} + +/// Build a `CString` for a stage or unit name without panicking. +/// +/// Lance core never emits interior NUL bytes, but the FFI contract is +/// panic-free, so any that appear are stripped rather than aborting the +/// process. The strip removes every NUL byte, so `CString::new` cannot fail; +/// the `map_err` below is defensive-only and unreachable. +fn progress_c_string(value: &str) -> Result { + let cleaned = if value.contains('\0') { + value.replace('\0', "") + } else { + value.to_owned() + }; + CString::new(cleaned) + .map_err(|_| Error::internal("index build progress stage/unit contained an interior NUL")) +} + +#[async_trait::async_trait] +impl IndexBuildProgress for SendIndexBuildProgressCallback { + async fn stage_start(&self, stage: &str, total: Option, unit: &str) -> Result<()> { + let stage = progress_c_string(stage)?; + let unit = progress_c_string(unit)?; + unsafe { + (self.callback)( + self.ctx, + IndexBuildProgressEvent::StageStart as i32, + stage.as_ptr(), + total.unwrap_or(0), + unit.as_ptr(), + 0, + ) + }; + Ok(()) + } + + async fn stage_progress(&self, stage: &str, completed: u64) -> Result<()> { + let stage = progress_c_string(stage)?; + unsafe { + (self.callback)( + self.ctx, + IndexBuildProgressEvent::StageProgress as i32, + stage.as_ptr(), + 0, + c"".as_ptr(), + completed, + ) + }; + Ok(()) + } + + async fn stage_complete(&self, stage: &str) -> Result<()> { + let stage = progress_c_string(stage)?; + unsafe { + (self.callback)( + self.ctx, + IndexBuildProgressEvent::StageComplete as i32, + stage.as_ptr(), + 0, + c"".as_ptr(), + 0, + ) + }; + Ok(()) + } +} + /// Opaque parsed index segment metadata. pub struct LanceIndexSegmentMetadata { metadata: IndexMetadata, @@ -317,6 +433,7 @@ unsafe fn new_scalar_builder_inner( }, fragment_ids: parsed.fragment_ids, index_uuid: parsed.index_uuid, + progress: None, executed: false, }))) } @@ -930,6 +1047,7 @@ unsafe fn new_vector_builder_inner( }, fragment_ids: parsed.fragment_ids, index_uuid: parsed.index_uuid, + progress: None, executed: false, }))) } @@ -988,6 +1106,9 @@ unsafe fn execute_uncommitted_inner( if let Some(index_uuid) = builder.index_uuid { core_builder = core_builder.index_uuid(index_uuid); } + if let Some(progress) = builder.progress.clone() { + core_builder = core_builder.progress(Arc::new(progress)); + } block_on(core_builder.execute_uncommitted())? } SegmentKind::Vector { @@ -1011,6 +1132,9 @@ unsafe fn execute_uncommitted_inner( if let Some(index_uuid) = builder.index_uuid { core_builder = core_builder.index_uuid(index_uuid); } + if let Some(progress) = builder.progress.clone() { + core_builder = core_builder.progress(Arc::new(progress)); + } // Core's train=false means "create an empty index". Model presence // itself controls whether IVF/PQ training is skipped. block_on(core_builder.train(true).execute_uncommitted())? @@ -1032,6 +1156,57 @@ unsafe fn execute_uncommitted_inner( Ok(0) } +/// Install (or replace) the index-build progress callback for a segment +/// builder. +/// +/// The callback is invoked from lance-c's internal tokio runtime worker +/// threads while `lance_index_segment_builder_execute_uncommitted` runs. +/// Lance core may deliver events from spawned worker tasks, and error paths +/// can detach them before they finish, so the callback and `callback_ctx` +/// must remain valid until the builder is freed rather than being retired +/// when `execute_uncommitted` returns. Certain stages report progress +/// concurrently from parallel worker tasks, so the callback must be +/// thread-safe and reentrant, must be non-blocking, and must not call back +/// into any `lance_*` function. It is invoked without a panic guard: it must +/// return normally, because unwinding or throwing across this boundary can +/// abort the host process. Progress reporting is advisory and cannot affect +/// the build outcome or abort the build. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_index_segment_builder_set_progress_callback( + builder: *mut LanceIndexSegmentBuilder, + callback: LanceIndexBuildProgressCallback, + callback_ctx: *mut c_void, +) -> i32 { + ffi_try!( + unsafe { set_progress_callback_inner(builder, callback, callback_ctx) }, + neg + ) +} + +unsafe fn set_progress_callback_inner( + builder: *mut LanceIndexSegmentBuilder, + callback: LanceIndexBuildProgressCallback, + callback_ctx: *mut c_void, +) -> Result { + if builder.is_null() { + return Err(invalid_input("builder must not be NULL")); + } + let Some(callback) = callback else { + return Err(invalid_input("progress callback must not be NULL")); + }; + let builder = unsafe { &mut *builder }; + if builder.executed { + return Err(invalid_input( + "progress callback must be set before the builder is executed", + )); + } + builder.progress = Some(SendIndexBuildProgressCallback { + callback, + ctx: callback_ctx, + }); + Ok(0) +} + /// Free bytes returned by `lance_index_segment_builder_execute_uncommitted`. #[unsafe(no_mangle)] pub unsafe extern "C" fn lance_free_bytes(bytes: *mut u8) { diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 7f4ea31..0524210 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -10,6 +10,7 @@ use std::ffi::{CString, c_char, c_void}; use std::process::Command; use std::ptr; use std::sync::Arc; +use std::sync::Mutex; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering as AtomicOrdering}; use arrow::ffi::from_ffi; @@ -178,6 +179,134 @@ unsafe extern "C" fn capture_scan_statistics_atomically( captured.calls.fetch_add(1, AtomicOrdering::SeqCst); } +// ─── Index build progress capture fixture ─── + +/// Records progress events plus the exact `callback_ctx` pointer each +/// invocation received, so tests can verify the context round-trips. +#[derive(Default)] +struct ProgressCapture { + events: Vec<(i32, String, u64, String, u64)>, + contexts: Vec<*mut c_void>, +} + +/// Heap-allocate a capture and return it as an opaque callback context. +fn new_progress_capture() -> *mut c_void { + let capture: Box> = Box::new(Mutex::new(ProgressCapture::default())); + Box::into_raw(capture).cast() +} + +/// Reclaim a capture created by `new_progress_capture` and return its contents. +fn take_progress_capture(callback_ctx: *mut c_void) -> ProgressCapture { + assert!(!callback_ctx.is_null()); + let capture = unsafe { Box::from_raw(callback_ctx.cast::>()) }; + capture.into_inner().unwrap() +} + +/// Progress callback that records every event (and the context pointer it was +/// invoked with) into the heap `ProgressCapture` passed as `callback_ctx`. +/// Tolerates a NULL context by ignoring the call. +unsafe extern "C" fn record_build_progress( + callback_ctx: *mut c_void, + event: i32, + stage: *const c_char, + total: u64, + unit: *const c_char, + completed: u64, +) { + if callback_ctx.is_null() { + return; + } + let capture = unsafe { &*callback_ctx.cast::>() }; + let stage = unsafe { std::ffi::CStr::from_ptr(stage) } + .to_string_lossy() + .into_owned(); + let unit = unsafe { std::ffi::CStr::from_ptr(unit) } + .to_string_lossy() + .into_owned(); + let mut guard = capture.lock().unwrap(); + guard.contexts.push(callback_ctx); + guard.events.push((event, stage, total, unit, completed)); +} + +/// Log of every raw `callback_ctx` a build invoked (as `usize` so the static +/// stays `Sync`), for round-trip checks that pass a sentinel or NULL context +/// instead of a capture. +static RECORDED_PROGRESS_CONTEXTS: Mutex> = Mutex::new(Vec::new()); + +/// Progress callback that records only the raw `callback_ctx` pointer, never +/// dereferencing it. Used for sentinel / NULL context round-trip checks. +unsafe extern "C" fn record_progress_ctx( + callback_ctx: *mut c_void, + _event: i32, + _stage: *const c_char, + _total: u64, + _unit: *const c_char, + _completed: u64, +) { + RECORDED_PROGRESS_CONTEXTS + .lock() + .unwrap() + .push(callback_ctx as usize); +} + +/// Assert the well-formedness invariants shared by every progress-capturing +/// build: event codes are only {START, PROGRESS, COMPLETE}, stage strings are +/// non-empty, the documented numeric mapping holds per event (PROGRESS +/// reports total == 0, START reports completed == 0, COMPLETE zeroes both, +/// and only START carries a unit), and per stage the first event is START, +/// the last is COMPLETE, and START/COMPLETE counts match (one active stage at +/// a time). +fn assert_progress_events_well_formed(capture: &ProgressCapture) { + use std::collections::HashMap; + for (event, stage, total, unit, completed) in &capture.events { + assert!( + *event == 0 || *event == 1 || *event == 2, + "unexpected progress event code {event}" + ); + assert!(!stage.is_empty(), "progress stage must be non-empty"); + if *event == 1 { + assert_eq!(*total, 0, "PROGRESS event must report total == 0"); + } + if *event == 0 { + assert_eq!(*completed, 0, "START event must report completed == 0"); + } + if *event == 2 { + assert_eq!(*total, 0, "COMPLETE event must report total == 0"); + assert_eq!(*completed, 0, "COMPLETE event must report completed == 0"); + } + if *event != 0 { + assert!(unit.is_empty(), "non-START event must report unit == \"\""); + } + } + let mut order: Vec<&str> = Vec::new(); + let mut by_stage: HashMap<&str, Vec> = HashMap::new(); + for (event, stage, ..) in &capture.events { + if !by_stage.contains_key(stage.as_str()) { + order.push(stage); + } + by_stage.entry(stage.as_str()).or_default().push(*event); + } + for stage in order { + let events = &by_stage[stage]; + assert_eq!( + events.first().copied(), + Some(0), + "stage {stage} must begin with START" + ); + assert_eq!( + events.last().copied(), + Some(2), + "stage {stage} must end with COMPLETE" + ); + let starts = events.iter().filter(|&&event| event == 0).count(); + let completes = events.iter().filter(|&&event| event == 2).count(); + assert_eq!( + starts, completes, + "stage {stage} must pair each START with a COMPLETE" + ); + } +} + /// Helper: build a tiny dataset whose `value` column is nullable AND contains /// at least one NULL. Used by tests that need to exercise upstream's /// nullability-tightening pre-scan failure path. @@ -4533,6 +4662,505 @@ fn test_vector_index_segment_trains_locally_for_fragment_subset() { } } +#[test] +fn test_vector_index_segment_progress_callback() { + let (_tmp, uri) = create_vector_dataset(256, 16); + let uri_c = c_str(&uri); + let dataset = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; + assert!(!dataset.is_null()); + let column = c_str("embedding"); + // Mirror test_create_vector_index_ivf_pq: IVF_PQ over 256 rows, dim 16, + // 8 partitions, 4 sub-vectors, driven here through the segment builder. + let params = LanceVectorIndexSegmentParams { + index_type: LanceVectorIndexType::IvfPq as i32, + metric: LanceMetricType::L2 as i32, + num_partitions: 8, + num_sub_vectors: 4, + num_bits: 8, + max_iterations: 2, + hnsw_m: 0, + hnsw_ef_construction: 0, + sample_rate: 16, + }; + let builder = unsafe { + lance_index_segment_builder_new_vector( + dataset, + column.as_ptr(), + ptr::null(), + ¶ms, + ptr::null(), + ) + }; + assert!(!builder.is_null(), "{}", take_last_error_message()); + + let capture_ctx = new_progress_capture(); + assert_eq!( + unsafe { + lance_index_segment_builder_set_progress_callback( + builder, + Some(record_build_progress), + capture_ctx, + ) + }, + 0, + "{}", + take_last_error_message() + ); + + let mut bytes = ptr::null_mut(); + let mut len = 0; + assert_eq!( + unsafe { lance_index_segment_builder_execute_uncommitted(builder, &mut bytes, &mut len) }, + 0, + "{}", + take_last_error_message() + ); + assert!(!bytes.is_null() && len > 0); + unsafe { + lance_free_bytes(bytes); + lance_index_segment_builder_free(builder); + lance_dataset_close(dataset); + } + + let capture = take_progress_capture(capture_ctx); + // The installed context pointer round-trips to every invocation. + assert!(!capture.contexts.is_empty(), "expected progress events"); + assert!( + capture.contexts.iter().all(|ctx| *ctx == capture_ctx), + "callback_ctx must round-trip unchanged" + ); + + assert_progress_events_well_formed(&capture); + + let stage_names: Vec<&str> = capture + .events + .iter() + .map(|(_, stage, ..)| stage.as_str()) + .collect(); + assert!( + stage_names.contains(&"shuffle"), + "expected a shuffle stage, saw {stage_names:?}" + ); + assert!( + stage_names.contains(&"merge_partitions"), + "expected a merge_partitions stage, saw {stage_names:?}" + ); + + // The shuffle stage must report at least one PROGRESS event whose + // completed count does not exceed the START total. + let shuffle_start = capture + .events + .iter() + .find(|(event, stage, ..)| *event == 0 && stage == "shuffle") + .expect("shuffle START must be present"); + let shuffle_total = shuffle_start.2; + // Shuffle counts rows (rust/lance/src/index/vector/builder.rs). + assert_eq!( + shuffle_start.3, "rows", + "shuffle START must report unit \"rows\"" + ); + assert!(shuffle_total > 0, "shuffle total must be positive"); + assert!( + capture + .events + .iter() + .any(|(event, stage, _, _, completed)| { + *event == 1 && stage == "shuffle" && *completed <= shuffle_total + }), + "expected shuffle PROGRESS with completed <= total ({shuffle_total})" + ); +} + +#[test] +fn test_scalar_index_segment_progress_callback_sees_load_data() { + let (_tmp, uri) = create_vector_dataset(256, 16); + let uri_c = c_str(&uri); + let dataset = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; + assert!(!dataset.is_null()); + let column = c_str("id"); + let builder = unsafe { + lance_index_segment_builder_new_scalar( + dataset, + column.as_ptr(), + ptr::null(), + LanceScalarIndexType::BTree as i32, + ptr::null(), + ptr::null(), + ) + }; + assert!(!builder.is_null(), "{}", take_last_error_message()); + + let capture_ctx = new_progress_capture(); + assert_eq!( + unsafe { + lance_index_segment_builder_set_progress_callback( + builder, + Some(record_build_progress), + capture_ctx, + ) + }, + 0, + "{}", + take_last_error_message() + ); + + let mut bytes = ptr::null_mut(); + let mut len = 0; + assert_eq!( + unsafe { lance_index_segment_builder_execute_uncommitted(builder, &mut bytes, &mut len) }, + 0, + "{}", + take_last_error_message() + ); + assert!(!bytes.is_null() && len > 0); + unsafe { + lance_free_bytes(bytes); + lance_index_segment_builder_free(builder); + lance_dataset_close(dataset); + } + + let capture = take_progress_capture(capture_ctx); + assert!(!capture.events.is_empty(), "expected progress events"); + assert!( + capture + .events + .iter() + .any(|(event, stage, ..)| { *event == 0 && stage == "load_data" }), + "expected load_data START, saw {:?}", + capture.events + ); + assert!( + capture + .events + .iter() + .any(|(event, stage, ..)| { *event == 2 && stage == "load_data" }), + "expected load_data COMPLETE, saw {:?}", + capture.events + ); +} + +#[test] +fn test_vector_index_segment_progress_callback_multi_fragment_subset() { + let (_tmp, uri) = create_multi_fragment_vector_dataset(2, 64, 8, false); + let uri_c = c_str(&uri); + let dataset = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; + assert!(!dataset.is_null()); + let mut fragment_ids = [0_u64; 2]; + assert_eq!( + unsafe { lance_dataset_fragment_ids(dataset, fragment_ids.as_mut_ptr()) }, + 0 + ); + let selected_fragment = fragment_ids[0] as u32; + let column = c_str("embedding"); + let params = LanceVectorIndexSegmentParams { + index_type: LanceVectorIndexType::IvfFlat as i32, + metric: LanceMetricType::L2 as i32, + num_partitions: 2, + num_sub_vectors: 0, + num_bits: 0, + max_iterations: 2, + hnsw_m: 0, + hnsw_ef_construction: 0, + sample_rate: 16, + }; + let options = LanceIndexSegmentBuildOptions { + fragment_ids: &selected_fragment, + fragment_count: 1, + index_uuid: ptr::null(), + ivf_centroids: ptr::null_mut(), + ivf_centroids_schema: ptr::null(), + pq_codebook: ptr::null_mut(), + pq_codebook_schema: ptr::null(), + mode: LanceIndexSegmentBuildMode::Auto as i32, + }; + let builder = unsafe { + lance_index_segment_builder_new_vector( + dataset, + column.as_ptr(), + ptr::null(), + ¶ms, + &options, + ) + }; + assert!(!builder.is_null(), "{}", take_last_error_message()); + + let capture_ctx = new_progress_capture(); + assert_eq!( + unsafe { + lance_index_segment_builder_set_progress_callback( + builder, + Some(record_build_progress), + capture_ctx, + ) + }, + 0, + "{}", + take_last_error_message() + ); + + let mut bytes = ptr::null_mut(); + let mut len = 0; + assert_eq!( + unsafe { lance_index_segment_builder_execute_uncommitted(builder, &mut bytes, &mut len) }, + 0, + "{}", + take_last_error_message() + ); + assert!(!bytes.is_null() && len > 0); + unsafe { + lance_free_bytes(bytes); + lance_index_segment_builder_free(builder); + lance_dataset_close(dataset); + } + + // The fragment-scoped build still succeeds and reports progress. + let capture = take_progress_capture(capture_ctx); + assert!(!capture.events.is_empty(), "expected progress events"); + assert_progress_events_well_formed(&capture); +} + +#[test] +fn test_index_segment_builder_progress_callback_edge_cases() { + // NULL builder is rejected and sets the error channel. + let capture_ctx = new_progress_capture(); + assert_eq!( + unsafe { + lance_index_segment_builder_set_progress_callback( + ptr::null_mut(), + Some(record_build_progress), + capture_ctx, + ) + }, + -1 + ); + assert_ne!(lance_last_error_code(), lance_c::LanceErrorCode::Ok); + take_progress_capture(capture_ctx); + + let (_tmp, uri) = create_vector_dataset(64, 8); + let uri_c = c_str(&uri); + let dataset = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; + assert!(!dataset.is_null()); + + // NULL callback is rejected. + let column = c_str("id"); + let builder = unsafe { + lance_index_segment_builder_new_scalar( + dataset, + column.as_ptr(), + ptr::null(), + LanceScalarIndexType::BTree as i32, + ptr::null(), + ptr::null(), + ) + }; + assert!(!builder.is_null(), "{}", take_last_error_message()); + assert_eq!( + unsafe { + lance_index_segment_builder_set_progress_callback(builder, None, ptr::null_mut()) + }, + -1 + ); + + // Setting with a NULL callback_ctx succeeds and the NULL context reaches + // the callback verbatim. + RECORDED_PROGRESS_CONTEXTS.lock().unwrap().clear(); + assert_eq!( + unsafe { + lance_index_segment_builder_set_progress_callback( + builder, + Some(record_progress_ctx), + ptr::null_mut(), + ) + }, + 0, + "{}", + take_last_error_message() + ); + let mut bytes = ptr::null_mut(); + let mut len = 0; + assert_eq!( + unsafe { lance_index_segment_builder_execute_uncommitted(builder, &mut bytes, &mut len) }, + 0, + "{}", + take_last_error_message() + ); + assert!(!bytes.is_null() && len > 0); + unsafe { lance_free_bytes(bytes) }; + assert!( + RECORDED_PROGRESS_CONTEXTS + .lock() + .unwrap() + .contains(&(ptr::null_mut::() as usize)), + "NULL callback_ctx must reach the callback" + ); + + // A distinctive sentinel context round-trips to the callback. + let sentinel = 0xC0FFEE_usize as *mut c_void; + RECORDED_PROGRESS_CONTEXTS.lock().unwrap().clear(); + let builder = unsafe { + lance_index_segment_builder_new_scalar( + dataset, + column.as_ptr(), + ptr::null(), + LanceScalarIndexType::BTree as i32, + ptr::null(), + ptr::null(), + ) + }; + assert!(!builder.is_null(), "{}", take_last_error_message()); + assert_eq!( + unsafe { + lance_index_segment_builder_set_progress_callback( + builder, + Some(record_progress_ctx), + sentinel, + ) + }, + 0, + "{}", + take_last_error_message() + ); + let mut bytes = ptr::null_mut(); + let mut len = 0; + assert_eq!( + unsafe { lance_index_segment_builder_execute_uncommitted(builder, &mut bytes, &mut len) }, + 0, + "{}", + take_last_error_message() + ); + assert!(!bytes.is_null() && len > 0); + unsafe { lance_free_bytes(bytes) }; + assert!( + RECORDED_PROGRESS_CONTEXTS + .lock() + .unwrap() + .contains(&(sentinel as usize)), + "sentinel callback_ctx must round-trip" + ); + + // Setting a callback after execution is rejected: the builder is single-use. + assert_eq!( + unsafe { + lance_index_segment_builder_set_progress_callback( + builder, + Some(record_progress_ctx), + sentinel, + ) + }, + -1 + ); + unsafe { lance_index_segment_builder_free(builder) }; + + // Setting a callback twice installs only the second one. + let first_ctx = new_progress_capture(); + let second_ctx = new_progress_capture(); + let builder = unsafe { + lance_index_segment_builder_new_scalar( + dataset, + column.as_ptr(), + ptr::null(), + LanceScalarIndexType::BTree as i32, + ptr::null(), + ptr::null(), + ) + }; + assert!(!builder.is_null(), "{}", take_last_error_message()); + assert_eq!( + unsafe { + lance_index_segment_builder_set_progress_callback( + builder, + Some(record_build_progress), + first_ctx, + ) + }, + 0 + ); + assert_eq!( + unsafe { + lance_index_segment_builder_set_progress_callback( + builder, + Some(record_build_progress), + second_ctx, + ) + }, + 0 + ); + let mut bytes = ptr::null_mut(); + let mut len = 0; + assert_eq!( + unsafe { lance_index_segment_builder_execute_uncommitted(builder, &mut bytes, &mut len) }, + 0, + "{}", + take_last_error_message() + ); + assert!(!bytes.is_null() && len > 0); + unsafe { + lance_free_bytes(bytes); + lance_index_segment_builder_free(builder); + lance_dataset_close(dataset); + } + let first = take_progress_capture(first_ctx); + let second = take_progress_capture(second_ctx); + assert!( + first.events.is_empty(), + "the replaced callback must not receive events" + ); + assert!( + !second.events.is_empty(), + "the replacement callback must receive events" + ); +} + +#[test] +fn test_index_segment_builder_progress_callback_success_clears_error() { + let (_tmp, uri) = create_vector_dataset(64, 8); + let uri_c = c_str(&uri); + let dataset = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; + assert!(!dataset.is_null()); + let column = c_str("id"); + let builder = unsafe { + lance_index_segment_builder_new_scalar( + dataset, + column.as_ptr(), + ptr::null(), + LanceScalarIndexType::BTree as i32, + ptr::null(), + ptr::null(), + ) + }; + assert!(!builder.is_null(), "{}", take_last_error_message()); + + // A failed set leaves a non-OK error code... + assert_eq!( + unsafe { + lance_index_segment_builder_set_progress_callback(builder, None, ptr::null_mut()) + }, + -1 + ); + assert_ne!(lance_last_error_code(), lance_c::LanceErrorCode::Ok); + + // ...and a successful set clears it back to OK. + let capture_ctx = new_progress_capture(); + assert_eq!( + unsafe { + lance_index_segment_builder_set_progress_callback( + builder, + Some(record_build_progress), + capture_ctx, + ) + }, + 0, + "{}", + take_last_error_message() + ); + assert_eq!(lance_last_error_code(), lance_c::LanceErrorCode::Ok); + take_progress_capture(capture_ctx); + unsafe { + lance_index_segment_builder_free(builder); + lance_dataset_close(dataset); + } +} + #[test] fn test_index_segment_options_reject_invalid_fragment_and_train_combinations() { let (_tmp, uri) = create_multi_fragment_vector_dataset(2, 16, 8, false); diff --git a/tests/cpp/test_c_api.c b/tests/cpp/test_c_api.c index d9539e9..5df9cde 100644 --- a/tests/cpp/test_c_api.c +++ b/tests/cpp/test_c_api.c @@ -69,6 +69,51 @@ static void capture_scan_statistics( captured->bytes_read = statistics->bytes_read; } +typedef struct { + uint64_t events; + uint64_t starts; + uint64_t completes; + int saw_shuffle_start; + int saw_shuffle_complete; + int invalid; + void *expected_ctx; + int ctx_mismatch; +} BuildProgressCapture; + +static void capture_build_progress( + void *callback_ctx, + int32_t event, + const char *stage, + uint64_t total, + const char *unit, + uint64_t completed) { + (void)total; + (void)completed; + if (callback_ctx == NULL) return; + BuildProgressCapture *captured = (BuildProgressCapture *)callback_ctx; + if (callback_ctx != captured->expected_ctx) { + captured->ctx_mismatch = 1; + } + if (stage == NULL || unit == NULL) { + captured->invalid = 1; + return; + } + /* Exercise strcmp on the borrowed stage string. */ + if (strcmp(stage, "shuffle") == 0) { + if (event == LANCE_INDEX_BUILD_PROGRESS_STAGE_START) + captured->saw_shuffle_start = 1; + if (event == LANCE_INDEX_BUILD_PROGRESS_STAGE_COMPLETE) + captured->saw_shuffle_complete = 1; + } + if (event == LANCE_INDEX_BUILD_PROGRESS_STAGE_START) + captured->starts += 1; + else if (event == LANCE_INDEX_BUILD_PROGRESS_STAGE_COMPLETE) + captured->completes += 1; + else if (event != LANCE_INDEX_BUILD_PROGRESS_STAGE_PROGRESS) + captured->invalid = 1; + captured->events += 1; +} + static void test_open_and_metadata(const char *uri) { printf(" test_open_and_metadata... "); @@ -996,6 +1041,59 @@ static void test_index_segment_builder(const char *uri) { printf("OK\n"); } +/* Runs a small vector segment build with a progress callback and verifies + * events are delivered with a readable stage name and a round-tripped + * callback context. */ +static void test_index_segment_builder_progress(const char *uri) { + printf(" test_index_segment_builder_progress... "); + LanceDataset *ds = lance_dataset_open(uri, NULL, 0); + ASSERT(ds != NULL, "open failed"); + uint64_t fragment_count = lance_dataset_fragment_count(ds); + ASSERT(fragment_count >= 2, "vector fixture must have two fragments"); + uint64_t all_ids[2] = {0, 0}; + ASSERT(lance_dataset_fragment_ids(ds, all_ids) == 0, + "fragment enumeration failed"); + uint32_t fragment_ids[2] = {(uint32_t)all_ids[0], (uint32_t)all_ids[1]}; + + LanceVectorIndexSegmentParams params = { + LANCE_INDEX_IVF_FLAT, LANCE_METRIC_L2, 2, 0, 0, 2, 0, 0, 16, + }; + LanceIndexSegmentBuildOptions options = {0}; + options.fragment_ids = fragment_ids; + options.fragment_count = 2; + options.mode = LANCE_INDEX_SEGMENT_BUILD_AUTO; + LanceIndexSegmentBuilder *builder = + lance_index_segment_builder_new_vector( + ds, "embedding", "c_progress_idx", ¶ms, &options); + ASSERT(builder != NULL, "vector segment builder failed"); + + BuildProgressCapture captured = {0}; + captured.expected_ctx = &captured; + ASSERT(lance_index_segment_builder_set_progress_callback( + builder, capture_build_progress, &captured) == 0, + "progress callback registration failed"); + + uint8_t *bytes = NULL; + size_t len = 0; + ASSERT(lance_index_segment_builder_execute_uncommitted( + builder, &bytes, &len) == 0, + "vector segment execution failed"); + ASSERT(bytes != NULL && len > 0, "empty segment metadata"); + + ASSERT(captured.events > 0, "expected progress events"); + ASSERT(captured.invalid == 0, "malformed progress event"); + ASSERT(captured.ctx_mismatch == 0, "callback_ctx must round-trip"); + ASSERT(captured.starts > 0 && captured.completes > 0, + "expected at least one START and one COMPLETE stage"); + ASSERT(captured.saw_shuffle_start && captured.saw_shuffle_complete, + "expected shuffle START and COMPLETE events"); + + lance_free_bytes(bytes); + lance_index_segment_builder_free(builder); + lance_dataset_close(ds); + printf("events=%llu... OK\n", (unsigned long long)captured.events); +} + static void test_vector_models_and_reusable_segments(const char *uri) { printf(" test_vector_models_and_reusable_segments... "); LanceDataset *ds = lance_dataset_open(uri, NULL, 0); @@ -1233,6 +1331,7 @@ int main(int argc, char **argv) { test_restore_to_current(uri); test_error_handling(); test_index_segment_builder(uri); + test_index_segment_builder_progress(uri); test_vector_models_and_reusable_segments(uri); test_commit_index_segments(uri); test_dataset_write_roundtrip(uri, write_uri); diff --git a/tests/cpp/test_cpp_api.cpp b/tests/cpp/test_cpp_api.cpp index 9f4ded0..0332d1a 100644 --- a/tests/cpp/test_cpp_api.cpp +++ b/tests/cpp/test_cpp_api.cpp @@ -71,6 +71,37 @@ static void capture_async_scan( captured->ready.notify_one(); } +struct BuildProgressCapture { + uint64_t events = 0; + uint64_t starts = 0; + uint64_t completes = 0; + bool invalid = false; +}; + +static void capture_build_progress( + void* callback_ctx, + int32_t event, + const char* stage, + uint64_t total, + const char* unit, + uint64_t completed) noexcept { + (void)total; + (void)completed; + if (!callback_ctx) return; + auto* captured = static_cast(callback_ctx); + if (!stage || !unit) { + captured->invalid = true; + return; + } + if (event == LANCE_INDEX_BUILD_PROGRESS_STAGE_START) + captured->starts += 1; + else if (event == LANCE_INDEX_BUILD_PROGRESS_STAGE_COMPLETE) + captured->completes += 1; + else if (event != LANCE_INDEX_BUILD_PROGRESS_STAGE_PROGRESS) + captured->invalid = true; + captured->events += 1; +} + static void test_dataset_open(const std::string& uri) { TEST(test_dataset_open); @@ -634,6 +665,35 @@ static void test_index_segment_builder(const std::string& uri) { PASS(); } +static void test_index_segment_builder_progress(const std::string& uri) { + TEST(test_index_segment_builder_progress); + auto ds = lance::Dataset::open(uri); + auto all_ids = ds.fragment_ids(); + assert(all_ids.size() >= 2); + std::vector fragment_ids; + for (auto id : all_ids) fragment_ids.push_back(static_cast(id)); + + LanceVectorIndexParams params = { + LANCE_INDEX_IVF_FLAT, LANCE_METRIC_L2, 2, 0, 0, 2, 0, 0, 16, + }; + LanceIndexSegmentBuildOptions options = {}; + options.fragment_ids = fragment_ids.data(); + options.fragment_count = fragment_ids.size(); + options.mode = LANCE_INDEX_SEGMENT_BUILD_AUTO; + + BuildProgressCapture captured; + auto builder = ds.new_vector_index_segment_builder( + "embedding", params, "cpp_progress_idx", &options); + builder.progress_callback(capture_build_progress, &captured); + auto bytes = builder.execute_uncommitted(); + + assert(!bytes.empty()); + assert(captured.events > 0); + assert(captured.starts > 0 && captured.completes > 0); + assert(!captured.invalid); + PASS(); +} + static void test_vector_models_and_reusable_segments(const std::string& uri) { TEST(test_vector_models_and_reusable_segments); auto ds = lance::Dataset::open(uri); @@ -1167,6 +1227,7 @@ int main(int argc, char** argv) { test_multivector_rejects_flat_column(uri); test_index_segments_smoke(uri); test_index_segment_builder(uri); + test_index_segment_builder_progress(uri); test_vector_models_and_reusable_segments(uri); test_commit_index_segments(uri); test_fts_smoke(uri); From 4a5f53bb8c4e048a0d72cd43812f5d7b5d4c405d Mon Sep 17 00:00:00 2001 From: kid Date: Sun, 20 Sep 2026 10:15:54 +0800 Subject: [PATCH 2/2] fix: retire index build progress callbacks before execute returns Review on #86 found the callback lifetime contract was unenforceable on error paths: lance core's inverted builder clones the progress handle into spawned tokenize_docs workers whose JoinHandles can be dropped early, so a detached clone could invoke the raw C callback/context after execute_uncommitted returned -- a use-after-free once the caller retired the context. Add a shared ProgressCallbackGate (Arc-shared across every clone core hands to worker tasks): execute_uncommitted retires it through a drop guard on the success, error, and panic exit paths, disabling new callback entries and draining in-flight invocations before returning; late detached calls become no-ops. The "invocations only occur while executing" contract is now enforced rather than advisory, and the callback/context only need to stay valid until execute_uncommitted returns (C and C++ docs updated). Also document that stage names are diagnostic-only and not a stable cross-version contract. Regression tests at the gate level: a detached task-owned clone calling after retirement is a no-op (the review reproducer, inverted), and retire blocks until an in-flight invocation exits. --- include/lance/lance.h | 22 ++-- include/lance/lance.hpp | 12 +- src/index_segment.rs | 249 +++++++++++++++++++++++++++++++++++++--- 3 files changed, 257 insertions(+), 26 deletions(-) diff --git a/include/lance/lance.h b/include/lance/lance.h index 136910a..7630913 100644 --- a/include/lance/lance.h +++ b/include/lance/lance.h @@ -1846,6 +1846,11 @@ typedef enum { * - LANCE_INDEX_BUILD_PROGRESS_STAGE_COMPLETE: `total` is 0, `unit` is "", * and `completed` is 0. * + * Stage names are index-type-specific (e.g. "train_ivf", "shuffle", + * "merge_partitions" for vector indices; "load_data" for scalar indices) and + * are diagnostic-only: they are not a stable cross-version contract, so + * consumers must treat them as opaque strings. + * * The callback is invoked from lance-c's internal tokio runtime worker * threads. Certain stages report progress concurrently from parallel worker * tasks, so the callback MUST be thread-safe and reentrant. It must be @@ -1875,14 +1880,15 @@ typedef void (*LanceIndexBuildProgressCallback)( * NULL. `callback_ctx` may be NULL and is passed through to the callback * opaquely. Setting a callback replaces any previously set callback. * - * `callback` and `callback_ctx` must remain valid and safe to invoke until the - * builder is released with lance_index_segment_builder_free. Invocations only - * occur while lance_index_segment_builder_execute_uncommitted is executing. - * Lance core may deliver events from spawned worker tasks, and error paths can - * detach them before they finish; keep callback and callback_ctx valid until - * the builder is released rather than retiring them when execute returns. See - * LanceIndexBuildProgressCallback for the full threading and reentrancy - * contract. + * Invocations occur only while lance_index_segment_builder_execute_uncommitted + * is executing, and this is enforced rather than contractual: lance-c + * disables the callback and drains in-flight invocations through a retirement + * gate before that call returns, including on error, so a worker task + * detached by lance core on an error path can never invoke the callback + * afterwards. `callback` and `callback_ctx` must therefore remain valid and + * safe to invoke until lance_index_segment_builder_execute_uncommitted + * returns. See LanceIndexBuildProgressCallback for the full threading and + * reentrancy contract. * * @return 0 on success, -1 on error. */ diff --git a/include/lance/lance.hpp b/include/lance/lance.hpp index 752b4e1..286724e 100644 --- a/include/lance/lance.hpp +++ b/include/lance/lance.hpp @@ -1286,11 +1286,13 @@ class IndexSegmentBuilder { /// from parallel worker tasks, so it must be thread-safe, non-blocking, /// and must not re-enter any lance_* function. It must return normally; /// unwinding or throwing across this boundary can abort the host process. - /// The callback and a non-null context must remain valid until the - /// builder is released: core may deliver events from spawned worker - /// tasks, and error paths can detach them before they finish, so do not - /// retire them when execute_uncommitted returns. Progress reporting is - /// advisory and cannot affect the build outcome. + /// Invocations occur only while execute_uncommitted runs, and this is + /// enforced: lance-c disables the callback and drains in-flight + /// invocations before that call returns, including on error, so a worker + /// task detached by core can never invoke it afterwards. The callback and + /// the context (if non-null) must remain valid until execute_uncommitted + /// returns. Progress reporting is advisory and cannot affect the build + /// outcome. IndexSegmentBuilder& progress_callback(LanceIndexBuildProgressCallback callback, void* callback_ctx) { if (lance_index_segment_builder_set_progress_callback(handle_.get(), callback, diff --git a/src/index_segment.rs b/src/index_segment.rs index e2acb36..5209321 100644 --- a/src/index_segment.rs +++ b/src/index_segment.rs @@ -8,6 +8,7 @@ use std::ffi::{CStr, CString, c_char, c_void}; use std::ptr; use std::slice; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema, from_ffi, to_ffi}; use arrow_array::{Array, ArrayRef, FixedSizeListArray, make_array}; @@ -162,6 +163,66 @@ pub type LanceIndexBuildProgressCallback = Option< ), >; +/// Shared retirement gate for the C progress callback bridge. +/// +/// Lance core clones the installed progress handle into spawned worker tasks +/// (e.g. the inverted index's `tokenize_docs` workers), and error paths can +/// drop their `JoinHandle`s before the workers finish, so a detached clone +/// can outlive the build that installed it. The gate turns callback +/// retirement from a documentation contract into an enforced boundary: once +/// `retire` returns, no callback invocation is in flight and every later +/// entry becomes a no-op, so a detached clone can never touch the raw +/// `callback`/`ctx` after the owning `execute_uncommitted` call has returned. +#[derive(Debug, Default)] +struct ProgressCallbackGate { + /// Set by `retire`; entries that observe it become no-ops. + retired: AtomicBool, + /// Number of C callback invocations currently executing. + in_flight: AtomicUsize, +} + +impl ProgressCallbackGate { + /// Admit one callback invocation, or report that the gate is retired. + /// + /// The increment is ordered before the `retired` check (both SeqCst), so + /// an entry racing with `retire` either observes `retired` and backs out, + /// or is counted in `in_flight` and therefore awaited by `retire`'s + /// drain loop. + fn enter(&self) -> bool { + self.in_flight.fetch_add(1, Ordering::SeqCst); + if self.retired.load(Ordering::SeqCst) { + self.in_flight.fetch_sub(1, Ordering::SeqCst); + return false; + } + true + } + + fn exit(&self) { + self.in_flight.fetch_sub(1, Ordering::SeqCst); + } + + /// Disable new entries and wait for in-flight invocations to drain. + /// + /// A callback that blocks (violating its documented non-blocking + /// contract) stalls this drain rather than causing a use-after-free. + fn retire(&self) { + self.retired.store(true, Ordering::SeqCst); + while self.in_flight.load(Ordering::SeqCst) != 0 { + std::thread::yield_now(); + } + } +} + +/// Retires the shared progress gate on drop, covering the success, error, +/// and panic exits of `execute_uncommitted_inner`. +struct ProgressRetireGuard(Arc); + +impl Drop for ProgressRetireGuard { + fn drop(&mut self) { + self.0.retire(); + } +} + #[derive(Debug, Clone)] struct SendIndexBuildProgressCallback { callback: unsafe extern "C" fn( @@ -173,13 +234,18 @@ struct SendIndexBuildProgressCallback { completed: u64, ), ctx: *mut c_void, + /// Shared with every clone core hands to worker tasks; retired before the + /// owning `execute_uncommitted` call returns. + gate: Arc, } // SAFETY: The C API requires the callback and its context to remain valid and -// safe to invoke until the segment builder is freed, and the builder is -// single-use, so no invocation can outlive the builder. Certain build stages -// report progress concurrently from parallel worker tasks, so the callback -// may be invoked concurrently; the C contract requires it to be thread-safe. +// safe to invoke until `execute_uncommitted` returns, and the shared gate +// guarantees no invocation is in flight — and no new one can start — once +// that boundary is reached, so no invocation can touch the raw context after +// retirement. Certain build stages report progress concurrently from parallel +// worker tasks, so the callback may be invoked concurrently; the C contract +// requires it to be thread-safe. unsafe impl Send for SendIndexBuildProgressCallback {} unsafe impl Sync for SendIndexBuildProgressCallback {} @@ -202,8 +268,14 @@ fn progress_c_string(value: &str) -> Result { #[async_trait::async_trait] impl IndexBuildProgress for SendIndexBuildProgressCallback { async fn stage_start(&self, stage: &str, total: Option, unit: &str) -> Result<()> { + // Strings are built before entering the gate so the defensive `?` + // paths cannot leak an in-flight count; after `enter` the only code + // is the FFI call, so `exit` always runs. let stage = progress_c_string(stage)?; let unit = progress_c_string(unit)?; + if !self.gate.enter() { + return Ok(()); + } unsafe { (self.callback)( self.ctx, @@ -214,11 +286,15 @@ impl IndexBuildProgress for SendIndexBuildProgressCallback { 0, ) }; + self.gate.exit(); Ok(()) } async fn stage_progress(&self, stage: &str, completed: u64) -> Result<()> { let stage = progress_c_string(stage)?; + if !self.gate.enter() { + return Ok(()); + } unsafe { (self.callback)( self.ctx, @@ -229,11 +305,15 @@ impl IndexBuildProgress for SendIndexBuildProgressCallback { completed, ) }; + self.gate.exit(); Ok(()) } async fn stage_complete(&self, stage: &str) -> Result<()> { let stage = progress_c_string(stage)?; + if !self.gate.enter() { + return Ok(()); + } unsafe { (self.callback)( self.ctx, @@ -244,6 +324,7 @@ impl IndexBuildProgress for SendIndexBuildProgressCallback { 0, ) }; + self.gate.exit(); Ok(()) } } @@ -1085,7 +1166,14 @@ unsafe fn execute_uncommitted_inner( } let columns = [builder.column.as_str()]; - let metadata = match &builder.kind { + // Retire the shared progress gate before this call returns — on success, + // error, and panic paths alike — so a detached core worker holding a + // progress clone can no longer enter the C callback afterwards. + let _progress_retire = builder + .progress + .as_ref() + .map(|progress| ProgressRetireGuard(progress.gate.clone())); + let result = match &builder.kind { SegmentKind::Scalar { scalar_type, params_json, @@ -1109,7 +1197,7 @@ unsafe fn execute_uncommitted_inner( if let Some(progress) = builder.progress.clone() { core_builder = core_builder.progress(Arc::new(progress)); } - block_on(core_builder.execute_uncommitted())? + block_on(core_builder.execute_uncommitted()) } SegmentKind::Vector { params, @@ -1137,9 +1225,10 @@ unsafe fn execute_uncommitted_inner( } // Core's train=false means "create an empty index". Model presence // itself controls whether IVF/PQ training is skipped. - block_on(core_builder.train(true).execute_uncommitted())? + block_on(core_builder.train(true).execute_uncommitted()) } }; + let metadata = result?; let bytes = pb::IndexMetadata::from(&metadata).encode_to_vec(); let allocation = unsafe { libc::malloc(bytes.len()) }.cast::(); if allocation.is_null() { @@ -1160,12 +1249,13 @@ unsafe fn execute_uncommitted_inner( /// builder. /// /// The callback is invoked from lance-c's internal tokio runtime worker -/// threads while `lance_index_segment_builder_execute_uncommitted` runs. -/// Lance core may deliver events from spawned worker tasks, and error paths -/// can detach them before they finish, so the callback and `callback_ctx` -/// must remain valid until the builder is freed rather than being retired -/// when `execute_uncommitted` returns. Certain stages report progress -/// concurrently from parallel worker tasks, so the callback must be +/// threads while `lance_index_segment_builder_execute_uncommitted` runs, and +/// only while it runs: lance-c disables and drains the callback through a +/// shared retirement gate before `execute_uncommitted` returns (including on +/// error), so a core worker task detached by an error path can never enter +/// the callback afterwards. The callback and `callback_ctx` must therefore +/// remain valid until `execute_uncommitted` returns. Certain stages report +/// progress concurrently from parallel worker tasks, so the callback must be /// thread-safe and reentrant, must be non-blocking, and must not call back /// into any `lance_*` function. It is invoked without a panic guard: it must /// return normally, because unwinding or throwing across this boundary can @@ -1203,6 +1293,9 @@ unsafe fn set_progress_callback_inner( builder.progress = Some(SendIndexBuildProgressCallback { callback, ctx: callback_ctx, + // A replacement gets a fresh gate: the previous callback was never + // shared with a build, so there is nothing to retire. + gate: Arc::new(ProgressCallbackGate::default()), }); Ok(0) } @@ -1680,3 +1773,133 @@ pub unsafe extern "C" fn lance_index_segment_metadata_free( }); } } + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::mpsc; + use std::time::Duration; + + struct LateCallState { + entered_after_retire: AtomicBool, + } + + unsafe extern "C" fn record_late_entry( + ctx: *mut c_void, + _event: i32, + _stage: *const c_char, + _total: u64, + _unit: *const c_char, + _completed: u64, + ) { + let state = unsafe { &*ctx.cast::() }; + state.entered_after_retire.store(true, Ordering::SeqCst); + } + + /// Regression test for the review finding that a progress clone held by a + /// detached core worker task (e.g. the inverted index's tokenize_docs + /// workers) must not enter the C callback once the owning build has + /// retired the gate at the end of `execute_uncommitted`. + #[test] + fn detached_clone_cannot_enter_callback_after_retire() { + let state = Box::new(LateCallState { + entered_after_retire: AtomicBool::new(false), + }); + let ctx = Box::into_raw(state); + let bridge = SendIndexBuildProgressCallback { + callback: record_late_entry, + ctx: ctx.cast(), + gate: Arc::new(ProgressCallbackGate::default()), + }; + let detached = Arc::new(bridge.clone()); + // Simulate the retirement boundary at the end of execute_uncommitted: + // the builder's own handle is gone and the gate is retired, while a + // detached worker still holds its clone. + drop(bridge); + detached.gate.retire(); + crate::runtime::block_on(async move { + tokio::task::spawn(async move { + detached + .stage_progress("tokenize_docs", 1) + .await + .expect("a late progress call must be a no-op, not an error"); + }) + .await + .unwrap(); + }); + let state = unsafe { Box::from_raw(ctx) }; + assert!( + !state.entered_after_retire.load(Ordering::SeqCst), + "detached clone entered the C callback after retirement" + ); + } + + struct BlockingState { + entered: mpsc::Sender<()>, + release: mpsc::Receiver<()>, + } + + unsafe extern "C" fn blocking_callback( + ctx: *mut c_void, + _event: i32, + _stage: *const c_char, + _total: u64, + _unit: *const c_char, + _completed: u64, + ) { + let state = unsafe { &*ctx.cast::() }; + state.entered.send(()).unwrap(); + state.release.recv().unwrap(); + } + + /// `retire` must disable new entries and then wait until every in-flight + /// invocation has exited before returning. + #[test] + fn retire_drains_in_flight_invocation() { + let (entered_tx, entered_rx) = mpsc::channel::<()>(); + let (release_tx, release_rx) = mpsc::channel::<()>(); + let state = Box::new(BlockingState { + entered: entered_tx, + release: release_rx, + }); + let ctx = Box::into_raw(state); + let bridge = SendIndexBuildProgressCallback { + callback: blocking_callback, + ctx: ctx.cast(), + gate: Arc::new(ProgressCallbackGate::default()), + }; + let gate = bridge.gate.clone(); + + // Park one invocation inside the C callback. + let caller = std::thread::spawn(move || { + crate::runtime::block_on(bridge.stage_progress("shuffle", 1)).unwrap(); + }); + entered_rx.recv().unwrap(); + + let retire_returned = Arc::new(AtomicBool::new(false)); + let retire_thread = { + let gate = gate.clone(); + let retire_returned = retire_returned.clone(); + std::thread::spawn(move || { + gate.retire(); + retire_returned.store(true, Ordering::SeqCst); + }) + }; + // Wait until retirement has actually disabled new entries... + while !gate.retired.load(Ordering::SeqCst) { + std::thread::yield_now(); + } + // ...then prove it is still draining the in-flight invocation. + std::thread::sleep(Duration::from_millis(50)); + assert!( + !retire_returned.load(Ordering::SeqCst), + "retire returned while a callback invocation was in flight" + ); + // Unblock the callback; the drain must now complete. + release_tx.send(()).unwrap(); + caller.join().unwrap(); + retire_thread.join().unwrap(); + assert!(retire_returned.load(Ordering::SeqCst)); + drop(unsafe { Box::from_raw(ctx) }); + } +}