From 71b9496f2706bbc8ae9f08d20b17678ae5984c28 Mon Sep 17 00:00:00 2001 From: Beinan Date: Tue, 14 Jul 2026 21:32:03 +0000 Subject: [PATCH] fix(master): backfill legacy rollout registry entries --- crates/lance-context-core/src/registry.rs | 78 ++++++++++++- crates/lance-context-master/src/discovery.rs | 110 +++++++++++++++++++ crates/lance-context-master/src/lib.rs | 1 + crates/lance-context-master/src/state.rs | 58 ++++++++-- 4 files changed, 239 insertions(+), 8 deletions(-) create mode 100644 crates/lance-context-master/src/discovery.rs diff --git a/crates/lance-context-core/src/registry.rs b/crates/lance-context-core/src/registry.rs index 452a7ed1..6320a327 100644 --- a/crates/lance-context-core/src/registry.rs +++ b/crates/lance-context-core/src/registry.rs @@ -17,7 +17,7 @@ //! //! It deliberately holds only cheap directory metadata, never rollout rows. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use arrow_array::{Int64Array, RecordBatch, RecordBatchIterator, StringArray}; @@ -150,6 +150,51 @@ impl RolloutRegistry { Ok(()) } + /// Insert every entry whose name is not already registered in one append. + /// + /// This is intended for migration/backfill jobs where many pre-existing + /// rollout datasets need directory entries. Existing rows are left + /// unchanged, duplicate names in `entries` are ignored, and the return + /// value is the number of rows inserted. + pub async fn insert_missing(&mut self, entries: &[(String, String)]) -> LanceResult { + if entries.is_empty() { + return Ok(0); + } + + self.reload().await?; + let mut known: HashSet = self + .list() + .await? + .into_iter() + .map(|entry| entry.name) + .collect(); + let missing: Vec<&(String, String)> = entries + .iter() + .filter(|(name, _)| known.insert(name.clone())) + .collect(); + if missing.is_empty() { + return Ok(0); + } + + let schema = Arc::new(registry_schema()); + let created_at = Utc::now().timestamp_millis(); + let names: Vec<&str> = missing.iter().map(|(name, _)| name.as_str()).collect(); + let uris: Vec<&str> = missing.iter().map(|(_, uri)| uri.as_str()).collect(); + let created = vec![created_at; missing.len()]; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(names)), + Arc::new(StringArray::from(uris)), + Arc::new(Int64Array::from(created)), + ], + )?; + let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema); + let params = Self::write_params(WriteMode::Append, self.storage_options.clone()); + self.dataset.append(reader, Some(params)).await?; + Ok(missing.len()) + } + /// Remove the directory entry for `name`, if present. No-op when absent. pub async fn remove(&mut self, name: &str) -> LanceResult<()> { self.reload().await?; @@ -333,4 +378,35 @@ mod tests { assert!(reader.contains("external").await.unwrap()); assert_eq!(reader.list().await.unwrap()[0].name, "external"); } + + #[tokio::test] + async fn insert_missing_batches_new_entries() { + let dir = TempDir::new().unwrap(); + let mut reg = new_registry(&dir).await; + reg.upsert("existing", "/data/existing.rollout.lance") + .await + .unwrap(); + + let entries = vec![ + ( + "existing".to_string(), + "/other/existing.rollout.lance".to_string(), + ), + ("new-a".to_string(), "/data/new-a.rollout.lance".to_string()), + ("new-b".to_string(), "/data/new-b.rollout.lance".to_string()), + ( + "new-a".to_string(), + "/duplicate/new-a.rollout.lance".to_string(), + ), + ]; + assert_eq!(reg.insert_missing(&entries).await.unwrap(), 2); + assert_eq!(reg.insert_missing(&entries).await.unwrap(), 0); + + let mut listed = reg.list().await.unwrap(); + listed.sort_by(|a, b| a.name.cmp(&b.name)); + assert_eq!(listed.len(), 3); + assert_eq!(listed[0].uri, "/data/existing.rollout.lance"); + assert_eq!(listed[1].uri, "/data/new-a.rollout.lance"); + assert_eq!(listed[2].uri, "/data/new-b.rollout.lance"); + } } diff --git a/crates/lance-context-master/src/discovery.rs b/crates/lance-context-master/src/discovery.rs new file mode 100644 index 00000000..4f77af8e --- /dev/null +++ b/crates/lance-context-master/src/discovery.rs @@ -0,0 +1,110 @@ +//! Startup discovery for rollout datasets created before the registry existed. + +use std::path::Path; + +use lance::io::ObjectStore; +use lance_context_core::RolloutRegistry; + +const ROLLOUT_SUFFIX: &str = ".rollout.lance"; + +/// Discover top-level rollout datasets under `data_dir` and add any missing +/// registry rows in one batch. Returns the number of rows inserted. +pub async fn backfill_registry( + data_dir: &str, + registry: &mut RolloutRegistry, +) -> lance::Result { + let (store, base_path) = ObjectStore::from_uri(data_dir).await?; + let mut children = store.read_dir(base_path).await?; + children.sort(); + + let entries: Vec<(String, String)> = children + .into_iter() + .filter_map(|child| { + let name = child.strip_suffix(ROLLOUT_SUFFIX)?; + if name.is_empty() || matches!(name, "_registry" | "_stats") { + return None; + } + Some((name.to_string(), join_uri(data_dir, &child))) + }) + .collect(); + + registry.insert_missing(&entries).await +} + +fn join_uri(base: &str, child: &str) -> String { + if base.contains("://") { + format!("{}/{}", base.trim_end_matches('/'), child) + } else { + Path::new(base).join(child).to_string_lossy().to_string() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use lance_context_core::RolloutStore; + use tempfile::TempDir; + + #[tokio::test] + async fn discovers_only_unregistered_rollout_datasets() { + let dir = TempDir::new().unwrap(); + let existing_uri = dir.path().join("existing.rollout.lance"); + let legacy_uri = dir.path().join("legacy.rollout.lance"); + RolloutStore::open(existing_uri.to_str().unwrap()) + .await + .unwrap(); + RolloutStore::open(legacy_uri.to_str().unwrap()) + .await + .unwrap(); + tokio::fs::create_dir(dir.path().join("context.lance")) + .await + .unwrap(); + tokio::fs::create_dir(dir.path().join("_stats.rollout.lance")) + .await + .unwrap(); + + let registry_uri = dir.path().join("_registry.rollout.lance"); + let mut registry = RolloutRegistry::open_or_create(registry_uri.to_str().unwrap(), None) + .await + .unwrap(); + registry + .upsert("existing", existing_uri.to_str().unwrap()) + .await + .unwrap(); + + assert_eq!( + backfill_registry(dir.path().to_str().unwrap(), &mut registry) + .await + .unwrap(), + 1 + ); + assert_eq!( + backfill_registry(dir.path().to_str().unwrap(), &mut registry) + .await + .unwrap(), + 0 + ); + + let mut names: Vec = registry + .list() + .await + .unwrap() + .into_iter() + .map(|entry| entry.name) + .collect(); + names.sort(); + assert_eq!(names, vec!["existing", "legacy"]); + } + + #[test] + fn joins_local_paths_and_object_store_uris() { + assert_eq!( + join_uri("/data/rollouts", "exp.rollout.lance"), + "/data/rollouts/exp.rollout.lance" + ); + assert_eq!( + join_uri("s3://bucket/prefix/", "exp.rollout.lance"), + "s3://bucket/prefix/exp.rollout.lance" + ); + } +} diff --git a/crates/lance-context-master/src/lib.rs b/crates/lance-context-master/src/lib.rs index 48c6f64d..505d6ae8 100644 --- a/crates/lance-context-master/src/lib.rs +++ b/crates/lance-context-master/src/lib.rs @@ -1,6 +1,7 @@ //! Control-plane (master) library surface. pub mod config; +pub mod discovery; pub mod error; pub mod routes; pub mod scanner; diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs index f603231d..d0393470 100644 --- a/crates/lance-context-master/src/state.rs +++ b/crates/lance-context-master/src/state.rs @@ -1,4 +1,4 @@ -//! Shared master state: durable registry (read-only) + stats table (read/write). +//! Shared master state: durable registry + stats table. use std::collections::HashMap; use std::path::PathBuf; @@ -9,16 +9,18 @@ use lance_context_core::RolloutRegistry; use tokio::sync::{mpsc, Mutex, RwLock}; use crate::config::MasterConfig; +use crate::discovery; use crate::stats_store::StatsStore; /// Shared state for the master process. /// -/// The `registry` is written only by the data-plane (store create/delete); the -/// master reads it to enumerate experiments. The `stats_store` is owned and -/// written exclusively by the master. Both are wrapped in locks because their -/// mutating methods take `&mut self`. +/// The data-plane owns steady-state registry writes (store create/delete); the +/// master only adds missing legacy rows during startup discovery, then reads it +/// to enumerate experiments. The `stats_store` is owned and written exclusively +/// by the master. Both are wrapped in locks because their mutating methods take +/// `&mut self`. pub struct MasterState { - /// Durable directory of which rollout stores exist (data-plane-owned). + /// Durable directory of which rollout stores exist. pub registry: RwLock, /// Periodically-refreshed per-experiment metrics (master-owned). pub stats: Mutex, @@ -48,7 +50,14 @@ impl MasterState { .join("_stats.rollout.lance") .to_string_lossy() .to_string(); - let registry = RolloutRegistry::open_or_create(®istry_uri, None).await?; + let mut registry = RolloutRegistry::open_or_create(®istry_uri, None).await?; + let backfilled = discovery::backfill_registry(&config.data_dir, &mut registry).await?; + if backfilled > 0 { + tracing::info!( + experiments = backfilled, + "backfilled rollout registry from data directory" + ); + } let stats = StatsStore::open_or_create(&stats_uri, None).await?; let (compact_tx, compact_rx) = mpsc::unbounded_channel(); Ok(Arc::new(Self { @@ -71,3 +80,38 @@ impl MasterState { .to_string() } } + +#[cfg(test)] +mod tests { + use super::*; + use lance_context_core::RolloutStore; + use tempfile::TempDir; + + fn test_config(dir: &TempDir) -> MasterConfig { + MasterConfig { + data_dir: dir.path().to_string_lossy().to_string(), + host: "127.0.0.1".to_string(), + port: 0, + stats_scan_interval_secs: 0, + scan_concurrency: 4, + compaction_interval_secs: 0, + min_fragments: 16, + target_rows_per_fragment: 1_048_576, + ui_dir: None, + } + } + + #[tokio::test] + async fn startup_backfills_legacy_rollout_datasets() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().join("legacy.rollout.lance"); + RolloutStore::open(uri.to_str().unwrap()).await.unwrap(); + + let state = MasterState::new(test_config(&dir)).await.unwrap(); + let entries = state.registry.write().await.list().await.unwrap(); + + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].name, "legacy"); + assert_eq!(entries[0].uri, uri.to_string_lossy().to_string()); + } +}