From d124d80c86dfb8f7b316db2c70c3e854b77da802 Mon Sep 17 00:00:00 2001 From: Lucy Ge Date: Sun, 9 Aug 2026 05:34:43 -0700 Subject: [PATCH 01/11] use ArcSwap for StorageBase's dataset instead of let StorageBase own Dataset by value --- Cargo.lock | 1 + crates/lance-context-core/Cargo.toml | 1 + .../lance-context-core/src/datagen_store.rs | 6 +- .../lance-context-core/src/generic_store.rs | 4 +- .../lance-context-core/src/rollout_store.rs | 87 ++++++------ crates/lance-context-core/src/store.rs | 52 +++---- crates/lance-context-core/src/store_base.rs | 129 ++++++++++++------ 7 files changed, 164 insertions(+), 116 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6928852..7bb5a46 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6344,6 +6344,7 @@ dependencies = [ name = "lance-context-core" version = "0.6.5" dependencies = [ + "arc-swap", "arrow-array 58.3.0", "arrow-ipc 58.3.0", "arrow-json 58.3.0", diff --git a/crates/lance-context-core/Cargo.toml b/crates/lance-context-core/Cargo.toml index d77daf2..c38b516 100644 --- a/crates/lance-context-core/Cargo.toml +++ b/crates/lance-context-core/Cargo.toml @@ -18,6 +18,7 @@ default = ["metrics"] metrics = ["dep:metrics"] [dependencies] +arc-swap = "1" base64 = "0.22" arrow-array = "58" arrow-ipc = "58" diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index b967a46..3d34ff1 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -126,7 +126,7 @@ impl DatagenStore { #[must_use] pub fn uri(&self) -> &str { - self.base.dataset.uri() + self.base.uri() } #[must_use] @@ -336,7 +336,7 @@ impl DatagenStore { } } - Ok(Self::get_blob_from_dataset(&self.base.dataset, event_id) + Ok(Self::get_blob_from_dataset(self.base.current_dataset().as_ref(), event_id) .await? .flatten()) } @@ -529,7 +529,7 @@ impl DatagenStore { fn non_blob_columns(&self) -> Vec { self.base - .dataset + .current_dataset() .schema() .fields .iter() diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index 4c55247..202af55 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -184,7 +184,7 @@ impl GenericStore { ) .await?; - let schema: Arc = Arc::new(base.dataset.schema().into()); + let schema: Arc = Arc::new(base.current_dataset().schema().into()); let persisted = spec_from_schema(&schema)?; // Reopening with a different schema would reinterpret existing data. @@ -406,7 +406,7 @@ impl GenericStore { /// Row count of the base table. Excludes rows still in unmerged /// generations or buffered in the writer. pub async fn count_base_rows(&self) -> LanceResult { - self.base.dataset.count_rows(None).await + self.base.current_dataset().count_rows(None).await } } diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 23c4149..90fd35a 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -687,12 +687,12 @@ impl RolloutStore { /// writer owns and merges its own WAL shard. pub async fn observe(&self) -> LanceResult { let shard_snapshots = self.wal_shard_snapshots().await?; - let base_rows = self.base.dataset.count_rows(None).await? as u64; + let base_rows = self.base.current_dataset().count_rows(None).await? as u64; let pending_rows = self.base.pending_wal_rows(&shard_snapshots).await?; let row_count = (base_rows + pending_rows) as i64; - let fragment_count = self.base.dataset.count_fragments() as i64; - let version = self.base.dataset.manifest.version; - let last_updated = self.base.dataset.manifest.timestamp().timestamp_millis(); + let fragment_count = self.base.current_dataset().count_fragments() as i64; + let version = self.base.current_dataset().manifest.version; + let last_updated = self.base.current_dataset().manifest.timestamp().timestamp_millis(); let pending_wal_generations = shard_snapshots .iter() .map(|snapshot| snapshot.flushed_generations.len() as i64) @@ -876,13 +876,13 @@ impl RolloutStore { "pagination offset exceeds i64::MAX".to_string(), )) })?; - let mut scanner = self.base.dataset.scan(); + let mut scanner = self.base.current_dataset().scan(); scanner.project(&refs)?; // Lance 7's late take path can panic on nested list columns. // Keep those early while deferring only potentially large text. scanner.materialization_style(MaterializationStyle::all_early_except( &PAGINATION_LATE_COLUMNS, - self.base.dataset.schema(), + self.base.current_dataset().schema(), )?); if let Some(filter) = &filter { scanner.filter(filter)?; @@ -978,14 +978,14 @@ impl RolloutStore { } let columns = Arc::new(self.non_blob_columns()); - let target_schema = Arc::new(projected_arrow_schema(&self.base.dataset, &columns)?); + let target_schema = Arc::new(projected_arrow_schema(self.base.current_dataset().as_ref(), &columns)?); let id_filter = Arc::new(format!("id IN ({})", sql_quoted_list(&page_ids))); let wanted: HashSet = page_ids.iter().cloned().collect(); let mut records_by_id = HashMap::with_capacity(page_ids.len()); if source == ListSource::All { for record in Self::take_page_rows_from_dataset( - self.base.dataset.clone(), + (*self.base.current_dataset()).clone(), id_filter.clone(), columns.clone(), target_schema.clone(), @@ -1125,7 +1125,7 @@ impl RolloutStore { let schema = match table_schema { Some(schema) => schema, None => { - let full: Schema = self.base.dataset.schema().into(); + let full: Schema = self.base.current_dataset().schema().into(); let projected: Vec = full .fields() .iter() @@ -1234,7 +1234,7 @@ impl RolloutStore { ) -> LanceResult>)>> { // Base table first — no manifest reads, no per-generation opens. if let Some(record) = self.scan_one_by_id(id, ListSource::Fragments).await? { - let payload = Self::get_blob_from_dataset(&self.base.dataset, id) + let payload = Self::get_blob_from_dataset(self.base.current_dataset().as_ref(), id) .await? .flatten(); return Ok(Some((record, payload))); @@ -1318,7 +1318,7 @@ impl RolloutStore { pub async fn get_blob(&self, id: &str) -> LanceResult>> { // Base-table-first: an already-merged row is found here with no MemWAL // manifest reads and no per-generation opens. - if let Some(payload) = Self::get_blob_from_dataset(&self.base.dataset, id).await? { + if let Some(payload) = Self::get_blob_from_dataset(self.base.current_dataset().as_ref(), id).await? { return Ok(payload); } @@ -1411,7 +1411,7 @@ impl RolloutStore { /// never materialize artifact bytes. fn non_blob_columns(&self) -> Vec { self.base - .dataset + .current_dataset() .schema() .fields .iter() @@ -1436,7 +1436,7 @@ impl RolloutStore { } fn records_to_batch(&self, records: &[RolloutRecord]) -> LanceResult { - let field_paths = self.base.dataset.schema().field_paths(); + let field_paths = self.base.current_dataset().schema().field_paths(); let has = |name: &str| field_paths.iter().any(|path| path == name); let include_relationships = has(RELATIONSHIPS_COLUMN); let include_metadata = has("metadata"); @@ -1691,7 +1691,7 @@ impl RolloutStore { arrays_by_name.insert("metadata".to_string(), Arc::new(metadata_builder.finish())); } - let schema: Arc = Arc::new(self.base.dataset.schema().into()); + let schema: Arc = Arc::new(self.base.current_dataset().schema().into()); let arrays = schema .fields() .iter() @@ -2286,7 +2286,11 @@ mod tests { vec![Ok::(base_batch)].into_iter(), base_schema, ); - store.base.dataset.append(base_reader, None).await.unwrap(); + { + let mut dataset = (*store.base.current_dataset()).clone(); + dataset.append(base_reader, None).await.unwrap(); + store.base.set_dataset(dataset); + } store.add(&[assistant_record("legacy-wal")]).await.unwrap(); store.flush().await.unwrap(); @@ -2300,16 +2304,18 @@ mod tests { .filter(|field| legacy_schema.field_with_name(field.name()).is_err()) .cloned() .collect::>(); - store - .base - .dataset - .add_columns( - NewColumnTransform::AllNulls(Arc::new(Schema::new(claim_check_fields))), - None, - None, - ) - .await - .unwrap(); + { + let mut dataset = (*store.base.current_dataset()).clone(); + dataset + .add_columns( + NewColumnTransform::AllNulls(Arc::new(Schema::new(claim_check_fields))), + None, + None, + ) + .await + .unwrap(); + store.base.set_dataset(dataset); + } } (dir, store) @@ -3050,8 +3056,8 @@ mod tests { /// Read the number of un-merged flushed generations recorded for a store's /// own write shard. Used by merge tests to assert the manifest drains. async fn flushed_generation_count(store: &RolloutStore) -> usize { - let object_store = store.base.dataset.object_store(None).await.unwrap(); - let branch_location = store.base.dataset.branch_location(); + let object_store = store.base.current_dataset().object_store(None).await.unwrap(); + let branch_location = store.base.current_dataset().branch_location(); let manifest_store = ShardManifestStore::new( object_store, &branch_location.path, @@ -3070,8 +3076,8 @@ mod tests { /// Used to assert the resident writer claims the epoch once instead of /// bumping it on every append. async fn shard_writer_epoch(store: &RolloutStore) -> u64 { - let object_store = store.base.dataset.object_store(None).await.unwrap(); - let branch_location = store.base.dataset.branch_location(); + let object_store = store.base.current_dataset().object_store(None).await.unwrap(); + let branch_location = store.base.current_dataset().branch_location(); let manifest_store = ShardManifestStore::new( object_store, &branch_location.path, @@ -3355,7 +3361,7 @@ mod tests { store.flush().await.unwrap(); store.maybe_merge_own_shard().await.unwrap(); - let before = store.base.dataset.count_fragments(); + let before = store.base.current_dataset().count_fragments(); assert!(before > 1, "expected several fragments, got {before}"); assert!(store.should_compact(&CompactionConfig { min_fragments: 2, @@ -3380,7 +3386,7 @@ mod tests { "one incremental pass must honor max_source_fragments" ); - let after = store.base.dataset.count_fragments(); + let after = store.base.current_dataset().count_fragments(); assert!( after < before, "compaction should reduce fragments: {before} -> {after}" @@ -3426,7 +3432,7 @@ mod tests { .unwrap(); store.add(&[assistant_record("a-0")]).await.unwrap(); - let frags = store.base.dataset.count_fragments(); + let frags = store.base.current_dataset().count_fragments(); // A threshold above the current fragment count says "don't compact". assert!(!store.should_compact(&CompactionConfig { min_fragments: frags + 1, @@ -3702,7 +3708,7 @@ mod tests { assert_eq!(store.cleanup_own_shard().await.unwrap(), 1); assert_eq!(flushed_generation_count(&store).await, 0); - let field_paths = store.base.dataset.schema().field_paths(); + let field_paths = store.base.current_dataset().schema().field_paths(); for column in CLAIM_CHECK_COLUMNS { assert!(field_paths.iter().any(|path| path == column)); } @@ -3743,18 +3749,17 @@ mod tests { let generation_batch = current_store.records_to_batch(&[record]).unwrap(); legacy_store.base.ensure_latest_schema().await.unwrap(); - let merge_schema: Arc = Arc::new(legacy_store.base.dataset.schema().into()); + let merge_schema: Arc = Arc::new(legacy_store.base.current_dataset().schema().into()); let aligned = align_batch_to_schema(generation_batch, merge_schema.clone()).unwrap(); let reader = RecordBatchIterator::new( vec![Ok::(aligned)].into_iter(), merge_schema, ); - legacy_store - .base - .dataset - .append(reader, None) - .await - .unwrap(); + { + let mut dataset = (*legacy_store.base.current_dataset()).clone(); + dataset.append(reader, None).await.unwrap(); + legacy_store.base.set_dataset(dataset); + } let merged = legacy_store .get_by_id_source("current-generation", ListSource::Fragments) @@ -3825,7 +3830,7 @@ mod tests { store.create_id_zonemap_index().await.unwrap(); let has_id_index = |s: &RolloutStore| { - let dataset = s.base.dataset.clone(); + let dataset = s.base.current_dataset(); async move { dataset .load_indices() diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 2dc504b..f18c839 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -659,7 +659,7 @@ impl ContextStore { /// URI of the underlying Lance dataset. #[must_use] pub fn uri(&self) -> &str { - self.base.dataset.uri() + self.base.uri() } /// Distance metric this context ranks vector-search results with. @@ -1374,7 +1374,7 @@ impl ContextStore { fn has_relationships_column(&self) -> bool { self.base - .dataset + .current_dataset() .schema() .field_paths() .iter() @@ -1383,7 +1383,7 @@ impl ContextStore { fn has_external_id_column(&self) -> bool { self.base - .dataset + .current_dataset() .schema() .field_paths() .iter() @@ -1405,10 +1405,11 @@ impl ContextStore { } let schema = Arc::new(Schema::new(vec![relationship_field()])); - self.base - .dataset + let mut dataset = (*self.base.current_dataset()).clone(); + dataset .add_columns(NewColumnTransform::AllNulls(schema), None, None) .await?; + self.base.set_dataset(dataset); self.base.clear_version_pin(); Ok(true) } @@ -1432,7 +1433,7 @@ impl ContextStore { /// Retrieve a single record by its unique ID. pub async fn get(&self, id: &str) -> LanceResult> { let escaped_id = id.replace('\'', "''"); - let mut scanner = self.base.dataset.scan(); + let mut scanner = self.base.current_dataset().scan(); scanner.filter(&format!("id = '{}'", escaped_id))?; scanner.limit(Some(1), None)?; @@ -1835,7 +1836,7 @@ impl ContextStore { /// filtering and metadata stay correct). fn projected_columns(&self, projection: ReadProjection) -> Vec { self.base - .dataset + .current_dataset() .schema() .fields .iter() @@ -1894,7 +1895,7 @@ impl ContextStore { info!( "Starting compaction: {} fragments", - self.base.dataset.count_fragments() + self.base.current_dataset().count_fragments() ); let start = std::time::Instant::now(); @@ -2003,7 +2004,7 @@ impl ContextStore { /// Check if compaction should run based on configuration thresholds. pub async fn should_compact(&self) -> LanceResult { - let fragment_count = self.base.dataset.count_fragments(); + let fragment_count = self.base.current_dataset().count_fragments(); if fragment_count < self.compaction_config.min_fragments { return Ok(false); @@ -2030,7 +2031,7 @@ impl ContextStore { let state = self.compaction_state.lock().await; Ok(CompactionStats { - total_fragments: self.base.dataset.count_fragments(), + total_fragments: self.base.current_dataset().count_fragments(), is_compacting: state.is_compacting, last_compaction: state.last_compaction, last_error: state.last_error.clone(), @@ -2044,7 +2045,7 @@ impl ContextStore { return Ok(()); } - let indices = self.base.dataset.load_indices().await?; + let indices = self.base.current_dataset().load_indices().await?; if indices.iter().any(|i| i.name == ID_INDEX_NAME) { return Ok(()); } @@ -2064,12 +2065,13 @@ impl ContextStore { let params = ScalarIndexParams::default(); - self.base - .dataset + let mut dataset = (*self.base.current_dataset()).clone(); + dataset .create_index_builder(&["id"], index_type, ¶ms) .name(ID_INDEX_NAME.to_string()) .replace(true) .await?; + self.base.set_dataset(dataset); // Reload through the base so the new index is visible to subsequent // reads, keeping the storage options and session (a bare @@ -2338,42 +2340,42 @@ impl ContextStore { fn records_to_batch(&self, entries: &[ContextRecord]) -> LanceResult { let include_external_id = self .base - .dataset + .current_dataset() .schema() .field_paths() .iter() .any(|path| path == "external_id"); let include_lifecycle = self .base - .dataset + .current_dataset() .schema() .field_paths() .iter() .any(|path| path == "expires_at"); let include_metadata = self .base - .dataset + .current_dataset() .schema() .field_paths() .iter() .any(|path| path == "metadata"); let include_tenant = self .base - .dataset + .current_dataset() .schema() .field_paths() .iter() .any(|path| path == "tenant"); let include_source = self .base - .dataset + .current_dataset() .schema() .field_paths() .iter() .any(|path| path == "source"); let include_external_reference = self .base - .dataset + .current_dataset() .schema() .field_paths() .iter() @@ -2698,7 +2700,7 @@ impl ContextStore { ]); } - let schema: Arc = Arc::new(self.base.dataset.schema().into()); + let schema: Arc = Arc::new(self.base.current_dataset().schema().into()); let arrays = schema .fields() .iter() @@ -5561,7 +5563,7 @@ mod tests { .unwrap(); // Index should be created eagerly on open - let indices = store.base.dataset.load_indices().await.unwrap(); + let indices = store.base.current_dataset().load_indices().await.unwrap(); assert!( indices.iter().any(|i| i.name == ID_INDEX_NAME), "btree index should be created on open" @@ -5577,7 +5579,7 @@ mod tests { store.compact(None).await.unwrap(); // Index should still exist after compaction - let indices = store.base.dataset.load_indices().await.unwrap(); + let indices = store.base.current_dataset().load_indices().await.unwrap(); assert!( indices.iter().any(|i| i.name == ID_INDEX_NAME), "btree index should persist after compaction" @@ -5601,7 +5603,7 @@ mod tests { .unwrap(); // Index should be created eagerly on open - let indices = store.base.dataset.load_indices().await.unwrap(); + let indices = store.base.current_dataset().load_indices().await.unwrap(); assert!( indices.iter().any(|i| i.name == ID_INDEX_NAME), "zonemap index should be created on open" @@ -5615,7 +5617,7 @@ mod tests { } store.compact(None).await.unwrap(); - let indices = store.base.dataset.load_indices().await.unwrap(); + let indices = store.base.current_dataset().load_indices().await.unwrap(); assert!( indices.iter().any(|i| i.name == ID_INDEX_NAME), "zonemap index should persist after compaction" @@ -5635,7 +5637,7 @@ mod tests { store.add(&[text_record("no-idx-1", 0.0)]).await.unwrap(); store.compact(None).await.unwrap(); - let indices = store.base.dataset.load_indices().await.unwrap(); + let indices = store.base.current_dataset().load_indices().await.unwrap(); assert!( !indices.iter().any(|i| i.name == ID_INDEX_NAME), "no id index should be created when IdIndexType::None" diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index b534a22..fde7fa7 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -38,6 +38,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; +use arc_swap::ArcSwap; use arrow_array::{new_null_array, RecordBatch, RecordBatchIterator, UInt32Array}; use arrow_schema::{ArrowError, Schema}; use arrow_select::take::take; @@ -264,7 +265,12 @@ pub(crate) struct StorageBaseOptions { pub(crate) struct StorageBase { /// The base table. `pub(crate)` because concrete stores build their own /// schema-specific scans and projections directly against it. - pub dataset: Dataset, + /// + /// Wrapped in [`ArcSwap`] so a merge/compact/reload can publish a new + /// handle without requiring exclusive `&mut` for every reader. + pub dataset: ArcSwap, + /// Dataset URI; stable for the lifetime of this handle. + uri: String, /// MemWAL shard this instance writes to (derived from `shard_id`). pub write_shard: Uuid, /// Object-store options, retained so a self-merge can re-append flushed @@ -402,8 +408,10 @@ impl StorageBase { .into()); } + let uri = dataset.uri().to_string(); let mut base = Self { - dataset, + dataset: ArcSwap::from_pointee(dataset), + uri, write_shard: derive_shard_id(shard_id.as_deref()), storage_options, session, @@ -427,18 +435,22 @@ impl StorageBase { /// URI of the underlying Lance dataset. #[must_use] pub fn uri(&self) -> &str { - self.dataset.uri() + &self.uri } /// Current base dataset manifest version. #[must_use] pub fn version(&self) -> u64 { - self.dataset.manifest.version + self.current_dataset().manifest.version } /// Check out a specific base dataset version (time travel). pub async fn checkout(&mut self, version_id: u64) -> LanceResult<()> { - self.dataset = self.dataset.checkout_version(version_id).await?; + let dataset = self + .current_dataset() + .checkout_version(version_id) + .await?; + self.set_dataset(dataset); self.pinned_version = Some(version_id); Ok(()) } @@ -456,7 +468,9 @@ impl StorageBase { /// WAL merges committed by another process become visible without paying the /// cost of reopening the dataset and rebuilding all session caches. pub async fn refresh_latest(&mut self) -> LanceResult<()> { - self.dataset.checkout_latest().await?; + let mut dataset = (*self.current_dataset()).clone(); + dataset.checkout_latest().await?; + self.set_dataset(dataset); self.pinned_version = None; Ok(()) } @@ -467,6 +481,18 @@ impl StorageBase { self.pinned_version = None; } + /// Current dataset snapshot (`Arc` clone; cheap). + #[inline] + pub(crate) fn current_dataset(&self) -> Arc { + self.dataset.load_full() + } + + /// Publish a replacement dataset handle after a mutating Lance op. + #[inline] + pub(crate) fn set_dataset(&self, dataset: Dataset) { + self.dataset.store(Arc::new(dataset)); + } + // ---------------------------------------------------------------- writes /// Durably append `batches` through this instance's MemWAL shard. @@ -560,7 +586,7 @@ impl StorageBase { ..Default::default() }; let writer = Arc::new( - self.dataset + self.current_dataset() .mem_wal_writer(self.write_shard, config) .await?, ); @@ -765,8 +791,9 @@ impl StorageBase { // Materialize anything buffered so it is eligible for this pass. self.flush().await?; } - let object_store = self.dataset.object_store(None).await?; - let branch_location = self.dataset.branch_location(); + let dataset = self.current_dataset(); + let object_store = dataset.object_store(None).await?; + let branch_location = dataset.branch_location(); let manifest_store = ShardManifestStore::new( object_store, &branch_location.path, @@ -939,8 +966,9 @@ impl StorageBase { /// only leaks one directory. async fn delete_merged_generation_dirs(&self, merged_paths: &[String]) -> LanceResult<()> { let phase = timer_start!(); - let object_store = self.dataset.object_store(None).await?; - let branch_path = self.dataset.branch_location().path.clone(); + let dataset = self.current_dataset(); + let object_store = dataset.object_store(None).await?; + let branch_path = dataset.branch_location().path.clone(); for path in merged_paths { let gen_dir = branch_path .clone() @@ -978,11 +1006,12 @@ impl StorageBase { &self, manifest: &ShardManifest, ) -> LanceResult<(HashSet, Vec, Vec, Arc)> { - let base_uri = self.dataset.uri().trim_end_matches('/').to_string(); + let dataset = self.current_dataset(); + let base_uri = dataset.uri().trim_end_matches('/').to_string(); let mut merged_generations: HashSet = HashSet::new(); let mut merged_paths: Vec = Vec::new(); let mut generation_batches: Vec<(u64, Vec)> = Vec::new(); - let merge_schema: Arc = Arc::new(self.dataset.schema().into()); + let merge_schema: Arc = Arc::new(dataset.schema().into()); // Read at most `merge_max_generations` generations per pass. // @@ -1058,13 +1087,13 @@ impl StorageBase { merge_schema, ); let mut builder = MergeInsertBuilder::try_new( - Arc::new(self.dataset.clone()), + self.current_dataset(), vec![self.key_column.clone()], )?; builder.when_matched(WhenMatched::UpdateAll); let job = builder.try_build()?; let (dataset, _) = job.execute_reader(reader).await?; - self.dataset = Arc::unwrap_or_clone(dataset); + self.set_dataset(Arc::unwrap_or_clone(dataset)); Ok(()) } @@ -1079,7 +1108,7 @@ impl StorageBase { }; self.refresh_latest().await?; - let base_schema: Arc = Arc::new(self.dataset.schema().into()); + let base_schema: Arc = Arc::new(self.current_dataset().schema().into()); align_batch_to_schema( RecordBatch::new_empty(base_schema.clone()), latest_schema.clone(), @@ -1092,13 +1121,15 @@ impl StorageBase { .cloned() .collect::>(); if !missing_fields.is_empty() { - self.dataset + let mut dataset = (*self.current_dataset()).clone(); + dataset .add_columns( NewColumnTransform::AllNulls(Arc::new(Schema::new(missing_fields))), None, None, ) .await?; + self.set_dataset(dataset); } Ok(()) } @@ -1147,17 +1178,19 @@ impl StorageBase { ..Default::default() }; + let mut dataset = (*self.current_dataset()).clone(); let result = match config.max_source_fragments { Some(max_source_fragments) => { compact_files_incremental( - &mut self.dataset, + &mut dataset, lance_options, max_source_fragments.max(1), ) .await } - None => compact_files(&mut self.dataset, lance_options, None).await, + None => compact_files(&mut dataset, lance_options, None).await, }; + self.set_dataset(dataset); match result { Ok(metrics) => { @@ -1199,7 +1232,8 @@ impl StorageBase { /// scan of those generations. pub async fn create_key_zonemap_index(&mut self) -> LanceResult<()> { info!(column = %self.key_column, "creating ZoneMap index on key column"); - self.dataset + let mut dataset = (*self.current_dataset()).clone(); + dataset .create_index_builder( &[self.key_column.as_str()], IndexType::ZoneMap, @@ -1208,6 +1242,7 @@ impl StorageBase { .name(ID_INDEX_NAME.to_string()) .replace(true) .await?; + self.set_dataset(dataset); // Reload the handle so subsequent reads on this instance observe the new // index (mirrors the reload done after `compact`). self.reload().await @@ -1219,7 +1254,7 @@ impl StorageBase { /// same config it would pass to [`Self::compact`]. #[must_use] pub fn should_compact(&self, config: &CompactionConfig) -> bool { - if self.dataset.count_fragments() < config.min_fragments { + if self.current_dataset().count_fragments() < config.min_fragments { return false; } if !config.quiet_hours.is_empty() { @@ -1242,7 +1277,7 @@ impl StorageBase { #[must_use] pub fn compaction_stats(&self) -> CompactionStats { CompactionStats { - total_fragments: self.dataset.count_fragments(), + total_fragments: self.current_dataset().count_fragments(), is_compacting: false, last_compaction: self.last_compaction, last_error: self.last_compaction_error.clone(), @@ -1253,10 +1288,13 @@ impl StorageBase { /// Reload the base dataset handle through [`Self::load_with_options`], so /// the shared session and storage options are never dropped. pub async fn reload(&mut self) -> LanceResult<()> { - let uri = self.dataset.uri().to_string(); - self.dataset = - Self::load_with_options(&uri, self.storage_options.clone(), self.session.clone()) - .await?; + let dataset = Self::load_with_options( + &self.uri, + self.storage_options.clone(), + self.session.clone(), + ) + .await?; + self.set_dataset(dataset); self.pinned_version = None; Ok(()) } @@ -1278,14 +1316,17 @@ impl StorageBase { if self.mem_wal_index_present().await? { return Ok(()); } - match self - .dataset + let mut dataset = (*self.current_dataset()).clone(); + match dataset .initialize_mem_wal() .unsharded() .execute() .await { - Ok(()) => Ok(()), + Ok(()) => { + self.set_dataset(dataset); + Ok(()) + } Err(err) => { // A concurrent first-writer may have created the index between // our check and our commit. Reload and accept it if so. @@ -1300,7 +1341,7 @@ impl StorageBase { } async fn mem_wal_index_present(&self) -> LanceResult { - let indices = self.dataset.load_indices().await?; + let indices = self.current_dataset().load_indices().await?; Ok(indices.iter().any(|i| i.name == MEM_WAL_INDEX_NAME)) } @@ -1309,14 +1350,14 @@ impl StorageBase { pub fn flushed_generation_uri(&self, shard_id: Uuid, path: &str) -> String { format!( "{}/_mem_wal/{shard_id}/{path}", - self.dataset.uri().trim_end_matches('/') + self.uri.trim_end_matches('/') ) } /// Open a flushed generation dataset, inheriting the base dataset's session /// and this store's storage options. pub async fn open_flushed_dataset(&self, uri: &str) -> LanceResult { - let mut builder = DatasetBuilder::from_uri(uri).with_session(self.dataset.session()); + let mut builder = DatasetBuilder::from_uri(uri).with_session(self.current_dataset().session()); if let Some(options) = self.storage_options.clone() { builder = builder.with_storage_options(options); } @@ -1327,9 +1368,10 @@ impl StorageBase { /// bounded-concurrent so stores with many writer instances do not pay one /// object-store round trip per shard serially. pub async fn wal_shard_snapshots(&self) -> LanceResult> { - let object_store = self.dataset.object_store(None).await?; - let branch_path = self.dataset.branch_location().path.clone(); - let shard_ids = self.dataset.list_mem_wal_latest_shard_ids().await?; + let dataset = self.current_dataset(); + let object_store = dataset.object_store(None).await?; + let branch_path = dataset.branch_location().path.clone(); + let shard_ids = dataset.list_mem_wal_latest_shard_ids().await?; let snapshots: Vec> = stream::iter(shard_ids) .map(|shard_id| { @@ -1385,7 +1427,7 @@ impl StorageBase { }) }) .collect(); - let session = self.dataset.session(); + let session = self.current_dataset().session(); let storage_options = self.storage_options.clone(); stream::iter(generation_paths) @@ -1431,22 +1473,19 @@ impl StorageBase { shard_snapshots: Vec, ) -> LsmScanner { let merge_key = vec![self.key_column.clone()]; + let dataset = self.current_dataset(); match source { - ListSource::Fragments => { - LsmScanner::new(Arc::new(self.dataset.clone()), Vec::new(), merge_key) - } - ListSource::All => { - LsmScanner::new(Arc::new(self.dataset.clone()), shard_snapshots, merge_key) - } + ListSource::Fragments => LsmScanner::new(dataset, Vec::new(), merge_key), + ListSource::All => LsmScanner::new(dataset, shard_snapshots, merge_key), ListSource::Wal => { - let arrow_schema: Schema = self.dataset.schema().into(); + let arrow_schema: Schema = dataset.schema().into(); LsmScanner::without_base_table( Arc::new(arrow_schema), - self.dataset.uri().trim_end_matches('/').to_string(), + dataset.uri().trim_end_matches('/').to_string(), shard_snapshots, merge_key, ) - .with_session(self.dataset.session()) + .with_session(dataset.session()) } } } From a4a42be4b808deb216e4169b339e4ad3d8e5f18d Mon Sep 17 00:00:00 2001 From: Lucy Ge Date: Sun, 9 Aug 2026 06:07:30 -0700 Subject: [PATCH 02/11] uri change --- crates/lance-context-core/src/datagen_store.rs | 2 +- crates/lance-context-core/src/generic_store.rs | 2 +- crates/lance-context-core/src/rollout_store.rs | 2 +- crates/lance-context-core/src/store.rs | 4 ++-- crates/lance-context-core/src/store_base.rs | 13 +++++-------- python/src/lib.rs | 2 +- 6 files changed, 11 insertions(+), 14 deletions(-) diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index 3d34ff1..c2cf860 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -125,7 +125,7 @@ impl DatagenStore { } #[must_use] - pub fn uri(&self) -> &str { + pub fn uri(&self) -> String { self.base.uri() } diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index 202af55..013693c 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -221,7 +221,7 @@ impl GenericStore { /// URI of the underlying Lance dataset. #[must_use] - pub fn uri(&self) -> &str { + pub fn uri(&self) -> String { self.base.uri() } diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 90fd35a..a2b4047 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -496,7 +496,7 @@ impl RolloutStore { } /// URI of the underlying Lance dataset. #[must_use] - pub fn uri(&self) -> &str { + pub fn uri(&self) -> String { self.base.uri() } diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index f18c839..5d22430 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -658,7 +658,7 @@ impl ContextStore { /// URI of the underlying Lance dataset. #[must_use] - pub fn uri(&self) -> &str { + pub fn uri(&self) -> String { self.base.uri() } @@ -2102,7 +2102,7 @@ impl ContextStore { // have exactly one owner. A second handle is the right model anyway -- // compaction only rewrites base-table fragments and takes `&mut`, so // sharing a handle with the write path would mean contending for it. - let uri = self.uri().to_string(); + let uri = self.uri(); let interval_secs = self.compaction_config.check_interval_secs; let options = ContextStoreOptions { storage_options: self.base.storage_options.clone(), diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index fde7fa7..43d53e7 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -269,8 +269,6 @@ pub(crate) struct StorageBase { /// Wrapped in [`ArcSwap`] so a merge/compact/reload can publish a new /// handle without requiring exclusive `&mut` for every reader. pub dataset: ArcSwap, - /// Dataset URI; stable for the lifetime of this handle. - uri: String, /// MemWAL shard this instance writes to (derived from `shard_id`). pub write_shard: Uuid, /// Object-store options, retained so a self-merge can re-append flushed @@ -408,10 +406,8 @@ impl StorageBase { .into()); } - let uri = dataset.uri().to_string(); let mut base = Self { dataset: ArcSwap::from_pointee(dataset), - uri, write_shard: derive_shard_id(shard_id.as_deref()), storage_options, session, @@ -434,8 +430,8 @@ impl StorageBase { /// URI of the underlying Lance dataset. #[must_use] - pub fn uri(&self) -> &str { - &self.uri + pub fn uri(&self) -> String { + self.current_dataset().uri().to_string() } /// Current base dataset manifest version. @@ -1288,8 +1284,9 @@ impl StorageBase { /// Reload the base dataset handle through [`Self::load_with_options`], so /// the shared session and storage options are never dropped. pub async fn reload(&mut self) -> LanceResult<()> { + let uri = self.uri(); let dataset = Self::load_with_options( - &self.uri, + &uri, self.storage_options.clone(), self.session.clone(), ) @@ -1350,7 +1347,7 @@ impl StorageBase { pub fn flushed_generation_uri(&self, shard_id: Uuid, path: &str) -> String { format!( "{}/_mem_wal/{shard_id}/{path}", - self.uri.trim_end_matches('/') + self.uri().trim_end_matches('/') ) } diff --git a/python/src/lib.rs b/python/src/lib.rs index 63504d1..d0c862d 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -790,7 +790,7 @@ impl Context { // two writers racing for one shard. A fork branches the in-memory // `Context` and shares the underlying dataset, which a fresh handle // gives it. - let uri = self.store.uri().to_string(); + let uri = self.store.uri(); let store = py.allow_threads(|| self.runtime.block_on(ContextStore::open(&uri))); Ok(Self { inner: self.inner.fork(branch_name), From 71c0718e967586eeff06c25b74a0bfc013b2db5c Mon Sep 17 00:00:00 2001 From: Lucy Ge Date: Sun, 9 Aug 2026 19:22:49 -0700 Subject: [PATCH 03/11] add merge_write_lock_is_only_held_for_short_commit test --- .../tests/wal_merge_concurrency.rs | 103 ++++++++++++++++++ 1 file changed, 103 insertions(+) diff --git a/crates/lance-context-core/tests/wal_merge_concurrency.rs b/crates/lance-context-core/tests/wal_merge_concurrency.rs index c2025fe..fc46819 100644 --- a/crates/lance-context-core/tests/wal_merge_concurrency.rs +++ b/crates/lance-context-core/tests/wal_merge_concurrency.rs @@ -414,3 +414,106 @@ async fn append_is_not_blocked_for_the_duration_of_a_merge() { ); assert_eq!(ids.len(), 26, "all rows readable exactly once"); } + +/// Issue #198 / the ~17s production stall, measured as lock phases. +/// +/// Old shape: one `RwLock::write()` spanned seal + generation reads + append + +/// drain, so every concurrent `add` waited for the whole merge (~17s on abfss). +/// New shape (same as the server sweeper): prepare under a *read* lock, write +/// lock only for the short commit. +/// +/// This test records those two durations and races an append against prepare. +/// It fails if the exclusive lock is held for most of the merge again. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn merge_write_lock_is_only_held_for_short_commit() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = Arc::new(RwLock::new( + RolloutStore::open_with_options(&uri, opts("solo")) + .await + .unwrap(), + )); + + // Fat generations so prepare (read every flushed gen) dominates wall time + // even on a fast local disk — otherwise the assertion becomes vacuous. + let fat = "x".repeat(64 * 1024); + for i in 0..20 { + let mut r: RolloutRecord = rec(&format!("bulk-{i}")); + r.content = Some(fat.clone()); + store.read().await.add(&[r]).await.unwrap(); + store.read().await.flush().await.unwrap(); + } + + let store_for_merge = store.clone(); + let merge_handle = tokio::spawn(async move { + let prepare_start = Instant::now(); + let prepared = { + let guard = store_for_merge.read().await; + guard.prepare_cleanup_merge().await.unwrap() + }; + let prepare_elapsed = prepare_start.elapsed(); + + let Some((manifest_store, manifest, prepared)) = prepared else { + panic!("expected pending generations to merge"); + }; + + let write_start = Instant::now(); + let reclaimed = { + let mut guard = store_for_merge.write().await; + guard + .commit_prepared_merge(&manifest_store, &manifest, prepared) + .await + .unwrap() + }; + let write_lock_elapsed = write_start.elapsed(); + (reclaimed, prepare_elapsed, write_lock_elapsed) + }); + + // While prepare should be holding only a *shared* lock, appends must land + // quickly — this is the user-visible half of the ~17s stall. + tokio::time::sleep(Duration::from_millis(20)).await; + let append_start = Instant::now(); + store + .read() + .await + .add(&[rec("during-prepare")]) + .await + .unwrap(); + let append_elapsed = append_start.elapsed(); + + let (reclaimed, prepare_elapsed, write_lock_elapsed) = merge_handle.await.unwrap(); + + eprintln!( + "merge phases: prepare={prepare_elapsed:?} write_lock={write_lock_elapsed:?} \ + append_during_prepare={append_elapsed:?} reclaimed={reclaimed}" + ); + + assert!(reclaimed > 0, "merge should reclaim generations"); + assert!( + prepare_elapsed > Duration::from_millis(5), + "prepare should do measurable work so the lock split is observable; \ + got prepare={prepare_elapsed:?}" + ); + // Exclusive lock must be the short phase. In the old bug it ≈ prepare. + assert!( + write_lock_elapsed * 3 < prepare_elapsed + || write_lock_elapsed < Duration::from_millis(100), + "write lock held {write_lock_elapsed:?} but prepare took {prepare_elapsed:?}; \ + exclusive lock appears to cover the expensive merge work again (#198 / ~17s stall)" + ); + assert!( + append_elapsed < prepare_elapsed + && (append_elapsed * 2 < prepare_elapsed + || append_elapsed < Duration::from_millis(200)), + "append during prepare took {append_elapsed:?} while prepare took {prepare_elapsed:?}; \ + append was blocked as if the merge held the write lock" + ); + + store.read().await.flush().await.unwrap(); + let ids = read_ids(&store).await; + assert!( + ids.contains(&"during-prepare".to_string()), + "row appended under the shared prepare lock must be readable" + ); + assert_eq!(ids.len(), 21, "all rows readable exactly once: {ids:?}"); +} From 0e66d79e22ef1561dcc399df9f602fb285a8c4f3 Mon Sep 17 00:00:00 2001 From: Lucy Ge Date: Mon, 10 Aug 2026 06:25:07 -0700 Subject: [PATCH 04/11] remove all places that uses &mut self and replace with &self for storagebase so ops like merge, compact & schema evolution doenst need mutex of store --- .../lance-context-core/src/datagen_store.rs | 12 +- .../lance-context-core/src/generic_store.rs | 12 +- .../lance-context-core/src/rollout_store.rs | 28 ++-- crates/lance-context-core/src/store.rs | 18 +-- crates/lance-context-core/src/store_base.rs | 125 ++++++++++-------- .../tests/wal_merge_concurrency.rs | 42 +++--- .../tests/wal_merge_generation_cleanup.rs | 2 +- .../src/routes/compact.rs | 2 +- .../src/routes/datagen.rs | 10 +- .../src/routes/generic.rs | 2 +- .../src/routes/records.rs | 6 +- .../src/routes/rollouts.rs | 10 +- .../src/routes/versions.rs | 2 +- crates/lance-context-server/src/state.rs | 10 +- crates/lance-context-server/src/sweeper.rs | 20 +-- 15 files changed, 156 insertions(+), 145 deletions(-) diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index c2cf860..f31239c 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -141,7 +141,7 @@ impl DatagenStore { } /// Refresh this handle to the latest base-table manifest. - pub async fn refresh_latest(&mut self) -> LanceResult<()> { + pub async fn refresh_latest(&self) -> LanceResult<()> { self.base.refresh_latest().await } @@ -150,7 +150,7 @@ impl DatagenStore { /// The supplied slice is persisted as one MemWAL generation. Callers should /// include FIELD_* events and the corresponding STEP_COMPLETED marker in /// the same call so a crash cannot expose a partially checkpointed step. - pub async fn append(&mut self, events: &[DatagenEvent]) -> LanceResult { + pub async fn append(&self, events: &[DatagenEvent]) -> LanceResult { if events.is_empty() { return Ok(self.base.version()); } @@ -189,7 +189,7 @@ impl DatagenStore { } /// Gracefully stop this store's resident MemWAL writer. - pub async fn close(&mut self) -> LanceResult<()> { + pub async fn close(&self) -> LanceResult<()> { self.base.close().await } @@ -365,7 +365,7 @@ impl DatagenStore { /// Merge every currently flushed generation owned by this writer into the /// base table. - pub async fn cleanup_own_shard(&mut self) -> LanceResult { + pub async fn cleanup_own_shard(&self) -> LanceResult { self.base.cleanup_own_shard().await } @@ -380,7 +380,7 @@ impl DatagenStore { /// the shared base table and Lance treats two concurrent `Rewrite` commits /// as a conflict. pub async fn compact( - &mut self, + &self, options: Option, ) -> LanceResult { self.base.compact(options).await @@ -402,7 +402,7 @@ impl DatagenStore { /// Build a ZoneMap scalar index on `event_id`, the table's key column. /// Idempotent. Datagen previously had no scalar index, so every point /// lookup by event id scanned. - pub async fn create_event_id_index(&mut self) -> LanceResult<()> { + pub async fn create_event_id_index(&self) -> LanceResult<()> { self.base.create_key_zonemap_index().await } diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index 013693c..3515553 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -238,7 +238,7 @@ impl GenericStore { } /// Refresh this handle to the latest base-table manifest. - pub async fn refresh_latest(&mut self) -> LanceResult<()> { + pub async fn refresh_latest(&self) -> LanceResult<()> { self.base.refresh_latest().await } @@ -356,19 +356,19 @@ impl GenericStore { } /// Close the resident writer, draining its background tasks. Idempotent. - pub async fn close(&mut self) -> LanceResult<()> { + pub async fn close(&self) -> LanceResult<()> { self.base.close().await } /// Merge flushed generations into the base table once the count trigger is /// met. Returns how many were reclaimed. - pub async fn maybe_merge_wal(&mut self) -> LanceResult { + pub async fn maybe_merge_wal(&self) -> LanceResult { self.base.maybe_merge_own_shard().await } /// Seal, then merge **every** pending generation into the base table — the /// time half of the "time OR count" trigger. - pub async fn cleanup_wal(&mut self) -> LanceResult { + pub async fn cleanup_wal(&self) -> LanceResult { self.base.cleanup_own_shard().await } @@ -380,7 +380,7 @@ impl GenericStore { /// Compact the base table's small fragments. Drive from a single external /// trigger, not per worker — see [`StorageBase::compact`]. pub async fn compact( - &mut self, + &self, options: Option, ) -> LanceResult { self.base.compact(options).await @@ -399,7 +399,7 @@ impl GenericStore { } /// Build a ZoneMap scalar index on `id`. Idempotent. - pub async fn create_id_index(&mut self) -> LanceResult<()> { + pub async fn create_id_index(&self) -> LanceResult<()> { self.base.create_key_zonemap_index().await } diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index a2b4047..a301aaf 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -508,7 +508,7 @@ impl RolloutStore { /// Checkout a specific dataset version — recovers the exact rollout set that /// trained a checkpoint (spec §3, reproducibility). - pub async fn checkout(&mut self, version_id: u64) -> LanceResult<()> { + pub async fn checkout(&self, version_id: u64) -> LanceResult<()> { self.base.checkout(version_id).await } @@ -527,7 +527,7 @@ impl RolloutStore { /// Long-lived read handles call this before a new request so compaction or /// WAL merges committed by another process become visible without paying /// the cost of reopening the dataset and rebuilding all session caches. - pub async fn refresh_latest(&mut self) -> LanceResult<()> { + pub async fn refresh_latest(&self) -> LanceResult<()> { self.base.refresh_latest().await } @@ -595,14 +595,14 @@ impl RolloutStore { /// Gracefully close the resident writer, draining its background tasks. /// Idempotent. See `StorageBase::close`. - pub async fn close(&mut self) -> LanceResult<()> { + pub async fn close(&self) -> LanceResult<()> { self.base.close().await } /// Merge this instance's flushed generations into the base table **if** the /// shard has accumulated at least `merge_after_generations` of them (the /// count trigger; `0` disables it). No-op otherwise. - pub async fn maybe_merge_own_shard(&mut self) -> LanceResult { + pub async fn maybe_merge_own_shard(&self) -> LanceResult { self.base.maybe_merge_own_shard().await } @@ -625,7 +625,7 @@ impl RolloutStore { /// Commit a merge prepared by [`Self::prepare_merge_if_ready`]. pub async fn commit_prepared_merge( - &mut self, + &self, manifest_store: &ShardManifestStore, manifest: &ShardManifest, prepared: PreparedMerge, @@ -639,7 +639,7 @@ impl RolloutStore { /// then fold **every** pending flushed generation into the base table. This /// is the *time* half of the "time OR count" trigger and is deliberately not /// gated by the count threshold. See `StorageBase::cleanup_own_shard`. - pub async fn cleanup_own_shard(&mut self) -> LanceResult { + pub async fn cleanup_own_shard(&self) -> LanceResult { self.base.cleanup_own_shard().await } @@ -649,7 +649,7 @@ impl RolloutStore { /// compaction rewrites the shared base table and two concurrent `Rewrite` /// commits conflict. See `StorageBase::compact`. pub async fn compact( - &mut self, + &self, options: Option, ) -> LanceResult { self.base.compact(options).await @@ -657,7 +657,7 @@ impl RolloutStore { /// Build a ZoneMap scalar index on the base table's `id` column. Idempotent. /// See `StorageBase::create_key_zonemap_index`. - pub async fn create_id_zonemap_index(&mut self) -> LanceResult<()> { + pub async fn create_id_zonemap_index(&self) -> LanceResult<()> { self.base.create_key_zonemap_index().await } @@ -3006,7 +3006,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { shard_id: Some("trajectory-test".to_string()), @@ -3158,7 +3158,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { storage_options: None, @@ -3332,7 +3332,7 @@ mod tests { let artifact_bytes = b"\x00\x01\x02compacted"; let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { storage_options: None, @@ -3464,7 +3464,7 @@ mod tests { .unwrap(); runtime.block_on(async { // Seed the base table via A with several fragments to compact. - let mut a = RolloutStore::open_with_options( + let a = RolloutStore::open_with_options( &uri, RolloutStoreOptions { storage_options: None, @@ -3568,7 +3568,7 @@ mod tests { let artifact_bytes = b"\x00\x01\x02merged-trace"; let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { storage_options: None, @@ -3625,7 +3625,7 @@ mod tests { let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { storage_options: None, diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 5d22430..4387d44 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -1399,7 +1399,7 @@ impl ContextStore { /// /// Existing rows are stored as null in the new column and read back as an /// empty relationship list. - pub async fn migrate_relationships_column(&mut self) -> LanceResult { + pub async fn migrate_relationships_column(&self) -> LanceResult { if self.has_relationships_column() { return Ok(false); } @@ -1415,7 +1415,7 @@ impl ContextStore { } /// Checkout a specific dataset version. - pub async fn checkout(&mut self, version_id: u64) -> LanceResult<()> { + pub async fn checkout(&self, version_id: u64) -> LanceResult<()> { self.base.checkout(version_id).await } @@ -1426,7 +1426,7 @@ impl ContextStore { } /// Refresh this handle to the latest base-table manifest. - pub async fn refresh_latest(&mut self) -> LanceResult<()> { + pub async fn refresh_latest(&self) -> LanceResult<()> { self.base.refresh_latest().await } @@ -1888,7 +1888,7 @@ impl ContextStore { /// Manually trigger compaction to merge small fragments. pub async fn compact( - &mut self, + &self, options: Option, ) -> LanceResult { let config = options.unwrap_or_else(|| self.compaction_config.clone()); @@ -1971,7 +1971,7 @@ impl ContextStore { /// Gracefully close the resident MemWAL writer, draining its background /// tasks and sealing whatever it still buffers. Idempotent. - pub async fn close(&mut self) -> LanceResult<()> { + pub async fn close(&self) -> LanceResult<()> { self.base.close().await } @@ -1985,14 +1985,14 @@ impl ContextStore { /// unioned all of them, so read cost grew without bound in the number of /// writes. Merging is what keeps that bounded — drive it from a sweeper, or /// use [`Self::cleanup_wal`] for the time-based trigger. - pub async fn maybe_merge_wal(&mut self) -> LanceResult { + pub async fn maybe_merge_wal(&self) -> LanceResult { self.base.maybe_merge_own_shard().await } /// Seal, then fold **every** pending flushed generation into the base table. /// The time half of the "time OR count" trigger, so deliberately not gated /// by the count threshold. Returns the number of generations reclaimed. - pub async fn cleanup_wal(&mut self) -> LanceResult { + pub async fn cleanup_wal(&self) -> LanceResult { self.base.cleanup_own_shard().await } @@ -2040,7 +2040,7 @@ impl ContextStore { } /// Ensure the configured id index exists on the dataset. - async fn ensure_id_index(&mut self) -> LanceResult<()> { + async fn ensure_id_index(&self) -> LanceResult<()> { if self.id_index_type == IdIndexType::None { return Ok(()); } @@ -2054,7 +2054,7 @@ impl ContextStore { } /// Create (or replace) the scalar index on the `id` column. - pub async fn create_id_index(&mut self) -> LanceResult<()> { + pub async fn create_id_index(&self) -> LanceResult<()> { let index_type = match self.id_index_type { IdIndexType::ZoneMap => IndexType::ZoneMap, IdIndexType::BTree => IndexType::BTree, diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index 43d53e7..34b063f 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -36,7 +36,8 @@ //! latest schema to evolve a base table to — via [`StorageBaseOptions`]. use std::collections::{HashMap, HashSet}; -use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; use arc_swap::ArcSwap; use arrow_array::{new_null_array, RecordBatch, RecordBatchIterator, UInt32Array}; @@ -183,7 +184,7 @@ pub enum ListSource { /// /// Produced by `StorageBase::prepare_merge_if_ready` under `&self` (so appends /// keep running while it reads object storage) and consumed by -/// `StorageBase::commit_prepared_merge` under `&mut self`. `PreparedMerge` is +/// `StorageBase::commit_prepared_merge` (also `&self`). `PreparedMerge` is /// public because it appears in [`RolloutStore`]'s prepare/commit split, but /// its fields are opaque. /// @@ -260,6 +261,14 @@ pub(crate) struct StorageBaseOptions { pub seal_on_put: bool, } +/// Per-handle compaction counters/timestamps. +#[derive(Debug, Default)] +struct CompactionState { + last_compaction: Option>, + total_compactions: u64, + last_error: Option, +} + /// Schema-agnostic Lance storage: dataset handle, MemWAL write path, WAL merge, /// compaction, indexing, and LSM reads. See the module docs. pub(crate) struct StorageBase { @@ -287,18 +296,17 @@ pub(crate) struct StorageBase { /// Self-merge threshold; `0` disables it. merge_after_generations: usize, merge_max_generations: usize, - /// Timestamp of the last successful [`Self::compact`] on this handle. - last_compaction: Option>, - /// Number of successful compactions performed by this handle. - total_compactions: u64, - /// Error message from the most recent failed compaction on this handle. - last_compaction_error: Option, + /// Compaction bookkeeping for this handle (interior-mutable so [`Self::compact`] can be `&self`). + compaction: Mutex, /// Explicit time-travel version selected by [`Self::checkout`]. /// /// A point-read miss may refresh an ordinary long-lived handle to avoid a /// false negative from a stale manifest, but must never advance a handle /// whose caller deliberately selected a historical version. - pinned_version: Option, + /// + /// `0` means unpinned; any other value is the pinned manifest version. + /// (Lance dataset versions are 1-based, so `0` is never a real pin.) + pinned_version: AtomicU64, /// Resident MemWAL writer for this instance's shard, wrapped for `&self` /// concurrent access. The [`tokio::sync::Mutex`] is held only to /// fetch-or-open and clone the `Arc` (see [`Self::resident_writer`]) and to @@ -406,7 +414,7 @@ impl StorageBase { .into()); } - let mut base = Self { + let base = Self { dataset: ArcSwap::from_pointee(dataset), write_shard: derive_shard_id(shard_id.as_deref()), storage_options, @@ -416,14 +424,12 @@ impl StorageBase { seal_on_put, merge_after_generations: merge_after_generations.unwrap_or(0), merge_max_generations: merge_max_generations.unwrap_or(DEFAULT_MERGE_MAX_GENERATIONS), - last_compaction: None, - total_compactions: 0, - last_compaction_error: None, - pinned_version: None, + compaction: Mutex::new(CompactionState::default()), + pinned_version: AtomicU64::new(0), write_writer: tokio::sync::Mutex::new(None), }; // `ensure_mem_wal` may reload the dataset on a concurrent first-writer - // race, which is why it must run here where we hold `&mut`. + // race; it publishes the new handle via ArcSwap. base.ensure_mem_wal().await?; Ok(base) } @@ -441,20 +447,28 @@ impl StorageBase { } /// Check out a specific base dataset version (time travel). - pub async fn checkout(&mut self, version_id: u64) -> LanceResult<()> { + /// + /// `version_id` must be non-zero (`0` is reserved to mean "unpinned"). + pub async fn checkout(&self, version_id: u64) -> LanceResult<()> { + if version_id == 0 { + return Err(ArrowError::InvalidArgumentError( + "dataset version 0 is reserved; pin a real manifest version (>= 1)".to_string(), + ) + .into()); + } let dataset = self .current_dataset() .checkout_version(version_id) .await?; self.set_dataset(dataset); - self.pinned_version = Some(version_id); + self.pinned_version.store(version_id, Ordering::Release); Ok(()) } /// Whether this handle was explicitly checked out to a historical version. #[must_use] pub fn is_version_pinned(&self) -> bool { - self.pinned_version.is_some() + self.pinned_version.load(Ordering::Acquire) != 0 } /// Refresh this handle to the latest base-table manifest while retaining its @@ -463,18 +477,18 @@ impl StorageBase { /// Long-lived read handles call this before a new request so compaction or /// WAL merges committed by another process become visible without paying the /// cost of reopening the dataset and rebuilding all session caches. - pub async fn refresh_latest(&mut self) -> LanceResult<()> { + pub async fn refresh_latest(&self) -> LanceResult<()> { let mut dataset = (*self.current_dataset()).clone(); dataset.checkout_latest().await?; self.set_dataset(dataset); - self.pinned_version = None; + self.clear_version_pin(); Ok(()) } /// Mark this handle as no longer pinned after a concrete store mutates the /// dataset directly. - pub(crate) fn clear_version_pin(&mut self) { - self.pinned_version = None; + pub(crate) fn clear_version_pin(&self) { + self.pinned_version.store(0, Ordering::Release); } /// Current dataset snapshot (`Arc` clone; cheap). @@ -663,9 +677,9 @@ impl StorageBase { /// by an explicit `close().await`. Call this before dropping a store on a /// path that can `await` (e.g. an LRU eviction that owns the last handle). /// Idempotent: a no-op when no writer is resident. - pub async fn close(&mut self) -> LanceResult<()> { - // `&mut self` gives exclusive access, so `get_mut` avoids an async lock. - if let Some(writer) = self.write_writer.get_mut().take() { + pub async fn close(&self) -> LanceResult<()> { + let writer = self.write_writer.lock().await.take(); + if let Some(writer) = writer { match Arc::try_unwrap(writer) { // Sole owner: drain the writer's background tasks gracefully. Ok(writer) => writer.close().await?, @@ -703,7 +717,7 @@ impl StorageBase { /// Merge this instance's flushed generations into the base table **if** the /// shard has accumulated at least `merge_after_generations` of them (the /// count trigger; `0` disables it). No-op otherwise. - pub async fn maybe_merge_own_shard(&mut self) -> LanceResult { + pub async fn maybe_merge_own_shard(&self) -> LanceResult { if self.merge_after_generations == 0 { return Ok(0); } @@ -726,7 +740,7 @@ impl StorageBase { /// would stay empty, so the threshold check would return `0` and never reach /// the merge — leaving rows durable but permanently invisible until a /// process restart replayed the WAL. - pub async fn cleanup_own_shard(&mut self) -> LanceResult { + pub async fn cleanup_own_shard(&self) -> LanceResult { self.flush().await?; // Threshold `1`: merge whenever at least one generation is pending. The // time trigger must not depend on the count threshold — that is what @@ -734,7 +748,7 @@ impl StorageBase { self.merge_own_shard_if_ready(1).await } - async fn merge_own_shard_if_ready(&mut self, threshold: usize) -> LanceResult { + async fn merge_own_shard_if_ready(&self, threshold: usize) -> LanceResult { let Some((manifest_store, manifest, prepared)) = self.prepare_merge_if_ready(threshold).await? else { @@ -758,7 +772,7 @@ impl StorageBase { /// ```ignore /// let prepared = { store.read().await.prepare_merge_if_ready(1).await? }; /// if let Some((manifest_store, manifest, prepared)) = prepared { - /// store.write().await.commit_prepared_merge(&manifest_store, &manifest, prepared).await?; + /// store.commit_prepared_merge(&manifest_store, &manifest, prepared).await?; /// } /// ``` pub async fn prepare_merge_if_ready( @@ -811,7 +825,7 @@ impl StorageBase { /// Commit a merge prepared by [`Self::prepare_merge_if_ready`]. pub async fn commit_prepared_merge( - &mut self, + &self, manifest_store: &ShardManifestStore, manifest: &ShardManifest, prepared: PreparedMerge, @@ -861,9 +875,8 @@ impl StorageBase { })) } - /// The `&mut self` half of a merge: append the prepared rows to the base - /// table, drain the merged generations from the manifest, and delete their - /// directories. + /// Append the prepared rows to the base table, drain the merged generations + /// from the manifest, and delete their directories. /// /// # Surgical drain, not blanket clear /// @@ -884,7 +897,7 @@ impl StorageBase { /// of appending physical duplicates. The next attempt can then drain the /// manifest without relying on a particular Lance read-plan shape. async fn commit_merge( - &mut self, + &self, manifest_store: &ShardManifestStore, manifest: &ShardManifest, prepared: PreparedMerge, @@ -917,7 +930,7 @@ impl StorageBase { "append", self.merge_prepared_batches(batches, merge_schema).await )?; - self.pinned_version = None; + self.clear_version_pin(); } // Reuse the shard's *current* epoch rather than claiming a new one: @@ -946,7 +959,7 @@ impl StorageBase { )?; self.delete_merged_generation_dirs(&merged_paths).await?; - self.pinned_version = None; + self.clear_version_pin(); Ok(true) } @@ -1074,7 +1087,7 @@ impl StorageBase { /// row per key. `UpdateAll` preserves normal LSM last-write-wins semantics /// while also making a retry after an interrupted manifest drain idempotent. async fn merge_prepared_batches( - &mut self, + &self, batches: Vec, merge_schema: Arc, ) -> LanceResult<()> { @@ -1098,7 +1111,7 @@ impl StorageBase { /// Missing nullable columns are added as all-null arrays. Existing unknown /// columns, type changes, and missing required columns remain hard errors. /// A no-op when the store declared no `latest_schema`. - pub async fn ensure_latest_schema(&mut self) -> LanceResult<()> { + pub async fn ensure_latest_schema(&self) -> LanceResult<()> { let Some(latest_schema) = self.latest_schema.clone() else { return Ok(()); }; @@ -1147,7 +1160,7 @@ impl StorageBase { /// to call while other workers are appending or WAL-merging: `Append` vs /// `Rewrite` is non-conflicting in Lance's matrix. pub async fn compact( - &mut self, + &self, options: Option, ) -> LanceResult { let config = options.unwrap_or_default(); @@ -1193,9 +1206,12 @@ impl StorageBase { // Reload the handle so the caller (and subsequent reads on this // instance) observe the compacted version. self.reload().await?; - self.last_compaction = Some(Utc::now()); - self.total_compactions += 1; - self.last_compaction_error = None; + { + let mut state = self.compaction.lock().unwrap_or_else(|e| e.into_inner()); + state.last_compaction = Some(Utc::now()); + state.total_compactions += 1; + state.last_error = None; + } info!( fragments_removed = metrics.fragments_removed, fragments_added = metrics.fragments_added, @@ -1205,7 +1221,10 @@ impl StorageBase { } Err(e) => { warn!(error = %e, "base-table compaction failed"); - self.last_compaction_error = Some(e.to_string()); + { + let mut state = self.compaction.lock().unwrap_or_else(|e| e.into_inner()); + state.last_error = Some(e.to_string()); + } Err(e) } } @@ -1226,7 +1245,7 @@ impl StorageBase { /// only ever needs to describe the base table's already-merged fragments — /// rows still living in unmerged WAL generations are found by the normal /// scan of those generations. - pub async fn create_key_zonemap_index(&mut self) -> LanceResult<()> { + pub async fn create_key_zonemap_index(&self) -> LanceResult<()> { info!(column = %self.key_column, "creating ZoneMap index on key column"); let mut dataset = (*self.current_dataset()).clone(); dataset @@ -1267,23 +1286,23 @@ impl StorageBase { /// Current compaction statistics for the base table. /// - /// `is_compacting` is always `false`: compaction runs synchronously under - /// the caller's `&mut self`, so a stats read cannot observe an in-flight - /// compaction on this handle. + /// `is_compacting` is always `false`: compaction runs synchronously on this + /// handle, so a stats read cannot observe an in-flight compaction here. #[must_use] pub fn compaction_stats(&self) -> CompactionStats { + let state = self.compaction.lock().unwrap_or_else(|e| e.into_inner()); CompactionStats { total_fragments: self.current_dataset().count_fragments(), is_compacting: false, - last_compaction: self.last_compaction, - last_error: self.last_compaction_error.clone(), - total_compactions: self.total_compactions, + last_compaction: state.last_compaction, + last_error: state.last_error.clone(), + total_compactions: state.total_compactions, } } /// Reload the base dataset handle through [`Self::load_with_options`], so /// the shared session and storage options are never dropped. - pub async fn reload(&mut self) -> LanceResult<()> { + pub async fn reload(&self) -> LanceResult<()> { let uri = self.uri(); let dataset = Self::load_with_options( &uri, @@ -1292,7 +1311,7 @@ impl StorageBase { ) .await?; self.set_dataset(dataset); - self.pinned_version = None; + self.clear_version_pin(); Ok(()) } @@ -1309,7 +1328,7 @@ impl StorageBase { /// `RetryableCommitConflict`. That is benign here — the winner created /// exactly the index we wanted — so we reload and treat "index now present" /// as success. Any other error propagates. - async fn ensure_mem_wal(&mut self) -> LanceResult<()> { + async fn ensure_mem_wal(&self) -> LanceResult<()> { if self.mem_wal_index_present().await? { return Ok(()); } diff --git a/crates/lance-context-core/tests/wal_merge_concurrency.rs b/crates/lance-context-core/tests/wal_merge_concurrency.rs index fc46819..798fb2e 100644 --- a/crates/lance-context-core/tests/wal_merge_concurrency.rs +++ b/crates/lance-context-core/tests/wal_merge_concurrency.rs @@ -25,18 +25,17 @@ use std::time::{Duration, Instant}; use lance_context_core::{RolloutRecord, RolloutStore, RolloutStoreOptions, ROLE_ASSISTANT}; use tokio::sync::RwLock; -/// Run one full merge exactly the way the server's sweepers do: seal + read the -/// generations under a **read** lock (so appends keep running), then take the -/// write lock only for the short commit. Returns generations reclaimed. +/// Run one full merge: seal + read generations, then commit. Both phases use +/// `&self` on the store (dataset handle is ArcSwap), so callers only need a +/// shared lock — appends are not blocked. Returns generations reclaimed. /// /// Every test drives merges through this helper so the lock discipline under -/// test is the same one production uses -- a test that merged under a single -/// exclusive lock would pass while the real stall persisted. +/// test matches production. async fn merge_like_sweeper(store: &Arc>) -> usize { let prepared = { store.read().await.prepare_cleanup_merge().await.unwrap() }; match prepared { Some((manifest_store, manifest, prepared)) => store - .write() + .read() .await .commit_prepared_merge(&manifest_store, &manifest, prepared) .await @@ -419,11 +418,11 @@ async fn append_is_not_blocked_for_the_duration_of_a_merge() { /// /// Old shape: one `RwLock::write()` spanned seal + generation reads + append + /// drain, so every concurrent `add` waited for the whole merge (~17s on abfss). -/// New shape (same as the server sweeper): prepare under a *read* lock, write -/// lock only for the short commit. +/// New shape: prepare and commit both use shared locks (`&self` + ArcSwap), so +/// appends are never blocked by merge. /// -/// This test records those two durations and races an append against prepare. -/// It fails if the exclusive lock is held for most of the merge again. +/// This test records prepare vs commit durations and races an append against +/// prepare. It fails if append appears serialized behind the merge again. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn merge_write_lock_is_only_held_for_short_commit() { let tmp = tempfile::tempdir().unwrap(); @@ -457,16 +456,16 @@ async fn merge_write_lock_is_only_held_for_short_commit() { panic!("expected pending generations to merge"); }; - let write_start = Instant::now(); + let commit_start = Instant::now(); let reclaimed = { - let mut guard = store_for_merge.write().await; + let guard = store_for_merge.read().await; guard .commit_prepared_merge(&manifest_store, &manifest, prepared) .await .unwrap() }; - let write_lock_elapsed = write_start.elapsed(); - (reclaimed, prepare_elapsed, write_lock_elapsed) + let commit_elapsed = commit_start.elapsed(); + (reclaimed, prepare_elapsed, commit_elapsed) }); // While prepare should be holding only a *shared* lock, appends must land @@ -481,10 +480,10 @@ async fn merge_write_lock_is_only_held_for_short_commit() { .unwrap(); let append_elapsed = append_start.elapsed(); - let (reclaimed, prepare_elapsed, write_lock_elapsed) = merge_handle.await.unwrap(); + let (reclaimed, prepare_elapsed, commit_elapsed) = merge_handle.await.unwrap(); eprintln!( - "merge phases: prepare={prepare_elapsed:?} write_lock={write_lock_elapsed:?} \ + "merge phases: prepare={prepare_elapsed:?} commit={commit_elapsed:?} \ append_during_prepare={append_elapsed:?} reclaimed={reclaimed}" ); @@ -494,19 +493,18 @@ async fn merge_write_lock_is_only_held_for_short_commit() { "prepare should do measurable work so the lock split is observable; \ got prepare={prepare_elapsed:?}" ); - // Exclusive lock must be the short phase. In the old bug it ≈ prepare. + // Commit must be the short phase. In the old bug exclusive work ≈ prepare. assert!( - write_lock_elapsed * 3 < prepare_elapsed - || write_lock_elapsed < Duration::from_millis(100), - "write lock held {write_lock_elapsed:?} but prepare took {prepare_elapsed:?}; \ - exclusive lock appears to cover the expensive merge work again (#198 / ~17s stall)" + commit_elapsed * 3 < prepare_elapsed || commit_elapsed < Duration::from_millis(100), + "commit took {commit_elapsed:?} but prepare took {prepare_elapsed:?}; \ + expensive merge work appears serialized again (#198 / ~17s stall)" ); assert!( append_elapsed < prepare_elapsed && (append_elapsed * 2 < prepare_elapsed || append_elapsed < Duration::from_millis(200)), "append during prepare took {append_elapsed:?} while prepare took {prepare_elapsed:?}; \ - append was blocked as if the merge held the write lock" + append was blocked as if the merge held an exclusive store lock" ); store.read().await.flush().await.unwrap(); diff --git a/crates/lance-context-core/tests/wal_merge_generation_cleanup.rs b/crates/lance-context-core/tests/wal_merge_generation_cleanup.rs index 0848448..ae8e505 100644 --- a/crates/lance-context-core/tests/wal_merge_generation_cleanup.rs +++ b/crates/lance-context-core/tests/wal_merge_generation_cleanup.rs @@ -83,7 +83,7 @@ async fn serial_merge_deletes_merged_generation_dirs() { ..Default::default() }; - let mut store = RolloutStore::open_with_options(&uri, opts.clone()) + let store = RolloutStore::open_with_options(&uri, opts.clone()) .await .unwrap(); diff --git a/crates/lance-context-server/src/routes/compact.rs b/crates/lance-context-server/src/routes/compact.rs index 0c2c4a9..d8f6902 100644 --- a/crates/lance-context-server/src/routes/compact.rs +++ b/crates/lance-context-server/src/routes/compact.rs @@ -28,7 +28,7 @@ pub async fn compact( None }; - let mut store = store_lock.write().await; + let store = store_lock.read().await; let metrics = store.compact(config).await.map_err(AppError::from_lance)?; Ok(Json(CompactResponse { diff --git a/crates/lance-context-server/src/routes/datagen.rs b/crates/lance-context-server/src/routes/datagen.rs index 98ec618..fdb1dda 100644 --- a/crates/lance-context-server/src/routes/datagen.rs +++ b/crates/lance-context-server/src/routes/datagen.rs @@ -174,7 +174,7 @@ async fn fold_datagen_item_refreshing_on_miss( } } - let mut store = store_lock.write().await; + let store = store_lock.read().await; if !store.is_version_pinned() { store.refresh_latest().await.map_err(AppError::from_lance)?; } @@ -197,7 +197,7 @@ async fn datagen_failures_refreshing_on_empty( } } - let mut store = store_lock.write().await; + let store = store_lock.read().await; if !store.is_version_pinned() { store.refresh_latest().await.map_err(AppError::from_lance)?; } @@ -220,7 +220,7 @@ async fn datagen_events_for_root_refreshing_on_empty( } } - let mut store = store_lock.write().await; + let store = store_lock.read().await; if !store.is_version_pinned() { store.refresh_latest().await.map_err(AppError::from_lance)?; } @@ -243,7 +243,7 @@ async fn datagen_root_statuses_refreshing_on_missing( } } - let mut store = store_lock.write().await; + let store = store_lock.read().await; if !store.is_version_pinned() { store.refresh_latest().await.map_err(AppError::from_lance)?; } @@ -266,7 +266,7 @@ async fn get_datagen_blob_refreshing_on_miss( } } - let mut store = store_lock.write().await; + let store = store_lock.read().await; if !store.is_version_pinned() { store.refresh_latest().await.map_err(AppError::from_lance)?; } diff --git a/crates/lance-context-server/src/routes/generic.rs b/crates/lance-context-server/src/routes/generic.rs index 87477a1..2ab0f75 100644 --- a/crates/lance-context-server/src/routes/generic.rs +++ b/crates/lance-context-server/src/routes/generic.rs @@ -215,7 +215,7 @@ async fn get_generic_row_refreshing_on_miss( } } - let mut store = store_lock.write().await; + let store = store_lock.read().await; if !store.is_version_pinned() { store.refresh_latest().await.map_err(AppError::from_lance)?; } diff --git a/crates/lance-context-server/src/routes/records.rs b/crates/lance-context-server/src/routes/records.rs index 2d21ae3..0a43441 100644 --- a/crates/lance-context-server/src/routes/records.rs +++ b/crates/lance-context-server/src/routes/records.rs @@ -32,7 +32,7 @@ async fn get_context_record_refreshing_on_miss( } } - let mut store = store_lock.write().await; + let store = store_lock.read().await; if !store.is_version_pinned() { store.refresh_latest().await.map_err(AppError::from_lance)?; } @@ -73,7 +73,7 @@ async fn fetch_context_payload_refreshing_on_miss( } } - let mut store = store_lock.write().await; + let store = store_lock.read().await; if !store.is_version_pinned() { store.refresh_latest().await.map_err(AppError::from_lance)?; } @@ -96,7 +96,7 @@ async fn get_context_by_external_id_refreshing_on_miss( } } - let mut store = store_lock.write().await; + let store = store_lock.read().await; if !store.is_version_pinned() { store.refresh_latest().await.map_err(AppError::from_lance)?; } diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index f9c7c5a..c4372bb 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -494,7 +494,7 @@ async fn get_rollout_refreshing_on_miss( } } - let mut store = store_lock.write().await; + let store = store_lock.read().await; if !store.is_version_pinned() { store.refresh_latest().await.map_err(AppError::from_lance)?; } @@ -516,7 +516,7 @@ async fn get_rollout_blob_refreshing_on_miss( } } - let mut store = store_lock.write().await; + let store = store_lock.read().await; if !store.is_version_pinned() { store.refresh_latest().await.map_err(AppError::from_lance)?; } @@ -575,7 +575,7 @@ pub async fn checkout_rollout( ) -> Result, AppError> { let store_lock = state.get_or_open_rollout_store(&name).await?; - let mut store = store_lock.write().await; + let store = store_lock.read().await; store .checkout(req.version) .await @@ -613,7 +613,7 @@ pub async fn compact_rollout( }; let lock_start = std::time::Instant::now(); - let mut store = store_lock.write().await; + let store = store_lock.read().await; ::metrics::histogram!("rollout_compaction_lock_wait_seconds") .record(lock_start.elapsed().as_secs_f64()); let compact_start = std::time::Instant::now(); @@ -692,7 +692,7 @@ pub async fn merge_wal( let merge_start = std::time::Instant::now(); let reclaimed = match prepared { Some((manifest_store, manifest, prepared)) => { - let mut store = store_lock.write().await; + let store = store_lock.read().await; match store .commit_prepared_merge(&manifest_store, &manifest, prepared) .await diff --git a/crates/lance-context-server/src/routes/versions.rs b/crates/lance-context-server/src/routes/versions.rs index 7bb06d2..f99e303 100644 --- a/crates/lance-context-server/src/routes/versions.rs +++ b/crates/lance-context-server/src/routes/versions.rs @@ -26,7 +26,7 @@ pub async fn checkout( ) -> Result, AppError> { let store_lock = state.get_or_open_context_store(&name).await?; - let mut store = store_lock.write().await; + let store = store_lock.read().await; store .checkout(req.version) .await diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index dce945e..ea4e2ea 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -787,7 +787,7 @@ impl AppState { .collect() }; for (name, store) in resident { - if let Err(e) = store.write().await.close().await { + if let Err(e) = store.read().await.close().await { tracing::warn!( store = %name, error = %e, @@ -808,7 +808,7 @@ impl AppState { .collect() }; for (name, store) in datagen { - if let Err(e) = store.write().await.close().await { + if let Err(e) = store.read().await.close().await { tracing::warn!( store = %name, error = %e, @@ -825,7 +825,7 @@ impl AppState { .collect() }; for (name, store) in generic { - if let Err(e) = store.write().await.close().await { + if let Err(e) = store.read().await.close().await { tracing::warn!( store = %name, error = %e, @@ -994,7 +994,7 @@ mod tests { // --- datagen: seals on append, so the merge pass is what it needs --- let datagen_uri = state.datagen_uri("d1"); - let mut datagen = DatagenStore::open(&datagen_uri).await.unwrap(); + let datagen = DatagenStore::open(&datagen_uri).await.unwrap(); datagen.append(&[datagen_event()]).await.unwrap(); assert!( datagen.pending_wal_generations().await.unwrap() > 0, @@ -1069,6 +1069,6 @@ mod tests { // The handle survives shutdown (shutdown only drains the writer); a // fresh close is still a no-op. - store.write().await.close().await.unwrap(); + store.read().await.close().await.unwrap(); } } diff --git a/crates/lance-context-server/src/sweeper.rs b/crates/lance-context-server/src/sweeper.rs index 0500e25..a5c2f01 100644 --- a/crates/lance-context-server/src/sweeper.rs +++ b/crates/lance-context-server/src/sweeper.rs @@ -25,11 +25,9 @@ use tokio::sync::{Mutex, RwLock}; /// A store the sweepers can maintain. /// /// Implemented on `Arc>` rather than on the store itself so each -/// kind decides its own locking. That is load-bearing for rollout, whose merge -/// deliberately splits into a shared-lock prepare (the expensive object-storage -/// reads, during which appends keep flowing) and a brief exclusive-lock commit. -/// A trait over `&mut Store` would have forced the exclusive lock across the -/// whole merge and quietly stalled the write path. +/// kind decides its own locking. Merge/flush/commit are `&self` on the store +/// (dataset handle is ArcSwap), so these impls only need a shared lock — +/// concurrent appends keep flowing. pub(crate) trait Sweepable: Send + Sync + 'static { /// Human-readable kind, for log and metric labels. fn kind() -> &'static str; @@ -48,14 +46,13 @@ impl Sweepable for Arc> { } async fn flush(&self) -> Result<(), String> { - // Read lock: `flush` is `&self`, so concurrent appends are not blocked. let guard = self.read().await; let result = guard.flush().await.map_err(|e| e.to_string()); if result.is_ok() { // The count-triggered merge rides this timer; it is a no-op unless // the threshold is configured and met. drop(guard); - let mut guard = self.write().await; + let guard = self.read().await; guard .maybe_merge_own_shard() .await @@ -65,9 +62,6 @@ impl Sweepable for Arc> { } async fn merge_wal(&self) -> Result { - // The prepare/commit split: seal and read every flushed generation - // under the *shared* lock so appends keep running, then take the - // exclusive lock only for the short commit. let prepared = { let guard = self.read().await; guard @@ -78,7 +72,7 @@ impl Sweepable for Arc> { let Some((manifest_store, manifest, prepared)) = prepared else { return Ok(0); }; - let mut guard = self.write().await; + let guard = self.read().await; guard .commit_prepared_merge(&manifest_store, &manifest, prepared) .await @@ -99,7 +93,7 @@ impl Sweepable for Arc> { } async fn merge_wal(&self) -> Result { - let mut guard = self.write().await; + let guard = self.read().await; guard.cleanup_own_shard().await.map_err(|e| e.to_string()) } } @@ -115,7 +109,7 @@ impl Sweepable for Arc> { } async fn merge_wal(&self) -> Result { - let mut guard = self.write().await; + let guard = self.read().await; guard.cleanup_wal().await.map_err(|e| e.to_string()) } } From d83976f509488a7b9c32f3226991930e9722ffc3 Mon Sep 17 00:00:00 2001 From: Lucy Ge Date: Mon, 10 Aug 2026 17:20:52 -0700 Subject: [PATCH 05/11] remove unnecessary outer mut handles --- .../lance-context-core/src/datagen_store.rs | 12 +++--- .../lance-context-core/src/generic_store.rs | 4 +- .../lance-context-core/src/rollout_store.rs | 42 +++++++++---------- crates/lance-context-core/src/store.rs | 14 +++---- 4 files changed, 36 insertions(+), 36 deletions(-) diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index f31239c..4387d15 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -433,7 +433,7 @@ impl DatagenStore { let Some(store) = weak.upgrade() else { return; }; - let mut guard = store.write().await; + let guard = store.write().await; match tokio::time::timeout(pass_timeout, guard.cleanup_own_shard()).await { Ok(Ok(0)) => {} Ok(Ok(reclaimed)) => info!( @@ -1323,7 +1323,7 @@ mod tests { let uri = directory.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut writer_a = DatagenStore::open_with_options( + let writer_a = DatagenStore::open_with_options( &uri, DatagenStoreOptions { storage_options: None, @@ -1333,7 +1333,7 @@ mod tests { ) .await .unwrap(); - let mut writer_b = DatagenStore::open_with_options( + let writer_b = DatagenStore::open_with_options( &uri, DatagenStoreOptions { storage_options: None, @@ -1424,7 +1424,7 @@ mod tests { let uri = directory.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = DatagenStore::open(&uri).await.unwrap(); + let store = DatagenStore::open(&uri).await.unwrap(); // Root item "7" fans out into one sub-item "7/solve_twice:0". let mut root_created = event("7", 0, "created-root", 0, DatagenEventType::ItemCreated); @@ -1478,7 +1478,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = DatagenStore::open(&uri).await.unwrap(); + let store = DatagenStore::open(&uri).await.unwrap(); // One cleanup pass appends one fragment, so merge after each append // to accumulate several -- this is exactly the growth pattern that @@ -1519,7 +1519,7 @@ mod tests { let uri = directory.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = DatagenStore::open_with_options( + let store = DatagenStore::open_with_options( &uri, DatagenStoreOptions { storage_options: None, diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index 3515553..f3d6f5b 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -699,7 +699,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let rt = tokio::runtime::Runtime::new().unwrap(); rt.block_on(async { - let mut store = GenericStore::open(&uri, spec(), sealing()).await.unwrap(); + let store = GenericStore::open(&uri, spec(), sealing()).await.unwrap(); let payload: Vec = (0..4 * 1024 * 1024).map(|i| (i % 251) as u8).collect(); store .add(&[row(json!({"id": "big", "payload": payload}))]) @@ -745,7 +745,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let rt = tokio::runtime::Runtime::new().unwrap(); rt.block_on(async { - let mut store = GenericStore::open(&uri, spec(), sealing()).await.unwrap(); + let store = GenericStore::open(&uri, spec(), sealing()).await.unwrap(); for i in 0..3 { store .add(&[row(json!({"id": format!("r{i}")}))]) diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index a301aaf..1bc8fad 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -2267,7 +2267,7 @@ mod tests { let legacy_schema = pre_claim_check_schema(); create_empty_dataset(&uri, legacy_schema.clone()).await; - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { shard_id: Some("pre-claim-check".to_string()), @@ -2580,7 +2580,7 @@ mod tests { let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { storage_options: None, @@ -2865,7 +2865,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { shard_id: Some("fragment-pagination".to_string()), @@ -2959,7 +2959,7 @@ mod tests { writer.add(&[assistant_record("row-0")]).await.unwrap(); writer.flush().await.unwrap(); - let mut cached_reader = + let cached_reader = RolloutStore::open_existing_with_options(&uri, RolloutStoreOptions::default()) .await .unwrap(); @@ -2989,7 +2989,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open(&uri).await.unwrap(); + let store = RolloutStore::open(&uri).await.unwrap(); assert!(!store.is_version_pinned()); store.checkout(store.version()).await.unwrap(); @@ -3198,7 +3198,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { storage_options: None, @@ -3505,11 +3505,11 @@ mod tests { // A compacts the base table; B merges its shard into it — concurrently. let (ca, mb) = tokio::join!( async { - let mut g = a.write().await; + let g = a.write().await; g.compact(None).await }, async { - let mut g = b.write().await; + let g = b.write().await; g.cleanup_own_shard().await }, ); @@ -3665,7 +3665,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { storage_options: None, @@ -3704,7 +3704,7 @@ mod tests { fn cleanup_merges_pre_claim_check_generations_after_schema_evolution() { let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let (_dir, mut store) = store_with_legacy_base_and_wal(false).await; + let (_dir, store) = store_with_legacy_base_and_wal(false).await; assert_eq!(store.cleanup_own_shard().await.unwrap(), 1); assert_eq!(flushed_generation_count(&store).await, 0); @@ -3735,7 +3735,7 @@ mod tests { let legacy_dir = TempDir::new().unwrap(); let legacy_uri = legacy_dir.path().to_string_lossy().to_string(); create_empty_dataset(&legacy_uri, pre_claim_check_schema()).await; - let mut legacy_store = RolloutStore::open(&legacy_uri).await.unwrap(); + let legacy_store = RolloutStore::open(&legacy_uri).await.unwrap(); let current_dir = TempDir::new().unwrap(); let current_uri = current_dir.path().to_string_lossy().to_string(); @@ -3820,7 +3820,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open(&uri).await.unwrap(); + let store = RolloutStore::open(&uri).await.unwrap(); // Add rows and fold them into the base table so there is data (and a // MemWAL index) present when we build the scalar index. store.add(&[assistant_record("a-0")]).await.unwrap(); @@ -3905,11 +3905,11 @@ mod tests { // Both merge into the shared base table concurrently. let (ra, rb) = tokio::join!( async { - let mut g = a.write().await; + let g = a.write().await; g.cleanup_own_shard().await }, async { - let mut g = b.write().await; + let g = b.write().await; g.cleanup_own_shard().await }, ); @@ -3946,7 +3946,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { storage_options: None, @@ -3996,7 +3996,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { storage_options: None, @@ -4063,7 +4063,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { shard_id: Some("rollout-pagination".to_string()), @@ -4164,7 +4164,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut writer = RolloutStore::open_with_options( + let writer = RolloutStore::open_with_options( &uri, RolloutStoreOptions { shard_id: Some("pagination-benchmark".to_string()), @@ -4343,7 +4343,7 @@ mod tests { runtime.block_on(async { // merge_after_generations = None: appended rows stay in the WAL, // un-merged, so this exercises the base-miss -> WAL-fallback path. - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { storage_options: None, @@ -4426,7 +4426,7 @@ mod tests { let bytes = b"base-version-bytes"; let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { storage_options: None, @@ -4465,7 +4465,7 @@ mod tests { let bytes = b"\x00\x01record-with-blob"; let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = RolloutStore::open_with_options( + let store = RolloutStore::open_with_options( &uri, RolloutStoreOptions { storage_options: None, diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 4387d44..b6f7f89 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -2133,7 +2133,7 @@ impl ContextStore { loop { interval.tick().await; - let mut store = match open_for_compaction(&uri, compaction_options.clone()).await { + let store = match open_for_compaction(&uri, compaction_options.clone()).await { Ok(store) => store, Err(e) => { error!("Background compaction could not open store: {}", e); @@ -4473,7 +4473,7 @@ mod tests { .await .unwrap(); - let mut store = ContextStore::open(&uri).await.unwrap(); + let store = ContextStore::open(&uri).await.unwrap(); assert!(!store.has_relationships_column()); let mut record = text_record("with-relationships", 0.0); @@ -5048,7 +5048,7 @@ mod tests { let uri = dir.path().to_string_lossy().to_string(); let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = ContextStore::open(&uri).await.unwrap(); + let store = ContextStore::open(&uri).await.unwrap(); for i in 0..3 { store .add(&[text_record(&format!("r{i}"), i as f32)]) @@ -5558,7 +5558,7 @@ mod tests { id_index_type: IdIndexType::BTree, ..Default::default() }; - let mut store = ContextStore::open_with_options(&uri, options) + let store = ContextStore::open_with_options(&uri, options) .await .unwrap(); @@ -5598,7 +5598,7 @@ mod tests { id_index_type: IdIndexType::ZoneMap, ..Default::default() }; - let mut store = ContextStore::open_with_options(&uri, options) + let store = ContextStore::open_with_options(&uri, options) .await .unwrap(); @@ -5632,7 +5632,7 @@ mod tests { let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = ContextStore::open(&uri).await.unwrap(); + let store = ContextStore::open(&uri).await.unwrap(); store.add(&[text_record("no-idx-1", 0.0)]).await.unwrap(); store.compact(None).await.unwrap(); @@ -5656,7 +5656,7 @@ mod tests { id_index_type: IdIndexType::BTree, ..Default::default() }; - let mut store = ContextStore::open_with_options(&uri, options) + let store = ContextStore::open_with_options(&uri, options) .await .unwrap(); From 98d5e6b8547db2334db00980192393ce1743b61e Mon Sep 17 00:00:00 2001 From: Lucy Ge Date: Mon, 10 Aug 2026 20:35:07 -0700 Subject: [PATCH 06/11] add merge lock to lock prepare + commit --- .../lance-context-core/src/rollout_store.rs | 7 +- crates/lance-context-core/src/store_base.rs | 73 ++++++++++++------- .../tests/wal_merge_concurrency.rs | 5 +- crates/lance-context-metrics/src/lib.rs | 4 +- .../src/routes/rollouts.rs | 11 +-- crates/lance-context-server/src/sweeper.rs | 4 +- 6 files changed, 64 insertions(+), 40 deletions(-) diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 1bc8fad..b5f9e76 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -606,8 +606,9 @@ impl RolloutStore { self.base.maybe_merge_own_shard().await } - /// The shared-lock half of a merge; see `StorageBase::prepare_merge_if_ready` - /// for the intended read-lock/write-lock split. + /// Prepare half of a merge; see `StorageBase::prepare_merge_if_ready`. + /// Merge exclusivity is the base's internal `merge_lock` (held inside the + /// returned [`PreparedMerge`]), not an outer store write lock. pub async fn prepare_merge_if_ready( &self, threshold: usize, @@ -624,6 +625,8 @@ impl RolloutStore { } /// Commit a merge prepared by [`Self::prepare_merge_if_ready`]. + /// + /// Consumes [`PreparedMerge`], releasing the base's merge lock when done. pub async fn commit_prepared_merge( &self, manifest_store: &ShardManifestStore, diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index 34b063f..78d7fa8 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -182,11 +182,12 @@ pub enum ListSource { /// Rows read out of the flushed generations, ready to be appended to the base /// table and drained from the shard manifest. /// -/// Produced by `StorageBase::prepare_merge_if_ready` under `&self` (so appends -/// keep running while it reads object storage) and consumed by -/// `StorageBase::commit_prepared_merge` (also `&self`). `PreparedMerge` is -/// public because it appears in [`RolloutStore`]'s prepare/commit split, but -/// its fields are opaque. +/// Produced by `StorageBase::prepare_merge_if_ready` / `prepare_cleanup_merge` +/// and consumed by `StorageBase::commit_prepared_merge`. Holding a value of this +/// type retains the store's internal merge lock until it is committed or +/// dropped, so callers of the prepare/commit split do not need to lock +/// externally. `PreparedMerge` is public because it appears in +/// [`RolloutStore`]'s prepare/commit split, but its fields are opaque. /// /// [`RolloutStore`]: crate::RolloutStore pub struct PreparedMerge { @@ -194,6 +195,9 @@ pub struct PreparedMerge { merged_paths: Vec, batches: Vec, merge_schema: Arc, + /// Serializes this prepare+commit against other merges on the same store. + /// Released when `PreparedMerge` is dropped (after commit or on abandon). + _merge_guard: tokio::sync::OwnedMutexGuard<()>, } impl PreparedMerge { @@ -298,6 +302,10 @@ pub(crate) struct StorageBase { merge_max_generations: usize, /// Compaction bookkeeping for this handle (interior-mutable so [`Self::compact`] can be `&self`). compaction: Mutex, + /// Serializes WAL→base merge (prepare through commit). Taken with + /// `try_lock_owned`: a loser no-ops (`Ok(0)` / `Ok(None)`). Not held by + /// `add`/`flush`, so appends keep running while a merge is in flight. + merge_lock: Arc>, /// Explicit time-travel version selected by [`Self::checkout`]. /// /// A point-read miss may refresh an ordinary long-lived handle to avoid a @@ -425,6 +433,7 @@ impl StorageBase { merge_after_generations: merge_after_generations.unwrap_or(0), merge_max_generations: merge_max_generations.unwrap_or(DEFAULT_MERGE_MAX_GENERATIONS), compaction: Mutex::new(CompactionState::default()), + merge_lock: Arc::new(tokio::sync::Mutex::new(())), pinned_version: AtomicU64::new(0), write_writer: tokio::sync::Mutex::new(None), }; @@ -761,16 +770,16 @@ impl StorageBase { Ok(if committed { pending } else { 0 }) } - /// The shared-lock half of a merge: decide whether one is due and read the - /// flushed generations into memory. + /// Prepare half of a merge: decide whether one is due and read the flushed + /// generations into memory. /// - /// Takes `&self`, so a caller holding a *read* lock can run the expensive - /// part while appends continue, then take the write lock only to hand the - /// result to [`Self::commit_prepared_merge`]. Returns `None` when nothing is - /// due. + /// Acquires the internal merge lock with `try_lock` (returned inside + /// [`PreparedMerge`]) so prepare+commit stay exclusive without blocking + /// `add`. Returns `None` when nothing is due **or** another merge already + /// holds the lock. /// /// ```ignore - /// let prepared = { store.read().await.prepare_merge_if_ready(1).await? }; + /// let prepared = store.prepare_merge_if_ready(1).await?; /// if let Some((manifest_store, manifest, prepared)) = prepared { /// store.commit_prepared_merge(&manifest_store, &manifest, prepared).await?; /// } @@ -797,6 +806,12 @@ impl StorageBase { threshold: usize, seal_first: bool, ) -> LanceResult> { + // Exclusive for the whole prepare→commit lifetime (guard lives in + // PreparedMerge). Losers no-op: the holder will drain current gens. + let Ok(merge_guard) = Arc::clone(&self.merge_lock).try_lock_owned() else { + return Ok(None); + }; + if seal_first { // Materialize anything buffered so it is eligible for this pass. self.flush().await?; @@ -817,13 +832,15 @@ impl StorageBase { if pending == 0 || pending < threshold.max(1) { return Ok(None); } - let Some(prepared) = self.prepare_merge(&manifest).await? else { + let Some(prepared) = self.prepare_merge(&manifest, merge_guard).await? else { return Ok(None); }; Ok(Some((manifest_store, manifest, prepared))) } /// Commit a merge prepared by [`Self::prepare_merge_if_ready`]. + /// + /// Consumes [`PreparedMerge`], releasing the merge lock when it returns. pub async fn commit_prepared_merge( &self, manifest_store: &ShardManifestStore, @@ -837,21 +854,21 @@ impl StorageBase { Ok(if committed { pending } else { 0 }) } - /// The `&self` half of a merge: everything that can run while appends - /// continue — sealing the memtable and reading every flushed generation - /// into memory. + /// Seal + read every flushed generation into a [`PreparedMerge`]. /// - /// # Concurrency: the expensive phase does not need exclusive access + /// # Concurrency /// - /// A merge only ever touches *sealed* generations — history — while a `put` - /// writes the active memtable at the WAL tail. They operate on disjoint - /// data, which is the whole point of an LSM, so a merge must not stop the - /// write path. Notably the merge does **not** `claim_epoch`: the epoch is an - /// *ownership* token, not a per-commit token, and + /// Caller already holds [`Self::merge_lock`] via `merge_guard`. A merge only + /// ever touches *sealed* generations — history — while a `put` writes the + /// active memtable at the WAL tail, so appends keep running. The merge does + /// **not** `claim_epoch`: the epoch is an *ownership* token, and /// [`ShardManifestStore::commit_update`] only rejects a writer whose epoch is - /// **older** than the stored one. Reusing the shard's current epoch commits - /// the drain and leaves the live writer untouched. - async fn prepare_merge(&self, manifest: &ShardManifest) -> LanceResult> { + /// **older** than the stored one. + async fn prepare_merge( + &self, + manifest: &ShardManifest, + merge_guard: tokio::sync::OwnedMutexGuard<()>, + ) -> LanceResult> { if manifest.flushed_generations.is_empty() { return Ok(None); } @@ -862,8 +879,8 @@ impl StorageBase { observe_phase!("seal", self.flush().await)?; // The expensive phase: pull every flushed generation out of object - // storage. Buffered in memory, so this is the part that must not hold an - // exclusive lock. + // storage. Runs under the merge lock so a second merge cannot prepare + // the same generations concurrently. let (merged_generations, merged_paths, batches, merge_schema) = observe_phase!("read", self.read_flushed_generations(manifest).await)?; @@ -872,6 +889,7 @@ impl StorageBase { merged_paths, batches, merge_schema, + _merge_guard: merge_guard, })) } @@ -907,6 +925,7 @@ impl StorageBase { merged_paths, batches, merge_schema, + _merge_guard, } = prepared; // Several sweepers can prepare the same immutable generations under a diff --git a/crates/lance-context-core/tests/wal_merge_concurrency.rs b/crates/lance-context-core/tests/wal_merge_concurrency.rs index 798fb2e..4a45113 100644 --- a/crates/lance-context-core/tests/wal_merge_concurrency.rs +++ b/crates/lance-context-core/tests/wal_merge_concurrency.rs @@ -254,8 +254,9 @@ async fn generation_sealed_during_merge_is_not_dropped() { /// Two merges racing must not append the same generations twice. /// -/// Merges are serialized by an internal mutex (not the store lock, which would -/// also exclude appends); the loser returns 0 rather than waiting. +/// Merges are serialized by `StorageBase`'s internal `merge_lock` (prepare +/// through commit). A `try_lock` loser gets `prepare_* -> None` / reclaim `0` +/// rather than waiting — appends never take this lock. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn concurrent_merges_do_not_duplicate_rows() { let tmp = tempfile::tempdir().unwrap(); diff --git a/crates/lance-context-metrics/src/lib.rs b/crates/lance-context-metrics/src/lib.rs index 3c54d9c..c0e6c52 100644 --- a/crates/lance-context-metrics/src/lib.rs +++ b/crates/lance-context-metrics/src/lib.rs @@ -175,12 +175,12 @@ fn describe_metrics() { describe_histogram!( "rollout_wal_merge_lock_wait_seconds", Unit::Seconds, - "Time waiting for the store write lock before a WAL merge (blocks all ingest)." + "Time waiting for the store shared lock before a WAL merge prepare." ); describe_histogram!( "rollout_compaction_lock_wait_seconds", Unit::Seconds, - "Time waiting for the store write lock before compaction." + "Time waiting for the store shared lock before compaction." ); // Master task lifecycle. diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index c4372bb..2ca7e9c 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -307,7 +307,7 @@ pub async fn add_rollouts( // A read lock: `add` is `&self` and MemWAL appends are internally // concurrent, so multiple ingest requests to the same store run in parallel. - // Mutating ops (merge, compact, checkout, close) still take the write lock. + // Merge exclusivity is `StorageBase::merge_lock`, not this outer write lock. let store = store_lock.read().await; // Times only the store work (`add` + optional `flush`), excluding body @@ -671,10 +671,11 @@ pub async fn merge_wal( Path(name): Path, ) -> Result, AppError> { let store_lock = state.get_or_open_rollout_store(&name).await?; - // Split by lock scope: seal + read every flushed generation under the - // *read* lock so ingest on this store keeps running, then take the write - // lock only for the short commit. Holding the write lock across the whole - // merge is what used to stall every concurrent append for its duration. + // Store `RwLock` is shared (read) only so we can call `&self` APIs while + // ingest keeps running. Merge exclusivity is `StorageBase::merge_lock` + // inside prepare→commit (`try_lock`; a loser returns `reclaimed: 0`), not + // this outer write lock — holding write across the whole merge is what used + // to stall every concurrent append. let lock_start = std::time::Instant::now(); let prepared = { let store = store_lock.read().await; diff --git a/crates/lance-context-server/src/sweeper.rs b/crates/lance-context-server/src/sweeper.rs index a5c2f01..183ac75 100644 --- a/crates/lance-context-server/src/sweeper.rs +++ b/crates/lance-context-server/src/sweeper.rs @@ -26,8 +26,8 @@ use tokio::sync::{Mutex, RwLock}; /// /// Implemented on `Arc>` rather than on the store itself so each /// kind decides its own locking. Merge/flush/commit are `&self` on the store -/// (dataset handle is ArcSwap), so these impls only need a shared lock — -/// concurrent appends keep flowing. +/// (dataset handle is ArcSwap; merge exclusivity is an internal `try_lock`), so +/// these impls only need a shared store lock — concurrent appends keep flowing. pub(crate) trait Sweepable: Send + Sync + 'static { /// Human-readable kind, for log and metric labels. fn kind() -> &'static str; From 3e93c6324dcf9df86c9ed7ae95129606f3287c4f Mon Sep 17 00:00:00 2001 From: lucyge2022 Date: Mon, 10 Aug 2026 22:08:07 -0700 Subject: [PATCH 07/11] add concurrency test for narrow mutex for merge around prepare + commit two phase --- .../tests/wal_merge_concurrency.rs | 103 +++++++++++++++++- 1 file changed, 101 insertions(+), 2 deletions(-) diff --git a/crates/lance-context-core/tests/wal_merge_concurrency.rs b/crates/lance-context-core/tests/wal_merge_concurrency.rs index 4a45113..d922cd4 100644 --- a/crates/lance-context-core/tests/wal_merge_concurrency.rs +++ b/crates/lance-context-core/tests/wal_merge_concurrency.rs @@ -15,7 +15,8 @@ //! //! 1. appends succeed while a merge runs, and no row is lost; //! 2. a generation sealed *during* a merge is not silently dropped by the drain; -//! 3. concurrent merges do not duplicate rows; +//! 3. concurrent merges do not duplicate rows, and `merge_lock` excludes a +//! second prepare while the first `PreparedMerge` is still live; //! 4. an interrupted merge loses nothing (rows stay readable exactly once); //! 5. `add` is not blocked for the merge's duration. @@ -285,11 +286,23 @@ async fn concurrent_merges_do_not_duplicate_rows() { async move { merge_like_sweeper(&store).await }, )); } + let mut reclaimed = Vec::new(); for h in handles { // None may error; a loser simply reports 0. - h.await.unwrap(); + reclaimed.push(h.await.unwrap()); } + let winners: Vec = reclaimed.iter().copied().filter(|&n| n > 0).collect(); + assert_eq!( + winners.len(), + 1, + "exactly one merge may reclaim; got reclaimed={reclaimed:?}" + ); + assert_eq!( + winners[0], 10, + "winner should reclaim every pending generation; got reclaimed={reclaimed:?}" + ); + let ids = read_ids(&store).await; let mut deduped = ids.clone(); deduped.dedup(); @@ -300,6 +313,92 @@ async fn concurrent_merges_do_not_duplicate_rows() { assert_eq!(ids.len(), 10, "all rows readable exactly once: {ids:?}"); } +/// Direct mutual-exclusion check for `merge_lock`: while one `PreparedMerge` +/// is live (prepare done, commit not yet), a second `prepare_*` must lose +/// `try_lock` and return `None` — even though both only hold the outer store +/// `RwLock` for shared/`read` access. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn second_merge_prepare_is_rejected_while_first_holds_prepared_merge() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = Arc::new(RwLock::new( + RolloutStore::open_with_options(&uri, opts("solo")) + .await + .unwrap(), + )); + + for i in 0..4 { + store + .read() + .await + .add(&[rec(&format!("row-{i}"))]) + .await + .unwrap(); + store.read().await.flush().await.unwrap(); + } + + let first = { + let guard = store.read().await; + guard.prepare_cleanup_merge().await.unwrap() + }; + let Some((manifest_store, manifest, prepared)) = first else { + panic!("expected pending generations for the first prepare"); + }; + + // Outer store lock is already dropped; only `_merge_guard` inside + // `prepared` serializes merges. A concurrent prepare must no-op. + let second = { + let guard = store.read().await; + guard.prepare_cleanup_merge().await.unwrap() + }; + assert!( + second.is_none(), + "second prepare must lose merge_lock try_lock while PreparedMerge is live" + ); + + // Appends must still flow under the held merge lock. + store + .read() + .await + .add(&[rec("during-held-merge")]) + .await + .unwrap(); + + let reclaimed = { + let guard = store.read().await; + guard + .commit_prepared_merge(&manifest_store, &manifest, prepared) + .await + .unwrap() + }; + assert_eq!(reclaimed, 4); + + // Lock released with PreparedMerge; nothing left to merge until a new seal. + let after_commit = { + let guard = store.read().await; + guard.prepare_cleanup_merge().await.unwrap() + }; + // prepare_cleanup_merge seals first, so the during-held-merge row becomes + // one pending generation — that prepare must succeed now that the lock is free. + assert!( + after_commit.is_some(), + "after commit, merge_lock must be free for a new prepare" + ); + let (manifest_store, manifest, prepared) = after_commit.unwrap(); + let reclaimed = { + let guard = store.read().await; + guard + .commit_prepared_merge(&manifest_store, &manifest, prepared) + .await + .unwrap() + }; + assert_eq!(reclaimed, 1); + + let ids = read_ids(&store).await; + assert_eq!(ids.len(), 5, "all rows readable exactly once: {ids:?}"); + assert!(ids.contains(&"during-held-merge".to_string())); +} + /// A merge abandoned partway (the sweeper's timeout does exactly this) must not /// lose data. A retry merge-inserts by id, so nothing may disappear or remain /// duplicated after a later merge converges. From c984c756022461df720a08ae75d514e3674c68e2 Mon Sep 17 00:00:00 2001 From: lucyge2022 Date: Tue, 11 Aug 2026 01:30:35 -0700 Subject: [PATCH 08/11] add test asserting compact does not block appends --- .../tests/wal_merge_concurrency.rs | 84 ++++++++++++++++++- 1 file changed, 83 insertions(+), 1 deletion(-) diff --git a/crates/lance-context-core/tests/wal_merge_concurrency.rs b/crates/lance-context-core/tests/wal_merge_concurrency.rs index d922cd4..ccbf95b 100644 --- a/crates/lance-context-core/tests/wal_merge_concurrency.rs +++ b/crates/lance-context-core/tests/wal_merge_concurrency.rs @@ -23,7 +23,9 @@ use std::sync::Arc; use std::time::{Duration, Instant}; -use lance_context_core::{RolloutRecord, RolloutStore, RolloutStoreOptions, ROLE_ASSISTANT}; +use lance_context_core::{ + CompactionConfig, RolloutRecord, RolloutStore, RolloutStoreOptions, ROLE_ASSISTANT, +}; use tokio::sync::RwLock; /// Run one full merge: seal + read generations, then commit. Both phases use @@ -615,3 +617,83 @@ async fn merge_write_lock_is_only_held_for_short_commit() { ); assert_eq!(ids.len(), 21, "all rows readable exactly once: {ids:?}"); } + +/// Compact is `&self` (ArcSwap dataset handle) and must not take the outer +/// store write lock — otherwise every concurrent `add` waits for the whole +/// rewrite. Assert append finishes within a timeout and remains readable; do +/// not compare wall times to compact (flaky on fast disks). +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn append_is_not_blocked_for_the_duration_of_a_compact() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = Arc::new(RwLock::new( + RolloutStore::open_with_options(&uri, opts("solo")) + .await + .unwrap(), + )); + + // Many small base-table fragments so compact has real rewrite work. + for i in 0..20 { + store + .read() + .await + .add(&[rec(&format!("bulk-{i}"))]) + .await + .unwrap(); + store.read().await.flush().await.unwrap(); + // Fold each generation into the base table as its own fragment. + assert!(merge_like_sweeper(&store).await > 0); + } + + let fragments_before = store.read().await.observe().await.unwrap().fragment_count; + assert!( + fragments_before > 1, + "need several base fragments to compact, got {fragments_before}" + ); + + let compactor = { + let store = store.clone(); + tokio::spawn(async move { + let guard = store.read().await; + guard + .compact(Some(CompactionConfig { + min_fragments: 2, + num_threads: Some(1), + batch_size: Some(1), + ..Default::default() + })) + .await + .unwrap() + }) + }; + + // Give compact a moment to enter the rewrite. + tokio::time::sleep(Duration::from_millis(10)).await; + + // Fail only on multi-second stalls — not on wall-clock ratios vs compact. + const APPEND_NOT_STALLED: Duration = Duration::from_secs(5); + tokio::time::timeout(APPEND_NOT_STALLED, async { + store + .read() + .await + .add(&[rec("during-compact")]) + .await + .unwrap(); + }) + .await + .expect("append during compact stalled; it appears blocked behind compact"); + + let metrics = compactor.await.unwrap(); + assert!( + metrics.fragments_removed > 0, + "compact should have rewritten fragments" + ); + + store.read().await.flush().await.unwrap(); + let ids = read_ids(&store).await; + assert!( + ids.contains(&"during-compact".to_string()), + "the row appended during compact must be readable" + ); + assert_eq!(ids.len(), 21, "all rows readable exactly once: {ids:?}"); +} From 565529cfbbb0855befbee0ebf4470c932d7cae1b Mon Sep 17 00:00:00 2001 From: lucyge2022 Date: Tue, 11 Aug 2026 01:31:55 -0700 Subject: [PATCH 09/11] remove unnecessary test --- .../tests/wal_merge_concurrency.rs | 102 ------------------ 1 file changed, 102 deletions(-) diff --git a/crates/lance-context-core/tests/wal_merge_concurrency.rs b/crates/lance-context-core/tests/wal_merge_concurrency.rs index ccbf95b..e45f3d2 100644 --- a/crates/lance-context-core/tests/wal_merge_concurrency.rs +++ b/crates/lance-context-core/tests/wal_merge_concurrency.rs @@ -516,108 +516,6 @@ async fn append_is_not_blocked_for_the_duration_of_a_merge() { assert_eq!(ids.len(), 26, "all rows readable exactly once"); } -/// Issue #198 / the ~17s production stall, measured as lock phases. -/// -/// Old shape: one `RwLock::write()` spanned seal + generation reads + append + -/// drain, so every concurrent `add` waited for the whole merge (~17s on abfss). -/// New shape: prepare and commit both use shared locks (`&self` + ArcSwap), so -/// appends are never blocked by merge. -/// -/// This test records prepare vs commit durations and races an append against -/// prepare. It fails if append appears serialized behind the merge again. -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn merge_write_lock_is_only_held_for_short_commit() { - let tmp = tempfile::tempdir().unwrap(); - let uri = tmp.path().to_string_lossy().to_string(); - let store = Arc::new(RwLock::new( - RolloutStore::open_with_options(&uri, opts("solo")) - .await - .unwrap(), - )); - - // Fat generations so prepare (read every flushed gen) dominates wall time - // even on a fast local disk — otherwise the assertion becomes vacuous. - let fat = "x".repeat(64 * 1024); - for i in 0..20 { - let mut r: RolloutRecord = rec(&format!("bulk-{i}")); - r.content = Some(fat.clone()); - store.read().await.add(&[r]).await.unwrap(); - store.read().await.flush().await.unwrap(); - } - - let store_for_merge = store.clone(); - let merge_handle = tokio::spawn(async move { - let prepare_start = Instant::now(); - let prepared = { - let guard = store_for_merge.read().await; - guard.prepare_cleanup_merge().await.unwrap() - }; - let prepare_elapsed = prepare_start.elapsed(); - - let Some((manifest_store, manifest, prepared)) = prepared else { - panic!("expected pending generations to merge"); - }; - - let commit_start = Instant::now(); - let reclaimed = { - let guard = store_for_merge.read().await; - guard - .commit_prepared_merge(&manifest_store, &manifest, prepared) - .await - .unwrap() - }; - let commit_elapsed = commit_start.elapsed(); - (reclaimed, prepare_elapsed, commit_elapsed) - }); - - // While prepare should be holding only a *shared* lock, appends must land - // quickly — this is the user-visible half of the ~17s stall. - tokio::time::sleep(Duration::from_millis(20)).await; - let append_start = Instant::now(); - store - .read() - .await - .add(&[rec("during-prepare")]) - .await - .unwrap(); - let append_elapsed = append_start.elapsed(); - - let (reclaimed, prepare_elapsed, commit_elapsed) = merge_handle.await.unwrap(); - - eprintln!( - "merge phases: prepare={prepare_elapsed:?} commit={commit_elapsed:?} \ - append_during_prepare={append_elapsed:?} reclaimed={reclaimed}" - ); - - assert!(reclaimed > 0, "merge should reclaim generations"); - assert!( - prepare_elapsed > Duration::from_millis(5), - "prepare should do measurable work so the lock split is observable; \ - got prepare={prepare_elapsed:?}" - ); - // Commit must be the short phase. In the old bug exclusive work ≈ prepare. - assert!( - commit_elapsed * 3 < prepare_elapsed || commit_elapsed < Duration::from_millis(100), - "commit took {commit_elapsed:?} but prepare took {prepare_elapsed:?}; \ - expensive merge work appears serialized again (#198 / ~17s stall)" - ); - assert!( - append_elapsed < prepare_elapsed - && (append_elapsed * 2 < prepare_elapsed - || append_elapsed < Duration::from_millis(200)), - "append during prepare took {append_elapsed:?} while prepare took {prepare_elapsed:?}; \ - append was blocked as if the merge held an exclusive store lock" - ); - - store.read().await.flush().await.unwrap(); - let ids = read_ids(&store).await; - assert!( - ids.contains(&"during-prepare".to_string()), - "row appended under the shared prepare lock must be readable" - ); - assert_eq!(ids.len(), 21, "all rows readable exactly once: {ids:?}"); -} - /// Compact is `&self` (ArcSwap dataset handle) and must not take the outer /// store write lock — otherwise every concurrent `add` waits for the whole /// rewrite. Assert append finishes within a timeout and remains readable; do From ef55e747ae8510e6c07e18ed1ca879f574481b76 Mon Sep 17 00:00:00 2001 From: lucyge2022 Date: Tue, 11 Aug 2026 22:29:42 -0700 Subject: [PATCH 10/11] format check fix --- .../lance-context-core/src/datagen_store.rs | 8 +++-- .../lance-context-core/src/rollout_store.rs | 33 +++++++++++++++---- crates/lance-context-core/src/store_base.rs | 32 +++++------------- 3 files changed, 41 insertions(+), 32 deletions(-) diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index 4387d15..0909972 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -336,9 +336,11 @@ impl DatagenStore { } } - Ok(Self::get_blob_from_dataset(self.base.current_dataset().as_ref(), event_id) - .await? - .flatten()) + Ok( + Self::get_blob_from_dataset(self.base.current_dataset().as_ref(), event_id) + .await? + .flatten(), + ) } /// Materialize a folded item's blob field by name, resolving the `event_id` for the caller. diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index b5f9e76..35ef646 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -695,7 +695,12 @@ impl RolloutStore { let row_count = (base_rows + pending_rows) as i64; let fragment_count = self.base.current_dataset().count_fragments() as i64; let version = self.base.current_dataset().manifest.version; - let last_updated = self.base.current_dataset().manifest.timestamp().timestamp_millis(); + let last_updated = self + .base + .current_dataset() + .manifest + .timestamp() + .timestamp_millis(); let pending_wal_generations = shard_snapshots .iter() .map(|snapshot| snapshot.flushed_generations.len() as i64) @@ -981,7 +986,10 @@ impl RolloutStore { } let columns = Arc::new(self.non_blob_columns()); - let target_schema = Arc::new(projected_arrow_schema(self.base.current_dataset().as_ref(), &columns)?); + let target_schema = Arc::new(projected_arrow_schema( + self.base.current_dataset().as_ref(), + &columns, + )?); let id_filter = Arc::new(format!("id IN ({})", sql_quoted_list(&page_ids))); let wanted: HashSet = page_ids.iter().cloned().collect(); @@ -1321,7 +1329,9 @@ impl RolloutStore { pub async fn get_blob(&self, id: &str) -> LanceResult>> { // Base-table-first: an already-merged row is found here with no MemWAL // manifest reads and no per-generation opens. - if let Some(payload) = Self::get_blob_from_dataset(self.base.current_dataset().as_ref(), id).await? { + if let Some(payload) = + Self::get_blob_from_dataset(self.base.current_dataset().as_ref(), id).await? + { return Ok(payload); } @@ -3059,7 +3069,12 @@ mod tests { /// Read the number of un-merged flushed generations recorded for a store's /// own write shard. Used by merge tests to assert the manifest drains. async fn flushed_generation_count(store: &RolloutStore) -> usize { - let object_store = store.base.current_dataset().object_store(None).await.unwrap(); + let object_store = store + .base + .current_dataset() + .object_store(None) + .await + .unwrap(); let branch_location = store.base.current_dataset().branch_location(); let manifest_store = ShardManifestStore::new( object_store, @@ -3079,7 +3094,12 @@ mod tests { /// Used to assert the resident writer claims the epoch once instead of /// bumping it on every append. async fn shard_writer_epoch(store: &RolloutStore) -> u64 { - let object_store = store.base.current_dataset().object_store(None).await.unwrap(); + let object_store = store + .base + .current_dataset() + .object_store(None) + .await + .unwrap(); let branch_location = store.base.current_dataset().branch_location(); let manifest_store = ShardManifestStore::new( object_store, @@ -3752,7 +3772,8 @@ mod tests { let generation_batch = current_store.records_to_batch(&[record]).unwrap(); legacy_store.base.ensure_latest_schema().await.unwrap(); - let merge_schema: Arc = Arc::new(legacy_store.base.current_dataset().schema().into()); + let merge_schema: Arc = + Arc::new(legacy_store.base.current_dataset().schema().into()); let aligned = align_batch_to_schema(generation_batch, merge_schema.clone()).unwrap(); let reader = RecordBatchIterator::new( vec![Ok::(aligned)].into_iter(), diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index 78d7fa8..fbedb34 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -465,10 +465,7 @@ impl StorageBase { ) .into()); } - let dataset = self - .current_dataset() - .checkout_version(version_id) - .await?; + let dataset = self.current_dataset().checkout_version(version_id).await?; self.set_dataset(dataset); self.pinned_version.store(version_id, Ordering::Release); Ok(()) @@ -1209,12 +1206,8 @@ impl StorageBase { let mut dataset = (*self.current_dataset()).clone(); let result = match config.max_source_fragments { Some(max_source_fragments) => { - compact_files_incremental( - &mut dataset, - lance_options, - max_source_fragments.max(1), - ) - .await + compact_files_incremental(&mut dataset, lance_options, max_source_fragments.max(1)) + .await } None => compact_files(&mut dataset, lance_options, None).await, }; @@ -1323,12 +1316,9 @@ impl StorageBase { /// the shared session and storage options are never dropped. pub async fn reload(&self) -> LanceResult<()> { let uri = self.uri(); - let dataset = Self::load_with_options( - &uri, - self.storage_options.clone(), - self.session.clone(), - ) - .await?; + let dataset = + Self::load_with_options(&uri, self.storage_options.clone(), self.session.clone()) + .await?; self.set_dataset(dataset); self.clear_version_pin(); Ok(()) @@ -1352,12 +1342,7 @@ impl StorageBase { return Ok(()); } let mut dataset = (*self.current_dataset()).clone(); - match dataset - .initialize_mem_wal() - .unsharded() - .execute() - .await - { + match dataset.initialize_mem_wal().unsharded().execute().await { Ok(()) => { self.set_dataset(dataset); Ok(()) @@ -1392,7 +1377,8 @@ impl StorageBase { /// Open a flushed generation dataset, inheriting the base dataset's session /// and this store's storage options. pub async fn open_flushed_dataset(&self, uri: &str) -> LanceResult { - let mut builder = DatasetBuilder::from_uri(uri).with_session(self.current_dataset().session()); + let mut builder = + DatasetBuilder::from_uri(uri).with_session(self.current_dataset().session()); if let Some(options) = self.storage_options.clone() { builder = builder.with_storage_options(options); } From 267b0c779bd96df9a67f07c747add433620ef2dd Mon Sep 17 00:00:00 2001 From: lucyge2022 Date: Fri, 21 Aug 2026 00:42:46 -0700 Subject: [PATCH 11/11] for ops that needs dataset RMW, such as 1) dataset.fn(&mut self...) 2) dataset.fn(&self) but needs replace StorageBase's dataset we fetch write_writer for exclusive access for ops that involves multiple dataset RMWs, we fetch write_writer covering multiple dataset RMWs, coupled with task local field for lock reentrant within entire scope --- .../lance-context-core/src/rollout_store.rs | 6 +- crates/lance-context-core/src/store.rs | 38 +- crates/lance-context-core/src/store_base.rs | 683 +++++++++++++++--- .../tests/wal_merge_concurrency.rs | 300 +++++++- 4 files changed, 893 insertions(+), 134 deletions(-) diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 35ef646..979f60d 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -506,8 +506,10 @@ impl RolloutStore { self.base.version() } - /// Checkout a specific dataset version — recovers the exact rollout set that - /// trained a checkpoint (spec §3, reproducibility). + /// Checkout a specific base-table version (time travel over the base table + /// only). Rollout training does not use this — reproduce a checkpoint by + /// filtering immutable rows (`policy_version`), not by pinning a dataset + /// version. See `docs/src/specs/rollout-deployment.md` §7. pub async fn checkout(&self, version_id: u64) -> LanceResult<()> { self.base.checkout(version_id).await } diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index b6f7f89..748ab11 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -1405,11 +1405,14 @@ impl ContextStore { } let schema = Arc::new(Schema::new(vec![relationship_field()])); - let mut dataset = (*self.base.current_dataset()).clone(); - dataset - .add_columns(NewColumnTransform::AllNulls(schema), None, None) + self.base + .call_dataset_mut_fn(|mut dataset| async move { + dataset + .add_columns(NewColumnTransform::AllNulls(schema), None, None) + .await?; + Ok(dataset) + }) .await?; - self.base.set_dataset(dataset); self.base.clear_version_pin(); Ok(true) } @@ -2065,18 +2068,21 @@ impl ContextStore { let params = ScalarIndexParams::default(); - let mut dataset = (*self.base.current_dataset()).clone(); - dataset - .create_index_builder(&["id"], index_type, ¶ms) - .name(ID_INDEX_NAME.to_string()) - .replace(true) - .await?; - self.base.set_dataset(dataset); - - // Reload through the base so the new index is visible to subsequent - // reads, keeping the storage options and session (a bare - // `Dataset::open` here silently dropped them). - self.base.reload().await + self.base + .with_exclusive_writer(|| async { + self.base + .call_dataset_mut_fn(|mut dataset| async move { + dataset + .create_index_builder(&["id"], index_type, ¶ms) + .name(ID_INDEX_NAME.to_string()) + .replace(true) + .await?; + Ok(dataset) + }) + .await?; + self.base.reload().await + }) + .await } /// Start background compaction task if enabled. diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index fbedb34..aa2cd4c 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -27,6 +27,20 @@ //! - **Every dataset open goes through [`StorageBase::load_with_options`]**, so //! storage options and the shared session are never silently dropped. //! +//! # Dataset-handle publish +//! +//! Readers `ArcSwap::load` a snapshot and never take a writer lock. Every +//! replacement of the handle takes [`StorageBase::write_writer`]: +//! +//! - **Single RMW** ([`StorageBase::call_dataset_mut_fn`], +//! [`StorageBase::call_dataset_with`]): lock → mutate → `set_dataset`. +//! - **Multi-step RMW** ([`StorageBase::with_exclusive_writer`]): one lock +//! covers several publishes so a concurrent `checkout` cannot land between +//! them. Nesting is task-local; another task always waits on the mutex. +//! +//! Steady-state MemWAL `put`s clone the resident `ShardWriter` and do **not** +//! hold this lock; only first-open / fence-reopen and handle publish do. +//! //! # What stays in the concrete store //! //! Anything that needs to know the schema: the Arrow schema itself, @@ -39,6 +53,20 @@ use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; +tokio::task_local! { + /// Set while *this task* holds [`StorageBase::write_writer`] for exclusive + /// handle publish / nested RMW. + /// + /// Must be task-local: a shared `AtomicBool` would let a *different* task + /// observe "held" and skip the mutex, defeating the critical section. + static WRITE_WRITER_HELD: (); +} + +#[inline] +fn exclusive_writer_held_here() -> bool { + WRITE_WRITER_HELD.try_with(|_| ()).is_ok() +} + use arc_swap::ArcSwap; use arrow_array::{new_null_array, RecordBatch, RecordBatchIterator, UInt32Array}; use arrow_schema::{ArrowError, Schema}; @@ -305,6 +333,10 @@ pub(crate) struct StorageBase { /// Serializes WAL→base merge (prepare through commit). Taken with /// `try_lock_owned`: a loser no-ops (`Ok(0)` / `Ok(None)`). Not held by /// `add`/`flush`, so appends keep running while a merge is in flight. + /// + /// Does **not** alone protect [`Self::set_dataset`]. Handle publish is + /// exclusive under [`Self::write_writer`]. Lock order: this lock, then + /// `write_writer`. Never acquire this while holding `write_writer`. merge_lock: Arc>, /// Explicit time-travel version selected by [`Self::checkout`]. /// @@ -316,10 +348,15 @@ pub(crate) struct StorageBase { /// (Lance dataset versions are 1-based, so `0` is never a real pin.) pinned_version: AtomicU64, /// Resident MemWAL writer for this instance's shard, wrapped for `&self` - /// concurrent access. The [`tokio::sync::Mutex`] is held only to - /// fetch-or-open and clone the `Arc` (see [`Self::resident_writer`]) and to - /// invalidate a fenced writer (see [`Self::invalidate_writer`]); it is - /// **never** held across `put`, so steady-state appends run concurrently. + /// concurrent access. + /// + /// Also the exclusive lock for every dataset-handle RMW (`checkout`, + /// `refresh_latest`, merge append, compact, reload, schema/index). Held + /// only to fetch-or-open / invalidate the `ShardWriter` (see + /// [`Self::resident_writer`]) or across a handle publish — **never** across + /// `put`, so steady-state appends run concurrently. Multi-step publishers + /// use [`Self::with_exclusive_writer`]; nesting is task-local so only the + /// holding task skips re-acquire. write_writer: tokio::sync::Mutex>>, } @@ -458,6 +495,11 @@ impl StorageBase { /// Check out a specific base dataset version (time travel). /// /// `version_id` must be non-zero (`0` is reserved to mean "unpinned"). + /// + /// Lance's `checkout_version` takes `&self` and returns a **new** `Dataset` + /// handle aimed at that manifest (same URI/session, different view) — it + /// does not mutate the caller's value in place. Publishing is exclusive + /// under [`Self::write_writer`] so an older version can actually be pinned. pub async fn checkout(&self, version_id: u64) -> LanceResult<()> { if version_id == 0 { return Err(ArrowError::InvalidArgumentError( @@ -465,8 +507,11 @@ impl StorageBase { ) .into()); } - let dataset = self.current_dataset().checkout_version(version_id).await?; - self.set_dataset(dataset); + self.call_dataset_with(|current| async move { + let dataset = current.checkout_version(version_id).await?; + Ok((dataset, ())) + }) + .await?; self.pinned_version.store(version_id, Ordering::Release); Ok(()) } @@ -484,9 +529,11 @@ impl StorageBase { /// WAL merges committed by another process become visible without paying the /// cost of reopening the dataset and rebuilding all session caches. pub async fn refresh_latest(&self) -> LanceResult<()> { - let mut dataset = (*self.current_dataset()).clone(); - dataset.checkout_latest().await?; - self.set_dataset(dataset); + self.call_dataset_mut_fn(|mut dataset| async move { + dataset.checkout_latest().await?; + Ok(dataset) + }) + .await?; self.clear_version_pin(); Ok(()) } @@ -504,11 +551,90 @@ impl StorageBase { } /// Publish a replacement dataset handle after a mutating Lance op. + /// + /// Unconditional `store`. Prefer [`Self::call_dataset_with`] / + /// [`Self::call_dataset_mut_fn`]. #[inline] pub(crate) fn set_dataset(&self, dataset: Dataset) { self.dataset.store(Arc::new(dataset)); } + /// Build a new handle from the current one (or ignore it — e.g. + /// [`Self::reload`]) and publish it under [`Self::write_writer`]. + /// + /// Nested publishes from [`Self::with_exclusive_writer`] reuse that section + /// instead of acquiring the lock again. + /// + /// `f` receives the current `Arc` and returns `(new_handle, out)`. + /// The handle is published only if `f` succeeds. Do not call + /// [`Self::reload`] (or anything else that takes `write_writer`) from + /// inside `f` — the mutex is not reentrant. + pub(crate) async fn call_dataset_with(&self, f: F) -> LanceResult + where + F: FnOnce(Arc) -> Fut, + Fut: std::future::Future>, + { + if exclusive_writer_held_here() { + return self.call_dataset_with_locked(f).await; + } + let _guard = self.write_writer.lock().await; + WRITE_WRITER_HELD + .scope((), self.call_dataset_with_locked(f)) + .await + } + + async fn call_dataset_with_locked(&self, f: F) -> LanceResult + where + F: FnOnce(Arc) -> Fut, + Fut: std::future::Future>, + { + let current = self.current_dataset(); + let (dataset, out) = f(current).await?; + self.before_publish_dataset().await; + self.set_dataset(dataset); + Ok(out) + } + + /// Clone the current handle, run `f`, and publish the result. + /// + /// Single-step RMW (`append`, `checkout_latest`, `add_columns`, …). For + /// several publishes in one critical section, wrap them in + /// [`Self::with_exclusive_writer`]. + pub(crate) async fn call_dataset_mut_fn(&self, f: F) -> LanceResult<()> + where + F: FnOnce(Dataset) -> Fut, + Fut: std::future::Future>, + { + self.call_dataset_with(|current| async move { + let dataset = f((*current).clone()).await?; + Ok((dataset, ())) + }) + .await + } + + /// Run `f` under one exclusive `write_writer` section so multiple handle + /// publishes compose. + /// + /// Nesting is allowed only for the **same task** (via + /// [`WRITE_WRITER_HELD`]); other tasks block on the mutex. + pub(crate) async fn with_exclusive_writer(&self, f: F) -> LanceResult + where + F: FnOnce() -> Fut, + Fut: std::future::Future>, + { + if exclusive_writer_held_here() { + return f().await; + } + let _guard = self.write_writer.lock().await; + WRITE_WRITER_HELD.scope((), f()).await + } + + /// Test-only pause point after modify, before the handle is published. + async fn before_publish_dataset(&self) { + #[cfg(test)] + dataset_rmw_test_hooks::await_before_publish().await; + } + // ---------------------------------------------------------------- writes /// Durably append `batches` through this instance's MemWAL shard. @@ -939,15 +1065,19 @@ impl StorageBase { return Ok(false); } - self.ensure_latest_schema().await?; + self.with_exclusive_writer(|| async { + self.ensure_latest_schema().await?; - if !batches.is_empty() { - observe_phase!( - "append", - self.merge_prepared_batches(batches, merge_schema).await - )?; - self.clear_version_pin(); - } + if !batches.is_empty() { + observe_phase!( + "append", + self.merge_prepared_batches(batches, merge_schema).await + )?; + self.clear_version_pin(); + } + Ok(()) + }) + .await?; // Reuse the shard's *current* epoch rather than claiming a new one: // claiming would fence our own live writer. `commit_update` still fails @@ -1111,15 +1241,15 @@ impl StorageBase { batches.into_iter().map(Ok::), merge_schema, ); - let mut builder = MergeInsertBuilder::try_new( - self.current_dataset(), - vec![self.key_column.clone()], - )?; - builder.when_matched(WhenMatched::UpdateAll); - let job = builder.try_build()?; - let (dataset, _) = job.execute_reader(reader).await?; - self.set_dataset(Arc::unwrap_or_clone(dataset)); - Ok(()) + let key_column = self.key_column.clone(); + self.call_dataset_with(|current| async move { + let mut builder = MergeInsertBuilder::try_new(current, vec![key_column])?; + builder.when_matched(WhenMatched::UpdateAll); + let job = builder.try_build()?; + let (dataset, _) = job.execute_reader(reader).await?; + Ok((Arc::unwrap_or_clone(dataset), ())) + }) + .await } /// Evolve an older base table to the store's latest additive schema. @@ -1131,32 +1261,37 @@ impl StorageBase { let Some(latest_schema) = self.latest_schema.clone() else { return Ok(()); }; - self.refresh_latest().await?; + self.with_exclusive_writer(|| async { + self.refresh_latest().await?; - let base_schema: Arc = Arc::new(self.current_dataset().schema().into()); - align_batch_to_schema( - RecordBatch::new_empty(base_schema.clone()), - latest_schema.clone(), - )?; + let base_schema: Arc = Arc::new(self.current_dataset().schema().into()); + align_batch_to_schema( + RecordBatch::new_empty(base_schema.clone()), + latest_schema.clone(), + )?; - let missing_fields = latest_schema - .fields() - .iter() - .filter(|field| base_schema.field_with_name(field.name()).is_err()) - .cloned() - .collect::>(); - if !missing_fields.is_empty() { - let mut dataset = (*self.current_dataset()).clone(); - dataset - .add_columns( - NewColumnTransform::AllNulls(Arc::new(Schema::new(missing_fields))), - None, - None, - ) + let missing_fields = latest_schema + .fields() + .iter() + .filter(|field| base_schema.field_with_name(field.name()).is_err()) + .cloned() + .collect::>(); + if !missing_fields.is_empty() { + self.call_dataset_mut_fn(|mut dataset| async move { + dataset + .add_columns( + NewColumnTransform::AllNulls(Arc::new(Schema::new(missing_fields))), + None, + None, + ) + .await?; + Ok(dataset) + }) .await?; - self.set_dataset(dataset); - } - Ok(()) + } + Ok(()) + }) + .await } // ------------------------------------------------- compaction & indexing @@ -1203,43 +1338,56 @@ impl StorageBase { ..Default::default() }; - let mut dataset = (*self.current_dataset()).clone(); - let result = match config.max_source_fragments { - Some(max_source_fragments) => { - compact_files_incremental(&mut dataset, lance_options, max_source_fragments.max(1)) - .await - } - None => compact_files(&mut dataset, lance_options, None).await, - }; - self.set_dataset(dataset); + // Compact then reload as one exclusive section so a checkout cannot + // land between the two publishes. Still publish the local handle even + // when compact returns Err (preserves prior behavior: Lance may have + // partially updated the in-memory view). + self.with_exclusive_writer(|| async { + let result = self + .call_dataset_with(|current| async move { + let mut dataset = (*current).clone(); + let result = match config.max_source_fragments { + Some(max_source_fragments) => { + compact_files_incremental( + &mut dataset, + lance_options, + max_source_fragments.max(1), + ) + .await + } + None => compact_files(&mut dataset, lance_options, None).await, + }; + Ok((dataset, result)) + }) + .await?; - match result { - Ok(metrics) => { - // Reload the handle so the caller (and subsequent reads on this - // instance) observe the compacted version. - self.reload().await?; - { - let mut state = self.compaction.lock().unwrap_or_else(|e| e.into_inner()); - state.last_compaction = Some(Utc::now()); - state.total_compactions += 1; - state.last_error = None; + match result { + Ok(metrics) => { + self.reload().await?; + { + let mut state = self.compaction.lock().unwrap_or_else(|e| e.into_inner()); + state.last_compaction = Some(Utc::now()); + state.total_compactions += 1; + state.last_error = None; + } + info!( + fragments_removed = metrics.fragments_removed, + fragments_added = metrics.fragments_added, + "base-table compaction completed" + ); + Ok(metrics) } - info!( - fragments_removed = metrics.fragments_removed, - fragments_added = metrics.fragments_added, - "base-table compaction completed" - ); - Ok(metrics) - } - Err(e) => { - warn!(error = %e, "base-table compaction failed"); - { - let mut state = self.compaction.lock().unwrap_or_else(|e| e.into_inner()); - state.last_error = Some(e.to_string()); + Err(e) => { + warn!(error = %e, "base-table compaction failed"); + { + let mut state = self.compaction.lock().unwrap_or_else(|e| e.into_inner()); + state.last_error = Some(e.to_string()); + } + Err(e) } - Err(e) } - } + }) + .await } /// Build a ZoneMap scalar index on the base table's key column. @@ -1259,20 +1407,24 @@ impl StorageBase { /// scan of those generations. pub async fn create_key_zonemap_index(&self) -> LanceResult<()> { info!(column = %self.key_column, "creating ZoneMap index on key column"); - let mut dataset = (*self.current_dataset()).clone(); - dataset - .create_index_builder( - &[self.key_column.as_str()], - IndexType::ZoneMap, - &ScalarIndexParams::default(), - ) - .name(ID_INDEX_NAME.to_string()) - .replace(true) + let key_column = self.key_column.clone(); + self.with_exclusive_writer(|| async { + self.call_dataset_mut_fn(|mut dataset| async move { + dataset + .create_index_builder( + &[key_column.as_str()], + IndexType::ZoneMap, + &ScalarIndexParams::default(), + ) + .name(ID_INDEX_NAME.to_string()) + .replace(true) + .await?; + Ok(dataset) + }) .await?; - self.set_dataset(dataset); - // Reload the handle so subsequent reads on this instance observe the new - // index (mirrors the reload done after `compact`). - self.reload().await + self.reload().await + }) + .await } /// Whether the base table has accumulated at least `min_fragments` @@ -1316,10 +1468,13 @@ impl StorageBase { /// the shared session and storage options are never dropped. pub async fn reload(&self) -> LanceResult<()> { let uri = self.uri(); - let dataset = - Self::load_with_options(&uri, self.storage_options.clone(), self.session.clone()) - .await?; - self.set_dataset(dataset); + let storage_options = self.storage_options.clone(); + let session = self.session.clone(); + self.call_dataset_with(move |_current| async move { + let dataset = Self::load_with_options(&uri, storage_options, session).await?; + Ok((dataset, ())) + }) + .await?; self.clear_version_pin(); Ok(()) } @@ -1341,23 +1496,27 @@ impl StorageBase { if self.mem_wal_index_present().await? { return Ok(()); } - let mut dataset = (*self.current_dataset()).clone(); - match dataset.initialize_mem_wal().unsharded().execute().await { - Ok(()) => { - self.set_dataset(dataset); - Ok(()) - } - Err(err) => { - // A concurrent first-writer may have created the index between - // our check and our commit. Reload and accept it if so. - self.reload().await?; - if self.mem_wal_index_present().await? { - Ok(()) - } else { - Err(err) + self.with_exclusive_writer(|| async { + let init_result = self + .call_dataset_with(|current| async move { + let mut dataset = (*current).clone(); + dataset.initialize_mem_wal().unsharded().execute().await?; + Ok((dataset, ())) + }) + .await; + match init_result { + Ok(()) => Ok(()), + Err(err) => { + self.reload().await?; + if self.mem_wal_index_present().await? { + Ok(()) + } else { + Err(err) + } } } - } + }) + .await } async fn mem_wal_index_present(&self) -> LanceResult { @@ -1771,3 +1930,299 @@ pub fn derive_shard_id(instance_id: Option<&str>) -> Uuid { let input = instance_id.unwrap_or("default"); Uuid::new_v5(&Uuid::NAMESPACE_OID, input.as_bytes()) } + +/// Test-only controls for proving ArcSwap handle lost-updates. +/// +/// Integration tests cannot see `cfg(test)` on this crate, so deterministic +/// RMW races live in unit tests below that use these hooks. +#[cfg(test)] +pub(crate) mod dataset_rmw_test_hooks { + use std::sync::{Arc, Mutex}; + + use tokio::sync::Notify; + + /// These hooks are process-global; RMW tests that use them must not overlap. + static SERIAL: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + static BEFORE_PUBLISH: Mutex>> = Mutex::new(None); + static READY_TO_PUBLISH: Mutex>> = Mutex::new(None); + + pub async fn serial() -> tokio::sync::MutexGuard<'static, ()> { + SERIAL.lock().await + } + + /// Install a barrier: the next `before_publish` notifies `ready`, then waits + /// on `proceed`. Returns `(ready, proceed)`. + pub fn install_before_publish_barrier() -> (Arc, Arc) { + let ready = Arc::new(Notify::new()); + let proceed = Arc::new(Notify::new()); + *BEFORE_PUBLISH.lock().unwrap() = Some(proceed.clone()); + *READY_TO_PUBLISH.lock().unwrap() = Some(ready.clone()); + (ready, proceed) + } + + pub fn clear() { + *BEFORE_PUBLISH.lock().unwrap() = None; + *READY_TO_PUBLISH.lock().unwrap() = None; + } + + pub async fn await_before_publish() { + let ready = READY_TO_PUBLISH.lock().unwrap().clone(); + let proceed = BEFORE_PUBLISH.lock().unwrap().clone(); + let (Some(ready), Some(proceed)) = (ready, proceed) else { + return; + }; + // One-shot: clear so only the first publisher hits the barrier. + *READY_TO_PUBLISH.lock().unwrap() = None; + *BEFORE_PUBLISH.lock().unwrap() = None; + ready.notify_one(); + proceed.notified().await; + } +} + +#[cfg(test)] +mod dataset_handle_rmw_tests { + use super::dataset_rmw_test_hooks as hooks; + use crate::{RolloutRecord, RolloutStore, RolloutStoreOptions, ROLE_ASSISTANT}; + use std::sync::Arc; + + fn rec(id: &str) -> RolloutRecord { + RolloutRecord { + id: id.to_string(), + rollout_id: "r".to_string(), + problem_id: "p".to_string(), + dataset: Some("d".to_string()), + sequence_order: 0, + role: ROLE_ASSISTANT.to_string(), + created_at: chrono::Utc::now(), + content: Some("x".to_string()), + content_type: "text/plain".to_string(), + model_input_string: None, + model_output_string: None, + rationale: None, + problem_text: None, + user_metadata: None, + input_tokens: None, + output_tokens: None, + num_input_tokens: None, + num_output_tokens: None, + output_logprobs: None, + input_logprobs: None, + ref_logprobs: None, + loss_mask: None, + advantage: None, + reward: None, + raw_reward: None, + grader_id: None, + score: None, + include_in_training: None, + exclude_reason: None, + policy_version: None, + relationships: vec![], + binary_payload: None, + payload_size: None, + payload_checksum: None, + artifact_type: None, + metadata: None, + } + } + + async fn merge_once(store: &RolloutStore) -> usize { + let prepared = store.prepare_cleanup_merge().await.unwrap(); + match prepared { + Some((ms, m, p)) => store.commit_prepared_merge(&ms, &m, p).await.unwrap(), + None => 0, + } + } + + /// Refresh holds `write_writer` across checkout_latest→store, so a concurrent + /// merge cannot publish in between. After refresh completes, merge proceeds + /// and the in-memory version never goes backwards. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn merge_waits_on_refresh_write_writer() { + use std::time::Duration; + + let _serial = hooks::serial().await; + hooks::clear(); + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = Arc::new( + RolloutStore::open_with_options( + &uri, + RolloutStoreOptions { + shard_id: Some("solo".into()), + merge_after_generations: None, + ..Default::default() + }, + ) + .await + .unwrap(), + ); + + for i in 0..4 { + store.add(&[rec(&format!("pre-{i}"))]).await.unwrap(); + store.flush().await.unwrap(); + } + assert!(merge_once(&store).await > 0); + + for i in 0..3 { + store.add(&[rec(&format!("pending-{i}"))]).await.unwrap(); + store.flush().await.unwrap(); + } + + hooks::clear(); + let (ready, proceed) = hooks::install_before_publish_barrier(); + + let refresher = { + let store = store.clone(); + tokio::spawn(async move { + store.refresh_latest().await.unwrap(); + }) + }; + + ready.notified().await; + let v_at_pause = store.version(); + + let merger = { + let store = store.clone(); + tokio::spawn(async move { merge_once(&store).await }) + }; + + tokio::time::sleep(Duration::from_millis(150)).await; + assert!( + !merger.is_finished(), + "merge must wait on write_writer held by refresh" + ); + + proceed.notify_one(); + refresher.await.unwrap(); + let reclaimed = merger.await.unwrap(); + hooks::clear(); + assert!(reclaimed > 0); + assert!( + store.version() >= v_at_pause, + "handle version must not roll back after refresh then merge" + ); + for i in 0..3 { + assert!(store + .get_by_id(&format!("pending-{i}")) + .await + .unwrap() + .is_some()); + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn merge_after_checkout_clears_pin_and_lands_on_tip() { + let _serial = hooks::serial().await; + hooks::clear(); + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = RolloutStore::open_with_options( + &uri, + RolloutStoreOptions { + shard_id: Some("solo".into()), + merge_after_generations: None, + ..Default::default() + }, + ) + .await + .unwrap(); + + for i in 0..3 { + store.add(&[rec(&format!("a-{i}"))]).await.unwrap(); + store.flush().await.unwrap(); + } + assert!(merge_once(&store).await > 0); + let pinned = store.version(); + store.checkout(pinned).await.unwrap(); + assert!(store.is_version_pinned()); + + for i in 0..3 { + store.add(&[rec(&format!("b-{i}"))]).await.unwrap(); + store.flush().await.unwrap(); + } + assert!(merge_once(&store).await > 0); + assert!( + !store.is_version_pinned(), + "rollout merge refresh should clear an explicit checkout pin" + ); + let tip = store.version(); + store.refresh_latest().await.unwrap(); + assert_eq!(store.version(), tip, "handle should already be at tip"); + for i in 0..3 { + assert!(store.get_by_id(&format!("b-{i}")).await.unwrap().is_some()); + } + } + + /// A shared AtomicBool "held" flag would let task B see task A's hold and + /// skip `write_writer` — defeating the critical section. Task-local nesting + /// must make B wait on the mutex instead. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn exclusive_writer_held_is_task_local() { + use arrow_schema::{DataType, Field, Schema}; + use std::time::Duration; + use tokio::sync::Notify; + + use super::{StorageBase, StorageBaseOptions}; + + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Utf8, false)])); + let base = Arc::new( + StorageBase::open( + &uri, + StorageBaseOptions { + storage_options: None, + shard_id: Some("solo".into()), + merge_after_generations: None, + session: None, + schema, + key_column: "id".into(), + latest_schema: None, + seal_on_put: true, + }, + true, + ) + .await + .unwrap(), + ); + + let ready = Arc::new(Notify::new()); + let proceed = Arc::new(Notify::new()); + + let holder = { + let base = Arc::clone(&base); + let ready = Arc::clone(&ready); + let proceed = Arc::clone(&proceed); + tokio::spawn(async move { + base.with_exclusive_writer(|| async { + ready.notify_one(); + proceed.notified().await; + Ok(()) + }) + .await + .unwrap(); + }) + }; + + ready.notified().await; + + let waiter = { + let base = Arc::clone(&base); + tokio::spawn(async move { + base.refresh_latest().await.unwrap(); + }) + }; + + // If held leaked across tasks, waiter would finish without waiting. + tokio::time::sleep(Duration::from_millis(150)).await; + assert!( + !waiter.is_finished(), + "another task must block on write_writer, not skip via a shared held flag" + ); + + proceed.notify_one(); + holder.await.unwrap(); + waiter.await.unwrap(); + } +} diff --git a/crates/lance-context-core/tests/wal_merge_concurrency.rs b/crates/lance-context-core/tests/wal_merge_concurrency.rs index e45f3d2..d0d3c85 100644 --- a/crates/lance-context-core/tests/wal_merge_concurrency.rs +++ b/crates/lance-context-core/tests/wal_merge_concurrency.rs @@ -1,3 +1,5 @@ +#![recursion_limit = "256"] + //! Concurrency tests for WAL self-merge: a merge must never block or corrupt //! concurrent appends. //! @@ -18,8 +20,16 @@ //! 3. concurrent merges do not duplicate rows, and `merge_lock` excludes a //! second prepare while the first `PreparedMerge` is still live; //! 4. an interrupted merge loses nothing (rows stay readable exactly once); -//! 5. `add` is not blocked for the merge's duration. - +//! 5. `add` is not blocked for the merge's duration; +//! 6. `add` is not blocked for a base-table compact's duration; +//! 7. concurrent `refresh_latest` cannot roll the in-memory dataset handle +//! backwards over a merge's published version (`write_writer` serializes +//! handle publish); +//! 8. same handle monotonicity under concurrent refresh vs compact; +//! 9. merge and compact can run together without losing rows; +//! 10. `get_by_id` does not flaky-miss merged rows under a refresh storm. + +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -595,3 +605,289 @@ async fn append_is_not_blocked_for_the_duration_of_a_compact() { ); assert_eq!(ids.len(), 21, "all rows readable exactly once: {ids:?}"); } + +/// Blind `ArcSwap::store` after load→modify→await lets `refresh_latest` publish +/// an older handle over a concurrent merge append. That rolls the in-memory +/// version backwards: object storage still has the merge, but `get_by_id` on +/// the base handle can flaky-miss until the next refresh (#234 class). +/// +/// Exclusive `write_writer` serializes handle publish so sampled versions never +/// decrease and merged rows stay visible immediately after commit. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn refresh_cannot_roll_back_dataset_handle_over_merge() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = Arc::new(RwLock::new( + RolloutStore::open_with_options(&uri, opts("solo")) + .await + .unwrap(), + )); + + for i in 0..12 { + store + .read() + .await + .add(&[rec(&format!("row-{i}"))]) + .await + .unwrap(); + store.read().await.flush().await.unwrap(); + } + + let stop = Arc::new(AtomicBool::new(false)); + let max_seen = Arc::new(AtomicU64::new(store.read().await.version())); + let dips = Arc::new(AtomicU64::new(0)); + let refresher = { + let store = store.clone(); + let stop = stop.clone(); + let max_seen = max_seen.clone(); + let dips = dips.clone(); + tokio::spawn(async move { + while !stop.load(Ordering::Acquire) { + store.read().await.refresh_latest().await.unwrap(); + let v = store.read().await.version(); + let prev_max = max_seen.fetch_max(v, Ordering::SeqCst); + if v < prev_max { + dips.fetch_add(1, Ordering::SeqCst); + } + } + }) + }; + + let mut extra = 0usize; + for _ in 0..3 { + for _ in 0..4 { + store + .read() + .await + .add(&[rec(&format!("extra-{extra}"))]) + .await + .unwrap(); + store.read().await.flush().await.unwrap(); + extra += 1; + } + let reclaimed = merge_like_sweeper(&store).await; + assert!(reclaimed > 0, "merge should reclaim pending generations"); + let v = store.read().await.version(); + let prev_max = max_seen.fetch_max(v, Ordering::SeqCst); + assert!( + v >= prev_max, + "merge published version {v} below previously seen max {prev_max}" + ); + } + + stop.store(true, Ordering::Release); + refresher.await.unwrap(); + + assert_eq!( + dips.load(Ordering::SeqCst), + 0, + "in-memory dataset version went backwards under concurrent refresh_latest" + ); + + let ids = read_ids(&store).await; + for i in 0..12 { + assert!( + ids.contains(&format!("row-{i}")), + "merged row-{i} missing from handle after refresh race: {ids:?}" + ); + } +} + +async fn track_version_dips( + store: Arc>, + stop: Arc, + max_seen: Arc, + dips: Arc, +) { + while !stop.load(Ordering::Acquire) { + // Compact can briefly make a concurrent checkout_latest miss a + // mid-rewrite manifest; retry rather than failing the storm. + if store.read().await.refresh_latest().await.is_err() { + continue; + } + let v = store.read().await.version(); + let prev_max = max_seen.fetch_max(v, Ordering::SeqCst); + if v < prev_max { + dips.fetch_add(1, Ordering::SeqCst); + } + } +} + +/// Compact publishes a new handle (then reloads); concurrent refresh must not +/// roll the in-memory version backwards over that publish. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn refresh_cannot_roll_back_dataset_handle_over_compact() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = Arc::new(RwLock::new( + RolloutStore::open_with_options(&uri, opts("solo")) + .await + .unwrap(), + )); + + for i in 0..12 { + store + .read() + .await + .add(&[rec(&format!("row-{i}"))]) + .await + .unwrap(); + store.read().await.flush().await.unwrap(); + assert!(merge_like_sweeper(&store).await > 0); + } + + let stop = Arc::new(AtomicBool::new(false)); + let max_seen = Arc::new(AtomicU64::new(store.read().await.version())); + let dips = Arc::new(AtomicU64::new(0)); + let refresher = tokio::spawn(track_version_dips( + store.clone(), + stop.clone(), + max_seen.clone(), + dips.clone(), + )); + + let metrics = { + let guard = store.read().await; + guard + .compact(Some(CompactionConfig { + min_fragments: 2, + num_threads: Some(1), + batch_size: Some(1), + ..Default::default() + })) + .await + .unwrap() + }; + assert!( + metrics.fragments_removed > 0, + "compact should rewrite fragments, got removed={} added={}", + metrics.fragments_removed, + metrics.fragments_added + ); + let v = store.read().await.version(); + let prev_max = max_seen.fetch_max(v, Ordering::SeqCst); + assert!( + v >= prev_max, + "compact published version {v} below previously seen max {prev_max}" + ); + + stop.store(true, Ordering::Release); + refresher.await.unwrap(); + assert_eq!( + dips.load(Ordering::SeqCst), + 0, + "in-memory dataset version went backwards under concurrent refresh vs compact" + ); + assert_eq!(read_ids(&store).await.len(), 12); +} + +/// Lance treats Append (WAL merge) vs Rewrite (compact) as non-conflicting; +/// both must succeed and conserve rows under shared store locks. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn merge_and_compact_concurrently_preserve_rows() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = Arc::new(RwLock::new( + RolloutStore::open_with_options(&uri, opts("solo")) + .await + .unwrap(), + )); + + for i in 0..8 { + store + .read() + .await + .add(&[rec(&format!("base-{i}"))]) + .await + .unwrap(); + store.read().await.flush().await.unwrap(); + assert!(merge_like_sweeper(&store).await > 0); + } + for i in 0..6 { + store + .read() + .await + .add(&[rec(&format!("wal-{i}"))]) + .await + .unwrap(); + store.read().await.flush().await.unwrap(); + } + + let merger = { + let store = store.clone(); + tokio::spawn(async move { merge_like_sweeper(&store).await }) + }; + let compactor = { + let store = store.clone(); + tokio::spawn(async move { + store + .read() + .await + .compact(Some(CompactionConfig { + min_fragments: 2, + num_threads: Some(1), + ..Default::default() + })) + .await + .unwrap() + }) + }; + + let reclaimed = merger.await.unwrap(); + let metrics = compactor.await.unwrap(); + assert!(reclaimed > 0, "merge should reclaim WAL generations"); + let _ = metrics; + + let ids = read_ids(&store).await; + assert_eq!(ids.len(), 14, "all rows readable exactly once: {ids:?}"); + for i in 0..8 { + assert!(ids.contains(&format!("base-{i}"))); + } + for i in 0..6 { + assert!(ids.contains(&format!("wal-{i}"))); + } +} + +/// Production symptom of handle rollback: merged id briefly missing from +/// `get_by_id` while refresh races merge. Exclusive handle publish keeps the +/// id visible once merge returns. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn get_by_id_sees_merged_rows_under_refresh_storm() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = Arc::new(RwLock::new( + RolloutStore::open_with_options(&uri, opts("solo")) + .await + .unwrap(), + )); + + let stop = Arc::new(AtomicBool::new(false)); + let refresher = { + let store = store.clone(); + let stop = stop.clone(); + tokio::spawn(async move { + while !stop.load(Ordering::Acquire) { + store.read().await.refresh_latest().await.unwrap(); + } + }) + }; + + store.read().await.add(&[rec("target")]).await.unwrap(); + store.read().await.flush().await.unwrap(); + assert!(merge_like_sweeper(&store).await > 0); + + for _ in 0..30 { + let hit = store + .read() + .await + .get_by_id("target") + .await + .unwrap() + .is_some(); + assert!(hit, "merged target must remain visible under refresh storm"); + tokio::time::sleep(Duration::from_millis(2)).await; + } + + stop.store(true, Ordering::Release); + refresher.await.unwrap(); +}