diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index bcb53c50..ca96c778 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -176,12 +176,37 @@ impl RolloutStore { } /// Open a rollout dataset with explicit storage and shard configuration. + /// Creates the dataset if it does not exist. pub async fn open_with_options(uri: &str, options: RolloutStoreOptions) -> LanceResult { + Self::open_inner(uri, options, true).await + } + + /// Open an **existing** rollout dataset. Unlike [`Self::open_with_options`], + /// this does **not** create the dataset when it is absent — it returns the + /// underlying [`LanceError::DatasetNotFound`] instead. + /// + /// This is the read/write path used by the server's lazy cache-fill: a + /// cache miss must load a store that genuinely exists on object storage + /// (create-on-absence would silently materialize an empty table for a + /// mistyped name, masking the 404). Store creation goes exclusively through + /// [`Self::open_with_options`] on the explicit `create` route. + pub async fn open_existing_with_options( + uri: &str, + options: RolloutStoreOptions, + ) -> LanceResult { + Self::open_inner(uri, options, false).await + } + + async fn open_inner( + uri: &str, + options: RolloutStoreOptions, + create_if_missing: bool, + ) -> LanceResult { let storage_options = options.storage_options.clone(); let write_shard = derive_shard_id(options.shard_id.as_deref()); let dataset = match Self::load_with_options(uri, storage_options.clone()).await { Ok(dataset) => dataset, - Err(LanceError::DatasetNotFound { .. }) => { + Err(LanceError::DatasetNotFound { .. }) if create_if_missing => { Self::create_with_options(uri, storage_options.clone()).await? } Err(err) => return Err(err), diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index ce67ec2c..5344ade4 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -337,7 +337,31 @@ impl ContextStore { } /// Open a dataset with explicit object store configuration (e.g. S3 credentials). + /// Creates the dataset if it does not exist. pub async fn open_with_options(uri: &str, options: ContextStoreOptions) -> LanceResult { + Self::open_inner(uri, options, true).await + } + + /// Open an **existing** context dataset. Unlike [`Self::open_with_options`], + /// this does **not** create the dataset when it is absent — it returns the + /// underlying [`LanceError::DatasetNotFound`] instead. + /// + /// Used by the server's lazy cache-fill path: a cache miss must load a store + /// that genuinely exists on object storage. Create-on-absence would silently + /// materialize an empty context for a mistyped name, masking the 404. Store + /// creation goes exclusively through [`Self::open_with_options`]. + pub async fn open_existing_with_options( + uri: &str, + options: ContextStoreOptions, + ) -> LanceResult { + Self::open_inner(uri, options, false).await + } + + async fn open_inner( + uri: &str, + options: ContextStoreOptions, + create_if_missing: bool, + ) -> LanceResult { // Validate blob_columns for col in &options.blob_columns { if !VALID_BLOB_COLUMNS.contains(&col.as_str()) { @@ -359,7 +383,7 @@ impl ContextStore { let blob_columns = options.blob_columns.clone(); let (dataset, created) = match Self::load_with_options(uri, storage_options.clone()).await { Ok(dataset) => (dataset, false), - Err(LanceError::DatasetNotFound { .. }) => { + Err(LanceError::DatasetNotFound { .. }) if create_if_missing => { let dataset = Self::create_with_options( uri, storage_options.clone(), diff --git a/crates/lance-context-server/src/routes/compact.rs b/crates/lance-context-server/src/routes/compact.rs index cfd7588c..0c2c4a90 100644 --- a/crates/lance-context-server/src/routes/compact.rs +++ b/crates/lance-context-server/src/routes/compact.rs @@ -13,12 +13,7 @@ pub async fn compact( Path(name): Path, Json(req): Json, ) -> Result, AppError> { - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let config = if req.target_rows_per_fragment.is_some() || req.materialize_deletions.is_some() { let mut c = CompactionConfig::default(); @@ -48,12 +43,7 @@ pub async fn compact_stats( State(state): State>, Path(name): Path, ) -> Result, AppError> { - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let store = store_lock.read().await; let stats = store diff --git a/crates/lance-context-server/src/routes/contexts.rs b/crates/lance-context-server/src/routes/contexts.rs index 56ece0ba..fbc45389 100644 --- a/crates/lance-context-server/src/routes/contexts.rs +++ b/crates/lance-context-server/src/routes/contexts.rs @@ -93,10 +93,7 @@ pub async fn get_context( State(state): State>, Path(name): Path, ) -> Result, AppError> { - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))?; + let store_lock = state.get_or_open_context_store(&name).await?; let store = store_lock.read().await; Ok(Json(ContextInfo { diff --git a/crates/lance-context-server/src/routes/records.rs b/crates/lance-context-server/src/routes/records.rs index c20a9ac8..98a453c5 100644 --- a/crates/lance-context-server/src/routes/records.rs +++ b/crates/lance-context-server/src/routes/records.rs @@ -32,12 +32,7 @@ pub async fn add_records( )); } - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let run_id = Uuid::new_v4().to_string(); let mut ids = Vec::with_capacity(req.records.len()); @@ -83,12 +78,7 @@ pub async fn upsert_record( )); } - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let record = record_from_add_request( &req.record, @@ -141,12 +131,7 @@ pub async fn upsert_records( } } - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let core_records: Vec = req .records @@ -191,12 +176,7 @@ pub async fn update_record( )); } - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let patch = patch_from_dto(&req.patch); let mut store = store_lock.write().await; @@ -236,12 +216,7 @@ pub async fn get_record( State(state): State>, Path((name, id)): Path<(String, String)>, ) -> Result, AppError> { - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + 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)?; @@ -260,12 +235,7 @@ pub async fn fetch_payload( State(state): State>, Path((name, id)): Path<(String, String)>, ) -> Result { - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let store = store_lock.read().await; let record = store @@ -306,12 +276,7 @@ pub async fn get_record_by_external_id( Path(name): Path, Query(params): Query, ) -> Result, AppError> { - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let store = store_lock.read().await; let record = store @@ -328,12 +293,7 @@ pub async fn delete_record( State(state): State>, Path((name, id)): Path<(String, String)>, ) -> Result, AppError> { - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let mut store = store_lock.write().await; let deleted = store @@ -350,12 +310,7 @@ pub async fn delete_record_by_external_id( Path(name): Path, Query(params): Query, ) -> Result, AppError> { - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let mut store = store_lock.write().await; let deleted = store @@ -396,12 +351,7 @@ pub async fn list_records( }) .transpose()?; - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let store = store_lock.read().await; let records = store @@ -435,12 +385,7 @@ pub async fn related_records( Path(name): Path, Query(params): Query, ) -> Result, AppError> { - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let store = store_lock.read().await; let records = store @@ -1255,4 +1200,53 @@ mod tests { .await; assert!(matches!(result, Err(AppError::InvalidRequest(_)))); } + + /// A second server instance (empty cache, shared data dir) must lazily open + /// a context created elsewhere instead of 404ing — the multi-replica bug. + #[tokio::test] + async fn second_instance_lazily_opens_context_created_elsewhere() { + let context_name = "ctx"; + let (state_a, dir) = test_state(context_name).await; + + // Instance A writes a record. + let _ = add_records( + State(state_a.clone()), + Path(context_name.to_string()), + Json(AddRecordsRequest { + records: vec![text_record("hello")], + }), + ) + .await + .expect("write on instance A"); + + // Instance B: fresh AppState over the same data dir, empty cache. + let state_b = Arc::new(AppState { + stores: RwLock::new(HashMap::new()), + rollout_stores: RwLock::new(HashMap::new()), + base_path: dir.path().to_path_buf(), + instance_id: None, + rollout_merge_after_generations: 0, + rollout_cleanup_interval_secs: 0, + rollout_cleanup_min_generations: 1, + }); + + let Json(list) = list_records( + State(state_b.clone()), + Path(context_name.to_string()), + Query(ListParams::default()), + ) + .await + .expect("instance B lazily opens the context instead of 404"); + assert_eq!(list.records.len(), 1); + + // Reading a context that was never created anywhere still 404s. + let err = list_records( + State(state_b), + Path("no-such-context".to_string()), + Query(ListParams::default()), + ) + .await + .unwrap_err(); + assert!(matches!(err, AppError::NotFound(_))); + } } diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index b74f8309..54d41a4f 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -96,10 +96,7 @@ pub async fn get_rollout_store( State(state): State>, Path(name): Path, ) -> Result, AppError> { - let stores = state.rollout_stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))?; + let store_lock = state.get_or_open_rollout_store(&name).await?; let store = store_lock.read().await; Ok(Json(RolloutStoreInfo { @@ -180,12 +177,7 @@ pub async fn add_rollouts( )); } - let stores = state.rollout_stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_rollout_store(&name).await?; let ids: Vec = records.iter().map(|r| r.id.clone()).collect(); let core_records: Vec = records @@ -310,12 +302,7 @@ pub async fn list_rollouts( Path(name): Path, Query(params): Query, ) -> Result, AppError> { - let stores = state.rollout_stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_rollout_store(&name).await?; let store = store_lock.read().await; let records = store @@ -332,12 +319,7 @@ pub async fn get_rollout( State(state): State>, Path((name, id)): Path<(String, String)>, ) -> Result, AppError> { - let stores = state.rollout_stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))? - .clone(); - drop(stores); + 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)?; @@ -354,12 +336,7 @@ pub async fn fetch_rollout_blob( State(state): State>, Path((name, id)): Path<(String, String)>, ) -> Result { - let stores = state.rollout_stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_rollout_store(&name).await?; let store = store_lock.read().await; let bytes = store @@ -379,12 +356,7 @@ pub async fn checkout_rollout( Path(name): Path, Json(req): Json, ) -> Result, AppError> { - let stores = state.rollout_stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_rollout_store(&name).await?; let mut store = store_lock.write().await; store @@ -408,12 +380,7 @@ pub async fn compact_rollout( Path(name): Path, Json(req): Json, ) -> Result, AppError> { - let stores = state.rollout_stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_rollout_store(&name).await?; let config = if req.target_rows_per_fragment.is_some() || req.materialize_deletions.is_some() { let mut c = CompactionConfig::default(); @@ -444,12 +411,7 @@ pub async fn compact_rollout_stats( State(state): State>, Path(name): Path, ) -> Result, AppError> { - let stores = state.rollout_stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Rollout store '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_rollout_store(&name).await?; let store = store_lock.read().await; let stats = store.compaction_stats(); @@ -819,4 +781,76 @@ mod tests { .unwrap_err(); assert!(matches!(err, AppError::InvalidRequest(_))); } + + /// A second server instance (distinct in-memory cache, shared data dir) + /// that never `create`d the store must still serve reads/writes for it by + /// lazily loading from storage — the multi-replica 404 regression. + #[tokio::test] + async fn second_instance_lazily_opens_store_created_elsewhere() { + // Pod A: create the store and write a row. + let (state_a, dir) = rollout_state().await; + let body = serde_json::to_vec(&AddRolloutsRequest { + records: vec![record_with_size("r0", None)], + }) + .unwrap(); + let _ = add_rollouts( + State(state_a.clone()), + Path("rl".to_string()), + json_request(body), + ) + .await + .expect("write on instance A"); + + // Pod B: a fresh AppState over the SAME data dir, with an empty cache and + // a different instance id (its own shard). It never saw the `create`. + let state_b = Arc::new(AppState { + stores: RwLock::new(HashMap::new()), + rollout_stores: RwLock::new(HashMap::new()), + base_path: dir.path().to_path_buf(), + instance_id: Some("rollout-1".to_string()), + rollout_merge_after_generations: 0, + rollout_cleanup_interval_secs: 0, + rollout_cleanup_min_generations: 1, + }); + + // Read routed to B must lazily open the store, not 404. + let Json(list) = list_rollouts( + State(state_b.clone()), + Path("rl".to_string()), + Query(RolloutListParams::default()), + ) + .await + .expect("instance B lazily opens the store instead of 404"); + assert_eq!(list.records.len(), 1); + assert_eq!(list.records[0].id, "r0"); + + // And a write on B lands in the same dataset (visible back on A). + let body = serde_json::to_vec(&AddRolloutsRequest { + records: vec![record_with_size("r1", None)], + }) + .unwrap(); + let _ = add_rollouts( + State(state_b.clone()), + Path("rl".to_string()), + json_request(body), + ) + .await + .expect("write on instance B"); + assert_eq!(count_rollouts(&state_a).await, 2); + } + + /// A read for a store that was never created anywhere must still 404 — lazy + /// open must not silently materialize an empty dataset for a bad name. + #[tokio::test] + async fn read_of_nonexistent_store_still_404s() { + let (state, _dir) = rollout_state().await; + let err = list_rollouts( + State(state.clone()), + Path("no-such-store".to_string()), + Query(RolloutListParams::default()), + ) + .await + .unwrap_err(); + assert!(matches!(err, AppError::NotFound(_))); + } } diff --git a/crates/lance-context-server/src/routes/search.rs b/crates/lance-context-server/src/routes/search.rs index 2c19bd31..e8bbe25a 100644 --- a/crates/lance-context-server/src/routes/search.rs +++ b/crates/lance-context-server/src/routes/search.rs @@ -24,12 +24,7 @@ pub async fn search( .transpose() .map_err(AppError::InvalidRequest)?; - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let store = store_lock.read().await; let results = store @@ -76,12 +71,7 @@ pub async fn retrieve( .transpose() .map_err(AppError::InvalidRequest)?; - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let store = store_lock.read().await; let results = store diff --git a/crates/lance-context-server/src/routes/versions.rs b/crates/lance-context-server/src/routes/versions.rs index 76af8aa8..7bb06d28 100644 --- a/crates/lance-context-server/src/routes/versions.rs +++ b/crates/lance-context-server/src/routes/versions.rs @@ -11,12 +11,7 @@ pub async fn get_version( State(state): State>, Path(name): Path, ) -> Result, AppError> { - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let store = store_lock.read().await; Ok(Json(VersionResponse { @@ -29,12 +24,7 @@ pub async fn checkout( Path(name): Path, Json(req): Json, ) -> Result, AppError> { - let stores = state.stores.read().await; - let store_lock = stores - .get(&name) - .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? - .clone(); - drop(stores); + let store_lock = state.get_or_open_context_store(&name).await?; let mut store = store_lock.write().await; store diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index aec854d6..5049e6a8 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -2,10 +2,11 @@ use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; -use lance_context_core::{ContextStore, RolloutStore}; +use lance_context_core::{ContextStore, ContextStoreOptions, RolloutStore, RolloutStoreOptions}; use tokio::sync::RwLock; use crate::config::ServerConfig; +use crate::error::AppError; pub struct AppState { pub stores: RwLock>>>, @@ -56,4 +57,86 @@ impl AppState { .to_string_lossy() .to_string() } + + /// Look up a rollout store by name, lazily loading it from object storage on + /// a local cache miss. + /// + /// The server caches each opened store in this process's memory. In a + /// multi-replica deployment a store `create`d on pod A only enters A's map, + /// so a read/write that the load balancer routes to pod B would otherwise + /// 404 even though the dataset exists on shared object storage. This helper + /// closes that gap: on a miss it *loads* (never creates — see + /// [`RolloutStore::open_existing_with_options`], which returns a + /// `DatasetNotFound`/404 for a genuinely-absent store rather than silently + /// materializing an empty table) and caches the handle for subsequent hits. + pub async fn get_or_open_rollout_store( + &self, + name: &str, + ) -> Result>, AppError> { + // Fast path: already cached in this process. + if let Some(store) = self.rollout_stores.read().await.get(name) { + return Ok(store.clone()); + } + + // Slow path: load from object storage WITHOUT holding the map lock, so a + // slow open does not block other stores' requests. + let uri = self.rollout_uri(name); + let options = RolloutStoreOptions { + // No request body on the read path: object-store credentials come + // from the pod's workload-identity environment, exactly as they do + // for the `create` route in production. + storage_options: None, + shard_id: self.instance_id.clone(), + merge_after_generations: (self.rollout_merge_after_generations > 0) + .then_some(self.rollout_merge_after_generations), + cleanup_interval_secs: (self.rollout_cleanup_interval_secs > 0) + .then_some(self.rollout_cleanup_interval_secs), + cleanup_min_generations: Some(self.rollout_cleanup_min_generations), + }; + let opened = RolloutStore::open_existing_with_options(&uri, options) + .await + .map_err(AppError::from_lance)?; + let opened = Arc::new(RwLock::new(opened)); + + // Insert under the write lock, re-checking for a store another request + // may have opened concurrently while we were loading. + let mut stores = self.rollout_stores.write().await; + if let Some(existing) = stores.get(name) { + return Ok(existing.clone()); + } + // A lazily-opened store must also get the background WAL-cleanup timer + // that the `create` route starts, or these stores would never merge. + let _cleanup = RolloutStore::spawn_periodic_cleanup(opened.clone()); + stores.insert(name.to_string(), opened.clone()); + Ok(opened) + } + + /// Look up a context store by name, lazily loading it from object storage on + /// a local cache miss. See [`Self::get_or_open_rollout_store`] for the + /// multi-replica rationale; the same in-memory-cache-as-existence-check bug + /// affects context stores. + pub async fn get_or_open_context_store( + &self, + name: &str, + ) -> Result>, AppError> { + if let Some(store) = self.stores.read().await.get(name) { + return Ok(store.clone()); + } + + let uri = self.context_uri(name); + // Load-only: existing schema (embedding dim, blob columns, distance + // metric) is read from the persisted dataset, so no create-time options + // are needed. Credentials come from the pod environment. + let opened = ContextStore::open_existing_with_options(&uri, ContextStoreOptions::default()) + .await + .map_err(AppError::from_lance)?; + let opened = Arc::new(RwLock::new(opened)); + + let mut stores = self.stores.write().await; + if let Some(existing) = stores.get(name) { + return Ok(existing.clone()); + } + stores.insert(name.to_string(), opened.clone()); + Ok(opened) + } }