From 948b9707d22a031308aca4d43bd81c8ac85b4339 Mon Sep 17 00:00:00 2001 From: Beinan Date: Thu, 10 Sep 2026 22:25:07 +0000 Subject: [PATCH 1/2] fix: harden storage consistency and blob admission --- .../lance-context-core/src/datagen_store.rs | 6 + .../lance-context-core/src/generic_store.rs | 6 + crates/lance-context-core/src/lib.rs | 4 +- .../lance-context-core/src/rollout_store.rs | 178 ++++++++--- crates/lance-context-core/src/storage.rs | 11 + crates/lance-context-core/src/store.rs | 17 ++ crates/lance-context-core/src/store_base.rs | 120 +++++++- .../tests/storage_reliability.rs | 280 ++++++++++++++++++ .../src/routes/contexts.rs | 20 +- .../src/routes/datagen.rs | 4 - .../src/routes/rollouts.rs | 244 +++++++++++++-- crates/lance-context-server/src/state.rs | 166 ++++++++--- docs/src/specs/rollout-deployment.md | 4 +- 13 files changed, 926 insertions(+), 134 deletions(-) create mode 100644 crates/lance-context-core/tests/storage_reliability.rs diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index cf78c7a7..b933b504 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -195,6 +195,12 @@ impl DatagenStore { self.base.close().await } + /// Stop this handle's writer and delete the dataset through its backend. + /// Remote writers must be quiesced by the caller before deletion. + pub async fn delete_data(&mut self) -> LanceResult<()> { + self.base.delete_data().await + } + /// Read one item's event history without materializing blob bytes. pub async fn events_for_item(&self, item_id: &str) -> LanceResult> { self.filtered_events(&format!("item_id = '{}'", escape_sql_literal(item_id))) diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index f568aabf..c835c5a6 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -362,6 +362,12 @@ impl GenericStore { self.base.close().await } + /// Stop this handle's writer and delete the dataset through its backend. + /// Remote writers must be quiesced by the caller before deletion. + pub async fn delete_data(&mut self) -> LanceResult<()> { + self.base.delete_data().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 { diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 21bcf48f..586fae01 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -67,7 +67,9 @@ pub use rollout_store::{ }; // Schema declaration lives in the API crate: it is part of the wire contract. pub use lance_context_api::{ColumnSpec, ColumnType, SchemaSpec, ID_COLUMN}; -pub use storage::{create_local_dir_if_needed, join_uri, validate_store_name, MAX_STORE_NAME_LEN}; +pub use storage::{ + create_local_dir_if_needed, join_uri, remove_dataset, validate_store_name, MAX_STORE_NAME_LEN, +}; pub use store::{ CompactionConfig, CompactionStats, ContextStore, ContextStoreOptions, DistanceMetric, IdIndexType, ReadProjection, diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 313f0c01..11cda7b6 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -471,7 +471,7 @@ impl RolloutStore { merge_max_bytes, session, } = options; - let base = StorageBase::open( + let mut base = StorageBase::open( uri, StorageBaseOptions { storage_options, @@ -495,6 +495,9 @@ impl RolloutStore { create_if_missing, ) .await?; + // Encoding follows the base schema. Evolve before accepting writes, + // otherwise values for newly supported fields are silently discarded. + base.ensure_latest_schema().await?; Ok(Self { base }) } /// URI of the underlying Lance dataset. @@ -509,8 +512,8 @@ impl RolloutStore { self.base.version() } - /// Checkout a specific dataset version — recovers the exact rollout set that - /// trained a checkpoint (spec §3, reproducibility). + /// Check out an immutable base-table version. Live WAL generations are + /// excluded; unmerged appends are not part of a base version. pub async fn checkout(&mut self, version_id: u64) -> LanceResult<()> { self.base.checkout(version_id).await } @@ -602,6 +605,12 @@ impl RolloutStore { self.base.close().await } + /// Stop this handle's writer and delete the dataset through its backend. + /// Remote writers must be quiesced by the caller before deletion. + pub async fn delete_data(&mut self) -> LanceResult<()> { + self.base.delete_data().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. @@ -813,17 +822,36 @@ impl RolloutStore { limit: usize, offset: usize, source: ListSource, + ) -> LanceResult { + self.base + .read_consistent(source, |dataset, snapshots| { + self.list_filtered_snapshot(filters, limit, offset, source, dataset, snapshots) + }) + .await + } + + async fn list_filtered_snapshot( + &self, + filters: &RolloutFilters, + limit: usize, + offset: usize, + source: ListSource, + dataset: Dataset, + shard_snapshots: Vec, ) -> LanceResult { let filter = filters.expression(); - let columns = self.non_blob_columns(); + let columns = Self::non_blob_columns_for(&dataset); let refs: Vec<&str> = columns.iter().map(String::as_str).collect(); let page_limit = limit.saturating_add(1); if matches!(source, ListSource::Wal | ListSource::All) { - let shard_snapshots = self.wal_shard_snapshots().await?; - let mut scanner = self - .lsm_scanner_for_source(source, shard_snapshots.clone()) - .project(&["id"])?; + let mut scanner = StorageBase::lsm_scanner_for_dataset( + &dataset, + "id", + source, + shard_snapshots.clone(), + ) + .project(&["id"])?; if let Some(filter) = &filter { scanner = scanner.filter(filter)?; } @@ -862,7 +890,7 @@ impl RolloutStore { let has_more = page_ids.len() > limit; page_ids.truncate(limit); let records = self - .take_lsm_page_rows(source, shard_snapshots, page_ids) + .take_lsm_page_rows(&dataset, source, shard_snapshots, page_ids) .await?; return Ok(RolloutPage { records, has_more }); } @@ -879,13 +907,13 @@ impl RolloutStore { "pagination offset exceeds i64::MAX".to_string(), )) })?; - let mut scanner = self.base.dataset.scan(); + let mut scanner = 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(), + dataset.schema(), )?); if let Some(filter) = &filter { scanner.filter(filter)?; @@ -906,13 +934,25 @@ impl RolloutStore { } async fn list_all_non_blob_records(&self) -> LanceResult> { - let columns = Arc::new(self.non_blob_columns()); - let target_schema = Arc::new(projected_arrow_schema(&self.base.dataset, &columns)?); + self.base + .read_consistent(ListSource::All, |dataset, snapshots| { + self.list_non_blob_snapshot(dataset, snapshots) + }) + .await + } + + async fn list_non_blob_snapshot( + &self, + dataset: Dataset, + snapshots: Vec, + ) -> LanceResult> { + let columns = Arc::new(Self::non_blob_columns_for(&dataset)); + let target_schema = Arc::new(projected_arrow_schema(&dataset, &columns)?); let mut records_by_id = HashMap::new(); let mut records = Vec::new(); Self::append_non_blob_records_from_dataset( - self.base.dataset.clone(), + dataset, columns.clone(), target_schema.clone(), &mut records_by_id, @@ -920,14 +960,10 @@ impl RolloutStore { ) .await?; - for snapshot in self.wal_shard_snapshots().await? { + for snapshot in snapshots { for generation in snapshot.flushed_generations { let uri = self.flushed_generation_uri(snapshot.shard_id, &generation.path); - let dataset = match self.open_flushed_dataset(&uri).await { - Ok(dataset) => dataset, - Err(err) if is_not_found_error(&err) => continue, - Err(err) => return Err(err), - }; + let dataset = self.open_flushed_dataset(&uri).await?; Self::append_non_blob_records_from_dataset( dataset, columns.clone(), @@ -972,6 +1008,7 @@ impl RolloutStore { /// take only those rows' non-blob columns. async fn take_lsm_page_rows( &self, + dataset: &Dataset, source: ListSource, shard_snapshots: Vec, page_ids: Vec, @@ -980,15 +1017,15 @@ impl RolloutStore { return Ok(Vec::new()); } - let columns = Arc::new(self.non_blob_columns()); - let target_schema = Arc::new(projected_arrow_schema(&self.base.dataset, &columns)?); + let columns = Arc::new(Self::non_blob_columns_for(dataset)); + let target_schema = Arc::new(projected_arrow_schema(dataset, &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(), + dataset.clone(), id_filter.clone(), columns.clone(), target_schema.clone(), @@ -1012,11 +1049,7 @@ impl RolloutStore { let target_schema = target_schema.clone(); let id_filter = id_filter.clone(); async move { - let dataset = match self.open_flushed_dataset(&uri).await { - Ok(dataset) => dataset, - Err(err) if is_not_found_error(&err) => return Ok(Vec::new()), - Err(err) => return Err(err), - }; + let dataset = self.open_flushed_dataset(&uri).await?; Self::take_page_rows_from_dataset(dataset, id_filter, columns, target_schema) .await } @@ -1413,8 +1446,11 @@ impl RolloutStore { /// Top-level column names excluding `binary_payload`, so list-style scans /// never materialize artifact bytes. fn non_blob_columns(&self) -> Vec { - self.base - .dataset + Self::non_blob_columns_for(&self.base.dataset) + } + + fn non_blob_columns_for(dataset: &Dataset) -> Vec { + dataset .schema() .fields .iter() @@ -2113,6 +2149,56 @@ mod tests { use serde_json::json; use tempfile::TempDir; + #[tokio::test] + async fn list_retries_when_a_generation_is_deleted_during_the_read() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let writer = RolloutStore::open_with_options( + uri, + RolloutStoreOptions { + shard_id: Some("merge-writer".into()), + ..Default::default() + }, + ) + .await + .unwrap(); + let reader = RolloutStore::open_with_options( + uri, + RolloutStoreOptions { + shard_id: Some("read-worker".into()), + ..Default::default() + }, + ) + .await + .unwrap(); + writer + .add(&[assistant_record("merged-during-read")]) + .await + .unwrap(); + writer.flush().await.unwrap(); + let writer = tokio::sync::Mutex::new(writer); + let attempts = AtomicUsize::new(0); + let rows = reader + .base + .read_consistent(ListSource::All, |dataset, snapshots| { + let writer = &writer; + let reader = &reader; + let attempts = &attempts; + async move { + if attempts.fetch_add(1, Ordering::SeqCst) == 0 { + writer.lock().await.cleanup_own_shard().await.unwrap(); + } + reader.list_non_blob_snapshot(dataset, snapshots).await + } + }) + .await + .unwrap(); + assert_eq!(attempts.load(Ordering::SeqCst), 2); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].id, "merged-during-read"); + } + #[test] fn rollout_filters_parse_supported_fields() { let filters = RolloutFilters::from_json_value(json!({ @@ -2270,16 +2356,28 @@ mod tests { let legacy_schema = pre_claim_check_schema(); create_empty_dataset(&uri, legacy_schema.clone()).await; - let mut store = RolloutStore::open_with_options( - &uri, - RolloutStoreOptions { - shard_id: Some("pre-claim-check".to_string()), - merge_after_generations: None, - ..Default::default() - }, - ) - .await - .unwrap(); + // Model a writer from the old release without running today's open- + // time migration, so the fixtures really contain the legacy schema. + let mut store = RolloutStore { + base: StorageBase::open( + &uri, + StorageBaseOptions { + storage_options: None, + shard_id: Some("pre-claim-check".to_string()), + merge_after_generations: None, + merge_max_generations: None, + merge_max_bytes: None, + session: None, + schema: legacy_schema.clone(), + key_column: "id".to_string(), + latest_schema: Some(Arc::new(rollout_schema())), + seal_on_put: false, + }, + false, + ) + .await + .unwrap(), + }; let base_batch = store .records_to_batch(&[assistant_record("legacy-base")]) diff --git a/crates/lance-context-core/src/storage.rs b/crates/lance-context-core/src/storage.rs index 17f4d355..e33ede18 100644 --- a/crates/lance-context-core/src/storage.rs +++ b/crates/lance-context-core/src/storage.rs @@ -3,6 +3,17 @@ use std::io; use std::path::Path; +/// Remove a dataset root through its URI's backend, without loading a manifest. +/// This also permits retrying a partially completed deletion. Callers must +/// stop local writers and quiesce remote writers before deleting a live store. +pub async fn remove_dataset(uri: &str) -> lance::Result<()> { + let (store, path) = lance::io::ObjectStore::from_uri(uri).await?; + match store.remove_dir_all(path).await { + Err(err) if crate::store_base::is_not_found_error(&err) => Ok(()), + result => result, + } +} + /// Maximum accepted length for a context or rollout store name. pub const MAX_STORE_NAME_LEN: usize = 128; diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index a1bef8f0..4d4e7ff0 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -237,6 +237,12 @@ pub struct ContextStore { /// one owner. Sharing the base is also what makes clones agree about the /// resident writer instead of each opening their own. base: StorageBase, + /// Keep uniqueness validation and the subsequent visible append atomic + /// with respect to other adds through this handle. Separate processes still + /// require application-level coordination for globally unique keys. + /// Trusted deferred-seal bulk writers opt out of read-back guarantees and + /// retain their concurrent append behavior. + add_lock: Option>, compaction_state: Arc>, pub compaction_config: CompactionConfig, blob_columns: HashSet, @@ -627,6 +633,7 @@ impl ContextStore { let mut store = Self { base, + add_lock: options.seal_on_add.then(|| Mutex::new(())), compaction_state: Arc::new(Mutex::new(CompactionState { background_task: None, is_compacting: false, @@ -681,6 +688,10 @@ impl ContextStore { return Ok(self.base.version()); } + let _guard = match &self.add_lock { + Some(lock) => Some(lock.lock().await), + None => None, + }; self.validate_unique_ids(entries).await?; self.write_entries(entries).await } @@ -2064,6 +2075,12 @@ impl ContextStore { self.base.close().await } + /// Stop this handle's writer and delete the dataset through its backend. + /// Remote writers must be quiesced by the caller before deletion. + pub async fn delete_data(&mut self) -> LanceResult<()> { + self.base.delete_data().await + } + /// Fold this instance's flushed MemWAL generations into the base table when /// it has accumulated [`ContextStoreOptions::merge_after_generations`] of /// them. Returns the number reclaimed; a no-op when the trigger is unset. diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index 3a622572..8433da4e 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -299,6 +299,8 @@ pub(crate) struct StorageBase { /// false negative from a stale manifest, but must never advance a handle /// whose caller deliberately selected a historical version. pinned_version: Option, + /// A deleted (or partially deleted) handle must never reopen a writer. + deleted: bool, /// 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 @@ -424,6 +426,7 @@ impl StorageBase { total_compactions: 0, last_compaction_error: None, pinned_version: None, + deleted: false, write_writer: tokio::sync::Mutex::new(None), }; // `ensure_mem_wal` may reload the dataset on a concurrent first-writer @@ -507,6 +510,11 @@ impl StorageBase { /// de-duplicated by the key column at read time, so a retried append can /// never double-count. pub async fn put(&self, batches: Vec) -> LanceResult<()> { + if self.deleted { + return Err(LanceError::invalid_input( + "store deletion has started; reopen after recreation", + )); + } if batches.is_empty() { return Ok(()); } @@ -663,6 +671,19 @@ impl StorageBase { Ok(()) } + /// Close this handle's writer, then remove the dataset using its configured + /// backend. Keep failed deletions retryable while refusing further appends. + pub async fn delete_data(&mut self) -> LanceResult<()> { + self.deleted = true; + self.close().await?; + let store = self.dataset.object_store(None).await?; + let path = self.dataset.branch_location().path.clone(); + match store.remove_dir_all(path).await { + Err(err) if is_not_found_error(&err) => Ok(()), + result => result, + } + } + /// Rows buffered in this instance's resident writer that have not yet been /// sealed into a flushed generation. /// @@ -670,6 +691,9 @@ impl StorageBase { /// no memtable, or a fenced writer all report `0` rather than erroring, so /// an observability read never fails because of writer state. pub async fn unflushed_rows(&self) -> i64 { + if self.is_version_pinned() { + return 0; + } let writer = { let guard = self.write_writer.lock().await; guard.as_ref().cloned() @@ -769,6 +793,10 @@ impl StorageBase { threshold: usize, seal_first: bool, ) -> LanceResult> { + // Periodic maintenance must not advance an explicitly selected snapshot. + if self.is_version_pinned() || self.deleted { + return Ok(None); + } if seal_first { // Materialize anything buffered so it is eligible for this pass. self.flush().await?; @@ -874,6 +902,10 @@ impl StorageBase { manifest: &ShardManifest, prepared: PreparedMerge, ) -> LanceResult { + // Checkout or deletion may have happened after shared-lock preparation. + if self.is_version_pinned() || self.deleted { + return Ok(false); + } let PreparedMerge { merged_generations, merged_paths, @@ -1089,13 +1121,28 @@ impl StorageBase { .cloned() .collect::>(); if !missing_fields.is_empty() { - self.dataset + if let Err(err) = self + .dataset .add_columns( NewColumnTransform::AllNulls(Arc::new(Schema::new(missing_fields))), None, None, ) - .await?; + .await + { + // Another opener may have committed the same additive upgrade. + // Accept that race only after validating the resulting schema. + self.refresh_latest().await?; + let schema: Arc = Arc::new(self.dataset.schema().into()); + if latest_schema + .fields() + .iter() + .any(|field| schema.field_with_name(field.name()).is_err()) + { + return Err(err); + } + align_batch_to_schema(RecordBatch::new_empty(schema), latest_schema)?; + } } Ok(()) } @@ -1324,6 +1371,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> { + // A checked-out version names a base-table snapshot, never the live WAL. + if self.is_version_pinned() { + return Ok(Vec::new()); + } 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?; @@ -1405,6 +1456,50 @@ impl StorageBase { // ------------------------------------------------------------ LSM reads + /// Read a fresh base/WAL view and retry if a merge invalidates it. A merge + /// commits base before draining/deleting its generations, so an unchanged + /// base version brackets a valid read. Never return a partial result after + /// a generation disappears. Ordinary WAL appends do not change base and do + /// not force retries. Pinned reads use only their selected base version. + pub async fn read_consistent( + &self, + source: ListSource, + mut read: F, + ) -> LanceResult + where + F: FnMut(Dataset, Vec) -> Fut, + Fut: std::future::Future>, + { + for _ in 0..3 { + let mut dataset = self.dataset.clone(); + if !self.is_version_pinned() { + dataset.checkout_latest().await?; + } + let version = dataset.manifest.version; + let snapshots = if source == ListSource::Fragments { + Vec::new() + } else { + self.wal_shard_snapshots().await? + }; + let result = read(dataset.clone(), snapshots).await; + if self.is_version_pinned() { + return result; + } + if let Err(err) = &result { + if !is_not_found_error(err) { + return result; + } + } + dataset.checkout_latest().await?; + if dataset.manifest.version == version { + return result; + } + } + Err(LanceError::io( + "dataset changed during read; retry the request", + )) + } + /// Build an LSM scanner over the base table unioned with every shard's /// flushed MemWAL generations, discovered from object storage. Because the /// snapshot is rebuilt from shard manifests on each call, one instance sees @@ -1427,23 +1522,32 @@ impl StorageBase { source: ListSource, shard_snapshots: Vec, ) -> LsmScanner { - let merge_key = vec![self.key_column.clone()]; + Self::lsm_scanner_for_dataset(&self.dataset, &self.key_column, source, shard_snapshots) + } + + pub fn lsm_scanner_for_dataset( + dataset: &Dataset, + key_column: &str, + source: ListSource, + shard_snapshots: Vec, + ) -> LsmScanner { + let merge_key = vec![key_column.to_string()]; match source { ListSource::Fragments => { - LsmScanner::new(Arc::new(self.dataset.clone()), Vec::new(), merge_key) + LsmScanner::new(Arc::new(dataset.clone()), Vec::new(), merge_key) } ListSource::All => { - LsmScanner::new(Arc::new(self.dataset.clone()), shard_snapshots, merge_key) + LsmScanner::new(Arc::new(dataset.clone()), 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()) } } } diff --git a/crates/lance-context-core/tests/storage_reliability.rs b/crates/lance-context-core/tests/storage_reliability.rs new file mode 100644 index 00000000..e91e6960 --- /dev/null +++ b/crates/lance-context-core/tests/storage_reliability.rs @@ -0,0 +1,280 @@ +use lance_context_api::AddRecordRequest; +use lance_context_core::{ + record_from_add_request, ContextStore, RolloutRecord, RolloutStore, RolloutStoreOptions, + ROLE_ASSISTANT, +}; + +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, + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn list_keeps_rows_after_peer_merge() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let mut writer = RolloutStore::open_with_options( + uri, + RolloutStoreOptions { + shard_id: Some("writer".into()), + ..Default::default() + }, + ) + .await + .unwrap(); + let reader = RolloutStore::open_with_options( + uri, + RolloutStoreOptions { + shard_id: Some("reader".into()), + ..Default::default() + }, + ) + .await + .unwrap(); + writer.add(&[rec("persisted")]).await.unwrap(); + writer.flush().await.unwrap(); + assert_eq!(reader.list(None, None).await.unwrap().len(), 1); + writer.cleanup_own_shard().await.unwrap(); + let page = reader + .list_filtered(&Default::default(), 10, 0) + .await + .unwrap(); + assert_eq!(page.records.len(), 1); + assert!(!page.has_more); + assert_eq!( + reader.list(None, None).await.unwrap().len(), + 1, + "a successful peer merge must not hide a previously visible row" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn checkout_excludes_future_wal() { + let dir = tempfile::tempdir().unwrap(); + let mut store = RolloutStore::open(dir.path().to_str().unwrap()) + .await + .unwrap(); + store.add(&[rec("old")]).await.unwrap(); + store.cleanup_own_shard().await.unwrap(); + let version = store.version(); + let mut future = rec("future"); + future.binary_payload = Some(vec![1, 2, 3]); + store.add(&[future]).await.unwrap(); + store.flush().await.unwrap(); + store.checkout(version).await.unwrap(); + assert!(store.get_by_id("future").await.unwrap().is_none()); + assert!(store.get_blob("future").await.unwrap().is_none()); + assert_eq!(store.observe().await.unwrap().row_count, 1); + assert_eq!(store.cleanup_own_shard().await.unwrap(), 0); + assert!(store.is_version_pinned()); + let ids: Vec<_> = store + .list(None, None) + .await + .unwrap() + .into_iter() + .map(|r| r.id) + .collect(); + assert_eq!( + ids, + vec!["old"], + "checkout must not include later WAL writes" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn concurrent_add_preserves_external_id_uniqueness() { + let dir = tempfile::tempdir().unwrap(); + let store = ContextStore::open(dir.path().to_str().unwrap()) + .await + .unwrap(); + let request: AddRecordRequest = serde_json::from_value(serde_json::json!({ + "text_payload":"hello", "external_id":"same-external-id" + })) + .unwrap(); + let left = [record_from_add_request(&request, "left".into(), "r".into())]; + let right = [record_from_add_request( + &request, + "right".into(), + "r".into(), + )]; + let (a, b) = tokio::join!(store.add(&left), store.add(&right)); + assert_eq!(store.list(None, None).await.unwrap().len(), 1); + assert_eq!( + usize::from(a.is_ok()) + usize::from(b.is_ok()), + 1, + "only one concurrent insertion of the same external_id may succeed" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn context_checkout_excludes_future_wal() { + let dir = tempfile::tempdir().unwrap(); + let mut store = ContextStore::open(dir.path().to_str().unwrap()) + .await + .unwrap(); + let request: AddRecordRequest = + serde_json::from_value(serde_json::json!({"text_payload":"hello"})).unwrap(); + let old = record_from_add_request(&request, "old".into(), "r".into()); + store.add(&[old]).await.unwrap(); + store.cleanup_wal().await.unwrap(); + let version = store.version(); + let future = record_from_add_request(&request, "future".into(), "r".into()); + store.add(&[future]).await.unwrap(); + store.checkout(version).await.unwrap(); + let ids: Vec<_> = store + .list(None, None) + .await + .unwrap() + .into_iter() + .map(|r| r.id) + .collect(); + assert_eq!( + ids, + vec!["old"], + "a pinned context snapshot must exclude subsequent writes" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn write_to_legacy_rollout_preserves_new_fields() { + use arrow_array::{RecordBatch, RecordBatchIterator}; + use arrow_schema::{ArrowError, Schema}; + use lance::Dataset; + use std::sync::Arc; + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let current = lance_context_core::rollout_schema(); + let legacy = Arc::new(Schema::new( + current + .fields() + .iter() + .filter(|f| { + ![ + "model_input_string", + "model_output_string", + "rationale", + "problem_text", + "user_metadata", + ] + .contains(&f.name().as_str()) + }) + .cloned() + .collect::>(), + )); + let empty = RecordBatch::new_empty(legacy.clone()); + let reader = RecordBatchIterator::new(vec![Ok::<_, ArrowError>(empty)].into_iter(), legacy); + Dataset::write(reader, uri, None).await.unwrap(); + let (left, right) = tokio::join!(RolloutStore::open(uri), RolloutStore::open(uri)); + let mut store = left.unwrap(); + let _peer = right.unwrap(); + let mut record = rec("new-format"); + record.model_input_string = Some("must survive".into()); + record.model_output_string = Some("output".into()); + record.rationale = Some("reason".into()); + record.problem_text = Some("question".into()); + record.user_metadata = Some("metadata".into()); + store + .add(&[record]) + .await + .expect("new writer should support an older additive schema"); + store.flush().await.unwrap(); + store.cleanup_own_shard().await.unwrap(); + let got = store.get_by_id("new-format").await.unwrap().unwrap(); + assert_eq!(got.model_input_string.as_deref(), Some("must survive")); + assert_eq!(got.model_output_string.as_deref(), Some("output")); + assert_eq!(got.rationale.as_deref(), Some("reason")); + assert_eq!(got.problem_text.as_deref(), Some("question")); + assert_eq!(got.user_metadata.as_deref(), Some("metadata")); + store.close().await.unwrap(); + let reopened = RolloutStore::open(uri).await.unwrap(); + assert_eq!( + reopened + .get_by_id("new-format") + .await + .unwrap() + .unwrap() + .model_input_string, + got.model_input_string + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn concurrent_add_preserves_internal_id_uniqueness() { + let dir = tempfile::tempdir().unwrap(); + let store = ContextStore::open(dir.path().to_str().unwrap()) + .await + .unwrap(); + let request: AddRecordRequest = + serde_json::from_value(serde_json::json!({"text_payload":"hello"})).unwrap(); + let records = [record_from_add_request( + &request, + "same-id".into(), + "r".into(), + )]; + let (left, right) = tokio::join!(store.add(&records), store.add(&records)); + assert_eq!(usize::from(left.is_ok()) + usize::from(right.is_ok()), 1); + assert_eq!(store.list(None, None).await.unwrap().len(), 1); +} + +#[tokio::test] +async fn checkout_cancels_an_already_prepared_merge() { + let dir = tempfile::tempdir().unwrap(); + let mut store = RolloutStore::open(dir.path().to_str().unwrap()) + .await + .unwrap(); + store.add(&[rec("old")]).await.unwrap(); + store.cleanup_own_shard().await.unwrap(); + let version = store.version(); + store.add(&[rec("future")]).await.unwrap(); + let (manifest_store, manifest, prepared) = + store.prepare_cleanup_merge().await.unwrap().unwrap(); + store.checkout(version).await.unwrap(); + assert_eq!( + store + .commit_prepared_merge(&manifest_store, &manifest, prepared) + .await + .unwrap(), + 0 + ); + assert_eq!(store.version(), version); + assert!(store.is_version_pinned()); + assert!(store.get_by_id("future").await.unwrap().is_none()); + store.refresh_latest().await.unwrap(); + assert_eq!(store.pending_wal_generations().await.unwrap(), 1); +} diff --git a/crates/lance-context-server/src/routes/contexts.rs b/crates/lance-context-server/src/routes/contexts.rs index d1dd1ef5..cb3529f5 100644 --- a/crates/lance-context-server/src/routes/contexts.rs +++ b/crates/lance-context-server/src/routes/contexts.rs @@ -109,18 +109,14 @@ pub async fn delete_context( Path(name): Path, ) -> Result { AppState::validate_name(&name)?; - let mut stores = state.stores.write().await; - if stores.remove(&name).is_none() { - return Err(AppError::NotFound(format!( - "Context '{}' does not exist", - name - ))); - } - - let uri = state.context_uri(&name); - if let Err(e) = tokio::fs::remove_dir_all(&uri).await { - tracing::warn!("Failed to remove context data at {}: {}", uri, e); - } + let store = state.get_or_open_context_store(&name).await?; + store + .write() + .await + .delete_data() + .await + .map_err(AppError::from_lance)?; + state.stores.write().await.remove(&name); Ok(axum::http::StatusCode::NO_CONTENT) } diff --git a/crates/lance-context-server/src/routes/datagen.rs b/crates/lance-context-server/src/routes/datagen.rs index f8bb19e9..d6dbd0be 100644 --- a/crates/lance-context-server/src/routes/datagen.rs +++ b/crates/lance-context-server/src/routes/datagen.rs @@ -111,10 +111,6 @@ pub async fn delete_datagen_store( name ))); } - let uri = state.datagen_uri(&name); - if let Err(e) = tokio::fs::remove_dir_all(&uri).await { - tracing::warn!("Failed to remove datagen data at {}: {}", uri, e); - } Ok(StatusCode::NO_CONTENT) } diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index c36583b0..bb90c094 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -60,15 +60,12 @@ fn blob_stream_body( Body::from_stream(stream) } -/// Parse the `Content-Length` header into a byte count, defaulting to `0` when -/// absent or unparsable (a chunked upload without a declared length reserves -/// nothing up-front; the body-size limit still caps it). -fn content_length(headers: &header::HeaderMap) -> usize { +/// Unknown lengths need exclusive admission before any body buffering. +fn content_length(headers: &header::HeaderMap) -> Option { headers .get(header::CONTENT_LENGTH) .and_then(|value| value.to_str().ok()) .and_then(|value| value.parse::().ok()) - .unwrap_or(0) } /// Whether the request asked for a synchronous flush via `?flush=true`. @@ -133,17 +130,37 @@ fn observe_blob_sizes(state: &AppState, experiment: &str, records: &[AddRolloutR /// budget is configured (`None`) this is a no-op that always admits. fn acquire_blob_budget( state: &AppState, - bytes: usize, + bytes: Option, ) -> Result, AppError> { match &state.blob_budget { None => Ok(None), - Some(budget) => budget.try_acquire(bytes).map(Some).ok_or_else(|| { - metrics::counter!("rollout_blob_budget_rejections_total").increment(1); - AppError::Overloaded( - "server is at its in-flight blob memory limit; retry shortly".to_string(), + Some(budget) => bytes + .map_or_else( + || budget.try_acquire_unknown(), + |bytes| budget.try_acquire(bytes), ) - }), + .map(Some) + .ok_or_else(|| { + metrics::counter!("rollout_blob_budget_rejections_total").increment(1); + AppError::Overloaded( + "server is at its in-flight blob memory limit; retry shortly".to_string(), + ) + }), + } +} + +/// Reserve before polling a payload read. Stored payload_size is optional and +/// caller supplied, so it cannot be used to admit an allocation safely. +async fn load_blob_with_budget( + state: &AppState, + load: impl std::future::Future>, AppError>>, +) -> Result<(Option>, Option), AppError> { + let mut reservation = acquire_blob_budget(state, None)?; + let bytes = load.await?; + if let Some(reservation) = reservation.as_mut() { + reservation.shrink_to(bytes.as_ref().map_or(0, Vec::len)); } + Ok((bytes, reservation)) } /// Response for the internal WAL-merge trigger: how many flushed generations @@ -257,11 +274,6 @@ pub async fn delete_rollout_store( ))); } - let uri = state.rollout_uri(&name); - if let Err(e) = tokio::fs::remove_dir_all(&uri).await { - tracing::warn!("Failed to remove rollout data at {}: {}", uri, e); - } - Ok(StatusCode::NO_CONTENT) } @@ -307,8 +319,8 @@ pub async fn add_rollouts( let flush = flush_requested(req.uri().query()); // Admit against the in-flight blob budget before buffering the body, using - // the declared Content-Length as the reservation size. Held for the whole - // handler so concurrent uploads cannot collectively exceed the budget and + // the declared Content-Length, or exclusive admission for an unknown size. + // Held for the whole handler so concurrent uploads cannot collectively exceed the budget and // OOM the worker; dropped when the request completes. let _budget = acquire_blob_budget(&state, content_length(req.headers()))?; @@ -592,15 +604,13 @@ pub async fn fetch_rollout_blob( ) -> Result { let store_lock = state.get_or_open_rollout_store(&name).await?; - let bytes = get_rollout_blob_refreshing_on_miss(&store_lock, &id) - .await? - .ok_or_else(|| AppError::NotFound(format!("Rollout '{}' has no payload", id)))?; - - // Reserve now that the payload size is known, and hold the reservation for - // the whole streamed send (moved into the body): a slow client keeps the - // blob resident until the last frame flushes, so the budget must account for - // it until then. Reject with 503 if the budget is currently exhausted. - let reservation = acquire_blob_budget(&state, bytes.len())?; + let (bytes, reservation) = load_blob_with_budget( + &state, + get_rollout_blob_refreshing_on_miss(&store_lock, &id), + ) + .await?; + let bytes = + bytes.ok_or_else(|| AppError::NotFound(format!("Rollout '{}' has no payload", id)))?; let len = bytes.len(); Response::builder() @@ -826,11 +836,11 @@ mod tests { } #[test] - fn content_length_parses_or_defaults_zero() { + fn content_length_distinguishes_unknown_from_empty() { let mut headers = header::HeaderMap::new(); - assert_eq!(content_length(&headers), 0); + assert_eq!(content_length(&headers), None); headers.insert(header::CONTENT_LENGTH, header::HeaderValue::from(4096)); - assert_eq!(content_length(&headers), 4096); + assert_eq!(content_length(&headers), Some(4096)); } /// The threshold is inclusive and `0` disables the log entirely, so an @@ -895,6 +905,180 @@ mod tests { assert_eq!(observed, vec![512.0, 2048.0]); } + #[tokio::test] + async fn chunked_upload_respects_exhausted_budget() { + let dir = TempDir::new().unwrap(); + let mut state = AppState::new_for_test(dir.path().to_path_buf()).await; + let budget = crate::state::BlobBudget::new(1024); + state.blob_budget = Some(budget.clone()); + let state = Arc::new(state); + let _ = create_rollout_store( + State(state.clone()), + Json(CreateRolloutStoreRequest { + name: "rl".to_string(), + storage_options: None, + }), + ) + .await + .unwrap(); + let _occupied = budget.try_acquire(1024).unwrap(); + assert!(budget.try_acquire(1).is_none()); + let body = serde_json::json!({"records":[{ + "id":"chunked-blob", "rollout_id":"r", "problem_id":"p", + "sequence_order":0, "role":"assistant", "content_type":"application/octet-stream", + "binary_payload":"eA==" + }]}) + .to_string(); + let known_length = Request::builder() + .method("POST") + .uri("/rollouts/rl/records?flush=true") + .header(header::CONTENT_TYPE, "application/json") + .header(header::CONTENT_LENGTH, body.len()) + .body(Body::from(body.clone())) + .unwrap(); + assert!(matches!( + add_rollouts(State(state.clone()), Path("rl".to_string()), known_length).await, + Err(AppError::Overloaded(_)) + )); + let request = Request::builder() + .method("POST") + .uri("/rollouts/rl/records?flush=true") + .header(header::CONTENT_TYPE, "application/json") + .header(header::TRANSFER_ENCODING, "chunked") + .body(Body::from(body)) + .unwrap(); + assert!(request.headers().get(header::CONTENT_LENGTH).is_none()); + let result = add_rollouts(State(state), Path("rl".to_string()), request).await; + assert!( + matches!(result, Err(AppError::Overloaded(_))), + "a blob upload must not bypass an exhausted global budget" + ); + } + + #[tokio::test] + async fn delete_uri_store_does_not_resurrect_rows() { + let dir = TempDir::new().unwrap(); + let base_uri = std::path::PathBuf::from(format!("file://{}", dir.path().display())); + let state = Arc::new(AppState::new_for_test(base_uri).await); + let _ = create_rollout_store( + State(state.clone()), + Json(CreateRolloutStoreRequest { + name: "rl".to_string(), + storage_options: None, + }), + ) + .await + .unwrap(); + let _ = add_rollouts( + State(state.clone()), + Path("rl".to_string()), + append_request("?flush=true", "deleted-row"), + ) + .await + .unwrap(); + let cached = state.get_or_open_rollout_store("rl").await.unwrap(); + cached.write().await.cleanup_own_shard().await.unwrap(); + assert_eq!( + delete_rollout_store(State(state.clone()), Path("rl".to_string())) + .await + .unwrap(), + StatusCode::NO_CONTENT + ); + let _ = create_rollout_store( + State(state.clone()), + Json(CreateRolloutStoreRequest { + name: "rl".to_string(), + storage_options: None, + }), + ) + .await + .unwrap(); + let reopened = state.get_or_open_rollout_store("rl").await.unwrap(); + let rows = reopened.read().await.list(None, None).await.unwrap(); + let ids: Vec<_> = rows.into_iter().map(|r| r.id).collect(); + assert!( + ids.is_empty(), + "deleted rows reappeared after recreation: {ids:?}" + ); + } + + #[tokio::test] + async fn blob_budget_admits_before_loading_and_releases_on_cancellation() { + let dir = TempDir::new().unwrap(); + let mut state = AppState::new_for_test(dir.path().to_path_buf()).await; + let budget = crate::state::BlobBudget::new(1024); + state.blob_budget = Some(budget.clone()); + let occupied = budget.try_acquire(1024).unwrap(); + let loaded = std::cell::Cell::new(false); + let result = load_blob_with_budget(&state, async { + loaded.set(true); + Ok(Some(vec![0; 256])) + }) + .await; + assert!(matches!(result, Err(AppError::Overloaded(_)))); + assert!( + !loaded.get(), + "admission must happen before payload allocation" + ); + drop(occupied); + + let (bytes, guard) = load_blob_with_budget(&state, async { Ok(Some(vec![0; 256])) }) + .await + .unwrap(); + assert_eq!(bytes.unwrap().len(), 256); + let remaining = budget.try_acquire(768).unwrap(); + assert!(budget.try_acquire_unknown().is_none()); + drop(remaining); + drop(guard); + + let mut pending = Box::pin(load_blob_with_budget(&state, std::future::pending())); + assert!(futures::poll!(pending.as_mut()).is_pending()); + assert!(budget.try_acquire(1).is_none()); + drop(pending); + assert!(budget.try_acquire(1024).is_some()); + } + + #[tokio::test] + async fn failed_deletion_retains_registry_and_can_be_retried() { + let (state, dir) = rollout_state().await; + let cached = state.get_or_open_rollout_store("rl").await.unwrap(); + cached.write().await.close().await.unwrap(); + let path = std::path::PathBuf::from(state.rollout_uri("rl")); + let saved = dir.path().join("saved-dataset"); + std::fs::rename(&path, &saved).unwrap(); + std::fs::write(&path, b"not a directory").unwrap(); + assert!( + delete_rollout_store(State(state.clone()), Path("rl".to_string())) + .await + .is_err() + ); + assert!(state + .rollout_registry + .write() + .await + .contains("rl") + .await + .unwrap()); + std::fs::remove_file(&path).unwrap(); + std::fs::rename(&saved, &path).unwrap(); + assert_eq!( + delete_rollout_store(State(state.clone()), Path("rl".to_string())) + .await + .unwrap(), + StatusCode::NO_CONTENT + ); + assert!(!path.exists()); + assert!(!state + .rollout_registry + .write() + .await + .contains("rl") + .await + .unwrap()); + let row = rollout_record_from_add_request(&record_with_size("late-write", None)); + assert!(cached.read().await.add(&[row]).await.is_err()); + } + async fn rollout_state() -> (Arc, TempDir) { let dir = TempDir::new().unwrap(); let state = Arc::new(AppState::new_for_test(dir.path().to_path_buf()).await); diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index e0a6e271..1c5d7e53 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -158,13 +158,14 @@ impl BlobBudget { loop { // Admit when it fits, or when the instance is otherwise idle (so a // single request bigger than the whole budget can still proceed). - let fits = current + bytes <= self.limit; + let next = current.checked_add(bytes)?; + let fits = next <= self.limit; if !fits && current != 0 { return None; } match self.used.compare_exchange_weak( current, - current + bytes, + next, Ordering::AcqRel, Ordering::Acquire, ) { @@ -178,6 +179,25 @@ impl BlobBudget { } } } + + /// Admit one unknown-size materialization only while the budget is idle. + /// Missing or client-supplied payload sizes cannot safely predict the + /// allocation. Holding the full budget excludes other blob allocations. + pub fn try_acquire_unknown(self: &Arc) -> Option { + self.try_acquire(self.limit.max(1)) + } +} + +impl BlobReservation { + /// After an exclusive load, return unused capacity for concurrent responses. + /// An oversized blob keeps the full reservation until it is released. + pub fn shrink_to(&mut self, bytes: usize) { + let retained = bytes.min(self.bytes); + self.budget + .used + .fetch_sub(self.bytes - retained, std::sync::atomic::Ordering::AcqRel); + self.bytes = retained; + } } impl Drop for BlobReservation { @@ -363,22 +383,30 @@ impl AppState { /// handle. Returns whether the store existed. pub async fn unregister_rollout(&self, name: &str) -> Result { Self::validate_name(name)?; - let existed = self - .rollout_registry - .write() - .await + // Keep the registry entry and cached handle until physical deletion + // succeeds, so failures can be retried without losing the store's URI. + let mut registry = self.rollout_registry.write().await; + if !registry .contains(name) .await - .map_err(AppError::from_lance)?; - if !existed { + .map_err(AppError::from_lance)? + { return Ok(false); } - self.rollout_registry - .write() - .await - .remove(name) - .await - .map_err(AppError::from_lance)?; + let cached = self.rollout_stores.lock().await.peek(name).cloned(); + if let Some(store) = cached { + store + .write() + .await + .delete_data() + .await + .map_err(AppError::from_lance)?; + } else { + lance_context_core::remove_dataset(&self.rollout_uri(name)) + .await + .map_err(AppError::from_lance)?; + } + registry.remove(name).await.map_err(AppError::from_lance)?; self.rollout_stores.lock().await.pop(name); Ok(true) } @@ -426,6 +454,18 @@ impl AppState { .await .map_err(AppError::from_lance)?; let opened = Arc::new(RwLock::new(opened)); + // Recheck after the slow open while excluding a concurrent deletion + // through registry removal and cache insertion. + let mut registry = self.rollout_registry.write().await; + if !registry + .contains(name) + .await + .map_err(AppError::from_lance)? + { + return Err(AppError::NotFound(format!( + "Store '{name}' was deleted while opening" + ))); + } // Insert under the lock, re-checking for a store another request may // have opened concurrently while we were loading. @@ -513,22 +553,30 @@ impl AppState { /// handle. Returns whether the store existed. pub async fn unregister_datagen(&self, name: &str) -> Result { Self::validate_name(name)?; - let existed = self - .datagen_registry - .write() - .await + // Keep the registry entry and cached handle until physical deletion + // succeeds, so failures can be retried without losing the store's URI. + let mut registry = self.datagen_registry.write().await; + if !registry .contains(name) .await - .map_err(AppError::from_lance)?; - if !existed { + .map_err(AppError::from_lance)? + { return Ok(false); } - self.datagen_registry - .write() - .await - .remove(name) - .await - .map_err(AppError::from_lance)?; + let cached = self.datagen_stores.lock().await.peek(name).cloned(); + if let Some(store) = cached { + store + .write() + .await + .delete_data() + .await + .map_err(AppError::from_lance)?; + } else { + lance_context_core::remove_dataset(&self.datagen_uri(name)) + .await + .map_err(AppError::from_lance)?; + } + registry.remove(name).await.map_err(AppError::from_lance)?; self.datagen_stores.lock().await.pop(name); Ok(true) } @@ -565,6 +613,18 @@ impl AppState { .await .map_err(AppError::from_lance)?; let opened = Arc::new(RwLock::new(opened)); + // Recheck after the slow open while excluding a concurrent deletion + // through registry removal and cache insertion. + let mut registry = self.datagen_registry.write().await; + if !registry + .contains(name) + .await + .map_err(AppError::from_lance)? + { + return Err(AppError::NotFound(format!( + "Store '{name}' was deleted while opening" + ))); + } let mut cache = self.datagen_stores.lock().await; if let Some(existing) = cache.get(name) { @@ -619,22 +679,30 @@ impl AppState { /// Returns whether the store existed. pub async fn unregister_generic(&self, name: &str) -> Result { Self::validate_name(name)?; - let existed = self - .generic_registry - .write() - .await + // Keep the registry entry and cached handle until physical deletion + // succeeds, so failures can be retried without losing the store's URI. + let mut registry = self.generic_registry.write().await; + if !registry .contains(name) .await - .map_err(AppError::from_lance)?; - if !existed { + .map_err(AppError::from_lance)? + { return Ok(false); } - self.generic_registry - .write() - .await - .remove(name) - .await - .map_err(AppError::from_lance)?; + let cached = self.generic_stores.lock().await.peek(name).cloned(); + if let Some(store) = cached { + store + .write() + .await + .delete_data() + .await + .map_err(AppError::from_lance)?; + } else { + lance_context_core::remove_dataset(&self.generic_uri(name)) + .await + .map_err(AppError::from_lance)?; + } + registry.remove(name).await.map_err(AppError::from_lance)?; self.generic_stores.lock().await.pop(name); Ok(true) } @@ -674,6 +742,18 @@ impl AppState { .await .map_err(AppError::from_lance)?; let opened = Arc::new(RwLock::new(opened)); + // Recheck after the slow open while excluding a concurrent deletion + // through registry removal and cache insertion. + let mut registry = self.generic_registry.write().await; + if !registry + .contains(name) + .await + .map_err(AppError::from_lance)? + { + return Err(AppError::NotFound(format!( + "Store '{name}' was deleted while opening" + ))); + } let mut cache = self.generic_stores.lock().await; if let Some(existing) = cache.get(name) { @@ -853,6 +933,16 @@ mod tests { use super::*; use tempfile::TempDir; + #[test] + fn blob_budget_rejects_counter_overflow() { + let budget = BlobBudget::new(1024); + let oversized = budget.try_acquire(usize::MAX).unwrap(); + assert!(budget.try_acquire(1).is_none()); + assert!(budget.try_acquire_unknown().is_none()); + drop(oversized); + assert!(budget.try_acquire_unknown().is_some()); + } + #[test] fn blob_budget_admits_until_full_then_releases_on_drop() { let budget = BlobBudget::new(1000); diff --git a/docs/src/specs/rollout-deployment.md b/docs/src/specs/rollout-deployment.md index ae9533ed..dbbd101e 100644 --- a/docs/src/specs/rollout-deployment.md +++ b/docs/src/specs/rollout-deployment.md @@ -179,6 +179,8 @@ The caps form an **OR**: whichever binds first ends the pass. Bytes are measured The manifest reclaims whole generations, so a generation is the smallest indivisible merge unit. A pass always finishes the generation that reaches the budget, even if its first generation alone is oversized. Thus the read buffer can exceed the budget by up to one generation; this is not a hard process-RSS limit, and scan, deduplication and commit allocations need additional headroom. +**Blob request admission.** `ROLLOUT_MAX_INFLIGHT_BLOB_BYTES` (default `0`, disabled) separately controls concurrent HTTP blob requests. Known-length uploads reserve their declared body size. Unknown-length uploads and blob downloads reserve the whole budget before buffering or loading; they are admitted only when it is idle. Download sizes cannot be trusted from optional, caller-supplied `payload_size` metadata. After loading, a smaller download releases unused capacity for known-length uploads and holds the remaining reservation through the response lifetime. A lone oversized request is still admitted. This conservative policy can reduce concurrency when the budget is enabled. + Only merged generation ids and directories are removed. Leftovers stay pending and drain on subsequent passes. The caps apply independently of the count-trigger threshold (`--rollout-merge-after-generations`), including when time-triggered cleanup merges at a threshold of one generation. Server-managed rollout, datagen and generic stores share these settings. This is the "external compactor" path that Lance's MemWAL LSM design explicitly anticipates. Two properties make it safe under the §2 deployment model: @@ -216,7 +218,7 @@ rows = await store.list(filters={ }) ``` -`checkout` remains available for base-table time travel, but is not the mechanism for per-checkpoint rollout reproducibility. (See schema-design §3 and §7 — and note that `learner_iteration` is likewise a column, not a dataset version.) +`checkout` selects a fixed base-table version and excludes live WAL generations, including unmerged writes made before checkout. Periodic WAL merging does not advance a pinned handle. It is not the mechanism for per-checkpoint rollout reproducibility. (See schema-design §3 and §7 — and note that `learner_iteration` is likewise a column, not a dataset version.) --- From b7d166940e079c401d419ff65be959fbd522d0b1 Mon Sep 17 00:00:00 2001 From: Beinan Date: Thu, 10 Sep 2026 23:04:04 +0000 Subject: [PATCH 2/2] fix: fence snapshot and evicted handles and retry WAL cleanup reads --- .../lance-context-core/src/rollout_store.rs | 122 ++++++- crates/lance-context-core/src/store.rs | 14 +- crates/lance-context-core/src/store_base.rs | 39 ++- .../tests/storage_reliability.rs | 106 ++++++ .../src/routes/datagen.rs | 5 +- .../src/routes/generic.rs | 5 +- .../src/routes/rollouts.rs | 50 ++- crates/lance-context-server/src/state.rs | 303 ++++++++++++++++-- docs/src/specs/rollout-deployment.md | 2 +- 9 files changed, 592 insertions(+), 54 deletions(-) diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 11cda7b6..cf5e6202 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -513,7 +513,8 @@ impl RolloutStore { } /// Check out an immutable base-table version. Live WAL generations are - /// excluded; unmerged appends are not part of a base version. + /// excluded; unmerged appends are not part of a base version. Writes are + /// rejected until `refresh_latest` is called. pub async fn checkout(&mut self, version_id: u64) -> LanceResult<()> { self.base.checkout(version_id).await } @@ -581,6 +582,7 @@ impl RolloutStore { /// visibility; call [`Self::flush`] instead. Retained only for API /// compatibility — see the module docs on reproducibility. pub async fn add(&self, records: &[RolloutRecord]) -> LanceResult { + self.base.ensure_writable()?; if records.is_empty() { return Ok(self.base.version()); } @@ -2199,6 +2201,124 @@ mod tests { assert_eq!(rows[0].id, "merged-during-read"); } + #[tokio::test] + async fn review_list_retries_when_merge_committed_before_read() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let mut writer = RolloutStore::open(uri).await.unwrap(); + let record = assistant_record("already-committed"); + writer.add(std::slice::from_ref(&record)).await.unwrap(); + writer.flush().await.unwrap(); + writer.close().await.unwrap(); + // Pause a merge after its base commit, before manifest drain and GC. + let batch = writer.records_to_batch(&[record]).unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema); + writer.base.dataset.append(batches, None).await.unwrap(); + let reader = RolloutStore::open(uri).await.unwrap(); + let attempts = AtomicUsize::new(0); + let result = reader + .base + .read_consistent(ListSource::All, |dataset, snapshots| { + let reader = &reader; + let attempts = &attempts; + async move { + if attempts.fetch_add(1, Ordering::SeqCst) == 0 { + let snapshot = &snapshots[0]; + let store = dataset.object_store(None).await.unwrap(); + let path = dataset.branch_location().path.clone(); + let manifests = ShardManifestStore::new( + store.clone(), + &path, + snapshot.shard_id, + DEFAULT_MANIFEST_SCAN_BATCH_SIZE, + ); + let current = manifests.read_latest().await.unwrap().unwrap(); + manifests + .commit_update(current.writer_epoch, |m| ShardManifest { + version: m.version + 1, + flushed_generations: vec![], + ..m.clone() + }) + .await + .unwrap(); + for generation in &snapshot.flushed_generations { + store + .remove_dir_all( + path.clone() + .join("_mem_wal") + .join(snapshot.shard_id.to_string().as_str()) + .join(generation.path.as_str()), + ) + .await + .unwrap(); + } + } + reader.list_non_blob_snapshot(dataset, snapshots).await + } + }) + .await; + assert!( + result.is_ok(), + "read failed instead of retrying: {result:?}; attempts={}", + attempts.load(Ordering::SeqCst) + ); + assert_eq!(result.unwrap().len(), 1); + assert_eq!(attempts.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn consistent_read_bounds_missing_generation_retries() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let dir = tempfile::tempdir().unwrap(); + let mut store = RolloutStore::open(dir.path().to_str().unwrap()) + .await + .unwrap(); + let attempts = AtomicUsize::new(0); + let missing = |_, _| { + attempts.fetch_add(1, Ordering::SeqCst); + async { + Err::<(), _>(LanceError::dataset_not_found( + "missing-generation", + "missing".into(), + )) + } + }; + let err = store + .base + .read_consistent(ListSource::All, missing) + .await + .unwrap_err(); + assert!(matches!(err, LanceError::DatasetNotFound { .. })); + assert_eq!(attempts.load(Ordering::SeqCst), 3); + + attempts.store(0, Ordering::SeqCst); + store.checkout(store.version()).await.unwrap(); + assert!(store + .base + .read_consistent(ListSource::All, missing) + .await + .is_err()); + assert_eq!( + attempts.load(Ordering::SeqCst), + 1, + "pinned reads must not retry against newer snapshots" + ); + + store.refresh_latest().await.unwrap(); + attempts.store(0, Ordering::SeqCst); + let result = store + .base + .read_consistent(ListSource::All, |_, _| { + attempts.fetch_add(1, Ordering::SeqCst); + async { Err::<(), _>(LanceError::invalid_input("bad filter")) } + }) + .await; + assert!(matches!(result, Err(LanceError::InvalidInput { .. }))); + assert_eq!(attempts.load(Ordering::SeqCst), 1); + } + #[test] fn rollout_filters_parse_supported_fields() { let filters = RolloutFilters::from_json_value(json!({ diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 4d4e7ff0..ca6c313f 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -684,6 +684,7 @@ impl ContextStore { /// Append context records to the store and return the new dataset version. pub async fn add(&self, entries: &[ContextRecord]) -> LanceResult { + self.base.ensure_writable()?; if entries.is_empty() { return Ok(self.base.version()); } @@ -749,6 +750,7 @@ impl ContextStore { /// then `add` a record whose `payload_uri` points at `uri`. Inline /// [`ContextRecord::binary_payload`] remains the small-payload path. pub async fn put_payload(&self, uri: &str, bytes: &[u8]) -> LanceResult { + self.base.ensure_writable()?; let registry = Arc::new(ObjectStoreRegistry::default()); let (store, path) = ObjectStore::from_uri_and_params(registry, uri, &self.payload_store_params()).await?; @@ -773,6 +775,7 @@ impl ContextStore { /// This writes a tombstone with the same primary key, preserving prior /// dataset versions while hiding the record from default reads. pub async fn delete_by_id(&mut self, id: &str) -> LanceResult { + self.base.ensure_writable()?; let Some(record) = self.get_by_id(id).await? else { return Ok(false); }; @@ -782,6 +785,7 @@ impl ContextStore { /// Logically forget a record by caller-supplied external id. pub async fn delete_by_external_id(&mut self, external_id: &str) -> LanceResult { + self.base.ensure_writable()?; let Some(record) = self.get_by_external_id(external_id).await? else { return Ok(false); }; @@ -800,6 +804,7 @@ impl ContextStore { &mut self, mut record: ContextRecord, ) -> LanceResult { + self.base.ensure_writable()?; let Some(external_id) = record.external_id.clone() else { return Err(ArrowError::InvalidArgumentError( "upsert_by_external_id requires external_id".to_string(), @@ -888,6 +893,7 @@ impl ContextStore { &mut self, mut records: Vec, ) -> LanceResult> { + self.base.ensure_writable()?; if records.is_empty() { return Ok(Vec::new()); } @@ -1077,6 +1083,7 @@ impl ContextStore { id: &str, patch: RecordPatch, ) -> LanceResult> { + self.base.ensure_writable()?; if id.is_empty() { return Err(ArrowError::InvalidArgumentError( "update_by_id requires a non-empty id".to_string(), @@ -1100,6 +1107,7 @@ impl ContextStore { external_id: &str, patch: RecordPatch, ) -> LanceResult> { + self.base.ensure_writable()?; if external_id.is_empty() { return Err(ArrowError::InvalidArgumentError( "update_by_external_id requires a non-empty external_id".to_string(), @@ -1405,6 +1413,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 { + self.base.ensure_writable()?; if self.has_relationships_column() { return Ok(false); } @@ -1414,11 +1423,10 @@ impl ContextStore { .dataset .add_columns(NewColumnTransform::AllNulls(schema), None, None) .await?; - self.base.clear_version_pin(); Ok(true) } - /// Checkout a specific dataset version. + /// Check out a read-only dataset version. Call `refresh_latest` to resume writes. pub async fn checkout(&mut self, version_id: u64) -> LanceResult<()> { self.base.checkout(version_id).await } @@ -1991,6 +1999,7 @@ impl ContextStore { &mut self, options: Option, ) -> LanceResult { + self.base.ensure_writable()?; let config = options.unwrap_or_else(|| self.compaction_config.clone()); info!( @@ -2161,6 +2170,7 @@ impl ContextStore { /// Create (or replace) the scalar index on the `id` column. pub async fn create_id_index(&mut self) -> LanceResult<()> { + self.base.ensure_writable()?; 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 8433da4e..0e75933c 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -447,7 +447,7 @@ impl StorageBase { self.dataset.manifest.version } - /// Check out a specific base dataset version (time travel). + /// Check out a read-only base dataset version. Call `refresh_latest` to resume writes. pub async fn checkout(&mut self, version_id: u64) -> LanceResult<()> { self.dataset = self.dataset.checkout_version(version_id).await?; self.pinned_version = Some(version_id); @@ -472,10 +472,19 @@ impl StorageBase { 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; + /// Reject mutations on historical snapshots and on handles being deleted. + pub(crate) fn ensure_writable(&self) -> LanceResult<()> { + if self.deleted { + return Err(LanceError::invalid_input( + "store deletion has started; reopen after recreation", + )); + } + if self.is_version_pinned() { + return Err(LanceError::invalid_input( + "checked-out store is read-only; call refresh_latest before writing", + )); + } + Ok(()) } // ---------------------------------------------------------------- writes @@ -510,11 +519,7 @@ impl StorageBase { /// de-duplicated by the key column at read time, so a retried append can /// never double-count. pub async fn put(&self, batches: Vec) -> LanceResult<()> { - if self.deleted { - return Err(LanceError::invalid_input( - "store deletion has started; reopen after recreation", - )); - } + self.ensure_writable()?; if batches.is_empty() { return Ok(()); } @@ -1103,6 +1108,7 @@ impl StorageBase { /// 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<()> { + self.ensure_writable()?; let Some(latest_schema) = self.latest_schema.clone() else { return Ok(()); }; @@ -1167,6 +1173,7 @@ impl StorageBase { &mut self, options: Option, ) -> LanceResult { + self.ensure_writable()?; let config = options.unwrap_or_default(); let lance_options = CompactionOptions { @@ -1242,6 +1249,7 @@ impl StorageBase { /// 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<()> { + self.ensure_writable()?; info!(column = %self.key_column, "creating ZoneMap index on key column"); self.dataset .create_index_builder( @@ -1458,7 +1466,8 @@ impl StorageBase { /// Read a fresh base/WAL view and retry if a merge invalidates it. A merge /// commits base before draining/deleting its generations, so an unchanged - /// base version brackets a valid read. Never return a partial result after + /// base version brackets a successful read. Retry missing generations even + /// when the base version is unchanged. Never return a partial result after /// a generation disappears. Ordinary WAL appends do not change base and do /// not force retries. Pinned reads use only their selected base version. pub async fn read_consistent( @@ -1470,7 +1479,7 @@ impl StorageBase { F: FnMut(Dataset, Vec) -> Fut, Fut: std::future::Future>, { - for _ in 0..3 { + for attempt in 0..3 { let mut dataset = self.dataset.clone(); if !self.is_version_pinned() { dataset.checkout_latest().await?; @@ -1486,9 +1495,13 @@ impl StorageBase { return result; } if let Err(err) = &result { - if !is_not_found_error(err) { + if !is_not_found_error(err) || attempt == 2 { return result; } + // A merge may already have committed before this read began. + // Its later WAL GC can invalidate our snapshot without changing + // the base version. Refresh both snapshots even in that case. + continue; } dataset.checkout_latest().await?; if dataset.manifest.version == version { diff --git a/crates/lance-context-core/tests/storage_reliability.rs b/crates/lance-context-core/tests/storage_reliability.rs index e91e6960..c052efa1 100644 --- a/crates/lance-context-core/tests/storage_reliability.rs +++ b/crates/lance-context-core/tests/storage_reliability.rs @@ -45,6 +45,112 @@ fn rec(id: &str) -> RolloutRecord { } } +#[tokio::test] +async fn review_pinned_add_preserves_external_id_uniqueness() { + let dir = tempfile::tempdir().unwrap(); + let mut store = ContextStore::open(dir.path().to_str().unwrap()) + .await + .unwrap(); + store.checkout(store.version()).await.unwrap(); + let request: AddRecordRequest = serde_json::from_value(serde_json::json!({ + "text_payload": "hello", "external_id": "same-key" + })) + .unwrap(); + let left = record_from_add_request(&request, "left".into(), "r".into()); + let right = record_from_add_request(&request, "right".into(), "r".into()); + for record in [&left, &right] { + let err = store.add(std::slice::from_ref(record)).await.unwrap_err(); + assert!(err.to_string().contains("read-only")); + } + assert!(store.is_version_pinned()); + store.refresh_latest().await.unwrap(); + assert!(store.list(None, None).await.unwrap().is_empty()); + store.add(&[left]).await.unwrap(); + assert!(store.add(&[right]).await.is_err()); + assert_eq!(store.list(None, None).await.unwrap().len(), 1); +} + +#[tokio::test] +async fn pinned_context_rejects_mutations_without_unpinning() { + let dir = tempfile::tempdir().unwrap(); + let mut store = ContextStore::open(dir.path().to_str().unwrap()) + .await + .unwrap(); + let request: AddRecordRequest = serde_json::from_value(serde_json::json!({ + "text_payload": "hello", "external_id": "same-key" + })) + .unwrap(); + let row = record_from_add_request(&request, "original".into(), "r".into()); + store.add(&[row]).await.unwrap(); + store.cleanup_wal().await.unwrap(); + let version = store.version(); + store.checkout(version).await.unwrap(); + let replacement = record_from_add_request(&request, "replacement".into(), "r".into()); + let patch = lance_context_core::RecordPatch { + source: Some("updated".into()), + ..Default::default() + }; + assert!(store + .upsert_by_external_id(replacement.clone()) + .await + .is_err()); + assert!(store + .upsert_many_by_external_id(vec![replacement]) + .await + .is_err()); + assert!(store.update_by_id("original", patch).await.is_err()); + let patch = lance_context_core::RecordPatch { + source: Some("updated".into()), + ..Default::default() + }; + assert!(store + .update_by_external_id("same-key", patch) + .await + .is_err()); + assert!(store.delete_by_id("original").await.is_err()); + assert!(store.delete_by_external_id("same-key").await.is_err()); + assert!(store.migrate_relationships_column().await.is_err()); + assert!(store.compact(None).await.is_err()); + assert!(store.create_id_index().await.is_err()); + let payload = dir.path().join("payload.bin"); + assert!(store + .put_payload(payload.to_str().unwrap(), b"payload") + .await + .is_err()); + assert!(!payload.exists()); + assert!(store.is_version_pinned()); + assert_eq!(store.version(), version); + store.refresh_latest().await.unwrap(); + assert_eq!(store.version(), version); + assert_eq!(store.list(None, None).await.unwrap()[0].id, "original"); + assert!(store.delete_by_id("original").await.unwrap()); +} + +#[tokio::test] +async fn pinned_rollout_rejects_writes_until_refresh() { + let dir = tempfile::tempdir().unwrap(); + let mut store = RolloutStore::open(dir.path().to_str().unwrap()) + .await + .unwrap(); + let version = store.version(); + store.checkout(version).await.unwrap(); + assert!(store + .add(&[rec("blocked")]) + .await + .unwrap_err() + .to_string() + .contains("read-only")); + assert!(store.compact(None).await.is_err()); + assert!(store.create_id_zonemap_index().await.is_err()); + assert!(store.is_version_pinned()); + store.refresh_latest().await.unwrap(); + assert_eq!(store.version(), version); + store.add(&[rec("allowed")]).await.unwrap(); + store.flush().await.unwrap(); + assert_eq!(store.list(None, None).await.unwrap()[0].id, "allowed"); + store.close().await.unwrap(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn list_keeps_rows_after_peer_merge() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/lance-context-server/src/routes/datagen.rs b/crates/lance-context-server/src/routes/datagen.rs index d6dbd0be..c14c7323 100644 --- a/crates/lance-context-server/src/routes/datagen.rs +++ b/crates/lance-context-server/src/routes/datagen.rs @@ -26,6 +26,7 @@ pub async fn create_datagen_store( Json(req): Json, ) -> Result<(StatusCode, Json), AppError> { AppState::validate_name(&req.name)?; + let mut handle = state.datagen_handles.lock(&req.name).await; if state .datagen_registry .write() @@ -55,7 +56,9 @@ pub async fn create_datagen_store( let version = store.version(); let store = Arc::new(RwLock::new(store)); - state.register_datagen(&req.name, &uri, store).await?; + state + .register_datagen(&req.name, &uri, store, &mut handle) + .await?; Ok(( StatusCode::CREATED, diff --git a/crates/lance-context-server/src/routes/generic.rs b/crates/lance-context-server/src/routes/generic.rs index c9c4f439..69ccda2d 100644 --- a/crates/lance-context-server/src/routes/generic.rs +++ b/crates/lance-context-server/src/routes/generic.rs @@ -33,6 +33,7 @@ pub async fn create_generic_store( Json(req): Json, ) -> Result<(StatusCode, Json), AppError> { AppState::validate_name(&req.name)?; + let mut handle = state.generic_handles.lock(&req.name).await; // Reject an invalid schema before creating anything, so a bad declaration // cannot leave an empty dataset behind. req.schema.validate().map_err(AppError::InvalidRequest)?; @@ -67,7 +68,9 @@ pub async fn create_generic_store( let version = store.version(); let store = Arc::new(RwLock::new(store)); - state.register_generic(&req.name, &uri, store).await?; + state + .register_generic(&req.name, &uri, store, &mut handle) + .await?; Ok(( StatusCode::CREATED, diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index bb90c094..0d27dba0 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -175,6 +175,7 @@ pub async fn create_rollout_store( Json(req): Json, ) -> Result<(StatusCode, Json), AppError> { AppState::validate_name(&req.name)?; + let mut handle = state.rollout_handles.lock(&req.name).await; // Existence is tracked durably in the registry, not by cache membership. if state .rollout_registry @@ -210,7 +211,9 @@ pub async fn create_rollout_store( // Record in the durable registry and the LRU. WAL cleanup is handled by the // single process-wide sweeper (see `AppState::spawn_global_sweeper`), so no // per-store timer is started here. - state.register_rollout(&req.name, &uri, store).await?; + state + .register_rollout(&req.name, &uri, store, &mut handle) + .await?; Ok(( StatusCode::CREATED, @@ -1079,6 +1082,51 @@ mod tests { assert!(cached.read().await.add(&[row]).await.is_err()); } + #[tokio::test] + async fn review_evicted_handle_cannot_write_after_delete() { + let (state, _dir) = rollout_state().await; + let in_flight = state.get_or_open_rollout_store("rl").await.unwrap(); + let row = rollout_record_from_add_request(&record_with_size("before-delete", None)); + in_flight.read().await.add(&[row]).await.unwrap(); + // Reopening an evicted but live store must reuse its writer. + state.rollout_stores.lock().await.pop("rl"); + let reopened = state.get_or_open_rollout_store("rl").await.unwrap(); + assert!(Arc::ptr_eq(&in_flight, &reopened)); + state.rollout_stores.lock().await.pop("rl"); + delete_rollout_store(State(state.clone()), Path("rl".into())) + .await + .unwrap(); + let row = rollout_record_from_add_request(&record_with_size("late-write", None)); + assert!(in_flight + .read() + .await + .add(std::slice::from_ref(&row)) + .await + .is_err()); + assert!(state.get_or_open_rollout_store("rl").await.is_err()); + + // Recreation gets a new writable handle; the old request stays fenced. + let _ = create_rollout_store( + State(state.clone()), + Json(CreateRolloutStoreRequest { + name: "rl".into(), + storage_options: None, + }), + ) + .await + .unwrap(); + let recreated = state.get_or_open_rollout_store("rl").await.unwrap(); + assert!(!Arc::ptr_eq(&in_flight, &recreated)); + assert!(in_flight + .read() + .await + .add(std::slice::from_ref(&row)) + .await + .is_err()); + recreated.read().await.add(&[row]).await.unwrap(); + recreated.write().await.close().await.unwrap(); + } + async fn rollout_state() -> (Arc, TempDir) { let dir = TempDir::new().unwrap(); let state = Arc::new(AppState::new_for_test(dir.path().to_path_buf()).await); diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index 1c5d7e53..9857e0d3 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -1,7 +1,8 @@ +use std::collections::HashMap; use std::num::NonZeroUsize; #[cfg(test)] use std::path::PathBuf; -use std::sync::Arc; +use std::sync::{Arc, Weak}; use std::time::Duration; use lance_context_core::{ @@ -10,7 +11,7 @@ use lance_context_core::{ RolloutStoreOptions, Session, }; use lru::LruCache; -use tokio::sync::{Mutex, RwLock}; +use tokio::sync::{Mutex, OwnedMutexGuard, RwLock}; use tokio::task::JoinHandle; use crate::config::ServerConfig; @@ -22,6 +23,58 @@ use crate::sweeper; /// number of experiments (which may be hundreds of thousands). pub const DEFAULT_ROLLOUT_CACHE_CAPACITY: usize = 2000; +type StoreSlot = Arc>>>; + +/// Weak handles survive LRU eviction without retaining datasets. The per-name +/// lock serializes create/open/delete across storage I/O, so an evicted handle +/// is reused and an old open cannot cross a delete/recreate boundary. +pub(crate) struct StoreHandles { + slots: std::sync::Mutex>, +} + +struct HandleSlots { + by_name: HashMap>, + prune_at: usize, +} + +impl Default for HandleSlots { + fn default() -> Self { + Self { + by_name: HashMap::new(), + prune_at: 64, + } + } +} + +impl Default for StoreHandles { + fn default() -> Self { + Self { + slots: Default::default(), + } + } +} + +impl StoreHandles { + pub(crate) async fn lock(&self, name: &str) -> OwnedMutexGuard>> { + let slot = { + let mut slots = self.slots.lock().unwrap(); + if !slots.by_name.contains_key(name) && slots.by_name.len() >= slots.prune_at { + // Amortize collection and keep slots with either a live store + // or an operation holding/waiting for the per-name lock. + slots.by_name.retain(|_, slot| { + Arc::strong_count(slot) > 1 + || slot + .try_lock() + .map_or(true, |handle| handle.strong_count() > 0) + }); + slots.prune_at = (slots.by_name.len() * 2).max(64); + } + slots.by_name.entry(name.to_string()).or_default().clone() + }; + slot.lock_owned().await + } +} + pub struct AppState { pub stores: RwLock>>>, /// Bounded LRU of resident rollout-store handles. @@ -47,6 +100,7 @@ pub struct AppState { /// the `Arc`. Both cases are logged (see that `Drop` impl) rather than /// silently stranding rows. pub rollout_stores: Mutex>>>, + pub(crate) rollout_handles: StoreHandles, /// Durable directory of which rollout stores exist. Consulted on a cache /// miss (existence check) and to back the list endpoint. Guarded by a lock /// because every operation refreshes the snapshot and therefore takes @@ -91,6 +145,7 @@ pub struct AppState { /// experiment, so the same residency bound and durable-registry existence /// model applies. pub datagen_stores: Mutex>>>, + pub(crate) datagen_handles: StoreHandles, /// Durable directory of which datagen stores exist. A separate registry /// dataset from [`Self::rollout_registry`] so the two store kinds never /// collide on a shared name. @@ -98,6 +153,7 @@ pub struct AppState { /// Bounded LRU of resident generic-store handles, mirroring /// [`Self::rollout_stores`]. pub generic_stores: Mutex>>>, + pub(crate) generic_handles: StoreHandles, /// Durable directory of which generic stores exist. /// /// Records only name and URI: a generic store's **schema is stored in its @@ -249,6 +305,7 @@ impl AppState { Ok(Self { stores: RwLock::new(std::collections::HashMap::new()), rollout_stores: Mutex::new(LruCache::new(capacity)), + rollout_handles: StoreHandles::default(), rollout_registry: RwLock::new(registry), base_uri, instance_id, @@ -261,8 +318,10 @@ impl AppState { rollout_large_blob_log_bytes: config.rollout_large_blob_log_bytes, rollout_session, datagen_stores: Mutex::new(LruCache::new(capacity)), + datagen_handles: StoreHandles::default(), datagen_registry: RwLock::new(datagen_registry), generic_stores: Mutex::new(LruCache::new(capacity)), + generic_handles: StoreHandles::default(), generic_registry: RwLock::new(generic_registry), }) } @@ -311,6 +370,7 @@ impl AppState { rollout_stores: Mutex::new(LruCache::new( NonZeroUsize::new(DEFAULT_ROLLOUT_CACHE_CAPACITY).unwrap(), )), + rollout_handles: StoreHandles::default(), rollout_registry: RwLock::new(registry), base_uri, instance_id, @@ -325,10 +385,12 @@ impl AppState { generic_stores: Mutex::new(LruCache::new( std::num::NonZeroUsize::new(DEFAULT_ROLLOUT_CACHE_CAPACITY).unwrap(), )), + generic_handles: StoreHandles::default(), generic_registry: RwLock::new(generic_registry), datagen_stores: Mutex::new(LruCache::new( NonZeroUsize::new(DEFAULT_ROLLOUT_CACHE_CAPACITY).unwrap(), )), + datagen_handles: StoreHandles::default(), datagen_registry: RwLock::new(datagen_registry), } } @@ -355,14 +417,21 @@ impl AppState { } /// Record that a rollout store exists, in both the durable registry and the - /// in-memory LRU. Called by the create route after the dataset is written. + /// in-memory LRU. The caller must hold this name's handle lock throughout + /// the existence check, dataset creation, and registration. pub async fn register_rollout( &self, name: &str, uri: &str, store: Arc>, + handle: &mut OwnedMutexGuard>>, ) -> Result<(), AppError> { Self::validate_name(name)?; + if handle.upgrade().is_some() { + return Err(AppError::AlreadyExists(format!( + "Store '{name}' already has a live handle" + ))); + } self.rollout_registry .write() .await @@ -372,6 +441,7 @@ impl AppState { // Insertion may evict the LRU's least-recently-used entry; dropping the // returned handle releases it (no per-store timer to abort — cleanup is // now global). + **handle = Arc::downgrade(&store); self.rollout_stores .lock() .await @@ -383,6 +453,7 @@ impl AppState { /// handle. Returns whether the store existed. pub async fn unregister_rollout(&self, name: &str) -> Result { Self::validate_name(name)?; + let mut handle = self.rollout_handles.lock(name).await; // Keep the registry entry and cached handle until physical deletion // succeeds, so failures can be retried without losing the store's URI. let mut registry = self.rollout_registry.write().await; @@ -393,8 +464,8 @@ impl AppState { { return Ok(false); } - let cached = self.rollout_stores.lock().await.peek(name).cloned(); - if let Some(store) = cached { + let live = handle.upgrade(); + if let Some(store) = live { store .write() .await @@ -408,6 +479,7 @@ impl AppState { } registry.remove(name).await.map_err(AppError::from_lance)?; self.rollout_stores.lock().await.pop(name); + *handle = Weak::new(); Ok(true) } @@ -432,6 +504,15 @@ impl AppState { } metrics::counter!("rollout_store_cache_misses_total").increment(1); + let mut handle = self.rollout_handles.lock(name).await; + if let Some(store) = handle.upgrade() { + self.rollout_stores + .lock() + .await + .put(name.to_string(), store.clone()); + return Ok(store); + } + // Existence is the registry's job, not the cache's. let exists = self .rollout_registry @@ -454,8 +535,7 @@ impl AppState { .await .map_err(AppError::from_lance)?; let opened = Arc::new(RwLock::new(opened)); - // Recheck after the slow open while excluding a concurrent deletion - // through registry removal and cache insertion. + // Also observe registry removal by another server during the open. let mut registry = self.rollout_registry.write().await; if !registry .contains(name) @@ -466,13 +546,10 @@ impl AppState { "Store '{name}' was deleted while opening" ))); } - - // Insert under the lock, re-checking for a store another request may - // have opened concurrently while we were loading. + // The per-name lock excludes deletion and concurrent opens through + // publication in both the weak registry and the resident cache. + *handle = Arc::downgrade(&opened); let mut cache = self.rollout_stores.lock().await; - if let Some(existing) = cache.get(name) { - return Ok(existing.clone()); - } cache.put(name.to_string(), opened.clone()); metrics::gauge!("rollout_stores_resident").set(cache.len() as f64); Ok(opened) @@ -528,20 +605,28 @@ impl AppState { } /// Record that a datagen store exists, in both the durable registry and the - /// in-memory LRU. Called by the create route after the dataset is written. + /// in-memory LRU. The caller must hold this name's handle lock throughout + /// the existence check, dataset creation, and registration. pub async fn register_datagen( &self, name: &str, uri: &str, store: Arc>, + handle: &mut OwnedMutexGuard>>, ) -> Result<(), AppError> { Self::validate_name(name)?; + if handle.upgrade().is_some() { + return Err(AppError::AlreadyExists(format!( + "Store '{name}' already has a live handle" + ))); + } self.datagen_registry .write() .await .upsert(name, uri) .await .map_err(AppError::from_lance)?; + **handle = Arc::downgrade(&store); self.datagen_stores .lock() .await @@ -553,6 +638,7 @@ impl AppState { /// handle. Returns whether the store existed. pub async fn unregister_datagen(&self, name: &str) -> Result { Self::validate_name(name)?; + let mut handle = self.datagen_handles.lock(name).await; // Keep the registry entry and cached handle until physical deletion // succeeds, so failures can be retried without losing the store's URI. let mut registry = self.datagen_registry.write().await; @@ -563,8 +649,8 @@ impl AppState { { return Ok(false); } - let cached = self.datagen_stores.lock().await.peek(name).cloned(); - if let Some(store) = cached { + let live = handle.upgrade(); + if let Some(store) = live { store .write() .await @@ -578,6 +664,7 @@ impl AppState { } registry.remove(name).await.map_err(AppError::from_lance)?; self.datagen_stores.lock().await.pop(name); + *handle = Weak::new(); Ok(true) } @@ -594,6 +681,15 @@ impl AppState { return Ok(store.clone()); } + let mut handle = self.datagen_handles.lock(name).await; + if let Some(store) = handle.upgrade() { + self.datagen_stores + .lock() + .await + .put(name.to_string(), store.clone()); + return Ok(store); + } + let exists = self .datagen_registry .write() @@ -613,8 +709,7 @@ impl AppState { .await .map_err(AppError::from_lance)?; let opened = Arc::new(RwLock::new(opened)); - // Recheck after the slow open while excluding a concurrent deletion - // through registry removal and cache insertion. + // Also observe registry removal by another server during the open. let mut registry = self.datagen_registry.write().await; if !registry .contains(name) @@ -625,11 +720,10 @@ impl AppState { "Store '{name}' was deleted while opening" ))); } - + // The per-name lock excludes deletion and concurrent opens through + // publication in both the weak registry and the resident cache. + *handle = Arc::downgrade(&opened); let mut cache = self.datagen_stores.lock().await; - if let Some(existing) = cache.get(name) { - return Ok(existing.clone()); - } cache.put(name.to_string(), opened.clone()); Ok(opened) } @@ -654,20 +748,27 @@ impl AppState { } /// Record that a generic store exists, in both the durable registry and the - /// in-memory LRU. + /// in-memory LRU. The caller holds the name's handle lock throughout creation. pub async fn register_generic( &self, name: &str, uri: &str, store: Arc>, + handle: &mut OwnedMutexGuard>>, ) -> Result<(), AppError> { Self::validate_name(name)?; + if handle.upgrade().is_some() { + return Err(AppError::AlreadyExists(format!( + "Store '{name}' already has a live handle" + ))); + } self.generic_registry .write() .await .upsert(name, uri) .await .map_err(AppError::from_lance)?; + **handle = Arc::downgrade(&store); self.generic_stores .lock() .await @@ -679,6 +780,7 @@ impl AppState { /// Returns whether the store existed. pub async fn unregister_generic(&self, name: &str) -> Result { Self::validate_name(name)?; + let mut handle = self.generic_handles.lock(name).await; // Keep the registry entry and cached handle until physical deletion // succeeds, so failures can be retried without losing the store's URI. let mut registry = self.generic_registry.write().await; @@ -689,8 +791,8 @@ impl AppState { { return Ok(false); } - let cached = self.generic_stores.lock().await.peek(name).cloned(); - if let Some(store) = cached { + let live = handle.upgrade(); + if let Some(store) = live { store .write() .await @@ -704,6 +806,7 @@ impl AppState { } registry.remove(name).await.map_err(AppError::from_lance)?; self.generic_stores.lock().await.pop(name); + *handle = Weak::new(); Ok(true) } @@ -721,6 +824,15 @@ impl AppState { return Ok(store.clone()); } + let mut handle = self.generic_handles.lock(name).await; + if let Some(store) = handle.upgrade() { + self.generic_stores + .lock() + .await + .put(name.to_string(), store.clone()); + return Ok(store); + } + let exists = self .generic_registry .write() @@ -742,8 +854,7 @@ impl AppState { .await .map_err(AppError::from_lance)?; let opened = Arc::new(RwLock::new(opened)); - // Recheck after the slow open while excluding a concurrent deletion - // through registry removal and cache insertion. + // Also observe registry removal by another server during the open. let mut registry = self.generic_registry.write().await; if !registry .contains(name) @@ -754,11 +865,10 @@ impl AppState { "Store '{name}' was deleted while opening" ))); } - + // The per-name lock excludes deletion and concurrent opens through + // publication in both the weak registry and the resident cache. + *handle = Arc::downgrade(&opened); let mut cache = self.generic_stores.lock().await; - if let Some(existing) = cache.get(name) { - return Ok(existing.clone()); - } cache.put(name.to_string(), opened.clone()); Ok(opened) } @@ -1083,7 +1193,12 @@ mod tests { let generic = Arc::new(RwLock::new(generic)); state - .register_generic("g1", &generic_uri, generic.clone()) + .register_generic( + "g1", + &generic_uri, + generic.clone(), + &mut state.generic_handles.lock("g1").await, + ) .await .unwrap(); @@ -1104,7 +1219,12 @@ mod tests { ); let datagen = Arc::new(RwLock::new(datagen)); state - .register_datagen("d1", &datagen_uri, datagen.clone()) + .register_datagen( + "d1", + &datagen_uri, + datagen.clone(), + &mut state.datagen_handles.lock("d1").await, + ) .await .unwrap(); @@ -1139,6 +1259,116 @@ mod tests { ); } + #[tokio::test] + async fn evicted_datagen_and_generic_handles_are_fenced_by_delete() { + use lance_context_api::{ColumnSpec, ColumnType, SchemaSpec, ID_COLUMN}; + let dir = TempDir::new().unwrap(); + let state = state_with_interval(&dir, 0).await; + let uri = state.datagen_uri("dg"); + let datagen = Arc::new(RwLock::new(DatagenStore::open(&uri).await.unwrap())); + state + .register_datagen( + "dg", + &uri, + datagen.clone(), + &mut state.datagen_handles.lock("dg").await, + ) + .await + .unwrap(); + datagen + .write() + .await + .append(&[datagen_event()]) + .await + .unwrap(); + state.datagen_stores.lock().await.pop("dg"); + assert!(Arc::ptr_eq( + &datagen, + &state.get_or_open_datagen_store("dg").await.unwrap() + )); + state.datagen_stores.lock().await.pop("dg"); + assert!(state.unregister_datagen("dg").await.unwrap()); + assert!(datagen + .write() + .await + .append(&[datagen_event()]) + .await + .is_err()); + assert!(state.get_or_open_datagen_store("dg").await.is_err()); + + let uri = state.generic_uri("g"); + let schema = SchemaSpec::new(vec![( + ID_COLUMN.to_string(), + ColumnSpec::required(ColumnType::String { large: false }), + )]); + let generic = Arc::new(RwLock::new( + GenericStore::open(&uri, schema, GenericStoreOptions::default()) + .await + .unwrap(), + )); + state + .register_generic( + "g", + &uri, + generic.clone(), + &mut state.generic_handles.lock("g").await, + ) + .await + .unwrap(); + let row = serde_json::from_value::>( + serde_json::json!({"id": "r1"}), + ) + .unwrap(); + generic + .read() + .await + .add(std::slice::from_ref(&row)) + .await + .unwrap(); + state.generic_stores.lock().await.pop("g"); + assert!(Arc::ptr_eq( + &generic, + &state.get_or_open_generic_store("g").await.unwrap() + )); + state.generic_stores.lock().await.pop("g"); + assert!(state.unregister_generic("g").await.unwrap()); + assert!(generic.read().await.add(&[row]).await.is_err()); + assert!(state.get_or_open_generic_store("g").await.is_err()); + } + + #[tokio::test] + async fn weak_handles_release_stores_and_preserve_active_name_locks() { + let handles = StoreHandles::<()>::default(); + let store = Arc::new(RwLock::new(())); + let weak = Arc::downgrade(&store); + *handles.lock("live").await = weak.clone(); + let held = handles.lock("opening").await; + // Trigger collection repeatedly while a same-name open/delete is held. + for n in 0..200 { + drop(handles.lock(&format!("cold-{n}")).await); + } + assert!(handles.lock("live").await.upgrade().is_some()); + let waiting = handles.lock("opening"); + tokio::pin!(waiting); + assert!(futures::poll!(&mut waiting).is_pending()); + drop(held); + drop(waiting.await); + drop(store); + assert!( + weak.upgrade().is_none(), + "the registry must not retain datasets" + ); + for n in 200..400 { + drop(handles.lock(&format!("cold-{n}")).await); + } + let slots = handles.slots.lock().unwrap(); + assert!( + slots.by_name.len() < 64, + "dead names must not grow without bound" + ); + assert!(!slots.by_name.contains_key("live")); + } + /// `shutdown` drains resident rollout writers by awaiting `close()` on each /// (idempotent: `close` is a no-op when no writer is resident), so /// `ShardWriter` background tasks are reclaimed deterministically instead of @@ -1163,7 +1393,12 @@ mod tests { .unwrap(); let store = Arc::new(RwLock::new(store)); state - .register_rollout("exp", &uri, store.clone()) + .register_rollout( + "exp", + &uri, + store.clone(), + &mut state.rollout_handles.lock("exp").await, + ) .await .unwrap(); diff --git a/docs/src/specs/rollout-deployment.md b/docs/src/specs/rollout-deployment.md index dbbd101e..7d7a70da 100644 --- a/docs/src/specs/rollout-deployment.md +++ b/docs/src/specs/rollout-deployment.md @@ -218,7 +218,7 @@ rows = await store.list(filters={ }) ``` -`checkout` selects a fixed base-table version and excludes live WAL generations, including unmerged writes made before checkout. Periodic WAL merging does not advance a pinned handle. It is not the mechanism for per-checkpoint rollout reproducibility. (See schema-design §3 and §7 — and note that `learner_iteration` is likewise a column, not a dataset version.) +`checkout` selects a fixed base-table version and excludes live WAL generations, including unmerged writes made before checkout. Pinned handles reject appends, compaction, index creation, and schema changes. Call Rust `refresh_latest()` to return the handle to the latest writable view. Periodic WAL merging does not advance a pinned handle. It is not the mechanism for per-checkpoint rollout reproducibility. (See schema-design §3 and §7 — and note that `learner_iteration` is likewise a column, not a dataset version.) ---