diff --git a/Cargo.lock b/Cargo.lock index f497b537..307d850b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4152,6 +4152,7 @@ dependencies = [ name = "lance-context-server" version = "0.6.7" dependencies = [ + "async-trait", "axum", "bytes", "chrono", diff --git a/crates/lance-context-server/Cargo.toml b/crates/lance-context-server/Cargo.toml index 169ec7c9..a3f6b07c 100644 --- a/crates/lance-context-server/Cargo.toml +++ b/crates/lance-context-server/Cargo.toml @@ -34,6 +34,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] } uuid = { version = "1", features = ["v4"] } [dev-dependencies] +async-trait = "0.1.92" tokio = { version = "1", features = ["test-util", "io-util"] } lance = "13.0.0" tower = { version = "0.5", features = ["util"] } diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 2b39f7ed..3b2f92c6 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -258,8 +258,8 @@ pub async fn get_rollout_store( let store = store_lock.read().await; Ok(Json(RolloutStoreInfo { - name: name.clone(), - uri: state.rollout_uri(&name), + name, + uri: store.uri().to_string(), version: Some(store.version()), })) } diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index 00d5984b..cebd07c5 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -554,23 +554,20 @@ impl AppState { // Keep the registry entry and cached handle until physical deletion // succeeds, so failures can be retried without losing the store's URI. let registry = &self.rollout_registry; - if !registry - .contains(name) - .await - .map_err(AppError::from_lance)? - { + let Some(entry) = registry.get(name).await.map_err(AppError::from_lance)? else { return Ok(false); - } + }; let live = handle.upgrade(); if let Some(store) = live { - store - .write() - .await - .delete_data() - .await - .map_err(AppError::from_lance)?; + let mut store = store.write().await; + if store.uri() != entry.uri { + return Err(AppError::Overloaded(format!( + "Rollout store '{name}' registry changed; refusing to delete a different dataset" + ))); + } + store.delete_data().await.map_err(AppError::from_lance)?; } else { - lance_context_core::remove_dataset(&self.rollout_uri(name)) + lance_context_core::remove_dataset(&entry.uri) .await .map_err(AppError::from_lance)?; } @@ -610,36 +607,40 @@ impl AppState { return Ok(store); } - // Existence is the registry's job, not the cache's. - let exists = self + // The registry owns both existence and the physical dataset location. + let entry = self .rollout_registry - .contains(name) + .get(name) .await - .map_err(AppError::from_lance)?; - if !exists { - return Err(AppError::NotFound(format!( - "Rollout store '{}' does not exist", - name - ))); - } + .map_err(AppError::from_lance)? + .ok_or_else(|| AppError::NotFound(format!("Rollout store '{name}' does not exist")))?; // Load from storage WITHOUT holding the LRU lock, so a slow open does // not block other stores' requests. - let uri = self.rollout_uri(name); - let opened = RolloutStore::open_existing_with_options(&uri, self.rollout_store_options()) - .await - .map_err(AppError::from_lance)?; + let opened = + RolloutStore::open_existing_with_options(&entry.uri, self.rollout_store_options()) + .await + .map_err(AppError::from_lance)?; let opened = Arc::new(RwLock::new(opened)); - // Also observe registry removal by another server during the open. - let registry = &self.rollout_registry; - if !registry - .contains(name) + // Reject a removal or replacement observed during this open. This is + // not a distributed lease: cached handles still need lifecycle fencing. + match self + .rollout_registry + .get(name) .await .map_err(AppError::from_lance)? { - return Err(AppError::NotFound(format!( - "Store '{name}' was deleted while opening" - ))); + Some(current) if current == entry => {} + Some(_) => { + return Err(AppError::Overloaded(format!( + "Rollout store '{name}' registry changed while opening; retry" + ))); + } + None => { + return Err(AppError::NotFound(format!( + "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. @@ -1166,6 +1167,9 @@ impl AppState { } } +#[cfg(test)] +mod registry_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/crates/lance-context-server/src/state/registry_tests.rs b/crates/lance-context-server/src/state/registry_tests.rs new file mode 100644 index 00000000..5776b3e1 --- /dev/null +++ b/crates/lance-context-server/src/state/registry_tests.rs @@ -0,0 +1,191 @@ +use super::*; +use axum::extract::{Path, State}; +use lance::Result as LanceResult; +use lance_context_core::RegistryEntry; +use tempfile::TempDir; + +use crate::routes::rollouts::get_rollout_store; + +async fn create_at(uri: &str) { + RolloutStore::open_with_options(uri, RolloutStoreOptions::default()) + .await + .unwrap() + .close() + .await + .unwrap(); +} + +#[tokio::test] +async fn registered_uri_controls_open_metadata_and_uncached_deletion() { + let dir = TempDir::new().unwrap(); + let state = Arc::new(AppState::new_for_test(dir.path().to_path_buf()).await); + let default_uri = state.rollout_uri("logical"); + let registered_uri = dir.path().join("physical-generation.lance"); + let registered_uri = registered_uri.to_str().unwrap(); + // A valid dataset at the derived path must not shadow the registered one. + create_at(&default_uri).await; + create_at(registered_uri).await; + state + .rollout_registry + .upsert("logical", registered_uri) + .await + .unwrap(); + + let opened = state.get_or_open_rollout_store("logical").await.unwrap(); + assert_eq!(opened.read().await.uri(), registered_uri); + let info = get_rollout_store(State(state.clone()), Path("logical".into())) + .await + .unwrap() + .0; + assert_eq!(info.uri, registered_uri); + opened.write().await.close().await.unwrap(); + drop(opened); + state.rollout_stores.lock().await.pop("logical"); + assert!(state + .rollout_handles + .lock("logical") + .await + .upgrade() + .is_none()); + + assert!(state.unregister_rollout("logical").await.unwrap()); + assert!(!std::path::Path::new(registered_uri).exists()); + assert!(std::path::Path::new(&default_uri).exists()); + assert!(!state.rollout_registry.contains("logical").await.unwrap()); + assert!(matches!( + state.get_or_open_rollout_store("logical").await, + Err(AppError::NotFound(_)) + )); +} + +#[tokio::test] +async fn registered_uri_mismatch_does_not_delete_live_or_replacement_data() { + let dir = TempDir::new().unwrap(); + let state = AppState::new_for_test(dir.path().to_path_buf()).await; + let first = dir.path().join("first.lance").to_str().unwrap().to_string(); + let second = dir + .path() + .join("second.lance") + .to_str() + .unwrap() + .to_string(); + create_at(&first).await; + create_at(&second).await; + state + .rollout_registry + .upsert("logical", &first) + .await + .unwrap(); + let opened = state.get_or_open_rollout_store("logical").await.unwrap(); + state + .rollout_registry + .upsert("logical", &second) + .await + .unwrap(); + assert!(matches!( + state.unregister_rollout("logical").await, + Err(AppError::Overloaded(_)) + )); + assert!(std::path::Path::new(&first).exists()); + assert!(std::path::Path::new(&second).exists()); + assert_eq!( + state + .rollout_registry + .get("logical") + .await + .unwrap() + .unwrap() + .uri, + second + ); + assert!(Arc::ptr_eq( + state.rollout_stores.lock().await.peek("logical").unwrap(), + &opened + )); + state.shutdown().await; +} + +// Serve a changed registry entry at the post-open observation. Only registry +// timing is controlled; each open still uses a real Lance dataset. +struct ChangingRegistry { + initial: RegistryEntry, + current: Option, + reads: std::sync::atomic::AtomicUsize, +} + +#[async_trait::async_trait] +impl StoreRegistry for ChangingRegistry { + async fn get(&self, _: &str) -> LanceResult> { + if self.reads.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 { + Ok(Some(self.initial.clone())) + } else { + Ok(self.current.clone()) + } + } + async fn contains(&self, _: &str) -> LanceResult { + unreachable!() + } + async fn list(&self) -> LanceResult> { + unreachable!() + } + async fn upsert(&self, _: &str, _: &str) -> LanceResult<()> { + unreachable!() + } + async fn remove(&self, _: &str) -> LanceResult<()> { + unreachable!() + } + async fn insert_missing(&self, _: &[(String, String)]) -> LanceResult { + unreachable!() + } + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +#[tokio::test] +async fn registry_change_during_open_never_installs_a_stale_handle() { + let dir = TempDir::new().unwrap(); + let mut state = AppState::new_for_test(dir.path().to_path_buf()).await; + let uri = dir.path().join("first.lance").to_str().unwrap().to_string(); + create_at(&uri).await; + let initial = RegistryEntry { + name: "logical".into(), + uri, + created_at: 1, + }; + let moved = RegistryEntry { + uri: dir + .path() + .join("second.lance") + .to_str() + .unwrap() + .to_string(), + ..initial.clone() + }; + create_at(&moved.uri).await; + let recreated = RegistryEntry { + created_at: 2, + ..initial.clone() + }; + for current in [None, Some(moved), Some(recreated)] { + let removed = current.is_none(); + state.rollout_registry = Arc::new(ChangingRegistry { + initial: initial.clone(), + current, + reads: Default::default(), + }); + let result = state.get_or_open_rollout_store("logical").await; + if removed { + assert!(matches!(result, Err(AppError::NotFound(_)))); + } else { + assert!(matches!(result, Err(AppError::Overloaded(_)))); + } + assert!(state.rollout_stores.lock().await.peek("logical").is_none()); + assert!(state + .rollout_handles + .lock("logical") + .await + .upgrade() + .is_none()); + } +}