From 1aea8f6e906485438620287f0c06a38453feab7e Mon Sep 17 00:00:00 2001
From: KaiqiJinWow <kaiqi.jin@databricks.com>
Date: Tue, 8 Sep 2026 22:05:08 +0000
Subject: [PATCH] fix(registry): reopen when another caller wins creation

---
 crates/lance-context-core/src/registry.rs | 55 ++++++++++++++++++++++-
 1 file changed, 54 insertions(+), 1 deletion(-)

diff --git a/crates/lance-context-core/src/registry.rs b/crates/lance-context-core/src/registry.rs
index d1681cb..d6e9bb3 100644
--- a/crates/lance-context-core/src/registry.rs
+++ b/crates/lance-context-core/src/registry.rs
@@ -69,7 +69,7 @@ impl RolloutRegistry {
         let dataset = match Self::load(uri, storage_options.clone()).await {
             Ok(dataset) => dataset,
             Err(LanceError::DatasetNotFound { .. }) => {
-                Self::create(uri, storage_options.clone()).await?
+                Self::create_or_load(uri, storage_options.clone()).await?
             }
             Err(err) => return Err(err),
         };
@@ -94,6 +94,18 @@ impl RolloutRegistry {
         }
     }
 
+    /// Create the registry, or load it if another caller won the creation race.
+    async fn create_or_load(
+        uri: &str,
+        storage_options: Option<HashMap<String, String>>,
+    ) -> LanceResult<Dataset> {
+        match Self::create(uri, storage_options.clone()).await {
+            Ok(dataset) => Ok(dataset),
+            Err(LanceError::DatasetAlreadyExists { .. }) => Self::load(uri, storage_options).await,
+            Err(err) => Err(err),
+        }
+    }
+
     async fn create(
         uri: &str,
         storage_options: Option<HashMap<String, String>>,
@@ -349,6 +361,47 @@ mod tests {
             .unwrap()
     }
 
+    #[tokio::test]
+    async fn create_or_load_recovers_when_another_caller_wins() {
+        let dir = TempDir::new().unwrap();
+        let uri = dir.path().join("_registry.rollout.lance");
+        let uri = uri.to_str().unwrap();
+
+        // Force the check-then-create interleaving: both callers observe that
+        // the registry is absent before either one attempts to create it.
+        assert!(matches!(
+            RolloutRegistry::load(uri, None).await,
+            Err(LanceError::DatasetNotFound { .. })
+        ));
+        assert!(matches!(
+            RolloutRegistry::load(uri, None).await,
+            Err(LanceError::DatasetNotFound { .. })
+        ));
+
+        let winner = RolloutRegistry::create(uri, None).await.unwrap();
+        let mut winner = RolloutRegistry {
+            dataset: winner,
+            uri: uri.to_string(),
+            storage_options: None,
+        };
+        winner
+            .upsert("winner", "/data/winner.rollout.lance")
+            .await
+            .unwrap();
+
+        let loser = RolloutRegistry::create_or_load(uri, None).await.unwrap();
+
+        // In this late-loser interleaving, the caller reopens the winner's
+        // dataset instead of losing its rows or propagating DatasetAlreadyExists.
+        assert_eq!(loser.manifest.version, winner.dataset.manifest.version);
+        let mut loser = RolloutRegistry {
+            dataset: loser,
+            uri: uri.to_string(),
+            storage_options: None,
+        };
+        assert!(loser.contains("winner").await.unwrap());
+    }
+
     #[tokio::test]
     async fn upsert_is_idempotent() {
         let dir = TempDir::new().unwrap();