fix: prevent racing dataset creation from overwriting data

crates/lance-context-core/src/merge_write_scope.rs

@@ -22,12 +22,19 @@ use tokio::sync::{oneshot, Notify};
 
 tokio::task_local! { static CURRENT: Arc<MergeWriteScope>; }
 tokio::task_local! { static EXPECTED_BASE_VERSION: u64; }
+tokio::task_local! { static INITIAL_DATASET_CREATE: (); }
 
 /// Pin the next publication without replacing a dataset's captured owner guard.
 pub(crate) async fn at_base_version<F: Future>(version: u64, work: F) -> F::Output {
     EXPECTED_BASE_VERSION.scope(version, work).await
 }
 
+/// A racing create may only publish version one. This is scoped to the create
+/// future so the returned dataset remains usable for subsequent normal writes.
+pub(crate) async fn at_initial_version<F: Future>(work: F) -> F::Output {
+    INITIAL_DATASET_CREATE.scope((), work).await
+}
+
 #[derive(Debug, Default)]
 struct Progress {
     closed: bool,
@@ -660,6 +667,18 @@ impl CommitHandler for GuardedCommit {
         transaction: Option<Transaction>,
     ) -> std::result::Result<ManifestLocation, CommitError> {
         let work = async {
+            if INITIAL_DATASET_CREATE
+                .try_with(|()| manifest.version != 1)
+                .unwrap_or(false)
+            {
+                // Lance can rebase a raced Create into an overwrite of the
+                // winner. Reject before the inner commit, without inviting
+                // another conflict retry or retaining a restriction on the
+                // returned dataset handle.
+                return Err(CommitError::OtherError(Error::invalid_input(
+                    "dataset creation cannot publish a successor manifest",
+                )));
+            }
             if EXPECTED_BASE_VERSION
                 .try_with(|version| manifest.version != *version)
                 .unwrap_or(false)
@@ -768,6 +787,10 @@ impl CommitHandler for GuardedCommit {
     }
 }
 
+#[cfg(test)]
+#[path = "merge_write_scope/create_tests.rs"]
+mod create_tests;
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/lance-context-core/src/merge_write_scope/create_tests.rs

@@ -0,0 +1,169 @@
+use crate::store_base::StorageBase;
+use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator};
+use arrow_schema::{DataType, Field, Schema};
+use async_trait::async_trait;
+use futures::{stream::BoxStream, StreamExt, TryStreamExt};
+use lance::Result;
+use lance_io::object_store::{
+    providers::local::FileStoreProvider, ObjectStore, ObjectStoreParams, ObjectStoreProvider,
+    ObjectStoreRegistry,
+};
+use object_store::{
+    path::Path, signer::Url, CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload,
+    ObjectMeta, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions,
+    Result as StoreResult,
+};
+use std::{
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
+use tokio::sync::Notify;
+
+#[derive(Debug, Default)]
+struct CreationRace {
+    listed: AtomicBool,
+    ready: Notify,
+    release: Notify,
+}
+
+#[derive(Debug)]
+struct DelayedProvider(Arc<CreationRace>);
+
+#[async_trait]
+impl ObjectStoreProvider for DelayedProvider {
+    async fn new_store(&self, url: Url, params: &ObjectStoreParams) -> Result<ObjectStore> {
+        let mut store = FileStoreProvider.new_store(url, params).await?;
+        store.inner = Arc::new(DelayedListing {
+            inner: store.inner,
+            race: self.0.clone(),
+        });
+        Ok(store)
+    }
+
+    fn extract_path(&self, url: &Url) -> Result<Path> {
+        FileStoreProvider.extract_path(url)
+    }
+}
+
+#[derive(Debug)]
+struct DelayedListing {
+    inner: Arc<dyn object_store::ObjectStore>,
+    race: Arc<CreationRace>,
+}
+
+impl std::fmt::Display for DelayedListing {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "DelayedListing({})", self.inner)
+    }
+}
+
+#[async_trait]
+impl object_store::ObjectStore for DelayedListing {
+    fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, StoreResult<ObjectMeta>> {
+        let listing = self.inner.list(prefix);
+        if prefix.and_then(Path::filename) != Some("_versions")
+            || self.race.listed.swap(true, Ordering::SeqCst)
+        {
+            return listing;
+        }
+        let race = self.race.clone();
+        futures::stream::once(async move {
+            let entries = listing.try_collect::<Vec<_>>().await?;
+            assert!(entries.is_empty(), "the initial dataset must be absent");
+            race.ready.notify_one();
+            race.release.notified().await;
+            Ok::<_, object_store::Error>(futures::stream::iter(entries.into_iter().map(Ok)))
+        })
+        .try_flatten()
+        .boxed()
+    }
+
+    async fn put_opts(&self, p: &Path, b: PutPayload, o: PutOptions) -> StoreResult<PutResult> {
+        self.inner.put_opts(p, b, o).await
+    }
+    async fn put_multipart_opts(
+        &self,
+        p: &Path,
+        o: PutMultipartOptions,
+    ) -> StoreResult<Box<dyn MultipartUpload>> {
+        self.inner.put_multipart_opts(p, o).await
+    }
+    async fn get_opts(&self, p: &Path, o: GetOptions) -> StoreResult<GetResult> {
+        self.inner.get_opts(p, o).await
+    }
+    fn delete_stream(
+        &self,
+        p: BoxStream<'static, StoreResult<Path>>,
+    ) -> BoxStream<'static, StoreResult<Path>> {
+        self.inner.delete_stream(p)
+    }
+    async fn list_with_delimiter(&self, p: Option<&Path>) -> StoreResult<ListResult> {
+        self.inner.list_with_delimiter(p).await
+    }
+    async fn copy_opts(&self, from: &Path, to: &Path, o: CopyOptions) -> StoreResult<()> {
+        self.inner.copy_opts(from, to, o).await
+    }
+    async fn rename_opts(&self, from: &Path, to: &Path, o: RenameOptions) -> StoreResult<()> {
+        self.inner.rename_opts(from, to, o).await
+    }
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn stale_create_cannot_erase_the_winners_append() -> Result<()> {
+    let root = tempfile::tempdir().unwrap();
+    let uri = root
+        .path()
+        .join("create.lance")
+        .to_str()
+        .unwrap()
+        .to_owned();
+    let schema = Arc::new(Schema::new(vec![Field::new(
+        "value",
+        DataType::Int32,
+        false,
+    )]));
+    let race = Arc::new(CreationRace::default());
+    let registry = ObjectStoreRegistry::default();
+    registry.insert("file", Arc::new(DelayedProvider(race.clone())));
+    let session = Arc::new(lance::session::Session::new(0, 0, Arc::new(registry)));
+    let mut creators = tokio::task::JoinSet::new();
+    let (delayed_uri, delayed_schema) = (uri.clone(), schema.clone());
+    creators.spawn(async move {
+        StorageBase::create_with_options(&delayed_uri, delayed_schema, None, Some(session)).await
+    });
+    // Hold the first creator's actual empty listing result. After this initial
+    // existence check, Lance's commit builder reopens the URI. The other writer
+    // must create and append before that second check; pausing at manifest
+    // authorization is too late to exercise the overwrite race.
+    tokio::time::timeout(Duration::from_secs(30), race.ready.notified())
+        .await
+        .expect("initial existence check did not arrive");
+    let mut winner = StorageBase::create_with_options(&uri, schema.clone(), None, None).await?;
+    assert_eq!(winner.version().version, 1);
+    let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![7]))])?;
+    winner
+        .append(RecordBatchIterator::new(vec![Ok(batch)], schema), None)
+        .await?;
+    assert_eq!(winner.version().version, 2);
+    race.release.notify_one();
+    let loser = tokio::time::timeout(Duration::from_secs(30), creators.join_next())
+        .await
+        .expect("stale creator did not resolve")
+        .unwrap()
+        .unwrap();
+    let reopened = lance::Dataset::open(&uri).await?;
+    assert_eq!(
+        reopened.count_rows(None).await?,
+        1,
+        "stale create erased committed data"
+    );
+    assert_eq!(reopened.version().version, 2);
+    assert!(
+        loser.is_err(),
+        "a stale create must not overwrite committed data"
+    );
+    Ok(())
+}

crates/lance-context-core/src/store_base.rs

@@ -2460,7 +2460,8 @@ impl StorageBase {
             handler,
         )));
 
-        Dataset::write(batches, uri, Some(params)).await
+        crate::merge_write_scope::at_initial_version(Dataset::write(batches, uri, Some(params)))
+            .await
     }
 }