diff --git a/crates/lance-context-core/src/merge_write_scope.rs b/crates/lance-context-core/src/merge_write_scope.rs index 7a8c833..3cc31ae 100644 --- a/crates/lance-context-core/src/merge_write_scope.rs +++ b/crates/lance-context-core/src/merge_write_scope.rs @@ -22,12 +22,19 @@ use tokio::sync::{oneshot, Notify}; tokio::task_local! { static CURRENT: Arc; } 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(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(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, ) -> std::result::Result { 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::*; diff --git a/crates/lance-context-core/src/merge_write_scope/create_tests.rs b/crates/lance-context-core/src/merge_write_scope/create_tests.rs new file mode 100644 index 0000000..dcb5b21 --- /dev/null +++ b/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); + +#[async_trait] +impl ObjectStoreProvider for DelayedProvider { + async fn new_store(&self, url: Url, params: &ObjectStoreParams) -> Result { + 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 { + FileStoreProvider.extract_path(url) + } +} + +#[derive(Debug)] +struct DelayedListing { + inner: Arc, + race: Arc, +} + +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> { + 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::>().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 { + self.inner.put_opts(p, b, o).await + } + async fn put_multipart_opts( + &self, + p: &Path, + o: PutMultipartOptions, + ) -> StoreResult> { + self.inner.put_multipart_opts(p, o).await + } + async fn get_opts(&self, p: &Path, o: GetOptions) -> StoreResult { + self.inner.get_opts(p, o).await + } + fn delete_stream( + &self, + p: BoxStream<'static, StoreResult>, + ) -> BoxStream<'static, StoreResult> { + self.inner.delete_stream(p) + } + async fn list_with_delimiter(&self, p: Option<&Path>) -> StoreResult { + 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(()) +} diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index c8604b5..6b372c3 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/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 } }