fix: honor registered rollout dataset locations

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"] }

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()),
     }))
 }

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::*;

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<RegistryEntry>,
+    reads: std::sync::atomic::AtomicUsize,
+}
+
+#[async_trait::async_trait]
+impl StoreRegistry for ChangingRegistry {
+    async fn get(&self, _: &str) -> LanceResult<Option<RegistryEntry>> {
+        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<bool> {
+        unreachable!()
+    }
+    async fn list(&self) -> LanceResult<Vec<RegistryEntry>> {
+        unreachable!()
+    }
+    async fn upsert(&self, _: &str, _: &str) -> LanceResult<()> {
+        unreachable!()
+    }
+    async fn remove(&self, _: &str) -> LanceResult<()> {
+        unreachable!()
+    }
+    async fn insert_missing(&self, _: &[(String, String)]) -> LanceResult<usize> {
+        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());
+    }
+}