fix(master): backfill legacy rollout registry entries

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<usize> {
+        if entries.is_empty() {
+            return Ok(0);
+        }
+
+        self.reload().await?;
+        let mut known: HashSet<String> = 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");
+    }
 }

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<usize> {
+    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<String> = 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"
+        );
+    }
+}

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;

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<RolloutRegistry>,
     /// Periodically-refreshed per-experiment metrics (master-owned).
     pub stats: Mutex<StatsStore>,
@@ -48,7 +50,14 @@ impl MasterState {
             .join("_stats.rollout.lance")
             .to_string_lossy()
             .to_string();
-        let registry = RolloutRegistry::open_or_create(&registry_uri, None).await?;
+        let mut registry = RolloutRegistry::open_or_create(&registry_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());
+    }
+}