From 8fcf80b04b3593999fb341c6ba1bd391f1278a42 Mon Sep 17 00:00:00 2001 From: dentinyhao Date: Mon, 10 Aug 2026 18:15:09 -0700 Subject: [PATCH 1/4] fix(tag): make tag creation atomic --- rust/lance/src/dataset/refs.rs | 68 ++++++++++++------- .../src/dataset/tests/dataset_versioning.rs | 56 +++++++++++++++ 2 files changed, 99 insertions(+), 25 deletions(-) diff --git a/rust/lance/src/dataset/refs.rs b/rust/lance/src/dataset/refs.rs index cff86638e6a..e534a7cc4c8 100644 --- a/rust/lance/src/dataset/refs.rs +++ b/rust/lance/src/dataset/refs.rs @@ -8,7 +8,7 @@ use futures::stream::{StreamExt, TryStreamExt}; use itertools::Itertools; use lance_io::object_store::ObjectStore; use lance_table::io::commit::CommitHandler; -use object_store::path::Path; +use object_store::{Error as ObjectStoreError, PutMode, PutOptions, path::Path}; use serde::{Deserialize, Serialize}; use std::sync::Arc; @@ -144,6 +144,34 @@ impl Branches<'_> { } } +async fn put_ref_if_absent( + object_store: &ObjectStore, + path: &Path, + contents: Vec, + conflict_message: String, +) -> Result<()> { + object_store + .inner + .put_opts( + path, + contents.into(), + PutOptions { + mode: PutMode::Create, + ..Default::default() + }, + ) + .await + .map(|_| ()) + .map_err(|error| match error { + ObjectStoreError::AlreadyExists { .. } | ObjectStoreError::Precondition { .. } => { + Error::RefConflict { + message: conflict_message, + } + } + error => error.into(), + }) +} + impl Tags<'_> { pub async fn fetch_tags(&self) -> Result> { let root_location = self.refs.root()?; @@ -217,23 +245,18 @@ impl Tags<'_> { let root_location = self.refs.root()?; let tag_file = tag_path(&root_location.path, tag); - if self.object_store().exists(&tag_file).await? { - return Err(Error::RefConflict { - message: format!("tag {} already exists", tag), - }); - } let now = utc_now(); let tag_contents = self .build_tag_content_by_ref(reference, Some(now), Some(now)) .await?; - self.object_store() - .put( - &tag_file, - serde_json::to_string_pretty(&tag_contents)?.as_bytes(), - ) - .await - .map(|_| ()) + put_ref_if_absent( + self.object_store(), + &tag_file, + serde_json::to_vec_pretty(&tag_contents)?, + format!("tag {} already exists", tag), + ) + .await } pub async fn delete(&self, tag: &str) -> Result<()> { @@ -452,11 +475,6 @@ impl Branches<'_> { let source_branch = source_branch.and_then(standardize_branch); let root_location = self.refs.root()?; let branch_file = branch_contents_path(&root_location.path, branch_name); - if self.object_store().exists(&branch_file).await? { - return Err(Error::RefConflict { - message: format!("branch {} already exists", branch_name), - }); - } let branch_location = self .refs @@ -507,13 +525,13 @@ impl Branches<'_> { metadata: HashMap::new(), }; - self.object_store() - .put( - &branch_file, - serde_json::to_string_pretty(&branch_contents)?.as_bytes(), - ) - .await - .map(|_| ()) + put_ref_if_absent( + self.object_store(), + &branch_file, + serde_json::to_vec_pretty(&branch_contents)?, + format!("branch {} already exists", branch_name), + ) + .await } pub async fn replace_metadata( diff --git a/rust/lance/src/dataset/tests/dataset_versioning.rs b/rust/lance/src/dataset/tests/dataset_versioning.rs index c95e3a6d250..8be5bfbf2a9 100644 --- a/rust/lance/src/dataset/tests/dataset_versioning.rs +++ b/rust/lance/src/dataset/tests/dataset_versioning.rs @@ -21,6 +21,7 @@ use lance_core::utils::tempfile::{TempDir, TempStdDir, TempStrDir}; use lance_datagen::{BatchCount, RowCount, array, gen_batch}; use lance_file::version::LanceFileVersion; use mock_instant::thread_local::MockClock; +use tokio::sync::Barrier; use crate::dataset::refs::branch_contents_path; use crate::utils::test::copy_test_data_to_tmp; @@ -508,6 +509,61 @@ async fn test_tag( assert_eq!(dataset.manifest.version, 1); } +#[tokio::test] +async fn test_concurrent_tag_creation_conflict() { + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::UInt32, + false, + )])); + let data = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(UInt32Array::from_iter_values(0..10))], + ) + .unwrap(); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(data)], schema), + &test_uri, + None, + ) + .await + .unwrap(); + dataset.delete("i >= 5").await.unwrap(); + + let dataset = Arc::new(dataset); + let concurrency = 32; + let barrier = Arc::new(Barrier::new(concurrency)); + let handles = (0..concurrency) + .map(|attempt| { + let dataset = dataset.clone(); + let barrier = barrier.clone(); + let version = (attempt % 2 + 1) as u64; + tokio::spawn(async move { + barrier.wait().await; + (version, dataset.tags().create("race", version).await) + }) + }) + .collect::>(); + + let mut successful_version = None; + let mut conflicts = 0; + for handle in handles { + let (version, result) = handle.await.unwrap(); + match result { + Ok(()) => successful_version = Some(version), + Err(Error::RefConflict { .. }) => conflicts += 1, + Err(error) => panic!("unexpected tag creation error: {error}"), + } + } + + assert_eq!(conflicts, concurrency - 1); + assert_eq!( + dataset.tags().get_version("race").await.unwrap(), + successful_version.unwrap() + ); +} + #[rstest] #[tokio::test] async fn test_fragment_id_zero_not_reused() { From 606c31f063eed3ee705cb0efdb8bcffb86a0ae6e Mon Sep 17 00:00:00 2001 From: dentinyhao Date: Mon, 10 Aug 2026 20:44:16 -0700 Subject: [PATCH 2/4] extract object store put-if-absent --- Cargo.lock | 1 + rust/lance-io/Cargo.toml | 1 + rust/lance-io/src/object_store.rs | 67 ++++++++++++++++++- rust/lance/src/dataset/mem_wal/manifest.rs | 76 +++++----------------- rust/lance/src/dataset/mem_wal/wal.rs | 59 ++++------------- rust/lance/src/dataset/refs.rs | 13 +--- 6 files changed, 97 insertions(+), 120 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c9c2ef8a2d0..0d5d82806ff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4942,6 +4942,7 @@ dependencies = [ "tracing", "tracing-mock", "url", + "uuid", ] [[package]] diff --git a/rust/lance-io/Cargo.toml b/rust/lance-io/Cargo.toml index 03c6780263a..9c055c57e82 100644 --- a/rust/lance-io/Cargo.toml +++ b/rust/lance-io/Cargo.toml @@ -39,6 +39,7 @@ serde = { workspace = true, features = ["derive"] } tokio.workspace = true tracing.workspace = true url.workspace = true +uuid.workspace = true path_abs.workspace = true rand.workspace = true tempfile.workspace = true diff --git a/rust/lance-io/src/object_store.rs b/rust/lance-io/src/object_store.rs index a8230a578ca..c3ca6795e9b 100644 --- a/rust/lance-io/src/object_store.rs +++ b/rust/lance-io/src/object_store.rs @@ -26,7 +26,10 @@ use object_store::ObjectStoreExt as OSObjectStoreExt; use object_store::aws::AwsCredentialProvider; #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] use object_store::{ClientOptions, HeaderMap, HeaderValue}; -use object_store::{ListResult, ObjectMeta, ObjectStore as OSObjectStore, path::Path}; +use object_store::{ + ListResult, ObjectMeta, ObjectStore as OSObjectStore, PutMode, PutOptions, PutPayload, + path::Path, +}; use providers::local::FileStoreProvider; use providers::memory::MemoryStoreProvider; use tokio::io::AsyncWriteExt; @@ -823,6 +826,45 @@ impl ObjectStore { Writer::shutdown(writer.as_mut()).await } + /// Atomically creates an object without replacing an existing object. + /// + /// Local stores publish a uniquely named staging object with a conditional + /// rename. Other stores use their conditional create operation. + pub async fn put_if_absent( + &self, + path: &Path, + content: PutPayload, + ) -> object_store::Result<()> { + if self.is_local() { + let staging_path = + Path::from(format!("{}.tmp.{}", path, uuid::Uuid::new_v4().simple())); + self.inner.put(&staging_path, content).await?; + let result = self.inner.rename_if_not_exists(&staging_path, path).await; + if result.is_err() + && let Err(error) = self.inner.delete(&staging_path).await + { + log::warn!( + "Failed to remove staging object {} after atomic create failed: {}", + staging_path, + error + ); + } + result + } else { + self.inner + .put_opts( + path, + content, + PutOptions { + mode: PutMode::Create, + ..Default::default() + }, + ) + .await + .map(|_| ()) + } + } + pub async fn delete(&self, path: &Path) -> Result<()> { self.inner.delete(path).await?; Ok(()) @@ -1297,6 +1339,29 @@ mod tests { Ok(contents) } + #[tokio::test] + async fn test_put_if_absent() { + let temp_dir = TempStrDir::default(); + let path = Path::from(format!("{}/atomic-create", temp_dir.as_str())); + let store = ObjectStore::local(); + store + .put_if_absent(&path, Bytes::from_static(b"first").into()) + .await + .unwrap(); + let error = store + .put_if_absent(&path, Bytes::from_static(b"second").into()) + .await + .unwrap_err(); + assert!(matches!( + error, + object_store::Error::AlreadyExists { .. } | object_store::Error::Precondition { .. } + )); + assert_eq!( + store.read_one_all(&path).await.unwrap(), + b"first".as_slice() + ); + } + #[test] fn test_io_parallelism_clamped_to_nonzero() { // `io_parallelism()` feeds `buffered`/`buffer_unordered` windows; a value of 0 makes those diff --git a/rust/lance/src/dataset/mem_wal/manifest.rs b/rust/lance/src/dataset/mem_wal/manifest.rs index acfcbcc3a7c..7b55bdd4813 100644 --- a/rust/lance/src/dataset/mem_wal/manifest.rs +++ b/rust/lance/src/dataset/mem_wal/manifest.rs @@ -39,8 +39,6 @@ use lance_index::mem_wal::{ShardManifest, ShardStatus}; use lance_io::object_store::ObjectStore; use lance_table::format::pb; use log::{info, warn}; -use object_store::PutMode; -use object_store::PutOptions; use object_store::path::Path; use prost::Message; use serde::{Deserialize, Serialize}; @@ -184,68 +182,26 @@ impl ShardManifestStore { let pb_manifest = pb::ShardManifest::from(manifest); let bytes = pb_manifest.encode_to_vec(); - if self.object_store.is_local() { - // Local storage: Use temp file + atomic rename for fencing - let temp_filename = format!("{}.tmp.{}", filename, uuid::Uuid::new_v4()); - let temp_path = self.manifest_dir.clone().join(temp_filename.as_str()); - - // Write to temp file - self.object_store - .inner - .put(&temp_path, Bytes::from(bytes).into()) - .await - .map_err(|e| Error::io(format!("Failed to write temp manifest: {}", e)))?; - - // Atomically rename to final path - match self - .object_store - .inner - .rename_if_not_exists(&temp_path, &path) - .await - { - Ok(()) => {} - Err(object_store::Error::AlreadyExists { .. }) => { - // Clean up temp file - let _ = self.object_store.delete(&temp_path).await; - return Err(Error::io(format!( + self.object_store + .put_if_absent(&path, Bytes::from(bytes).into()) + .await + .map_err(|error| { + if matches!( + error, + object_store::Error::AlreadyExists { .. } + | object_store::Error::Precondition { .. } + ) { + Error::io(format!( "Manifest version {} already exists for shard {}", version, self.shard_id - ))); - } - Err(e) => { - // Clean up temp file - let _ = self.object_store.delete(&temp_path).await; - return Err(Error::io(format!( + )) + } else { + Error::io(format!( "Failed to write manifest version {} for shard {}: {}", - version, self.shard_id, e - ))); + version, self.shard_id, error + )) } - } - } else { - // Cloud storage: Use PUT-IF-NOT-EXISTS - let put_opts = PutOptions { - mode: PutMode::Create, - ..Default::default() - }; - - self.object_store - .inner - .put_opts(&path, Bytes::from(bytes).into(), put_opts) - .await - .map_err(|e| { - if matches!(e, object_store::Error::AlreadyExists { .. }) { - Error::io(format!( - "Manifest version {} already exists for shard {}", - version, self.shard_id - )) - } else { - Error::io(format!( - "Failed to write manifest version {} for shard {}: {}", - version, self.shard_id, e - )) - } - })?; - } + })?; // Best-effort update version hint (failures are logged as warnings) self.write_version_hint(version).await; diff --git a/rust/lance/src/dataset/mem_wal/wal.rs b/rust/lance/src/dataset/mem_wal/wal.rs index 68eb41a03c2..1360534c570 100644 --- a/rust/lance/src/dataset/mem_wal/wal.rs +++ b/rust/lance/src/dataset/mem_wal/wal.rs @@ -25,7 +25,6 @@ use lance_core::{Error, FenceReason, Result}; use lance_io::object_store::ObjectStore; use object_store::ObjectStoreExt; use object_store::path::Path; -use object_store::{PutMode, PutOptions}; use tokio::sync::{Mutex, mpsc, watch}; use tracing::instrument; @@ -1579,53 +1578,17 @@ async fn atomic_put( bytes: Bytes, ) -> std::result::Result<(), AtomicPutError> { let path = dir.clone().join(filename); - if object_store.is_local() { - let temp = dir - .clone() - .join(format!("{}.tmp.{}", filename, Uuid::new_v4())); - object_store - .inner - .put(&temp, bytes.into()) - .await - .map_err(|e| { - AtomicPutError::Other(Error::io(format!("failed to write temp file: {}", e))) - })?; - match object_store.inner.rename_if_not_exists(&temp, &path).await { - Ok(()) => Ok(()), - Err(object_store::Error::AlreadyExists { .. }) => { - let _ = object_store.delete(&temp).await; - Err(AtomicPutError::AlreadyExists) - } - Err(e) => { - let _ = object_store.delete(&temp).await; - Err(AtomicPutError::Other(Error::io(format!( - "failed to create {} atomically: {}", - path, e - )))) - } - } - } else { - object_store - .inner - .put_opts( - &path, - bytes.into(), - PutOptions { - mode: PutMode::Create, - ..Default::default() - }, - ) - .await - .map_err(|e| match e { - object_store::Error::AlreadyExists { .. } - | object_store::Error::Precondition { .. } => AtomicPutError::AlreadyExists, - _ => AtomicPutError::Other(Error::io(format!( - "failed to create {} atomically: {}", - path, e - ))), - })?; - Ok(()) - } + object_store + .put_if_absent(&path, bytes.into()) + .await + .map_err(|error| match error { + object_store::Error::AlreadyExists { .. } + | object_store::Error::Precondition { .. } => AtomicPutError::AlreadyExists, + _ => AtomicPutError::Other(Error::io(format!( + "failed to create {} atomically: {}", + path, error + ))), + }) } /// Probe forward from a hint position to find the next unwritten position. diff --git a/rust/lance/src/dataset/refs.rs b/rust/lance/src/dataset/refs.rs index e534a7cc4c8..79380a0799c 100644 --- a/rust/lance/src/dataset/refs.rs +++ b/rust/lance/src/dataset/refs.rs @@ -8,7 +8,7 @@ use futures::stream::{StreamExt, TryStreamExt}; use itertools::Itertools; use lance_io::object_store::ObjectStore; use lance_table::io::commit::CommitHandler; -use object_store::{Error as ObjectStoreError, PutMode, PutOptions, path::Path}; +use object_store::{Error as ObjectStoreError, path::Path}; use serde::{Deserialize, Serialize}; use std::sync::Arc; @@ -151,17 +151,8 @@ async fn put_ref_if_absent( conflict_message: String, ) -> Result<()> { object_store - .inner - .put_opts( - path, - contents.into(), - PutOptions { - mode: PutMode::Create, - ..Default::default() - }, - ) + .put_if_absent(path, contents.into()) .await - .map(|_| ()) .map_err(|error| match error { ObjectStoreError::AlreadyExists { .. } | ObjectStoreError::Precondition { .. } => { Error::RefConflict { From 2f6b2f835789392e54ace2bfbe735c8a844ff79e Mon Sep 17 00:00:00 2001 From: dentinyhao Date: Mon, 10 Aug 2026 20:48:23 -0700 Subject: [PATCH 3/4] explicitly fail COS --- rust/lance-io/src/object_store.rs | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/rust/lance-io/src/object_store.rs b/rust/lance-io/src/object_store.rs index c3ca6795e9b..5c536ede2e1 100644 --- a/rust/lance-io/src/object_store.rs +++ b/rust/lance-io/src/object_store.rs @@ -829,12 +829,24 @@ impl ObjectStore { /// Atomically creates an object without replacing an existing object. /// /// Local stores publish a uniquely named staging object with a conditional - /// rename. Other stores use their conditional create operation. + /// rename. Other stores use their conditional create operation. Tencent COS + /// is rejected because it can silently ignore conditional create requests. + /// + /// Returns [`object_store::Error::NotSupported`] without writing when the + /// backend cannot reliably provide put-if-absent semantics. pub async fn put_if_absent( &self, path: &Path, content: PutPayload, ) -> object_store::Result<()> { + if self.scheme == "cos" { + return Err(object_store::Error::NotSupported { + source: "Tencent COS does not reliably enforce put-if-absent after bucket \ + versioning has ever been enabled" + .into(), + }); + } + if self.is_local() { let staging_path = Path::from(format!("{}.tmp.{}", path, uuid::Uuid::new_v4().simple())); @@ -1362,6 +1374,21 @@ mod tests { ); } + #[tokio::test] + async fn test_put_if_absent_rejects_cos() { + let mut store = ObjectStore::memory(); + store.scheme = "cos".to_string(); + let path = Path::from("atomic-create"); + + let error = store + .put_if_absent(&path, Bytes::from_static(b"value").into()) + .await + .unwrap_err(); + + assert!(matches!(error, object_store::Error::NotSupported { .. })); + assert!(!store.exists(&path).await.unwrap()); + } + #[test] fn test_io_parallelism_clamped_to_nonzero() { // `io_parallelism()` feeds `buffered`/`buffer_unordered` windows; a value of 0 makes those From 03192e8dc6a8603359536dcbc6a613d11bc64925 Mon Sep 17 00:00:00 2001 From: dentinyhao Date: Mon, 10 Aug 2026 21:14:24 -0700 Subject: [PATCH 4/4] fix binding crate --- java/lance-jni/Cargo.lock | 1 + python/Cargo.lock | 1 + 2 files changed, 2 insertions(+) diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index 41c592bfd6f..f013f55a898 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -4108,6 +4108,7 @@ dependencies = [ "tokio", "tracing", "url", + "uuid", ] [[package]] diff --git a/python/Cargo.lock b/python/Cargo.lock index aacc676663e..cf9a4cf78f6 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -4436,6 +4436,7 @@ dependencies = [ "tokio", "tracing", "url", + "uuid", ] [[package]]