From 59f506111269d65b42d9228849f74e0cf4c66f85 Mon Sep 17 00:00:00 2001 From: Beinan Date: Tue, 14 Jul 2026 21:32:03 +0000 Subject: [PATCH] fix(registry): reload latest manifest before access --- crates/lance-context-core/src/registry.rs | 46 +++++++++++++++---- crates/lance-context-master/src/routes.rs | 26 +++++++++-- crates/lance-context-master/src/scanner.rs | 2 +- .../src/routes/rollouts.rs | 4 +- crates/lance-context-server/src/state.rs | 7 +-- 5 files changed, 66 insertions(+), 19 deletions(-) diff --git a/crates/lance-context-core/src/registry.rs b/crates/lance-context-core/src/registry.rs index ce20852a..452a7ed1 100644 --- a/crates/lance-context-core/src/registry.rs +++ b/crates/lance-context-core/src/registry.rs @@ -41,9 +41,10 @@ pub struct RegistryEntry { /// Durable directory of rollout stores, backed by a single Lance dataset. /// -/// All mutations (`upsert`, `remove`) take `&mut self` and are expected to be -/// serialized by the caller (the server wraps this in a lock), so the registry -/// itself does no internal locking. Reads (`list`, `contains`) take `&self`. +/// All operations take `&mut self` because Lance dataset handles are snapshots: +/// reads and writes first check out the latest manifest so commits made by +/// another process are visible. Callers are expected to serialize access (the +/// server and master wrap this in a lock). pub struct RolloutRegistry { dataset: Dataset, uri: String, @@ -132,6 +133,7 @@ impl RolloutRegistry { /// `create` retried after a crash (dataset on disk, registry row possibly /// present) converges to exactly one row. Callers must serialize mutations. pub async fn upsert(&mut self, name: &str, uri: &str) -> LanceResult<()> { + self.reload().await?; self.delete_row(name).await?; let schema = Arc::new(registry_schema()); let batch = RecordBatch::try_new( @@ -150,6 +152,7 @@ impl RolloutRegistry { /// Remove the directory entry for `name`, if present. No-op when absent. pub async fn remove(&mut self, name: &str) -> LanceResult<()> { + self.reload().await?; self.delete_row(name).await } @@ -161,8 +164,14 @@ impl RolloutRegistry { Ok(()) } - /// Whether a store named `name` exists in the directory. - pub async fn contains(&self, name: &str) -> LanceResult { + /// Refresh this handle to the registry's latest committed version. + pub async fn reload(&mut self) -> LanceResult<()> { + self.dataset.checkout_latest().await + } + + /// Whether a store named `name` exists in the latest registry version. + pub async fn contains(&mut self, name: &str) -> LanceResult { + self.reload().await?; let escaped = name.replace('\'', "''"); let mut scanner = self.dataset.scan(); scanner.project(&["name"])?; @@ -177,10 +186,12 @@ impl RolloutRegistry { Ok(false) } - /// All directory entries, ordered as stored (unspecified). The registry is a - /// narrow three-column table, so even hundreds of thousands of rows scan - /// quickly; pagination can be layered on later if needed. - pub async fn list(&self) -> LanceResult> { + /// All directory entries from the latest registry version, ordered as + /// stored (unspecified). The registry is a narrow three-column table, so + /// even hundreds of thousands of rows scan quickly; pagination can be + /// layered on later if needed. + pub async fn list(&mut self) -> LanceResult> { + self.reload().await?; let mut scanner = self.dataset.scan(); scanner.project(&["name", "uri", "created_at"])?; let mut stream = scanner.try_into_stream().await?; @@ -304,7 +315,22 @@ mod tests { .unwrap(); } // Reopen the same path: the entry must still be present. - let reg = new_registry(&dir).await; + let mut reg = new_registry(&dir).await; assert!(reg.contains("persist").await.unwrap()); } + + #[tokio::test] + async fn reads_see_commits_from_another_handle() { + let dir = TempDir::new().unwrap(); + let mut reader = new_registry(&dir).await; + let mut writer = new_registry(&dir).await; + + writer + .upsert("external", "/data/external.rollout.lance") + .await + .unwrap(); + + assert!(reader.contains("external").await.unwrap()); + assert_eq!(reader.list().await.unwrap()[0].name, "external"); + } } diff --git a/crates/lance-context-master/src/routes.rs b/crates/lance-context-master/src/routes.rs index 2c5b3b7a..0371b4aa 100644 --- a/crates/lance-context-master/src/routes.rs +++ b/crates/lance-context-master/src/routes.rs @@ -68,7 +68,7 @@ pub async fn get_experiment( // side effect, then read it back. let entries = state .registry - .read() + .write() .await .list() .await @@ -114,7 +114,7 @@ pub async fn compact_experiment( // Only enqueue known experiments. let exists = state .registry - .read() + .write() .await .contains(&name) .await @@ -159,7 +159,7 @@ pub fn api_router() -> Router> { mod tests { use super::*; use crate::config::MasterConfig; - use lance_context_core::RolloutStore; + use lance_context_core::{RolloutRegistry, RolloutStore}; use tempfile::TempDir; fn test_config(dir: &TempDir) -> MasterConfig { @@ -229,6 +229,26 @@ mod tests { assert_eq!(one.experiments[0].name, "exp-1"); } + #[tokio::test] + async fn scan_sees_registry_commits_from_another_handle() { + let dir = TempDir::new().unwrap(); + let state = MasterState::new(test_config(&dir)).await.unwrap(); + let name = "external"; + let uri = state.rollout_uri(name); + RolloutStore::open(&uri).await.unwrap(); + + let registry_uri = dir.path().join("_registry.rollout.lance"); + let mut worker_registry = + RolloutRegistry::open_or_create(registry_uri.to_str().unwrap(), None) + .await + .unwrap(); + worker_registry.upsert(name, &uri).await.unwrap(); + + let scanned = scanner::scan_once(&state).await.unwrap(); + assert_eq!(scanned, 1); + assert!(state.stats.lock().await.get(name).await.unwrap().is_some()); + } + #[tokio::test] async fn scan_reconciles_removed_experiments() { let dir = TempDir::new().unwrap(); diff --git a/crates/lance-context-master/src/scanner.rs b/crates/lance-context-master/src/scanner.rs index d4eaa570..ddaa3a78 100644 --- a/crates/lance-context-master/src/scanner.rs +++ b/crates/lance-context-master/src/scanner.rs @@ -27,7 +27,7 @@ const OBSERVE_TIMEOUT: Duration = Duration::from_secs(30); /// experiments successfully observed. pub async fn scan_once(state: &Arc) -> lance::Result { let scan_start = std::time::Instant::now(); - let entries = state.registry.read().await.list().await?; + let entries = state.registry.write().await.list().await?; let live: HashSet = entries.iter().map(|e| e.name.clone()).collect(); let concurrency = state.config.scan_concurrency.max(1); diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 80a9a979..ccdae359 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -33,7 +33,7 @@ pub async fn create_rollout_store( // Existence is tracked durably in the registry, not by cache membership. if state .rollout_registry - .read() + .write() .await .contains(&req.name) .await @@ -85,7 +85,7 @@ pub async fn list_rollout_stores( // opening every dataset. let entries = state .rollout_registry - .read() + .write() .await .list() .await diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index 70e74b74..913dd8c6 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -31,7 +31,8 @@ pub struct AppState { pub rollout_stores: Mutex>>>, /// 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 mutations take `&mut`. + /// because every operation refreshes the snapshot and therefore takes + /// `&mut`. pub rollout_registry: RwLock, pub base_path: PathBuf, /// Stable identity of this server instance, used as the MemWAL shard key for @@ -171,7 +172,7 @@ impl AppState { pub async fn unregister_rollout(&self, name: &str) -> Result { let existed = self .rollout_registry - .read() + .write() .await .contains(name) .await @@ -212,7 +213,7 @@ impl AppState { // Existence is the registry's job, not the cache's. let exists = self .rollout_registry - .read() + .write() .await .contains(name) .await