From 52f6911e41448e08aaa17ec67194f734d48a75e6 Mon Sep 17 00:00:00 2001 From: Beinan Date: Thu, 6 Aug 2026 21:30:25 +0000 Subject: [PATCH] fix(server): refresh stale point lookups across stores --- .../lance-context-core/src/datagen_store.rs | 11 + .../lance-context-core/src/generic_store.rs | 11 + .../lance-context-core/src/rollout_store.rs | 26 ++ crates/lance-context-core/src/store.rs | 16 +- crates/lance-context-core/src/store_base.rs | 32 +- .../src/routes/datagen.rs | 295 ++++++++++++++++-- .../src/routes/generic.rs | 73 ++++- .../src/routes/records.rs | 169 ++++++++-- .../src/routes/rollouts.rs | 165 +++++++++- 9 files changed, 728 insertions(+), 70 deletions(-) diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index dcc6064..a1abaf9 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -125,6 +125,17 @@ impl DatagenStore { self.base.version() } + /// Whether this handle was explicitly checked out to a dataset version. + #[must_use] + pub fn is_version_pinned(&self) -> bool { + self.base.is_version_pinned() + } + + /// Refresh this handle to the latest base-table manifest. + pub async fn refresh_latest(&mut self) -> LanceResult<()> { + self.base.refresh_latest().await + } + /// Append one or more complete checkpoint batches. /// /// The supplied slice is persisted as one MemWAL generation. Callers should diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index 15896be..a9e8fdf 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -222,6 +222,17 @@ impl GenericStore { self.base.version() } + /// Whether this handle was explicitly checked out to a dataset version. + #[must_use] + pub fn is_version_pinned(&self) -> bool { + self.base.is_version_pinned() + } + + /// Refresh this handle to the latest base-table manifest. + pub async fn refresh_latest(&mut self) -> LanceResult<()> { + self.base.refresh_latest().await + } + /// Append rows. /// /// Rows are matched to columns by name; an undeclared key is an error, and diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 2961478..f1419df 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -465,6 +465,15 @@ impl RolloutStore { self.base.checkout(version_id).await } + /// Whether this handle was explicitly checked out to a dataset version. + /// + /// Serving layers use this to preserve time-travel reads instead of + /// automatically advancing a pinned handle when a point lookup misses. + #[must_use] + pub fn is_version_pinned(&self) -> bool { + self.base.is_version_pinned() + } + /// Refresh this handle to the latest base-table manifest while retaining /// its session and metadata caches. /// @@ -2843,6 +2852,23 @@ mod tests { }); } + #[test] + fn explicit_checkout_pins_until_latest_refresh() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = RolloutStore::open(&uri).await.unwrap(); + assert!(!store.is_version_pinned()); + + store.checkout(store.version()).await.unwrap(); + assert!(store.is_version_pinned()); + + store.refresh_latest().await.unwrap(); + assert!(!store.is_version_pinned()); + }); + } + #[test] fn trajectory_rows_are_filtered_and_sorted_across_fragments() { let dir = TempDir::new().unwrap(); diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 00229fb..ab6618e 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -1310,14 +1310,24 @@ impl ContextStore { .dataset .add_columns(NewColumnTransform::AllNulls(schema), None, None) .await?; + self.base.clear_version_pin(); Ok(true) } /// Checkout a specific dataset version. pub async fn checkout(&mut self, version_id: u64) -> LanceResult<()> { - let dataset = self.base.dataset.checkout_version(version_id).await?; - self.base.dataset = dataset; - Ok(()) + self.base.checkout(version_id).await + } + + /// Whether this handle was explicitly checked out to a dataset version. + #[must_use] + pub fn is_version_pinned(&self) -> bool { + self.base.is_version_pinned() + } + + /// Refresh this handle to the latest base-table manifest. + pub async fn refresh_latest(&mut self) -> LanceResult<()> { + self.base.refresh_latest().await } /// Retrieve a single record by its unique ID. diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index 5c8bcdd..33d3740 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -259,6 +259,12 @@ pub(crate) struct StorageBase { total_compactions: u64, /// Error message from the most recent failed compaction on this handle. last_compaction_error: Option, + /// Explicit time-travel version selected by [`Self::checkout`]. + /// + /// A point-read miss may refresh an ordinary long-lived handle to avoid a + /// false negative from a stale manifest, but must never advance a handle + /// whose caller deliberately selected a historical version. + pinned_version: Option, /// 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 @@ -375,6 +381,7 @@ impl StorageBase { last_compaction: None, total_compactions: 0, last_compaction_error: None, + pinned_version: None, write_writer: tokio::sync::Mutex::new(None), }; // `ensure_mem_wal` may reload the dataset on a concurrent first-writer @@ -398,9 +405,16 @@ impl StorageBase { /// Check out a specific base dataset version (time travel). pub async fn checkout(&mut self, version_id: u64) -> LanceResult<()> { self.dataset = self.dataset.checkout_version(version_id).await?; + self.pinned_version = Some(version_id); Ok(()) } + /// Whether this handle was explicitly checked out to a historical version. + #[must_use] + pub fn is_version_pinned(&self) -> bool { + self.pinned_version.is_some() + } + /// Refresh this handle to the latest base-table manifest while retaining its /// session and metadata caches. /// @@ -408,7 +422,15 @@ impl StorageBase { /// WAL merges committed by another process become visible without paying the /// cost of reopening the dataset and rebuilding all session caches. pub async fn refresh_latest(&mut self) -> LanceResult<()> { - self.dataset.checkout_latest().await + self.dataset.checkout_latest().await?; + self.pinned_version = None; + 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; } // ---------------------------------------------------------------- writes @@ -822,6 +844,7 @@ impl StorageBase { "append", self.append_merged_batches(batches, merge_schema).await )?; + self.pinned_version = None; } // Reuse the shard's *current* epoch rather than claiming a new one: @@ -849,7 +872,9 @@ impl StorageBase { .await )?; - self.delete_merged_generation_dirs(&merged_paths).await + self.delete_merged_generation_dirs(&merged_paths).await?; + self.pinned_version = None; + Ok(()) } /// Delete the merged generations' directories now that no manifest @@ -966,7 +991,7 @@ impl StorageBase { let Some(latest_schema) = self.latest_schema.clone() else { return Ok(()); }; - self.dataset.checkout_latest().await?; + self.refresh_latest().await?; let base_schema: Arc = Arc::new(self.dataset.schema().into()); align_batch_to_schema( @@ -1146,6 +1171,7 @@ impl StorageBase { self.dataset = Self::load_with_options(&uri, self.storage_options.clone(), self.session.clone()) .await?; + self.pinned_version = None; Ok(()) } diff --git a/crates/lance-context-server/src/routes/datagen.rs b/crates/lance-context-server/src/routes/datagen.rs index b4ad1bc..b80ca2b 100644 --- a/crates/lance-context-server/src/routes/datagen.rs +++ b/crates/lance-context-server/src/routes/datagen.rs @@ -4,8 +4,9 @@ use axum::extract::{Path, Query, State}; use axum::http::StatusCode; use axum::Json; use lance_context_api::{ - AddDatagenEventsRequest, AddDatagenEventsResponse, CreateDatagenStoreRequest, DatagenStoreApi, - DatagenStoreInfo, GetFoldedDatagenItemResponse, ListDatagenFailuresResponse, + AddDatagenEventsRequest, AddDatagenEventsResponse, CreateDatagenStoreRequest, + DatagenRootItemStatusesResponse, DatagenStoreApi, DatagenStoreInfo, + GetFoldedDatagenItemResponse, ListDatagenEventsResponse, ListDatagenFailuresResponse, ListDatagenStoresResponse, }; use lance_context_core::{DatagenStore, DatagenStoreOptions}; @@ -157,16 +158,130 @@ pub struct FoldParams { pub load_blobs: bool, } +async fn fold_datagen_item_refreshing_on_miss( + store_lock: &RwLock, + item_id: &str, + load_blobs: bool, +) -> Result, AppError> { + { + let store = store_lock.read().await; + let item = DatagenStoreApi::fold_item_with_blobs(&*store, item_id, load_blobs) + .await + .map_err(AppError::from_context)?; + if item.is_some() || store.is_version_pinned() { + return Ok(item); + } + } + + let mut store = store_lock.write().await; + if !store.is_version_pinned() { + store.refresh_latest().await.map_err(AppError::from_lance)?; + } + DatagenStoreApi::fold_item_with_blobs(&*store, item_id, load_blobs) + .await + .map_err(AppError::from_context) +} + +async fn datagen_failures_refreshing_on_empty( + store_lock: &RwLock, + item_id: &str, +) -> Result, AppError> { + { + let store = store_lock.read().await; + let failures = DatagenStoreApi::item_failures(&*store, item_id) + .await + .map_err(AppError::from_context)?; + if !failures.is_empty() || store.is_version_pinned() { + return Ok(failures); + } + } + + let mut store = store_lock.write().await; + if !store.is_version_pinned() { + store.refresh_latest().await.map_err(AppError::from_lance)?; + } + DatagenStoreApi::item_failures(&*store, item_id) + .await + .map_err(AppError::from_context) +} + +async fn datagen_events_for_root_refreshing_on_empty( + store_lock: &RwLock, + root_item_id: &str, +) -> Result, AppError> { + { + let store = store_lock.read().await; + let events = DatagenStoreApi::events_for_root(&*store, root_item_id) + .await + .map_err(AppError::from_context)?; + if !events.is_empty() || store.is_version_pinned() { + return Ok(events); + } + } + + let mut store = store_lock.write().await; + if !store.is_version_pinned() { + store.refresh_latest().await.map_err(AppError::from_lance)?; + } + DatagenStoreApi::events_for_root(&*store, root_item_id) + .await + .map_err(AppError::from_context) +} + +async fn datagen_root_statuses_refreshing_on_missing( + store_lock: &RwLock, + ids: &[String], +) -> Result { + { + let store = store_lock.read().await; + let statuses = DatagenStoreApi::root_item_statuses(&*store, ids) + .await + .map_err(AppError::from_context)?; + if statuses.statuses.len() == ids.len() || store.is_version_pinned() { + return Ok(statuses); + } + } + + let mut store = store_lock.write().await; + if !store.is_version_pinned() { + store.refresh_latest().await.map_err(AppError::from_lance)?; + } + DatagenStoreApi::root_item_statuses(&*store, ids) + .await + .map_err(AppError::from_context) +} + +async fn get_datagen_blob_refreshing_on_miss( + store_lock: &RwLock, + event_id: &str, +) -> Result>, AppError> { + { + let store = store_lock.read().await; + let bytes = DatagenStoreApi::get_blob(&*store, event_id) + .await + .map_err(AppError::from_context)?; + if bytes.is_some() || store.is_version_pinned() { + return Ok(bytes); + } + } + + let mut store = store_lock.write().await; + if !store.is_version_pinned() { + store.refresh_latest().await.map_err(AppError::from_lance)?; + } + DatagenStoreApi::get_blob(&*store, event_id) + .await + .map_err(AppError::from_context) +} + pub async fn fold_datagen_item( State(state): State>, Path((name, item_id)): Path<(String, String)>, Query(params): Query, ) -> Result, AppError> { let store_lock = state.get_or_open_datagen_store(&name).await?; - let store = store_lock.read().await; - let item = DatagenStoreApi::fold_item_with_blobs(&*store, &item_id, params.load_blobs) - .await - .map_err(AppError::from_context)?; + let item = + fold_datagen_item_refreshing_on_miss(&store_lock, &item_id, params.load_blobs).await?; Ok(Json(GetFoldedDatagenItemResponse { item })) } @@ -189,10 +304,7 @@ pub async fn datagen_item_failures( Path((name, item_id)): Path<(String, String)>, ) -> Result, AppError> { let store_lock = state.get_or_open_datagen_store(&name).await?; - let store = store_lock.read().await; - let failures = DatagenStoreApi::item_failures(&*store, &item_id) - .await - .map_err(AppError::from_context)?; + let failures = datagen_failures_refreshing_on_empty(&store_lock, &item_id).await?; Ok(Json(ListDatagenFailuresResponse { failures })) } @@ -203,13 +315,8 @@ pub async fn datagen_events_for_root( Path((name, root_item_id)): Path<(String, String)>, ) -> Result, AppError> { let store_lock = state.get_or_open_datagen_store(&name).await?; - let store = store_lock.read().await; - let events = DatagenStoreApi::events_for_root(&*store, &root_item_id) - .await - .map_err(AppError::from_context)?; - Ok(Json(lance_context_api::ListDatagenEventsResponse { - events, - })) + let events = datagen_events_for_root_refreshing_on_empty(&store_lock, &root_item_id).await?; + Ok(Json(ListDatagenEventsResponse { events })) } #[derive(Debug, Default, serde::Deserialize)] @@ -236,10 +343,7 @@ pub async fn datagen_root_item_statuses( .unwrap_or_default(); let store_lock = state.get_or_open_datagen_store(&name).await?; - let store = store_lock.read().await; - let resp = DatagenStoreApi::root_item_statuses(&*store, &ids) - .await - .map_err(AppError::from_context)?; + let resp = datagen_root_statuses_refreshing_on_missing(&store_lock, &ids).await?; Ok(Json(resp)) } @@ -254,11 +358,152 @@ pub async fn fetch_datagen_blob( use axum::response::IntoResponse; let store_lock = state.get_or_open_datagen_store(&name).await?; - let store = store_lock.read().await; - let bytes = DatagenStoreApi::get_blob(&*store, &event_id) - .await - .map_err(AppError::from_context)? + let bytes = get_datagen_blob_refreshing_on_miss(&store_lock, &event_id) + .await? .ok_or_else(|| AppError::NotFound(format!("Datagen event '{}' has no blob", event_id)))?; Ok(([(header::CONTENT_TYPE, "application/octet-stream")], bytes).into_response()) } + +#[cfg(test)] +mod tests { + use chrono::Utc; + use lance_context_api::CreateDatagenStoreRequest; + use lance_context_core::{ + datagen_event_id, DatagenBlobValue, DatagenEvent, DatagenEventType, DatagenItemStatus, + DatagenStepKind, DatagenValue, DATAGEN_SCHEMA_VERSION, + }; + use tempfile::TempDir; + + use super::*; + + async fn test_state() -> (Arc, TempDir) { + let dir = TempDir::new().unwrap(); + let state = Arc::new(AppState::new_for_test(dir.path().to_path_buf()).await); + (state, dir) + } + + fn event( + item_id: &str, + seq: i64, + checkpoint_id: &str, + ordinal: u32, + event_type: DatagenEventType, + ) -> DatagenEvent { + DatagenEvent { + event_id: datagen_event_id(item_id, checkpoint_id, ordinal), + item_id: item_id.to_string(), + root_item_id: item_id.to_string(), + parent_item_id: None, + item_seq: seq, + checkpoint_id: checkpoint_id.to_string(), + event_type, + step_name: None, + step_kind: None, + step_index: None, + enclosing_step: None, + selector_step: None, + attempt: 0, + run_id: "external-run".to_string(), + writer_epoch: "external-writer".to_string(), + field_name: None, + field_type: None, + codec_version: None, + value: None, + query_tags: None, + status: Some(DatagenItemStatus::Running), + error_type: None, + error_dump: None, + traceback: None, + event_ts: Utc::now(), + schema_version: DATAGEN_SCHEMA_VERSION, + } + } + + #[tokio::test] + async fn point_reads_refresh_a_base_advanced_by_an_external_writer() { + let (state, _dir) = test_state().await; + for name in ["fold-store", "blob-store"] { + let _ = create_datagen_store( + State(state.clone()), + Json(CreateDatagenStoreRequest { + name: name.to_string(), + storage_options: None, + }), + ) + .await + .unwrap(); + } + + let cached_fold = state.get_or_open_datagen_store("fold-store").await.unwrap(); + let mut fold_writer = DatagenStore::open_existing_with_options( + &state.datagen_uri("fold-store"), + DatagenStoreOptions { + shard_id: Some("external-fold-writer".to_string()), + ..Default::default() + }, + ) + .await + .unwrap(); + let created = event( + "merged-item", + 0, + "created", + 0, + DatagenEventType::ItemCreated, + ); + fold_writer.append(&[created]).await.unwrap(); + assert_eq!(fold_writer.cleanup_own_shard().await.unwrap(), 1); + assert_eq!(fold_writer.pending_wal_generations().await.unwrap(), 0); + assert!(fold_writer.version() > cached_fold.read().await.version()); + + let Json(found) = fold_datagen_item( + State(state.clone()), + Path(("fold-store".to_string(), "merged-item".to_string())), + Query(FoldParams { load_blobs: false }), + ) + .await + .unwrap(); + assert_eq!(found.item.unwrap().item_id, "merged-item"); + assert_eq!(cached_fold.read().await.version(), fold_writer.version()); + + let cached_blob = state.get_or_open_datagen_store("blob-store").await.unwrap(); + let mut blob_writer = DatagenStore::open_existing_with_options( + &state.datagen_uri("blob-store"), + DatagenStoreOptions { + shard_id: Some("external-blob-writer".to_string()), + ..Default::default() + }, + ) + .await + .unwrap(); + let payload = b"externally merged blob".to_vec(); + let mut blob_event = event("blob-item", 0, "blob", 0, DatagenEventType::FieldSet); + blob_event.step_name = Some("capture".to_string()); + blob_event.step_kind = Some(DatagenStepKind::Leaf); + blob_event.step_index = Some(0); + blob_event.field_name = Some("artifact".to_string()); + blob_event.field_type = Some("blob".to_string()); + blob_event.codec_version = Some(1); + blob_event.value = Some(DatagenValue::Blob(DatagenBlobValue { + bytes: Some(payload.clone()), + size: payload.len() as i64, + checksum: None, + })); + blob_event.status = None; + let event_id = blob_event.event_id.clone(); + blob_writer.append(&[blob_event]).await.unwrap(); + assert_eq!(blob_writer.cleanup_own_shard().await.unwrap(), 1); + assert_eq!(blob_writer.pending_wal_generations().await.unwrap(), 0); + assert!(blob_writer.version() > cached_blob.read().await.version()); + + let response = fetch_datagen_blob(State(state), Path(("blob-store".to_string(), event_id))) + .await + .unwrap(); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + assert_eq!(bytes.as_ref(), payload.as_slice()); + assert_eq!(cached_blob.read().await.version(), blob_writer.version()); + } +} diff --git a/crates/lance-context-server/src/routes/generic.rs b/crates/lance-context-server/src/routes/generic.rs index 4d50a28..2e56580 100644 --- a/crates/lance-context-server/src/routes/generic.rs +++ b/crates/lance-context-server/src/routes/generic.rs @@ -201,6 +201,26 @@ pub struct GetRowQuery { pub columns: Option, } +async fn get_generic_row_refreshing_on_miss( + store_lock: &RwLock, + id: &str, + columns: Option<&[String]>, +) -> Result>, AppError> { + { + let store = store_lock.read().await; + let row = store.get(id, columns).await.map_err(AppError::from_lance)?; + if row.is_some() || store.is_version_pinned() { + return Ok(row); + } + } + + let mut store = store_lock.write().await; + if !store.is_version_pinned() { + store.refresh_latest().await.map_err(AppError::from_lance)?; + } + store.get(id, columns).await.map_err(AppError::from_lance) +} + /// `GET /api/v1/generic/{name}/rows/{id}` pub async fn get_row( State(state): State>, @@ -208,7 +228,6 @@ pub async fn get_row( Query(query): Query, ) -> Result>, AppError> { let store = state.get_or_open_generic_store(&name).await?; - let guard = store.read().await; let columns: Option> = query.columns.as_ref().map(|raw| { raw.split(',') @@ -218,10 +237,8 @@ pub async fn get_row( .collect() }); - let row = guard - .get(&id, columns.as_deref()) - .await - .map_err(AppError::from_lance)? + let row = get_generic_row_refreshing_on_miss(&store, &id, columns.as_deref()) + .await? .ok_or_else(|| AppError::NotFound(format!("Row '{id}' does not exist")))?; Ok(Json(row)) } @@ -291,6 +308,52 @@ mod tests { value.as_object().unwrap().clone() } + #[tokio::test] + async fn point_read_refreshes_a_base_advanced_by_an_external_writer() { + let (state, _dir) = test_state().await; + let _ = create_generic_store( + State(state.clone()), + Json(CreateGenericStoreRequest { + name: "s1".to_string(), + schema: spec(), + storage_options: None, + seal_on_add: true, + }), + ) + .await + .unwrap(); + let cached = state.get_or_open_generic_store("s1").await.unwrap(); + let mut writer = GenericStore::open_existing( + &state.generic_uri("s1"), + GenericStoreOptions { + shard_id: Some("external-writer".to_string()), + ..Default::default() + }, + ) + .await + .unwrap(); + + writer + .add(&[row( + serde_json::json!({"id": "merged-row", "user": "external"}), + )]) + .await + .unwrap(); + assert_eq!(writer.cleanup_wal().await.unwrap(), 1); + assert_eq!(writer.pending_wal_generations().await.unwrap(), 0); + assert!(writer.version() > cached.read().await.version()); + + let Json(found) = get_row( + State(state), + Path(("s1".to_string(), "merged-row".to_string())), + Query(GetRowQuery { columns: None }), + ) + .await + .unwrap(); + assert_eq!(found["user"], serde_json::json!("external")); + assert_eq!(cached.read().await.version(), writer.version()); + } + #[tokio::test] async fn create_list_add_and_read_round_trip() { let (state, _dir) = test_state().await; diff --git a/crates/lance-context-server/src/routes/records.rs b/crates/lance-context-server/src/routes/records.rs index 2698377..2d21ae3 100644 --- a/crates/lance-context-server/src/routes/records.rs +++ b/crates/lance-context-server/src/routes/records.rs @@ -11,14 +11,101 @@ use lance_context_api::{ UpsertRecordResponse, UpsertRecordsRequest, UpsertRecordsResponse, UpsertResultDto, }; use lance_context_core::{ - patch_from_dto, record_from_add_request, record_to_dto, ContextRecord, LifecycleQueryOptions, - RecordFilters, + patch_from_dto, record_from_add_request, record_to_dto, ContextRecord, ContextStore, + LifecycleQueryOptions, RecordFilters, }; +use tokio::sync::RwLock; use uuid::Uuid; use crate::error::AppError; use crate::state::AppState; +async fn get_context_record_refreshing_on_miss( + store_lock: &RwLock, + id: &str, +) -> Result, AppError> { + { + let store = store_lock.read().await; + let record = store.get(id).await.map_err(AppError::from_lance)?; + if record.is_some() || store.is_version_pinned() { + return Ok(record); + } + } + + let mut store = store_lock.write().await; + if !store.is_version_pinned() { + store.refresh_latest().await.map_err(AppError::from_lance)?; + } + store.get(id).await.map_err(AppError::from_lance) +} + +async fn materialize_context_payload( + store: &ContextStore, + record: Option, + id: &str, +) -> Result)>, AppError> { + let Some(record) = record else { + return Ok(None); + }; + if record.payload_uri.is_none() { + return Err(AppError::InvalidRequest(format!( + "record '{}' has no external payload reference to fetch", + id + ))); + } + let bytes = store + .fetch_payload(id) + .await + .map_err(AppError::from_lance)? + .ok_or_else(|| AppError::NotFound(format!("Record '{}' does not exist", id)))?; + Ok(Some((record, bytes))) +} + +async fn fetch_context_payload_refreshing_on_miss( + store_lock: &RwLock, + id: &str, +) -> Result)>, AppError> { + { + let store = store_lock.read().await; + let record = store.get_by_id(id).await.map_err(AppError::from_lance)?; + if record.is_some() || store.is_version_pinned() { + return materialize_context_payload(&store, record, id).await; + } + } + + let mut store = store_lock.write().await; + if !store.is_version_pinned() { + store.refresh_latest().await.map_err(AppError::from_lance)?; + } + let record = store.get_by_id(id).await.map_err(AppError::from_lance)?; + materialize_context_payload(&store, record, id).await +} + +async fn get_context_by_external_id_refreshing_on_miss( + store_lock: &RwLock, + external_id: &str, +) -> Result, AppError> { + { + let store = store_lock.read().await; + let record = store + .get_by_external_id(external_id) + .await + .map_err(AppError::from_lance)?; + if record.is_some() || store.is_version_pinned() { + return Ok(record); + } + } + + let mut store = store_lock.write().await; + if !store.is_version_pinned() { + store.refresh_latest().await.map_err(AppError::from_lance)?; + } + store + .get_by_external_id(external_id) + .await + .map_err(AppError::from_lance) +} + pub async fn add_records( State(state): State>, Path(name): Path, @@ -219,8 +306,7 @@ pub async fn get_record( ) -> Result, AppError> { let store_lock = state.get_or_open_context_store(&name).await?; - let store = store_lock.read().await; - let record = store.get(&id).await.map_err(AppError::from_lance)?; + let record = get_context_record_refreshing_on_miss(&store_lock, &id).await?; Ok(Json(GetRecordResponse { record: record.map(record_to_dto), @@ -238,22 +324,8 @@ pub async fn fetch_payload( ) -> Result { let store_lock = state.get_or_open_context_store(&name).await?; - let store = store_lock.read().await; - let record = store - .get_by_id(&id) - .await - .map_err(AppError::from_lance)? - .ok_or_else(|| AppError::NotFound(format!("Record '{}' does not exist", id)))?; - if record.payload_uri.is_none() { - return Err(AppError::InvalidRequest(format!( - "record '{}' has no external payload reference to fetch", - id - ))); - } - let bytes = store - .fetch_payload(&id) - .await - .map_err(AppError::from_lance)? + let (record, bytes) = fetch_context_payload_refreshing_on_miss(&store_lock, &id) + .await? .ok_or_else(|| AppError::NotFound(format!("Record '{}' does not exist", id)))?; let content_type = if record.content_type.is_empty() { @@ -279,11 +351,8 @@ pub async fn get_record_by_external_id( ) -> Result, AppError> { let store_lock = state.get_or_open_context_store(&name).await?; - let store = store_lock.read().await; - let record = store - .get_by_external_id(¶ms.external_id) - .await - .map_err(AppError::from_lance)?; + let record = + get_context_by_external_id_refreshing_on_miss(&store_lock, ¶ms.external_id).await?; Ok(Json(GetRecordResponse { record: record.map(record_to_dto), @@ -415,7 +484,7 @@ mod tests { AddRecordRequest, AddRecordsRequest, RecordPatchDto, RelationshipDto, UpdateRecordRequest, UpsertRecordRequest, UpsertRecordsRequest, }; - use lance_context_core::ContextStore; + use lance_context_core::{ContextStore, ContextStoreOptions}; use tempfile::TempDir; use tokio::sync::RwLock; @@ -448,6 +517,54 @@ mod tests { } } + #[tokio::test] + async fn point_reads_refresh_a_base_advanced_by_an_external_writer() { + let context_name = "ctx"; + let (state, _dir) = test_state(context_name).await; + let cached = state.get_or_open_context_store(context_name).await.unwrap(); + let mut writer = ContextStore::open_existing_with_options( + &state.context_uri(context_name), + ContextStoreOptions { + shard_id: Some("external-writer".to_string()), + ..Default::default() + }, + ) + .await + .unwrap(); + + let mut request = text_record("externally merged"); + request.external_id = Some("external-record".to_string()); + let record = record_from_add_request( + &request, + "merged-context-record".to_string(), + "external-run".to_string(), + ); + writer.add(std::slice::from_ref(&record)).await.unwrap(); + assert_eq!(writer.cleanup_wal().await.unwrap(), 1); + assert_eq!(writer.pending_wal_generations().await.unwrap(), 0); + assert!(writer.version() > cached.read().await.version()); + + let Json(found) = get_record( + State(state.clone()), + Path((context_name.to_string(), record.id.clone())), + ) + .await + .unwrap(); + assert_eq!(found.record.unwrap().id, record.id); + assert_eq!(cached.read().await.version(), writer.version()); + + let Json(found) = get_record_by_external_id( + State(state), + Path(context_name.to_string()), + Query(ExternalIdParams { + external_id: "external-record".to_string(), + }), + ) + .await + .unwrap(); + assert_eq!(found.record.unwrap().id, record.id); + } + #[tokio::test] async fn fetch_payload_returns_bytes_404_and_400() { let context_name = "ctx"; diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index d0ebc9d..feab6be 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -478,14 +478,57 @@ pub async fn list_rollouts( })) } +/// Run the normal shared-lock point lookup first, then refresh and retry under +/// the exclusive lock only when it would otherwise return a false miss from a +/// stale cached base dataset. +async fn get_rollout_refreshing_on_miss( + store_lock: &RwLock, + id: &str, +) -> Result, AppError> { + { + let store = store_lock.read().await; + let record = store.get_by_id(id).await.map_err(AppError::from_lance)?; + if record.is_some() || store.is_version_pinned() { + return Ok(record); + } + } + + let mut store = store_lock.write().await; + if !store.is_version_pinned() { + store.refresh_latest().await.map_err(AppError::from_lance)?; + } + store.get_by_id(id).await.map_err(AppError::from_lance) +} + +/// Blob equivalent of [`get_rollout_refreshing_on_miss`]. A successful payload +/// hit stays on the shared-lock fast path; only `None` pays the manifest refresh +/// and retry. +async fn get_rollout_blob_refreshing_on_miss( + store_lock: &RwLock, + id: &str, +) -> Result>, AppError> { + { + let store = store_lock.read().await; + let payload = store.get_blob(id).await.map_err(AppError::from_lance)?; + if payload.is_some() || store.is_version_pinned() { + return Ok(payload); + } + } + + let mut store = store_lock.write().await; + if !store.is_version_pinned() { + store.refresh_latest().await.map_err(AppError::from_lance)?; + } + store.get_blob(id).await.map_err(AppError::from_lance) +} + pub async fn get_rollout( State(state): State>, Path((name, id)): Path<(String, String)>, ) -> Result, AppError> { let store_lock = state.get_or_open_rollout_store(&name).await?; - let store = store_lock.read().await; - let record = store.get_by_id(&id).await.map_err(AppError::from_lance)?; + let record = get_rollout_refreshing_on_miss(&store_lock, &id).await?; Ok(Json(GetRolloutResponse { record: record.map(rollout_record_to_dto), @@ -506,11 +549,8 @@ pub async fn fetch_rollout_blob( ) -> Result { let store_lock = state.get_or_open_rollout_store(&name).await?; - let store = store_lock.read().await; - let bytes = store - .get_blob(&id) - .await - .map_err(AppError::from_lance)? + 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 @@ -518,7 +558,6 @@ pub async fn fetch_rollout_blob( // 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())?; - drop(store); let len = bytes.len(); Response::builder() @@ -766,6 +805,116 @@ mod tests { (state, dir) } + #[tokio::test] + async fn point_lookup_routes_refresh_stale_base_only_after_a_miss() { + let (state, _dir) = rollout_state().await; + let cached = state.get_or_open_rollout_store("rl").await.unwrap(); + let uri = state.rollout_uri("rl"); + let mut writer = RolloutStore::open_existing_with_options( + &uri, + RolloutStoreOptions { + shard_id: Some("external-writer".to_string()), + ..Default::default() + }, + ) + .await + .unwrap(); + + let first = rollout_record_from_add_request(&record_with_size("merged-record", None)); + writer.add(std::slice::from_ref(&first)).await.unwrap(); + writer.flush().await.unwrap(); + assert_eq!(writer.cleanup_own_shard().await.unwrap(), 1); + assert_eq!(writer.pending_wal_generations().await.unwrap(), 0); + assert!(writer.version() > cached.read().await.version()); + + let Json(found) = get_rollout( + State(state.clone()), + Path(("rl".to_string(), first.id.clone())), + ) + .await + .unwrap(); + assert_eq!(found.record.unwrap().id, first.id); + assert_eq!(cached.read().await.version(), writer.version()); + + let cached_version = cached.read().await.version(); + let payload = b"externally-merged-blob".to_vec(); + let mut blob_request = record_with_size("merged-blob", Some(payload.len() as i64)); + blob_request.binary_payload = Some(payload.clone()); + let blob_record = rollout_record_from_add_request(&blob_request); + writer + .add(std::slice::from_ref(&blob_record)) + .await + .unwrap(); + writer.flush().await.unwrap(); + assert_eq!(writer.cleanup_own_shard().await.unwrap(), 1); + assert!(writer.version() > cached_version); + + // A hit in the cached base remains on the shared-lock fast path and + // does not advance the handle just because a newer version exists. + let Json(found) = get_rollout( + State(state.clone()), + Path(("rl".to_string(), first.id.clone())), + ) + .await + .unwrap(); + assert_eq!(found.record.unwrap().id, first.id); + assert_eq!(cached.read().await.version(), cached_version); + + let response = fetch_rollout_blob( + State(state.clone()), + Path(("rl".to_string(), blob_record.id)), + ) + .await + .unwrap(); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + assert_eq!(bytes.as_ref(), payload.as_slice()); + assert_eq!(cached.read().await.version(), writer.version()); + } + + #[tokio::test] + async fn point_lookup_miss_preserves_explicit_checkout() { + let (state, _dir) = rollout_state().await; + let cached = state.get_or_open_rollout_store("rl").await.unwrap(); + let pinned_version = cached.read().await.version(); + let uri = state.rollout_uri("rl"); + let mut writer = RolloutStore::open_existing_with_options( + &uri, + RolloutStoreOptions { + shard_id: Some("external-writer".to_string()), + ..Default::default() + }, + ) + .await + .unwrap(); + let record = rollout_record_from_add_request(&record_with_size("future-record", None)); + writer.add(std::slice::from_ref(&record)).await.unwrap(); + writer.flush().await.unwrap(); + assert_eq!(writer.cleanup_own_shard().await.unwrap(), 1); + assert!(writer.version() > pinned_version); + + let Json(version) = checkout_rollout( + State(state.clone()), + Path("rl".to_string()), + Json(CheckoutRequest { + version: pinned_version, + }), + ) + .await + .unwrap(); + assert_eq!(version.version, pinned_version); + assert!(cached.read().await.is_version_pinned()); + + let Json(missing) = get_rollout(State(state), Path(("rl".to_string(), record.id))) + .await + .unwrap(); + assert!(missing.record.is_none()); + let cached = cached.read().await; + assert_eq!(cached.version(), pinned_version); + assert!(cached.is_version_pinned()); + } + /// Build a JSON append request for `add_rollouts`, with an optional query. fn append_request(query: &str, id: &str) -> Request { let body = serde_json::json!({