From 4fdf4256f1bc191b4a71d2c0504e9d2aa4924149 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 14:39:49 +0000 Subject: [PATCH 01/13] fix: protect active dataset versions from cleanup --- python/python/lance/__init__.py | 2 + python/python/lance/dataset.py | 28 ++ python/python/lance/lance/__init__.pyi | 13 +- python/python/tests/test_dataset.py | 20 ++ python/src/dataset.rs | 10 + python/src/dataset/version_lease.rs | 94 +++++++ python/src/lib.rs | 2 + rust/lance/src/dataset.rs | 2 + rust/lance/src/dataset/cleanup.rs | 179 +++++++++++- rust/lance/src/dataset/version_lease.rs | 359 ++++++++++++++++++++++++ 10 files changed, 693 insertions(+), 16 deletions(-) create mode 100644 python/src/dataset/version_lease.rs create mode 100644 rust/lance/src/dataset/version_lease.rs diff --git a/python/python/lance/__init__.py b/python/python/lance/__init__.py index 33f4ce00344..9d8582488fd 100644 --- a/python/python/lance/__init__.py +++ b/python/python/lance/__init__.py @@ -32,6 +32,7 @@ MergeInsertBuilder, Session, Transaction, + VersionLease, __version__, batch_udf, write_dataset, @@ -107,6 +108,7 @@ "MergeInsertBuilder", "ScanStatistics", "Transaction", + "VersionLease", "__version__", "batch_udf", "bytes_read_counter", diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index d387880be48..d1b50caf6c9 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -64,6 +64,7 @@ LanceSchema, PySearchFilter, ScanStatistics, + VersionLease, _Dataset, _format_field_path, _MergeInsertBuilder, @@ -3089,6 +3090,33 @@ def checkout_version( ds._ds = self._ds.checkout_version(version) return ds + def acquire_version_lease(self, ttl: timedelta) -> VersionLease: + """Protect this dataset version from cleanup with a renewable lease. + + Acquire the lease before starting a long-running read and renew it before + :attr:`VersionLease.expires_at`. Cleanup may remove the version after the + lease expires. A context manager releases the lease on exit; a process + that exits without releasing it stops protecting the version at expiry. + + Parameters + ---------- + ttl : timedelta + How long the lease protects this version. Must be positive. + + Returns + ------- + VersionLease + A renewable lease for this dataset version. + + Examples + -------- + >>> historical = dataset.checkout_version(1) + >>> with historical.acquire_version_lease(timedelta(minutes=5)) as lease: + ... table = historical.to_table() + ... lease.renew(timedelta(minutes=5)) + """ + return self._ds.acquire_version_lease(td_to_micros(ttl)) + def restore(self): """ Restore the currently checked out version as the latest version of the dataset. diff --git a/python/python/lance/lance/__init__.pyi b/python/python/lance/lance/__init__.pyi index 6a240f42860..e687ad4556d 100644 --- a/python/python/lance/lance/__init__.pyi +++ b/python/python/lance/lance/__init__.pyi @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from datetime import timedelta +from datetime import datetime, timedelta from pathlib import Path from typing import ( Any, @@ -137,6 +137,16 @@ class CleanupExplanation: referenced_branches: List[CleanupReferencedBranch] warnings: List[str] +class VersionLease: + @property + def version(self) -> int: ... + @property + def expires_at(self) -> datetime: ... + def renew(self, ttl: timedelta) -> None: ... + def release(self) -> None: ... + def __enter__(self) -> VersionLease: ... + def __exit__(self, exc_type, exc_value, traceback) -> bool: ... + class LanceFileWriteSummary: num_rows: int size_bytes: int @@ -460,6 +470,7 @@ class _Dataset: def checkout_version( self, version: int | str | Tuple[Optional[str], Optional[int]] ) -> _Dataset: ... + def acquire_version_lease(self, ttl_micros: int) -> VersionLease: ... def checkout_latest(self) -> _Dataset: ... def shallow_clone( self, diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 67c74663f2d..db7feb9d7a7 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -1592,6 +1592,26 @@ def test_cleanup_old_versions(tmp_path): assert stats.old_versions == 1 +def test_cleanup_retains_active_version_lease(tmp_path: Path): + table = pa.table({"value": range(10)}) + dataset = lance.write_dataset(table, tmp_path / "leased", mode="create") + historical = dataset.checkout_version(1) + dataset = lance.write_dataset(table, tmp_path / "leased", mode="overwrite") + + with historical.acquire_version_lease(timedelta(minutes=1)) as lease: + assert lease.version == 1 + assert lease.expires_at > datetime.now(tz=lease.expires_at.tzinfo) + lease.renew(timedelta(minutes=1)) + stats = dataset.cleanup_old_versions(retain_versions=1) + assert stats.old_versions == 0 + assert historical.to_table() == table + + stats = dataset.cleanup_old_versions(retain_versions=1) + assert stats.old_versions == 1 + with pytest.raises(OSError): + dataset.checkout_version(1) + + def test_explain_cleanup_old_versions(tmp_path): table = pa.Table.from_pydict({"a": range(100), "b": range(100)}) base_dir = tmp_path / "test" diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 69cdc3c8130..b12c840d195 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -114,6 +114,7 @@ use self::cleanup::{ }; use self::commit::PyCommitLock; use self::io_stats::IoStats; +use self::version_lease::PyVersionLease; pub mod blob; pub mod cleanup; @@ -121,6 +122,7 @@ pub mod commit; pub mod io_stats; pub mod optimize; pub mod stats; +pub mod version_lease; const DEFAULT_NPROBES: usize = 1; const LANCE_COMMIT_MESSAGE_KEY: &str = "__lance_commit_message"; @@ -2078,6 +2080,14 @@ impl Dataset { self._checkout_version(reference) } + fn acquire_version_lease(&self, py: Python<'_>, ttl_micros: i64) -> PyResult { + let ttl = PyVersionLease::ttl(ttl_micros)?; + let lease = rt() + .block_on(Some(py), self.ds.acquire_version_lease(ttl))? + .infer_error()?; + Ok(PyVersionLease::new(lease)) + } + /// Restore the current version #[pyo3(signature = (target_path, reference, storage_options=None))] fn shallow_clone( diff --git a/python/src/dataset/version_lease.rs b/python/src/dataset/version_lease.rs new file mode 100644 index 00000000000..bb6f5dc9e11 --- /dev/null +++ b/python/src/dataset/version_lease.rs @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use pyo3::{exceptions::PyValueError, prelude::*, types::PyAny}; + +use crate::{error::PythonErrorExt, rt}; + +/// Python wrapper for a renewable dataset version lease. +#[pyclass(name = "VersionLease", module = "_lib", skip_from_py_object)] +pub struct PyVersionLease { + inner: Option, +} + +impl PyVersionLease { + pub fn new(inner: lance::dataset::VersionLease) -> Self { + Self { inner: Some(inner) } + } + + pub(crate) fn ttl(ttl_micros: i64) -> PyResult { + let ttl_micros = u64::try_from(ttl_micros).map_err(|_| { + PyValueError::new_err(format!( + "version lease TTL must be greater than zero, got {ttl_micros} microseconds" + )) + })?; + if ttl_micros == 0 { + return Err(PyValueError::new_err( + "version lease TTL must be greater than zero, got 0 microseconds", + )); + } + Ok(Duration::from_micros(ttl_micros)) + } + + fn inner(&self) -> PyResult<&lance::dataset::VersionLease> { + self.inner + .as_ref() + .ok_or_else(|| PyValueError::new_err("version lease has been released")) + } +} + +#[pymethods] +impl PyVersionLease { + /// The dataset version protected by this lease. + #[getter] + fn version(&self) -> PyResult { + Ok(self.inner()?.version()) + } + + /// The time after which cleanup may remove the protected version. + #[getter] + fn expires_at(&self) -> PyResult> { + Ok(self.inner()?.expires_at()) + } + + /// Renew this lease for the given duration from now. + fn renew(&mut self, py: Python<'_>, ttl: Duration) -> PyResult<()> { + if ttl.is_zero() { + return Err(PyValueError::new_err( + "version lease TTL must be greater than zero", + )); + } + let lease = self + .inner + .as_mut() + .ok_or_else(|| PyValueError::new_err("version lease has been released"))?; + rt().block_on(Some(py), lease.renew(ttl))?.infer_error()?; + Ok(()) + } + + /// Release this lease before its TTL expires. + fn release(&mut self, py: Python<'_>) -> PyResult<()> { + if let Some(lease) = self.inner.take() { + rt().block_on(Some(py), lease.release())?.infer_error()?; + } + Ok(()) + } + + fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn __exit__( + mut slf: PyRefMut<'_, Self>, + py: Python<'_>, + _exc_type: &Bound<'_, PyAny>, + _exc_value: &Bound<'_, PyAny>, + _traceback: &Bound<'_, PyAny>, + ) -> PyResult { + slf.release(py)?; + Ok(false) + } +} diff --git a/python/src/lib.rs b/python/src/lib.rs index 069b591effc..9c4009d2cf5 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -45,6 +45,7 @@ use dataset::io_stats::IoStats; use dataset::optimize::{ PyCompaction, PyCompactionMetrics, PyCompactionPlan, PyCompactionTask, PyRewriteResult, }; +use dataset::version_lease::PyVersionLease; use dataset::{DatasetBasePath, MergeInsertBuilder, PyFullTextQuery, PySearchFilter}; use env_logger::{Builder, Env}; use file::{ @@ -284,6 +285,7 @@ fn lance(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index b98a2d1ee11..443a0ab668b 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -93,6 +93,7 @@ pub mod transaction; pub mod udtf; pub mod updater; mod utils; +mod version_lease; pub(crate) mod versions; pub mod write; @@ -135,6 +136,7 @@ pub use schema_evolution::{ }; pub use take::TakeBuilder; use uuid::Uuid; +pub use version_lease::VersionLease; pub use write::merge_insert::{ MergeInsertBuilder, MergeInsertJob, MergeStats, UncommittedMergeInsert, WhenMatched, WhenNotMatched, WhenNotMatchedBySource, diff --git a/rust/lance/src/dataset/cleanup.rs b/rust/lance/src/dataset/cleanup.rs index 5628f9d7e1f..d91605b44e8 100644 --- a/rust/lance/src/dataset/cleanup.rs +++ b/rust/lance/src/dataset/cleanup.rs @@ -34,6 +34,7 @@ //! happening at the same time) use super::refs::TagContents; +use super::version_lease::VersionLeaseStore; use crate::dataset::TRANSACTIONS_DIR; use crate::{Dataset, utils::temporal::utc_now}; use chrono::{DateTime, TimeDelta, Utc}; @@ -53,7 +54,7 @@ use lance_core::{ use lance_table::{ format::{IndexMetadata, Manifest}, io::{ - commit::ManifestLocation, + commit::{ManifestLocation, ManifestNamingScheme}, deletion::deletion_file_path, manifest::{read_manifest, read_manifest_indexes}, }, @@ -64,7 +65,7 @@ use std::fmt::Debug; use std::{ collections::{HashMap, HashSet}, future, - sync::{Mutex, MutexGuard}, + sync::Mutex, time::Duration, }; use tokio::time::{MissedTickBehavior, interval}; @@ -467,7 +468,14 @@ impl<'a> CleanupTask<'a> { .map(|tag_content| tag_content.version) .collect(); - let mut inspection = self.process_manifests(&tagged_versions).await?; + let version_lease_store = VersionLeaseStore::for_dataset(self.dataset).await?; + let leased_versions = version_lease_store + .active_versions(self.action.deletes_files()) + .await?; + + let mut inspection = self + .process_manifests(&tagged_versions, &leased_versions) + .await?; if self.policy.error_if_tagged_old_versions && !inspection.tagged_old_versions.is_empty() { return Err(tagged_old_versions_cleanup_error( @@ -487,6 +495,10 @@ impl<'a> CleanupTask<'a> { .await? }; + inspection = self + .fence_old_versions_and_retain_new_leases(inspection, &version_lease_store) + .await?; + final_result.merge( self.delete_unreferenced_files(inspection).await?, candidate_file_limit, @@ -498,6 +510,7 @@ impl<'a> CleanupTask<'a> { async fn process_manifests( &'a self, tagged_versions: &HashSet, + leased_versions: &HashSet, ) -> Result { let inspection = Mutex::new(CleanupInspection::default()); self.dataset @@ -505,7 +518,7 @@ impl<'a> CleanupTask<'a> { .list_manifest_locations(&self.dataset.base, &self.dataset.object_store, false) .try_filter(|location| future::ready(!self.ignored_manifests.contains(&location.path))) .try_for_each_concurrent(self.dataset.object_store.io_parallelism(), |location| { - self.process_manifest_file(location, &inspection, tagged_versions) + self.process_manifest_file(location, &inspection, tagged_versions, leased_versions) }) .await?; Ok(inspection.into_inner().unwrap()) @@ -516,6 +529,7 @@ impl<'a> CleanupTask<'a> { location: ManifestLocation, inspection: &Mutex, tagged_versions: &HashSet, + leased_versions: &HashSet, ) -> Result<()> { // TODO: We can't cleanup invalid manifests. There is no way to distinguish // between an invalid manifest and a temporary I/O error. It's also not safe @@ -551,7 +565,9 @@ impl<'a> CleanupTask<'a> { // version. These are either in-progress or newly added since we started. let is_latest = self.read_version <= manifest.version; let is_tagged = tagged_versions.contains(&manifest.version); - let in_working_set = is_latest || !self.policy.should_clean(&manifest) || is_tagged; + let is_leased = leased_versions.contains(&manifest.version); + let in_working_set = + is_latest || !self.policy.should_clean(&manifest) || is_tagged || is_leased; let mut inspection = inspection.lock().unwrap(); // Track tagged old versions in case we want to return a `CleanupError` later. @@ -566,16 +582,77 @@ impl<'a> CleanupTask<'a> { .old_manifests .insert(location.path.clone(), manifest.version); } else { - let commit_ts = manifest.timestamp(); - if let Some(ts) = inspection.earliest_retained_manifest_time { - if commit_ts < ts { - inspection.earliest_retained_manifest_time = Some(commit_ts); - } - } else { + Self::note_retained_manifest_time(&mut inspection, &manifest); + } + Ok(()) + } + + async fn fence_old_versions_and_retain_new_leases( + &self, + mut inspection: CleanupInspection, + version_lease_store: &VersionLeaseStore, + ) -> Result { + if !self.action.deletes_files() { + return Ok(inspection); + } + + let versions_to_delete: HashSet = inspection.old_manifests.values().copied().collect(); + version_lease_store + .fence_versions(&versions_to_delete) + .await?; + + // A lease acquired between the initial list and the fence either sees + // the fence and fails, or appears here and conservatively retains this + // version for the current cleanup pass. + let leased_versions = version_lease_store.active_versions(true).await?; + for version in versions_to_delete.intersection(&leased_versions) { + self.retain_version(&mut inspection, *version).await?; + } + Ok(inspection) + } + + async fn retain_version(&self, inspection: &mut CleanupInspection, version: u64) -> Result<()> { + let manifest_paths: Vec = inspection + .old_manifests + .iter() + .filter(|(_, manifest_version)| **manifest_version == version) + .map(|(path, _)| path.clone()) + .collect(); + + for path in manifest_paths { + let filename = path.filename().ok_or_else(|| { + Error::internal(format!("manifest path {} has no filename", path)) + })?; + let naming_scheme = ManifestNamingScheme::detect_scheme(filename).ok_or_else(|| { + Error::internal(format!("invalid manifest filename: '{filename}'")) + })?; + let location = ManifestLocation { + path: path.clone(), + version, + size: None, + naming_scheme, + e_tag: None, + }; + let manifest = + read_manifest(&self.dataset.object_store, &location.path, location.size).await?; + let indexes = + read_manifest_indexes(&self.dataset.object_store, &location, &manifest).await?; + self.process_manifest(&manifest, &indexes, true, inspection)?; + inspection.old_manifests.remove(&path); + Self::note_retained_manifest_time(inspection, &manifest); + } + Ok(()) + } + + fn note_retained_manifest_time(inspection: &mut CleanupInspection, manifest: &Manifest) { + let commit_ts = manifest.timestamp(); + if let Some(ts) = inspection.earliest_retained_manifest_time { + if commit_ts < ts { inspection.earliest_retained_manifest_time = Some(commit_ts); } + } else { + inspection.earliest_retained_manifest_time = Some(commit_ts); } - Ok(()) } fn process_manifest( @@ -583,7 +660,7 @@ impl<'a> CleanupTask<'a> { manifest: &Manifest, indexes: &Vec, in_working_set: bool, - inspection: &mut MutexGuard, + inspection: &mut CleanupInspection, ) -> Result<()> { // If this part of our working set then update referenced_files. Otherwise, just mark the // file as verified. @@ -2136,11 +2213,77 @@ mod tests { old_manifest, &Mutex::new(CleanupInspection::default()), &HashSet::new(), + &HashSet::new(), ) .await .unwrap(); } + #[tokio::test] + async fn cleanup_retains_version_with_active_lease() { + let fixture = MockDatasetFixture::try_new().unwrap(); + fixture.create_some_data().await.unwrap(); + let historical = fixture.load().await.unwrap(); + let mut lease = historical + .acquire_version_lease(Duration::from_secs(30 * 24 * 60 * 60)) + .await + .unwrap(); + lease + .renew(Duration::from_secs(30 * 24 * 60 * 60)) + .await + .unwrap(); + fixture.overwrite_some_data().await.unwrap(); + MockClock::set_system_time(TimeDelta::try_days(10).unwrap().to_std().unwrap()); + + let removed = fixture + .run_cleanup(utc_now() - TimeDelta::try_days(8).unwrap()) + .await + .unwrap(); + assert_eq!(removed.old_versions, 0); + assert_eq!(removed.data_files_removed, 0); + historical.scan().try_into_batch().await.unwrap(); + + lease.release().await.unwrap(); + let removed = fixture + .run_cleanup(utc_now() - TimeDelta::try_days(8).unwrap()) + .await + .unwrap(); + assert_eq!(removed.old_versions, 1); + assert!( + fixture + .load() + .await + .unwrap() + .checkout_version(1) + .await + .is_err() + ); + } + + #[tokio::test] + async fn expired_version_lease_does_not_block_cleanup() { + let fixture = MockDatasetFixture::try_new().unwrap(); + fixture.create_some_data().await.unwrap(); + let historical = fixture.load().await.unwrap(); + let _lease = historical + .acquire_version_lease(Duration::from_secs(1)) + .await + .unwrap(); + fixture.overwrite_some_data().await.unwrap(); + MockClock::set_system_time(TimeDelta::try_days(10).unwrap().to_std().unwrap()); + + let removed = fixture + .run_cleanup(utc_now() - TimeDelta::try_days(8).unwrap()) + .await + .unwrap(); + assert_eq!(removed.old_versions, 1); + let error = historical + .acquire_version_lease(Duration::from_secs(60)) + .await + .unwrap_err(); + assert!(error.to_string().contains("fenced for cleanup"), "{error}"); + } + #[tokio::test] async fn explain_cleanup_does_not_delete_files() { let fixture = MockDatasetFixture::try_new().unwrap(); @@ -2817,7 +2960,10 @@ mod tests { .build(), CleanupAction::Execute, ); - let inspection = task.process_manifests(&HashSet::new()).await.unwrap(); + let inspection = task + .process_manifests(&HashSet::new(), &HashSet::new()) + .await + .unwrap(); let referenced_branches = task.find_referenced_branches().await.unwrap(); let inspection = task .retain_branch_lineage_files(inspection, &referenced_branches, &HashSet::new()) @@ -2881,7 +3027,10 @@ mod tests { .build(), CleanupAction::Execute, ); - let inspection = task.process_manifests(&HashSet::new()).await.unwrap(); + let inspection = task + .process_manifests(&HashSet::new(), &HashSet::new()) + .await + .unwrap(); let kept: HashSet<&Path> = inspection .referenced_files .data_paths diff --git a/rust/lance/src/dataset/version_lease.rs b/rust/lance/src/dataset/version_lease.rs new file mode 100644 index 00000000000..6a1c4b59c3d --- /dev/null +++ b/rust/lance/src/dataset/version_lease.rs @@ -0,0 +1,359 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Advisory leases that protect dataset versions from cleanup. + +use std::{collections::HashSet, sync::Arc, time::Duration}; + +use bytes::Bytes; +use chrono::{DateTime, TimeDelta, Utc}; +use futures::{StreamExt, TryStreamExt, stream}; +use lance_io::object_store::ObjectStore; +use object_store::{PutMode, PutOptions, path::Path}; +use uuid::Uuid; + +use super::{Dataset, refs::MAIN_BRANCH}; +use crate::{Error, Result, utils::temporal::utc_now}; + +const LEASES_DIR: &str = "_refs/version_leases"; +const LEASE_GC_MARKERS_DIR: &str = "_refs/version_lease_gc"; +const LEASE_FILE_SUFFIX: &str = ".lease"; + +/// A renewable advisory lease that protects one dataset version from cleanup. +/// +/// The lease must be renewed before [`Self::expires_at`]. Once it expires, +/// [`Dataset::cleanup_old_versions`](Dataset::cleanup_old_versions) may delete +/// the protected version. Dropping a lease releases it on a best-effort basis; +/// if a process exits or cannot release the lease, it stops protecting the +/// version when its TTL expires. +#[derive(Debug)] +pub struct VersionLease { + store: VersionLeaseStore, + path: Path, + version: u64, + expires_at: DateTime, + released: bool, +} + +impl VersionLease { + /// The dataset version protected by this lease. + pub fn version(&self) -> u64 { + self.version + } + + /// The time after which cleanup may remove the protected version. + pub fn expires_at(&self) -> DateTime { + self.expires_at + } + + /// Renew this lease for `ttl` from the current time. + /// + /// Renewal must complete before the current expiration. It fails if the + /// lease has expired or cleanup has fenced the version for deletion. + pub async fn renew(&mut self, ttl: Duration) -> Result<()> { + let now = utc_now(); + if now >= self.expires_at { + return Err(Error::invalid_input(format!( + "cannot renew expired version lease for version {}: expired at {}", + self.version, self.expires_at + ))); + } + let expires_at = expiration_from_ttl(now, ttl)?; + self.store.ensure_not_fenced(self.version).await?; + + // Publish the replacement before removing the old file so cleanup never + // observes a gap between a timely renewal and its predecessor. + let path = self + .store + .create_lease_file(self.version, expires_at) + .await?; + if let Err(error) = self.store.ensure_not_fenced(self.version).await { + let _ = self.store.object_store.delete(&path).await; + return Err(error); + } + + if let Err(error) = self.store.object_store.delete(&self.path).await + && !error.is_not_found() + { + tracing::warn!( + path = %self.path, + error = %error, + "Failed to remove superseded version lease" + ); + } + self.path = path; + self.expires_at = expires_at; + Ok(()) + } + + /// Release this lease before its TTL expires. + pub async fn release(mut self) -> Result<()> { + match self.store.object_store.delete(&self.path).await { + Ok(()) => { + self.released = true; + Ok(()) + } + Err(error) if error.is_not_found() => { + self.released = true; + Ok(()) + } + Err(error) => Err(error), + } + } +} + +impl Drop for VersionLease { + fn drop(&mut self) { + if self.released { + return; + } + let object_store = Arc::clone(&self.store.object_store); + let path = self.path.clone(); + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + runtime.spawn(async move { + let _ = object_store.delete(&path).await; + }); + } + } +} + +#[derive(Clone, Debug)] +pub(super) struct VersionLeaseStore { + object_store: Arc, + leases_path: Path, + markers_path: Path, +} + +impl VersionLeaseStore { + pub(super) async fn for_dataset(dataset: &Dataset) -> Result { + let root = dataset.refs.root()?; + let branch = dataset.manifest.branch.as_deref(); + let branch_identifier = dataset.branches().get_identifier(branch).await?; + let namespace = if branch.is_none() { + MAIN_BRANCH + } else { + branch_identifier + .version_mapping + .last() + .map(|(_, id)| id.as_str()) + .ok_or_else(|| { + Error::internal(format!( + "branch {} has no branch identifier", + branch.unwrap_or_default() + )) + })? + }; + + Ok(Self { + object_store: Arc::clone(&dataset.object_store), + leases_path: root.path.clone().join(LEASES_DIR).join(namespace), + markers_path: root.path.join(LEASE_GC_MARKERS_DIR).join(namespace), + }) + } + + async fn acquire(&self, version: u64, ttl: Duration) -> Result { + self.ensure_not_fenced(version).await?; + let expires_at = expiration_from_ttl(utc_now(), ttl)?; + let path = self.create_lease_file(version, expires_at).await?; + if let Err(error) = self.ensure_not_fenced(version).await { + let _ = self.object_store.delete(&path).await; + return Err(error); + } + Ok(VersionLease { + store: self.clone(), + path, + version, + expires_at, + released: false, + }) + } + + async fn create_lease_file(&self, version: u64, expires_at: DateTime) -> Result { + let path = self.leases_path.clone().join(format!( + "{}-{}-{}{}", + version, + expires_at.timestamp_micros(), + Uuid::new_v4().simple(), + LEASE_FILE_SUFFIX + )); + self.object_store + .inner + .put_opts( + &path, + Bytes::new().into(), + PutOptions { + mode: PutMode::Create, + ..Default::default() + }, + ) + .await?; + Ok(path) + } + + pub(super) async fn active_versions(&self, remove_expired: bool) -> Result> { + let lease_files = match self.object_store.read_dir(self.leases_path.clone()).await { + Ok(files) => files, + Err(error) if error.is_not_found() => return Ok(HashSet::new()), + Err(error) => return Err(error), + }; + let now_micros = utc_now().timestamp_micros(); + let mut active_versions = HashSet::new(); + let mut expired_paths = Vec::new(); + + for file_name in lease_files { + let (version, expires_at_micros) = + parse_lease_file_name(&file_name).map_err(|error| { + Error::corrupt_file(self.leases_path.clone().join(file_name.as_str()), error) + })?; + if expires_at_micros > now_micros { + active_versions.insert(version); + } else if remove_expired { + expired_paths.push(self.leases_path.clone().join(file_name)); + } + } + + if remove_expired { + stream::iter(expired_paths) + .map(Ok) + .try_for_each_concurrent(self.object_store.io_parallelism(), |path| async move { + match self.object_store.delete(&path).await { + Ok(()) => Ok(()), + Err(error) if error.is_not_found() => Ok(()), + Err(error) => Err(error), + } + }) + .await?; + } + Ok(active_versions) + } + + pub(super) async fn fence_versions(&self, versions: &HashSet) -> Result<()> { + stream::iter(versions.iter().copied()) + .map(Ok) + .try_for_each_concurrent(self.object_store.io_parallelism(), |version| async move { + let path = self.markers_path.clone().join(version.to_string()); + match self + .object_store + .inner + .put_opts( + &path, + Bytes::new().into(), + PutOptions { + mode: PutMode::Create, + ..Default::default() + }, + ) + .await + { + Ok(_) => Ok(()), + Err(object_store::Error::AlreadyExists { .. }) + | Err(object_store::Error::Precondition { .. }) => Ok(()), + Err(error) => Err(Error::from(error)), + } + }) + .await + } + + async fn ensure_not_fenced(&self, version: u64) -> Result<()> { + let marker = self.markers_path.clone().join(version.to_string()); + if self.object_store.exists(&marker).await? { + Err(Error::VersionNotFound { + message: format!( + "version {} has been fenced for cleanup and cannot be leased", + version + ), + }) + } else { + Ok(()) + } + } +} + +impl Dataset { + /// Acquire a renewable advisory lease for this dataset version. + /// + /// Cleanup retains the version until the lease expires. Acquire the lease + /// before starting a long-running read and renew it before expiration. + /// + /// # Example + /// + /// ``` + /// # use std::time::Duration; + /// # use lance::{Dataset, Result}; + /// # async fn read_historical(dataset: &Dataset) -> Result<()> { + /// let historical = dataset.checkout_version(1).await?; + /// let lease = historical + /// .acquire_version_lease(Duration::from_secs(60)) + /// .await?; + /// let _batch = historical.scan().try_into_batch().await?; + /// lease.release().await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn acquire_version_lease(&self, ttl: Duration) -> Result { + VersionLeaseStore::for_dataset(self) + .await? + .acquire(self.version().version, ttl) + .await + } +} + +fn expiration_from_ttl(now: DateTime, ttl: Duration) -> Result> { + if ttl.is_zero() { + return Err(Error::invalid_input( + "version lease TTL must be greater than zero", + )); + } + let ttl = TimeDelta::from_std(ttl).map_err(|error| { + Error::invalid_input(format!( + "version lease TTL {ttl:?} is out of range: {error}" + )) + })?; + now.checked_add_signed(ttl).ok_or_else(|| { + Error::invalid_input(format!( + "version lease TTL {ttl:?} overflows its expiration" + )) + }) +} + +fn parse_lease_file_name(file_name: &str) -> std::result::Result<(u64, i64), String> { + let stem = file_name.strip_suffix(LEASE_FILE_SUFFIX).ok_or_else(|| { + format!("version lease file name '{file_name}' must end with {LEASE_FILE_SUFFIX}") + })?; + let mut parts = stem.splitn(3, '-'); + let version = parts + .next() + .ok_or_else(|| format!("version lease file name '{file_name}' has no version"))? + .parse::() + .map_err(|error| { + format!("version lease file name '{file_name}' has invalid version: {error}") + })?; + let expires_at_micros = parts + .next() + .ok_or_else(|| format!("version lease file name '{file_name}' has no expiration"))? + .parse::() + .map_err(|error| { + format!("version lease file name '{file_name}' has invalid expiration: {error}") + })?; + let lease_id = parts + .next() + .ok_or_else(|| format!("version lease file name '{file_name}' has no lease id"))?; + Uuid::parse_str(lease_id).map_err(|error| { + format!("version lease file name '{file_name}' has invalid lease id: {error}") + })?; + Ok((version, expires_at_micros)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_lease_file_name() { + let lease_id = Uuid::nil().simple(); + assert_eq!( + parse_lease_file_name(&format!("42-123-{lease_id}.lease")).unwrap(), + (42, 123) + ); + } +} From abe62f17d62c48d1523e5de093c484a996fdadc1 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 16:07:28 +0000 Subject: [PATCH 02/13] fix: make version lease retirement race-safe --- python/python/lance/dataset.py | 13 +- rust/lance/src/dataset/cleanup.rs | 134 +++++- rust/lance/src/dataset/refs.rs | 16 +- rust/lance/src/dataset/version_lease.rs | 599 ++++++++++++++++++++---- 4 files changed, 650 insertions(+), 112 deletions(-) diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index d1b50caf6c9..065a129218c 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3110,10 +3110,15 @@ def acquire_version_lease(self, ttl: timedelta) -> VersionLease: Examples -------- - >>> historical = dataset.checkout_version(1) - >>> with historical.acquire_version_lease(timedelta(minutes=5)) as lease: - ... table = historical.to_table() - ... lease.renew(timedelta(minutes=5)) + >>> import lance + >>> import pyarrow as pa + >>> from datetime import timedelta + >>> from tempfile import TemporaryDirectory + >>> with TemporaryDirectory() as tmp: + ... historical = lance.write_dataset(pa.table({"id": [1]}), tmp) + ... with historical.acquire_version_lease(timedelta(minutes=5)) as lease: + ... table = historical.to_table() + ... lease.renew(timedelta(minutes=5)) """ return self._ds.acquire_version_lease(td_to_micros(ttl)) diff --git a/rust/lance/src/dataset/cleanup.rs b/rust/lance/src/dataset/cleanup.rs index d91605b44e8..3cf36883805 100644 --- a/rust/lance/src/dataset/cleanup.rs +++ b/rust/lance/src/dataset/cleanup.rs @@ -34,7 +34,7 @@ //! happening at the same time) use super::refs::TagContents; -use super::version_lease::VersionLeaseStore; +use super::version_lease::{RetirementGuard, VersionLeaseStore}; use crate::dataset::TRANSACTIONS_DIR; use crate::{Dataset, utils::temporal::utc_now}; use chrono::{DateTime, TimeDelta, Utc}; @@ -315,6 +315,12 @@ struct CleanupInspection { earliest_retained_manifest_time: Option>, } +#[derive(Debug)] +struct CleanupRetirement { + guard: RetirementGuard, + manifest_paths: HashMap>, +} + /// If a file cannot be verified then it will only be deleted if it is at least /// this many days old. const UNVERIFIED_THRESHOLD_DAYS: i64 = 7; @@ -469,9 +475,14 @@ impl<'a> CleanupTask<'a> { .collect(); let version_lease_store = VersionLeaseStore::for_dataset(self.dataset).await?; - let leased_versions = version_lease_store - .active_versions(self.action.deletes_files()) - .await?; + // Execute establishes storage-clock retirement fences before deciding + // which leases are active. Explain remains read-only and conservatively + // treats every published lease as active. + let leased_versions = if self.action.deletes_files() { + HashSet::new() + } else { + version_lease_store.all_lease_versions().await? + }; let mut inspection = self .process_manifests(&tagged_versions, &leased_versions) @@ -495,14 +506,18 @@ impl<'a> CleanupTask<'a> { .await? }; - inspection = self + let (inspection, mut retirement) = self .fence_old_versions_and_retain_new_leases(inspection, &version_lease_store) .await?; - final_result.merge( - self.delete_unreferenced_files(inspection).await?, - candidate_file_limit, - ); + let cleanup_result = self.delete_unreferenced_files(inspection).await?; + if let Some(retirement) = retirement.as_mut() { + retirement + .guard + .finalize(&retirement.manifest_paths) + .await?; + } + final_result.merge(cleanup_result, candidate_file_limit); Ok(final_result) } @@ -591,24 +606,81 @@ impl<'a> CleanupTask<'a> { &self, mut inspection: CleanupInspection, version_lease_store: &VersionLeaseStore, - ) -> Result { + ) -> Result<(CleanupInspection, Option)> { if !self.action.deletes_files() { - return Ok(inspection); + return Ok((inspection, None)); } let versions_to_delete: HashSet = inspection.old_manifests.values().copied().collect(); - version_lease_store + let mut guard = version_lease_store .fence_versions(&versions_to_delete) .await?; - // A lease acquired between the initial list and the fence either sees - // the fence and fails, or appears here and conservatively retains this - // version for the current cleanup pass. - let leased_versions = version_lease_store.active_versions(true).await?; - for version in versions_to_delete.intersection(&leased_versions) { - self.retain_version(&mut inspection, *version).await?; + let result = async { + // Draining blocks new acquisitions but permits leases admitted before + // the marker to renew. Storage timestamps on the marker and lease + // provide a clock-skew-independent liveness comparison. + let leased_versions = version_lease_store + .active_versions_at(&guard.observed_at(), true) + .await?; + let retained_versions: HashSet<_> = versions_to_delete + .intersection(&leased_versions) + .copied() + .collect(); + for version in &retained_versions { + self.retain_version(&mut inspection, *version).await?; + } + guard.cancel_versions(&retained_versions).await?; + + let versions_to_seal: HashSet = + inspection.old_manifests.values().copied().collect(); + guard.seal_versions(&versions_to_seal).await?; + + // Recheck after sealing. A lease published across the draining + // transition either appears here and cancels this retirement, or + // observes the seal and fails before cleanup begins deletion. + let leased_versions = version_lease_store + .active_versions_at(&guard.observed_at(), true) + .await?; + let retained_versions: HashSet<_> = versions_to_seal + .intersection(&leased_versions) + .copied() + .collect(); + for version in &retained_versions { + self.retain_version(&mut inspection, *version).await?; + } + guard.cancel_versions(&retained_versions).await?; + + let mut manifest_paths = HashMap::>::new(); + for (path, version) in &inspection.old_manifests { + manifest_paths + .entry(*version) + .or_default() + .push(path.clone()); + } + Ok::<_, Error>((inspection, manifest_paths)) + } + .await; + + match result { + Ok((inspection, _)) if guard.is_empty() => Ok((inspection, None)), + Ok((inspection, manifest_paths)) => Ok(( + inspection, + Some(CleanupRetirement { + guard, + manifest_paths, + }), + )), + Err(error) => { + if let Err(cancel_error) = guard.cancel_all().await { + warn!( + error = %cancel_error, + "Failed to cancel version retirement before deletion" + ); + } + Err(error) + } } - Ok(inspection) } async fn retain_version(&self, inspection: &mut CleanupInspection, version: u64) -> Result<()> { @@ -2241,6 +2313,10 @@ mod tests { .unwrap(); assert_eq!(removed.old_versions, 0); assert_eq!(removed.data_files_removed, 0); + lease + .renew(Duration::from_secs(30 * 24 * 60 * 60)) + .await + .unwrap(); historical.scan().try_into_batch().await.unwrap(); lease.release().await.unwrap(); @@ -2249,6 +2325,21 @@ mod tests { .await .unwrap(); assert_eq!(removed.old_versions, 1); + let marker_prefix = historical + .refs + .root() + .unwrap() + .path + .join("_refs/version_lease_gc/main/1"); + assert!( + historical + .object_store + .list(Some(marker_prefix)) + .try_next() + .await + .unwrap() + .is_none() + ); assert!( fixture .load() @@ -2266,9 +2357,10 @@ mod tests { fixture.create_some_data().await.unwrap(); let historical = fixture.load().await.unwrap(); let _lease = historical - .acquire_version_lease(Duration::from_secs(1)) + .acquire_version_lease(Duration::from_millis(10)) .await .unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; fixture.overwrite_some_data().await.unwrap(); MockClock::set_system_time(TimeDelta::try_days(10).unwrap().to_std().unwrap()); @@ -2281,7 +2373,7 @@ mod tests { .acquire_version_lease(Duration::from_secs(60)) .await .unwrap_err(); - assert!(error.to_string().contains("fenced for cleanup"), "{error}"); + assert!(error.to_string().contains("no longer exists"), "{error}"); } #[tokio::test] diff --git a/rust/lance/src/dataset/refs.rs b/rust/lance/src/dataset/refs.rs index 0d3f65f7959..0f9e9ba4465 100644 --- a/rust/lance/src/dataset/refs.rs +++ b/rust/lance/src/dataset/refs.rs @@ -14,6 +14,7 @@ use std::sync::Arc; use crate::dataset::branch_location::BranchLocation; use crate::dataset::refs::Ref::{Tag, Version, VersionNumber}; +use crate::dataset::version_lease::remove_branch_state; use crate::utils::temporal::utc_now; use crate::{Error, Result}; use serde::de::DeserializeOwned; @@ -555,7 +556,7 @@ impl Branches<'_> { let branch_id = all_branches .get(branch) .map(|contents| contents.identifier.clone()); - if let Some(branch_id) = branch_id { + if let Some(branch_id) = branch_id.as_ref() { let referenced_versions = branch_id.collect_referenced_versions(&all_branches); if !referenced_versions.is_empty() { return Err(Error::RefConflict { @@ -579,8 +580,17 @@ impl Branches<'_> { self.object_store().delete(&branch_file).await?; } - // Clean up branch directories - self.cleanup_branch_directories(branch).await + // Clean up branch directories and operational state scoped to this + // branch incarnation. The identifier prevents a recreated branch with + // the same name from inheriting leases or retirement markers. + self.cleanup_branch_directories(branch).await?; + if let Some(namespace) = branch_id + .as_ref() + .and_then(|identifier| identifier.version_mapping.last().map(|(_, id)| id.as_str())) + { + remove_branch_state(self.object_store(), &root_location.path, namespace).await?; + } + Ok(()) } pub async fn list_ordered( diff --git a/rust/lance/src/dataset/version_lease.rs b/rust/lance/src/dataset/version_lease.rs index 6a1c4b59c3d..9a688a73e0f 100644 --- a/rust/lance/src/dataset/version_lease.rs +++ b/rust/lance/src/dataset/version_lease.rs @@ -3,21 +3,27 @@ //! Advisory leases that protect dataset versions from cleanup. -use std::{collections::HashSet, sync::Arc, time::Duration}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, + time::Duration, +}; use bytes::Bytes; use chrono::{DateTime, TimeDelta, Utc}; use futures::{StreamExt, TryStreamExt, stream}; use lance_io::object_store::ObjectStore; -use object_store::{PutMode, PutOptions, path::Path}; +use object_store::{ObjectMeta, ObjectStoreExt, PutMode, PutOptions, path::Path}; use uuid::Uuid; use super::{Dataset, refs::MAIN_BRANCH}; -use crate::{Error, Result, utils::temporal::utc_now}; +use crate::{Error, Result}; const LEASES_DIR: &str = "_refs/version_leases"; const LEASE_GC_MARKERS_DIR: &str = "_refs/version_lease_gc"; const LEASE_FILE_SUFFIX: &str = ".lease"; +const DRAINING_MARKER_SUFFIX: &str = ".draining"; +const SEALED_MARKER_SUFFIX: &str = ".sealed"; /// A renewable advisory lease that protects one dataset version from cleanup. /// @@ -48,27 +54,24 @@ impl VersionLease { /// Renew this lease for `ttl` from the current time. /// - /// Renewal must complete before the current expiration. It fails if the - /// lease has expired or cleanup has fenced the version for deletion. + /// Renewal must complete before the current expiration. A lease admitted + /// before cleanup starts draining its version may renew while draining; + /// renewal fails once cleanup seals the version for deletion. pub async fn renew(&mut self, ttl: Duration) -> Result<()> { - let now = utc_now(); - if now >= self.expires_at { - return Err(Error::invalid_input(format!( - "cannot renew expired version lease for version {}: expired at {}", - self.version, self.expires_at - ))); + if !self.store.object_store.exists(&self.path).await? { + return Err(expired_lease_error(self.version, self.expires_at)); } - let expires_at = expiration_from_ttl(now, ttl)?; - self.store.ensure_not_fenced(self.version).await?; + self.store.ensure_not_sealed(self.version).await?; // Publish the replacement before removing the old file so cleanup never // observes a gap between a timely renewal and its predecessor. - let path = self - .store - .create_lease_file(self.version, expires_at) - .await?; - if let Err(error) = self.store.ensure_not_fenced(self.version).await { - let _ = self.store.object_store.delete(&path).await; + let lease_file = self.store.create_lease_file(self.version, ttl).await?; + if lease_file.created_at >= self.expires_at { + let _ = self.store.object_store.delete(&lease_file.path).await; + return Err(expired_lease_error(self.version, self.expires_at)); + } + if let Err(error) = self.store.ensure_not_sealed(self.version).await { + let _ = self.store.object_store.delete(&lease_file.path).await; return Err(error); } @@ -81,8 +84,8 @@ impl VersionLease { "Failed to remove superseded version lease" ); } - self.path = path; - self.expires_at = expires_at; + self.path = lease_file.path; + self.expires_at = lease_file.expires_at; Ok(()) } @@ -122,6 +125,33 @@ pub(super) struct VersionLeaseStore { object_store: Arc, leases_path: Path, markers_path: Path, + manifest_path: Option, +} + +#[derive(Debug)] +struct LeaseFile { + path: Path, + created_at: DateTime, + expires_at: DateTime, +} + +#[derive(Debug)] +struct RetirementFence { + draining_path: Option, + sealed_path: Option, + observed_at: DateTime, +} + +/// Per-cleanup retirement markers for versions selected for deletion. +/// +/// Draining markers reject new leases while allowing already-admitted leases +/// to renew. Sealed markers reject both acquisition and renewal immediately +/// before deletion. Unique operation markers keep cancellation and finalization +/// safe when multiple cleanups overlap. +#[derive(Debug)] +pub(super) struct RetirementGuard { + store: VersionLeaseStore, + fences: HashMap, } impl VersionLeaseStore { @@ -148,31 +178,32 @@ impl VersionLeaseStore { object_store: Arc::clone(&dataset.object_store), leases_path: root.path.clone().join(LEASES_DIR).join(namespace), markers_path: root.path.join(LEASE_GC_MARKERS_DIR).join(namespace), + manifest_path: Some(dataset.manifest_location.path.clone()), }) } async fn acquire(&self, version: u64, ttl: Duration) -> Result { - self.ensure_not_fenced(version).await?; - let expires_at = expiration_from_ttl(utc_now(), ttl)?; - let path = self.create_lease_file(version, expires_at).await?; - if let Err(error) = self.ensure_not_fenced(version).await { - let _ = self.object_store.delete(&path).await; + self.ensure_version_available(version).await?; + let lease_file = self.create_lease_file(version, ttl).await?; + if let Err(error) = self.ensure_version_available(version).await { + let _ = self.object_store.delete(&lease_file.path).await; return Err(error); } Ok(VersionLease { store: self.clone(), - path, + path: lease_file.path, version, - expires_at, + expires_at: lease_file.expires_at, released: false, }) } - async fn create_lease_file(&self, version: u64, expires_at: DateTime) -> Result { + async fn create_lease_file(&self, version: u64, ttl: Duration) -> Result { + let ttl_micros = ttl_micros(ttl)?; let path = self.leases_path.clone().join(format!( "{}-{}-{}{}", version, - expires_at.timestamp_micros(), + ttl_micros, Uuid::new_v4().simple(), LEASE_FILE_SUFFIX )); @@ -187,28 +218,47 @@ impl VersionLeaseStore { }, ) .await?; - Ok(path) + let metadata = match self.object_store.inner.head(&path).await { + Ok(metadata) => metadata, + Err(error) => { + let _ = self.object_store.delete(&path).await; + return Err(error.into()); + } + }; + let expires_at = expiration_from_ttl(metadata.last_modified, ttl)?; + Ok(LeaseFile { + path, + created_at: metadata.last_modified, + expires_at, + }) } - pub(super) async fn active_versions(&self, remove_expired: bool) -> Result> { - let lease_files = match self.object_store.read_dir(self.leases_path.clone()).await { - Ok(files) => files, - Err(error) if error.is_not_found() => return Ok(HashSet::new()), - Err(error) => return Err(error), - }; - let now_micros = utc_now().timestamp_micros(); + pub(super) async fn all_lease_versions(&self) -> Result> { + self.lease_metadata() + .await? + .into_iter() + .map(|metadata| parse_lease_metadata(&metadata).map(|(version, _)| version)) + .collect() + } + + pub(super) async fn active_versions_at( + &self, + observed_at: &HashMap>, + remove_expired: bool, + ) -> Result> { let mut active_versions = HashSet::new(); let mut expired_paths = Vec::new(); - for file_name in lease_files { - let (version, expires_at_micros) = - parse_lease_file_name(&file_name).map_err(|error| { - Error::corrupt_file(self.leases_path.clone().join(file_name.as_str()), error) - })?; - if expires_at_micros > now_micros { + for metadata in self.lease_metadata().await? { + let (version, ttl) = parse_lease_metadata(&metadata)?; + let Some(reference_time) = observed_at.get(&version) else { + continue; + }; + let expires_at = expiration_from_ttl(metadata.last_modified, ttl)?; + if expires_at > *reference_time { active_versions.insert(version); } else if remove_expired { - expired_paths.push(self.leases_path.clone().join(file_name)); + expired_paths.push(metadata.location); } } @@ -227,46 +277,303 @@ impl VersionLeaseStore { Ok(active_versions) } - pub(super) async fn fence_versions(&self, versions: &HashSet) -> Result<()> { - stream::iter(versions.iter().copied()) - .map(Ok) - .try_for_each_concurrent(self.object_store.io_parallelism(), |version| async move { - let path = self.markers_path.clone().join(version.to_string()); - match self - .object_store - .inner - .put_opts( - &path, - Bytes::new().into(), - PutOptions { - mode: PutMode::Create, - ..Default::default() - }, - ) - .await - { - Ok(_) => Ok(()), - Err(object_store::Error::AlreadyExists { .. }) - | Err(object_store::Error::Precondition { .. }) => Ok(()), - Err(error) => Err(Error::from(error)), + pub(super) async fn fence_versions(&self, versions: &HashSet) -> Result { + let operation_id = Uuid::new_v4().simple().to_string(); + let results = stream::iter(versions.iter().copied()) + .map(|version| { + let path = self + .marker_version_path(version) + .join(format!("{operation_id}{DRAINING_MARKER_SUFFIX}")); + async move { + let metadata = self.create_marker(&path).await?; + Ok::<_, Error>((version, path, metadata.last_modified)) } }) - .await + .buffer_unordered(self.object_store.io_parallelism()) + .collect::>() + .await; + + let mut fences = HashMap::with_capacity(results.len()); + let mut first_error = None; + for result in results { + match result { + Ok((version, path, observed_at)) => { + fences.insert( + version, + RetirementFence { + draining_path: Some(path), + sealed_path: None, + observed_at, + }, + ); + } + Err(error) if first_error.is_none() => first_error = Some(error), + Err(_) => {} + } + } + let mut guard = RetirementGuard { + store: self.clone(), + fences, + }; + if let Some(error) = first_error { + guard.cancel_all().await?; + return Err(error); + } + Ok(guard) } - async fn ensure_not_fenced(&self, version: u64) -> Result<()> { - let marker = self.markers_path.clone().join(version.to_string()); - if self.object_store.exists(&marker).await? { + async fn ensure_version_available(&self, version: u64) -> Result<()> { + if let Some(manifest_path) = &self.manifest_path + && !self.object_store.exists(manifest_path).await? + { Err(Error::VersionNotFound { - message: format!( - "version {} has been fenced for cleanup and cannot be leased", - version - ), + message: format!("version {version} no longer exists and cannot be leased"), }) + } else if self.has_marker(version, None).await? { + Err(retiring_version_error(version)) } else { Ok(()) } } + + async fn ensure_not_sealed(&self, version: u64) -> Result<()> { + if self.has_marker(version, Some(SEALED_MARKER_SUFFIX)).await? { + Err(retiring_version_error(version)) + } else { + Ok(()) + } + } + + async fn has_marker(&self, version: u64, suffix: Option<&str>) -> Result { + let mut markers = self + .object_store + .list(Some(self.marker_version_path(version))); + while let Some(metadata) = markers.try_next().await? { + if suffix.is_none_or(|suffix| { + metadata + .location + .filename() + .is_some_and(|name| name.ends_with(suffix)) + }) { + return Ok(true); + } + } + Ok(false) + } + + async fn lease_metadata(&self) -> Result> { + self.object_store + .list(Some(self.leases_path.clone())) + .try_collect() + .await + } + + fn marker_version_path(&self, version: u64) -> Path { + self.markers_path.clone().join(version.to_string()) + } + + async fn create_marker(&self, path: &Path) -> Result { + self.object_store + .inner + .put_opts( + path, + Bytes::new().into(), + PutOptions { + mode: PutMode::Create, + ..Default::default() + }, + ) + .await?; + match self.object_store.inner.head(path).await { + Ok(metadata) => Ok(metadata), + Err(error) => { + let _ = self.object_store.delete(path).await; + Err(error.into()) + } + } + } +} + +impl RetirementGuard { + pub(super) fn observed_at(&self) -> HashMap> { + self.fences + .iter() + .map(|(version, fence)| (*version, fence.observed_at)) + .collect() + } + + pub(super) fn is_empty(&self) -> bool { + self.fences.is_empty() + } + + pub(super) async fn seal_versions(&mut self, versions: &HashSet) -> Result<()> { + let marker_paths = versions + .iter() + .filter_map(|version| { + self.fences.get(version).and_then(|fence| { + fence.draining_path.as_ref().and_then(|draining_path| { + let file_name = draining_path.filename()?; + let operation_id = file_name.strip_suffix(DRAINING_MARKER_SUFFIX)?; + Some(( + *version, + self.store + .marker_version_path(*version) + .join(format!("{operation_id}{SEALED_MARKER_SUFFIX}")), + )) + }) + }) + }) + .collect::>(); + + let store = self.store.clone(); + let results = stream::iter(marker_paths) + .map(move |(version, path)| { + let store = store.clone(); + async move { + let metadata = store.create_marker(&path).await?; + Ok::<_, Error>((version, path, metadata.last_modified)) + } + }) + .buffer_unordered(self.store.object_store.io_parallelism()) + .collect::>() + .await; + + let mut first_error = None; + for result in results { + match result { + Ok((version, sealed_path, observed_at)) => { + let fence = self.fences.get_mut(&version).ok_or_else(|| { + Error::internal(format!("missing draining fence for version {version}")) + })?; + fence.sealed_path = Some(sealed_path); + fence.observed_at = observed_at; + } + Err(error) if first_error.is_none() => first_error = Some(error), + Err(_) => {} + } + } + if let Some(error) = first_error { + return Err(error); + } + + let draining_paths = versions + .iter() + .filter_map(|version| { + self.fences + .get(version) + .and_then(|fence| fence.draining_path.clone()) + }) + .collect::>(); + self.store.delete_paths(draining_paths).await?; + for version in versions { + if let Some(fence) = self.fences.get_mut(version) { + fence.draining_path = None; + } + } + Ok(()) + } + + pub(super) async fn cancel_versions(&mut self, versions: &HashSet) -> Result<()> { + let paths = versions + .iter() + .filter_map(|version| self.fences.get(version)) + .flat_map(|fence| { + [fence.draining_path.clone(), fence.sealed_path.clone()] + .into_iter() + .flatten() + }) + .collect::>(); + self.store.delete_paths(paths).await?; + self.fences.retain(|version, _| !versions.contains(version)); + Ok(()) + } + + pub(super) async fn cancel_all(&mut self) -> Result<()> { + let versions = self.fences.keys().copied().collect(); + self.cancel_versions(&versions).await + } + + pub(super) async fn finalize( + &mut self, + manifest_paths: &HashMap>, + ) -> Result<()> { + for version in self.fences.keys() { + let paths = manifest_paths.get(version).ok_or_else(|| { + Error::internal(format!( + "missing manifest identity for retiring version {version}" + )) + })?; + for path in paths { + if self.store.object_store.exists(path).await? { + return Err(Error::internal(format!( + "cannot finalize retirement for version {version}: manifest {path} still exists" + ))); + } + } + } + + let versions: HashSet<_> = self.fences.keys().copied().collect(); + let marker_prefixes = versions + .iter() + .map(|version| self.store.marker_version_path(*version)) + .collect::>(); + let marker_streams = marker_prefixes + .into_iter() + .map(|prefix| self.store.object_store.list(Some(prefix))); + let mut paths = stream::iter(marker_streams) + .flatten() + .map_ok(|metadata| metadata.location) + .try_collect::>() + .await?; + for metadata in self.store.lease_metadata().await? { + let (version, _) = parse_lease_metadata(&metadata)?; + if versions.contains(&version) { + paths.push(metadata.location); + } + } + self.store.delete_paths(paths).await?; + self.fences.clear(); + Ok(()) + } +} + +impl VersionLeaseStore { + async fn delete_paths(&self, paths: Vec) -> Result<()> { + stream::iter(paths) + .map(Ok) + .try_for_each_concurrent(self.object_store.io_parallelism(), |path| async move { + match self.object_store.delete(&path).await { + Ok(()) => Ok(()), + Err(error) if error.is_not_found() => Ok(()), + Err(error) => Err(error), + } + }) + .await + } +} + +pub(super) async fn remove_branch_state( + object_store: &ObjectStore, + root_path: &Path, + namespace: &str, +) -> Result<()> { + for path in [ + root_path + .clone() + .join(LEASES_DIR) + .join(namespace.to_string()), + root_path + .clone() + .join(LEASE_GC_MARKERS_DIR) + .join(namespace.to_string()), + ] { + match object_store.remove_dir_all(path).await { + Ok(()) => {} + Err(error) if error.is_not_found() => {} + Err(error) => return Err(error), + } + } + Ok(()) } impl Dataset { @@ -298,7 +605,7 @@ impl Dataset { } } -fn expiration_from_ttl(now: DateTime, ttl: Duration) -> Result> { +fn ttl_micros(ttl: Duration) -> Result { if ttl.is_zero() { return Err(Error::invalid_input( "version lease TTL must be greater than zero", @@ -309,14 +616,38 @@ fn expiration_from_ttl(now: DateTime, ttl: Duration) -> Result, ttl: Duration) -> Result> { + let ttl = TimeDelta::from_std(ttl).map_err(|error| { + Error::invalid_input(format!( + "version lease TTL {ttl:?} is out of range: {error}" + )) + })?; + observed_at.checked_add_signed(ttl).ok_or_else(|| { Error::invalid_input(format!( "version lease TTL {ttl:?} overflows its expiration" )) }) } -fn parse_lease_file_name(file_name: &str) -> std::result::Result<(u64, i64), String> { +fn parse_lease_metadata(metadata: &ObjectMeta) -> Result<(u64, Duration)> { + let file_name = metadata.location.filename().ok_or_else(|| { + Error::corrupt_file( + metadata.location.clone(), + "version lease path has no filename", + ) + })?; + parse_lease_file_name(file_name) + .map_err(|error| Error::corrupt_file(metadata.location.clone(), error)) +} + +fn parse_lease_file_name(file_name: &str) -> std::result::Result<(u64, Duration), String> { let stem = file_name.strip_suffix(LEASE_FILE_SUFFIX).ok_or_else(|| { format!("version lease file name '{file_name}' must end with {LEASE_FILE_SUFFIX}") })?; @@ -328,32 +659,132 @@ fn parse_lease_file_name(file_name: &str) -> std::result::Result<(u64, i64), Str .map_err(|error| { format!("version lease file name '{file_name}' has invalid version: {error}") })?; - let expires_at_micros = parts + let ttl_micros = parts .next() - .ok_or_else(|| format!("version lease file name '{file_name}' has no expiration"))? - .parse::() + .ok_or_else(|| format!("version lease file name '{file_name}' has no TTL"))? + .parse::() .map_err(|error| { - format!("version lease file name '{file_name}' has invalid expiration: {error}") + format!("version lease file name '{file_name}' has invalid TTL: {error}") })?; + if ttl_micros == 0 { + return Err(format!( + "version lease file name '{file_name}' has a zero TTL" + )); + } let lease_id = parts .next() .ok_or_else(|| format!("version lease file name '{file_name}' has no lease id"))?; Uuid::parse_str(lease_id).map_err(|error| { format!("version lease file name '{file_name}' has invalid lease id: {error}") })?; - Ok((version, expires_at_micros)) + Ok((version, Duration::from_micros(ttl_micros))) +} + +fn expired_lease_error(version: u64, expires_at: DateTime) -> Error { + Error::invalid_input(format!( + "cannot renew expired version lease for version {version}: expired at {expires_at}" + )) +} + +fn retiring_version_error(version: u64) -> Error { + Error::VersionNotFound { + message: format!("version {version} is retiring and cannot accept a new lease"), + } } #[cfg(test)] mod tests { + use mock_instant::thread_local::MockClock; + use super::*; + fn memory_store() -> VersionLeaseStore { + VersionLeaseStore { + object_store: Arc::new(ObjectStore::memory()), + leases_path: Path::from("leases"), + markers_path: Path::from("markers"), + manifest_path: None, + } + } + #[test] fn parses_lease_file_name() { let lease_id = Uuid::nil().simple(); assert_eq!( parse_lease_file_name(&format!("42-123-{lease_id}.lease")).unwrap(), - (42, 123) + (42, Duration::from_micros(123)) + ); + } + + #[tokio::test] + async fn draining_lease_remains_renewable() { + let store = memory_store(); + let mut lease = store.acquire(42, Duration::from_secs(60)).await.unwrap(); + let mut guard = store.fence_versions(&HashSet::from([42])).await.unwrap(); + + assert!( + store + .active_versions_at(&guard.observed_at(), false) + .await + .unwrap() + .contains(&42) + ); + lease.renew(Duration::from_secs(60)).await.unwrap(); + guard.cancel_all().await.unwrap(); + } + + #[tokio::test] + async fn lease_expiry_uses_storage_clock() { + MockClock::set_system_time(Duration::from_secs(100)); + let store = memory_store(); + let _lease = store.acquire(42, Duration::from_secs(60)).await.unwrap(); + + // Moving the cleanup host clock ahead does not affect the storage + // timestamps used for lease liveness. + MockClock::set_system_time(Duration::from_secs(160)); + let mut guard = store.fence_versions(&HashSet::from([42])).await.unwrap(); + assert!( + store + .active_versions_at(&guard.observed_at(), false) + .await + .unwrap() + .contains(&42) ); + guard.cancel_all().await.unwrap(); + } + + #[tokio::test] + async fn sealed_version_rejects_renewal() { + let store = memory_store(); + let mut lease = store.acquire(42, Duration::from_secs(60)).await.unwrap(); + let mut guard = store.fence_versions(&HashSet::from([42])).await.unwrap(); + guard.seal_versions(&HashSet::from([42])).await.unwrap(); + + let error = lease.renew(Duration::from_secs(60)).await.unwrap_err(); + assert!(matches!(error, Error::VersionNotFound { .. })); + assert!(error.to_string().contains("retiring"), "{error}"); + guard.cancel_all().await.unwrap(); + } + + #[tokio::test] + async fn branch_deletion_removes_incarnation_state() { + let store = memory_store(); + let root = Path::from("dataset"); + let namespace = "branch-id"; + let lease_path = root.clone().join(LEASES_DIR).join(namespace).join("lease"); + let marker_path = root + .clone() + .join(LEASE_GC_MARKERS_DIR) + .join(namespace) + .join("marker"); + store.object_store.put(&lease_path, &[]).await.unwrap(); + store.object_store.put(&marker_path, &[]).await.unwrap(); + + remove_branch_state(&store.object_store, &root, namespace) + .await + .unwrap(); + + assert!(!store.object_store.exists(&lease_path).await.unwrap()); + assert!(!store.object_store.exists(&marker_path).await.unwrap()); } } From c6f89e998a081d9109cdcb8b73db1620da575572 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 17:36:51 +0000 Subject: [PATCH 03/13] fix: recover interrupted version retirement --- rust/lance/src/dataset/cleanup.rs | 165 ++++++- rust/lance/src/dataset/version_lease.rs | 588 +++++++++++++++++++++--- 2 files changed, 668 insertions(+), 85 deletions(-) diff --git a/rust/lance/src/dataset/cleanup.rs b/rust/lance/src/dataset/cleanup.rs index 3cf36883805..340936a65da 100644 --- a/rust/lance/src/dataset/cleanup.rs +++ b/rust/lance/src/dataset/cleanup.rs @@ -65,7 +65,7 @@ use std::fmt::Debug; use std::{ collections::{HashMap, HashSet}, future, - sync::Mutex, + sync::{Arc, Mutex}, time::Duration, }; use tokio::time::{MissedTickBehavior, interval}; @@ -475,6 +475,11 @@ impl<'a> CleanupTask<'a> { .collect(); let version_lease_store = VersionLeaseStore::for_dataset(self.dataset).await?; + let forced_retirement_versions = if self.action.deletes_files() { + version_lease_store.clone().recover_retirements().await? + } else { + HashSet::new() + }; // Execute establishes storage-clock retirement fences before deciding // which leases are active. Explain remains read-only and conservatively // treats every published lease as active. @@ -485,7 +490,11 @@ impl<'a> CleanupTask<'a> { }; let mut inspection = self - .process_manifests(&tagged_versions, &leased_versions) + .process_manifests( + &tagged_versions, + &leased_versions, + &forced_retirement_versions, + ) .await?; if self.policy.error_if_tagged_old_versions && !inspection.tagged_old_versions.is_empty() { @@ -502,12 +511,21 @@ impl<'a> CleanupTask<'a> { .cloned() .collect(); inspection = self - .retain_branch_lineage_files(inspection, &referenced_branches, &ignored_manifests) + .retain_branch_lineage_files( + inspection, + &referenced_branches, + &ignored_manifests, + &forced_retirement_versions, + ) .await? }; let (inspection, mut retirement) = self - .fence_old_versions_and_retain_new_leases(inspection, &version_lease_store) + .fence_old_versions_and_retain_new_leases( + inspection, + &version_lease_store, + &forced_retirement_versions, + ) .await?; let cleanup_result = self.delete_unreferenced_files(inspection).await?; @@ -526,6 +544,7 @@ impl<'a> CleanupTask<'a> { &'a self, tagged_versions: &HashSet, leased_versions: &HashSet, + forced_retirement_versions: &HashSet, ) -> Result { let inspection = Mutex::new(CleanupInspection::default()); self.dataset @@ -533,7 +552,13 @@ impl<'a> CleanupTask<'a> { .list_manifest_locations(&self.dataset.base, &self.dataset.object_store, false) .try_filter(|location| future::ready(!self.ignored_manifests.contains(&location.path))) .try_for_each_concurrent(self.dataset.object_store.io_parallelism(), |location| { - self.process_manifest_file(location, &inspection, tagged_versions, leased_versions) + self.process_manifest_file( + location, + &inspection, + tagged_versions, + leased_versions, + forced_retirement_versions, + ) }) .await?; Ok(inspection.into_inner().unwrap()) @@ -545,6 +570,7 @@ impl<'a> CleanupTask<'a> { inspection: &Mutex, tagged_versions: &HashSet, leased_versions: &HashSet, + forced_retirement_versions: &HashSet, ) -> Result<()> { // TODO: We can't cleanup invalid manifests. There is no way to distinguish // between an invalid manifest and a temporary I/O error. It's also not safe @@ -581,13 +607,14 @@ impl<'a> CleanupTask<'a> { let is_latest = self.read_version <= manifest.version; let is_tagged = tagged_versions.contains(&manifest.version); let is_leased = leased_versions.contains(&manifest.version); - let in_working_set = - is_latest || !self.policy.should_clean(&manifest) || is_tagged || is_leased; + let is_forced_retirement = forced_retirement_versions.contains(&manifest.version); + let in_working_set = !is_forced_retirement + && (is_latest || !self.policy.should_clean(&manifest) || is_tagged || is_leased); let mut inspection = inspection.lock().unwrap(); // Track tagged old versions in case we want to return a `CleanupError` later. // Only track tagged when it is old. - if is_tagged && !is_latest && self.policy.should_clean(&manifest) { + if is_tagged && !is_forced_retirement && !is_latest && self.policy.should_clean(&manifest) { inspection.tagged_old_versions.insert(manifest.version); } @@ -606,14 +633,25 @@ impl<'a> CleanupTask<'a> { &self, mut inspection: CleanupInspection, version_lease_store: &VersionLeaseStore, + forced_retirement_versions: &HashSet, ) -> Result<(CleanupInspection, Option)> { if !self.action.deletes_files() { return Ok((inspection, None)); } - let versions_to_delete: HashSet = inspection.old_manifests.values().copied().collect(); + let mut manifest_paths_to_delete = HashMap::>::new(); + for (path, version) in &inspection.old_manifests { + manifest_paths_to_delete + .entry(*version) + .or_default() + .push(path.clone()); + } + let versions_to_delete = manifest_paths_to_delete + .keys() + .copied() + .collect::>(); let mut guard = version_lease_store - .fence_versions(&versions_to_delete) + .fence_versions(&manifest_paths_to_delete) .await?; let result = async { @@ -625,6 +663,7 @@ impl<'a> CleanupTask<'a> { .await?; let retained_versions: HashSet<_> = versions_to_delete .intersection(&leased_versions) + .filter(|version| !forced_retirement_versions.contains(version)) .copied() .collect(); for version in &retained_versions { @@ -644,6 +683,7 @@ impl<'a> CleanupTask<'a> { .await?; let retained_versions: HashSet<_> = versions_to_seal .intersection(&leased_versions) + .filter(|version| !forced_retirement_versions.contains(version)) .copied() .collect(); for version in &retained_versions { @@ -1304,28 +1344,42 @@ impl<'a> CleanupTask<'a> { inspection: CleanupInspection, referenced_branches: &[(String, u64)], removed_branch_manifests: &HashSet, + forced_retirement_versions: &HashSet, ) -> Result { - let inspection = Mutex::new(inspection); - for (branch, root_version_number) in referenced_branches { + let inspection = Arc::new(Mutex::new(inspection)); + let removed_branch_manifests = Arc::new(removed_branch_manifests.clone()); + let forced_retirement_versions = Arc::new(forced_retirement_versions.clone()); + for (branch, referenced_version) in referenced_branches.iter().cloned() { // Use find_branch to get the branch path directly without checkout. // This avoids creating a dataset instance and prevents manifest deletion // during the retain operation. - let branch_location = self.dataset.branch_location().find_branch(Some(branch))?; + let branch_location = self.dataset.branch_location().find_branch(Some(&branch))?; + let removed_branch_manifests = Arc::clone(&removed_branch_manifests); + let forced_retirement_versions = Arc::clone(&forced_retirement_versions); + let branch_inspection = Arc::clone(&inspection); self.dataset .commit_handler .list_manifest_locations(&branch_location.path, &self.dataset.object_store, false) - .try_filter(|location| { + .try_filter(move |location| { future::ready(!removed_branch_manifests.contains(&location.path)) }) .try_for_each_concurrent(self.dataset.object_store.io_parallelism(), |location| { - self.process_branch_referenced_manifests( - location, - *root_version_number, - &inspection, - ) + let forced_retirement_versions = Arc::clone(&forced_retirement_versions); + let branch_inspection = Arc::clone(&branch_inspection); + async move { + self.process_branch_referenced_manifests( + location, + referenced_version, + branch_inspection.as_ref(), + forced_retirement_versions.as_ref(), + ) + .await + } }) .await?; } + let inspection = Arc::try_unwrap(inspection) + .map_err(|_| Error::internal("branch retention inspection still has active owners"))?; Ok(inspection.into_inner().unwrap()) } @@ -1334,6 +1388,7 @@ impl<'a> CleanupTask<'a> { location: ManifestLocation, referenced_version: u64, inspection: &Mutex, + forced_retirement_versions: &HashSet, ) -> Result<()> { let manifest = read_manifest(&self.dataset.object_store, &location.path, location.size).await?; @@ -1403,7 +1458,7 @@ impl<'a> CleanupTask<'a> { } } } - if is_referenced { + if is_referenced && !forced_retirement_versions.contains(&referenced_version) { inspection .old_manifests .retain(|_path, version_number| *version_number != referenced_version); @@ -2286,6 +2341,7 @@ mod tests { &Mutex::new(CleanupInspection::default()), &HashSet::new(), &HashSet::new(), + &HashSet::new(), ) .await .unwrap(); @@ -2360,7 +2416,7 @@ mod tests { .acquire_version_lease(Duration::from_millis(10)) .await .unwrap(); - tokio::time::sleep(Duration::from_millis(50)).await; + tokio::time::sleep(Duration::from_millis(1_100)).await; fixture.overwrite_some_data().await.unwrap(); MockClock::set_system_time(TimeDelta::try_days(10).unwrap().to_std().unwrap()); @@ -2376,6 +2432,62 @@ mod tests { assert!(error.to_string().contains("no longer exists"), "{error}"); } + #[tokio::test] + async fn cleanup_resumes_sealed_retirement() { + let fixture = MockDatasetFixture::try_new().unwrap(); + fixture.create_some_data().await.unwrap(); + let historical = fixture.load().await.unwrap(); + fixture.overwrite_some_data().await.unwrap(); + + let version_lease_store = VersionLeaseStore::for_dataset(&historical).await.unwrap(); + let manifest_paths = HashMap::from([( + historical.version().version, + vec![historical.manifest_location.path.clone()], + )]); + let mut guard = version_lease_store + .fence_versions(&manifest_paths) + .await + .unwrap(); + guard + .seal_versions(&HashSet::from([historical.version().version])) + .await + .unwrap(); + drop(guard); + + // This policy would normally retain version 1. The durable sealed state + // must resume the already-started retirement instead. + let removed = fixture + .run_cleanup_with_policy(CleanupPolicy { + before_version: Some(1), + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(removed.old_versions, 1); + assert!( + !historical + .object_store + .exists(&manifest_paths[&1][0]) + .await + .unwrap() + ); + let marker_prefix = historical + .refs + .root() + .unwrap() + .path + .join("_refs/version_lease_gc/main/1"); + assert!( + historical + .object_store + .list(Some(marker_prefix)) + .try_next() + .await + .unwrap() + .is_none() + ); + } + #[tokio::test] async fn explain_cleanup_does_not_delete_files() { let fixture = MockDatasetFixture::try_new().unwrap(); @@ -3053,12 +3165,17 @@ mod tests { CleanupAction::Execute, ); let inspection = task - .process_manifests(&HashSet::new(), &HashSet::new()) + .process_manifests(&HashSet::new(), &HashSet::new(), &HashSet::new()) .await .unwrap(); let referenced_branches = task.find_referenced_branches().await.unwrap(); let inspection = task - .retain_branch_lineage_files(inspection, &referenced_branches, &HashSet::new()) + .retain_branch_lineage_files( + inspection, + &referenced_branches, + &HashSet::new(), + &HashSet::new(), + ) .await .unwrap(); @@ -3120,7 +3237,7 @@ mod tests { CleanupAction::Execute, ); let inspection = task - .process_manifests(&HashSet::new(), &HashSet::new()) + .process_manifests(&HashSet::new(), &HashSet::new(), &HashSet::new()) .await .unwrap(); let kept: HashSet<&Path> = inspection diff --git a/rust/lance/src/dataset/version_lease.rs b/rust/lance/src/dataset/version_lease.rs index 9a688a73e0f..a8aca681562 100644 --- a/rust/lance/src/dataset/version_lease.rs +++ b/rust/lance/src/dataset/version_lease.rs @@ -5,15 +5,17 @@ use std::{ collections::{HashMap, HashSet}, - sync::Arc, + sync::{Arc, LazyLock}, time::Duration, }; use bytes::Bytes; use chrono::{DateTime, TimeDelta, Utc}; +use dashmap::DashSet; use futures::{StreamExt, TryStreamExt, stream}; use lance_io::object_store::ObjectStore; use object_store::{ObjectMeta, ObjectStoreExt, PutMode, PutOptions, path::Path}; +use serde::{Deserialize, Serialize}; use uuid::Uuid; use super::{Dataset, refs::MAIN_BRANCH}; @@ -24,6 +26,21 @@ const LEASE_GC_MARKERS_DIR: &str = "_refs/version_lease_gc"; const LEASE_FILE_SUFFIX: &str = ".lease"; const DRAINING_MARKER_SUFFIX: &str = ".draining"; const SEALED_MARKER_SUFFIX: &str = ".sealed"; +const STORAGE_CLOCK_PATH: &str = "_clock"; + +/// HTTP `Last-Modified`, used by supported cloud stores, has whole-second precision. +/// Adding one interval guarantees a lease is never shortened by timestamp truncation. +const STORAGE_TIMESTAMP_PRECISION: Duration = Duration::from_secs(1); + +/// Draining does no deletion and must complete within this ownership window. +/// A cleaner that exceeds the window fails before sealing; another actor may then +/// ignore and remove the abandoned drain without racing a live deletion. +const DRAINING_OWNERSHIP_TIMEOUT: Duration = Duration::from_secs(15 * 60); + +/// Drains whose in-process owner was dropped before sealing. The operation UUID +/// makes each path unique across stores, while process-wide scope lets a fresh +/// dataset handle immediately disregard a drain that can no longer delete. +static LOCALLY_ABANDONED_DRAINS: LazyLock> = LazyLock::new(DashSet::new); /// A renewable advisory lease that protects one dataset version from cleanup. /// @@ -140,6 +157,20 @@ struct RetirementFence { draining_path: Option, sealed_path: Option, observed_at: DateTime, + manifest_paths: Vec, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +struct RetirementMarker { + manifest_paths: Vec, +} + +#[derive(Debug)] +struct RetirementMarkerMetadata { + version: u64, + operation_id: String, + is_sealed: bool, + metadata: ObjectMeta, } /// Per-cleanup retirement markers for versions selected for deletion. @@ -277,16 +308,24 @@ impl VersionLeaseStore { Ok(active_versions) } - pub(super) async fn fence_versions(&self, versions: &HashSet) -> Result { + pub(super) async fn fence_versions( + &self, + manifest_paths: &HashMap>, + ) -> Result { let operation_id = Uuid::new_v4().simple().to_string(); - let results = stream::iter(versions.iter().copied()) - .map(|version| { + let version_manifests = manifest_paths + .iter() + .map(|(version, paths)| (*version, paths.clone())) + .collect::>(); + let results = stream::iter(version_manifests) + .map(|(version, manifest_paths)| { let path = self .marker_version_path(version) .join(format!("{operation_id}{DRAINING_MARKER_SUFFIX}")); async move { - let metadata = self.create_marker(&path).await?; - Ok::<_, Error>((version, path, metadata.last_modified)) + let payload = retirement_marker_payload(&manifest_paths)?; + let metadata = self.create_marker(&path, payload).await?; + Ok::<_, Error>((version, path, metadata.last_modified, manifest_paths)) } }) .buffer_unordered(self.object_store.io_parallelism()) @@ -297,13 +336,14 @@ impl VersionLeaseStore { let mut first_error = None; for result in results { match result { - Ok((version, path, observed_at)) => { + Ok((version, path, observed_at, manifest_paths)) => { fences.insert( version, RetirementFence { draining_path: Some(path), sealed_path: None, observed_at, + manifest_paths, }, ); } @@ -329,7 +369,7 @@ impl VersionLeaseStore { Err(Error::VersionNotFound { message: format!("version {version} no longer exists and cannot be leased"), }) - } else if self.has_marker(version, None).await? { + } else if self.has_active_retirement_marker(version).await? { Err(retiring_version_error(version)) } else { Ok(()) @@ -337,30 +377,242 @@ impl VersionLeaseStore { } async fn ensure_not_sealed(&self, version: u64) -> Result<()> { - if self.has_marker(version, Some(SEALED_MARKER_SUFFIX)).await? { + if self.has_sealed_marker(version).await? { Err(retiring_version_error(version)) } else { Ok(()) } } - async fn has_marker(&self, version: u64, suffix: Option<&str>) -> Result { - let mut markers = self - .object_store - .list(Some(self.marker_version_path(version))); - while let Some(metadata) = markers.try_next().await? { - if suffix.is_none_or(|suffix| { - metadata - .location - .filename() - .is_some_and(|name| name.ends_with(suffix)) - }) { + async fn has_sealed_marker(&self, version: u64) -> Result { + Ok(self + .version_marker_metadata(version) + .await? + .iter() + .any(|marker| marker.is_sealed)) + } + + async fn has_active_retirement_marker(&self, version: u64) -> Result { + let markers = self.version_marker_metadata(version).await?; + if markers.iter().any(|marker| marker.is_sealed) { + return Ok(true); + } + + let mut draining_markers = Vec::new(); + for marker in markers { + if LOCALLY_ABANDONED_DRAINS.contains(&marker.metadata.location) { + match self.object_store.delete(&marker.metadata.location).await { + Ok(()) => { + LOCALLY_ABANDONED_DRAINS.remove(&marker.metadata.location); + } + Err(error) if error.is_not_found() => { + LOCALLY_ABANDONED_DRAINS.remove(&marker.metadata.location); + } + Err(error) => { + tracing::warn!( + path = %marker.metadata.location, + error = %error, + "Failed to remove locally abandoned version retirement drain" + ); + } + } + } else { + draining_markers.push(marker); + } + } + if draining_markers.is_empty() { + return Ok(false); + } + + let observed_at = self.storage_observed_at().await?; + let mut stale_paths = Vec::new(); + for marker in draining_markers { + if draining_owner_is_active(marker.metadata.last_modified, observed_at)? { return Ok(true); } + stale_paths.push(marker.metadata.location); } + self.delete_paths(stale_paths).await?; Ok(false) } + async fn version_marker_metadata(&self, version: u64) -> Result> { + let metadata = self + .object_store + .list(Some(self.marker_version_path(version))) + .try_collect::>() + .await?; + metadata + .into_iter() + .map(|metadata| self.parse_retirement_marker_metadata(metadata)) + .collect() + } + + async fn all_marker_metadata(&self) -> Result> { + let metadata = self + .object_store + .list(Some(self.markers_path.clone())) + .try_collect::>() + .await?; + metadata + .into_iter() + .filter(|metadata| metadata.location != self.storage_clock_path()) + .map(|metadata| self.parse_retirement_marker_metadata(metadata)) + .collect() + } + + fn parse_retirement_marker_metadata( + &self, + metadata: ObjectMeta, + ) -> Result { + let relative_parts = metadata + .location + .prefix_match(&self.markers_path) + .ok_or_else(|| { + Error::corrupt_file( + metadata.location.clone(), + "retirement marker is outside its namespace", + ) + })? + .map(|part| part.as_ref().to_string()) + .collect::>(); + if relative_parts.len() != 2 { + return Err(Error::corrupt_file( + metadata.location, + "retirement marker path must contain a version and operation filename", + )); + } + let version = relative_parts[0].parse::().map_err(|error| { + Error::corrupt_file( + metadata.location.clone(), + format!("retirement marker has invalid version: {error}"), + ) + })?; + let file_name = &relative_parts[1]; + let (operation_id, is_sealed) = + if let Some(operation_id) = file_name.strip_suffix(DRAINING_MARKER_SUFFIX) { + (operation_id, false) + } else if let Some(operation_id) = file_name.strip_suffix(SEALED_MARKER_SUFFIX) { + (operation_id, true) + } else { + return Err(Error::corrupt_file( + metadata.location, + "retirement marker filename has an unknown state suffix", + )); + }; + Uuid::parse_str(operation_id).map_err(|error| { + Error::corrupt_file( + metadata.location.clone(), + format!("retirement marker has invalid operation id: {error}"), + ) + })?; + Ok(RetirementMarkerMetadata { + version, + operation_id: operation_id.to_string(), + is_sealed, + metadata, + }) + } + + async fn storage_observed_at(&self) -> Result> { + let path = self.storage_clock_path(); + self.object_store + .inner + .put(&path, Bytes::new().into()) + .await?; + Ok(self.object_store.inner.head(&path).await?.last_modified) + } + + fn storage_clock_path(&self) -> Path { + self.markers_path.clone().join(STORAGE_CLOCK_PATH) + } + + pub(super) async fn recover_retirements(self) -> Result> { + let markers = self.all_marker_metadata().await?; + if markers.is_empty() { + return Ok(HashSet::new()); + } + + let sealed_operations: HashSet<_> = markers + .iter() + .filter(|marker| marker.is_sealed) + .map(|marker| (marker.version, marker.operation_id.clone())) + .collect(); + let observed_at = self.storage_observed_at().await?; + let sealed_observation = markers + .iter() + .filter(|marker| marker.is_sealed) + .map(|marker| (marker.version, observed_at)) + .collect::>(); + let actively_leased_versions = self.active_versions_at(&sealed_observation, true).await?; + let mut stale_drains = Vec::new(); + let mut sealed_manifest_paths = HashMap::>::new(); + for marker in markers { + if marker.is_sealed { + let manifest_paths = self.read_retirement_marker(&marker.metadata).await?; + sealed_manifest_paths + .entry(marker.version) + .or_default() + .extend(manifest_paths); + } else if sealed_operations.contains(&(marker.version, marker.operation_id)) + || LOCALLY_ABANDONED_DRAINS.contains(&marker.metadata.location) + || !draining_owner_is_active(marker.metadata.last_modified, observed_at)? + { + stale_drains.push(marker.metadata.location); + } + } + self.delete_paths(stale_drains.clone()).await?; + for path in stale_drains { + LOCALLY_ABANDONED_DRAINS.remove(&path); + } + + let mut terminal_manifests = HashMap::new(); + let mut versions_to_resume = HashSet::new(); + for (version, manifest_paths) in sealed_manifest_paths { + let manifest_paths = manifest_paths.into_iter().collect::>(); + let mut has_existing_manifest = false; + for path in &manifest_paths { + has_existing_manifest |= self.object_store.exists(path).await?; + } + if has_existing_manifest { + // A seal is the renewal cutoff. If its owner disappeared before + // confirming the final lease scan, wait out every lease that was + // still entitled to its published TTL before resuming deletion. + if !actively_leased_versions.contains(&version) { + versions_to_resume.insert(version); + } + } else { + terminal_manifests.insert(version, manifest_paths); + } + } + self.finalize_versions(&terminal_manifests).await?; + Ok(versions_to_resume) + } + + async fn read_retirement_marker(&self, metadata: &ObjectMeta) -> Result> { + let bytes = self + .object_store + .inner + .get(&metadata.location) + .await? + .bytes() + .await?; + let marker: RetirementMarker = serde_json::from_slice(&bytes) + .map_err(|error| Error::corrupt_file(metadata.location.clone(), error.to_string()))?; + if marker.manifest_paths.is_empty() { + return Err(Error::corrupt_file( + metadata.location.clone(), + "retirement marker has no manifest identity", + )); + } + marker + .manifest_paths + .into_iter() + .map(Path::parse) + .collect::, _>>() + .map_err(Error::from) + } + async fn lease_metadata(&self) -> Result> { self.object_store .list(Some(self.leases_path.clone())) @@ -372,12 +624,12 @@ impl VersionLeaseStore { self.markers_path.clone().join(version.to_string()) } - async fn create_marker(&self, path: &Path) -> Result { + async fn create_marker(&self, path: &Path, payload: Bytes) -> Result { self.object_store .inner .put_opts( path, - Bytes::new().into(), + payload.into(), PutOptions { mode: PutMode::Create, ..Default::default() @@ -407,30 +659,46 @@ impl RetirementGuard { } pub(super) async fn seal_versions(&mut self, versions: &HashSet) -> Result<()> { - let marker_paths = versions - .iter() - .filter_map(|version| { - self.fences.get(version).and_then(|fence| { - fence.draining_path.as_ref().and_then(|draining_path| { - let file_name = draining_path.filename()?; - let operation_id = file_name.strip_suffix(DRAINING_MARKER_SUFFIX)?; - Some(( - *version, - self.store - .marker_version_path(*version) - .join(format!("{operation_id}{SEALED_MARKER_SUFFIX}")), - )) - }) - }) - }) - .collect::>(); + let mut marker_paths = Vec::with_capacity(versions.len()); + for version in versions { + let fence = self.fences.get(version).ok_or_else(|| { + Error::internal(format!("missing draining fence for version {version}")) + })?; + let draining_path = fence.draining_path.as_ref().ok_or_else(|| { + Error::internal(format!("version {version} has no draining marker")) + })?; + let file_name = draining_path.filename().ok_or_else(|| { + Error::internal(format!("draining marker {draining_path} has no filename")) + })?; + let operation_id = file_name + .strip_suffix(DRAINING_MARKER_SUFFIX) + .ok_or_else(|| { + Error::internal(format!( + "draining marker {draining_path} has an invalid suffix" + )) + })?; + marker_paths.push(( + *version, + self.store + .marker_version_path(*version) + .join(format!("{operation_id}{SEALED_MARKER_SUFFIX}")), + retirement_marker_payload(&fence.manifest_paths)?, + fence.observed_at, + )); + } let store = self.store.clone(); let results = stream::iter(marker_paths) - .map(move |(version, path)| { + .map(move |(version, path, payload, draining_observed_at)| { let store = store.clone(); async move { - let metadata = store.create_marker(&path).await?; + let metadata = store.create_marker(&path, payload).await?; + if !draining_owner_is_active(draining_observed_at, metadata.last_modified)? { + let _ = store.object_store.delete(&path).await; + return Err(Error::internal(format!( + "version {version} retirement ownership expired before sealing" + ))); + } Ok::<_, Error>((version, path, metadata.last_modified)) } }) @@ -497,14 +765,75 @@ impl RetirementGuard { &mut self, manifest_paths: &HashMap>, ) -> Result<()> { - for version in self.fences.keys() { - let paths = manifest_paths.get(version).ok_or_else(|| { - Error::internal(format!( + let guarded_manifest_paths = self + .fences + .keys() + .map(|version| { + manifest_paths + .get(version) + .map(|paths| (*version, paths.clone())) + .ok_or_else(|| { + Error::internal(format!( + "missing manifest identity for retiring version {version}" + )) + }) + }) + .collect::>>()?; + self.store + .finalize_versions(&guarded_manifest_paths) + .await?; + self.fences.clear(); + Ok(()) + } +} + +impl Drop for RetirementGuard { + fn drop(&mut self) { + let draining_paths = self + .fences + .values() + .filter_map(|fence| fence.draining_path.clone()) + .collect::>(); + if draining_paths.is_empty() { + return; + } + for path in &draining_paths { + LOCALLY_ABANDONED_DRAINS.insert(path.clone()); + } + let store = self.store.clone(); + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + runtime.spawn(async move { + match store.delete_paths(draining_paths.clone()).await { + Ok(()) => { + for path in draining_paths { + LOCALLY_ABANDONED_DRAINS.remove(&path); + } + } + Err(error) => { + tracing::warn!( + error = %error, + "Failed to remove abandoned version retirement drains" + ); + } + } + }); + } + } +} + +impl VersionLeaseStore { + async fn finalize_versions(&self, manifest_paths: &HashMap>) -> Result<()> { + if manifest_paths.is_empty() { + return Ok(()); + } + for (version, paths) in manifest_paths { + if paths.is_empty() { + return Err(Error::internal(format!( "missing manifest identity for retiring version {version}" - )) - })?; + ))); + } for path in paths { - if self.store.object_store.exists(path).await? { + if self.object_store.exists(path).await? { return Err(Error::internal(format!( "cannot finalize retirement for version {version}: manifest {path} still exists" ))); @@ -512,32 +841,28 @@ impl RetirementGuard { } } - let versions: HashSet<_> = self.fences.keys().copied().collect(); + let versions: HashSet<_> = manifest_paths.keys().copied().collect(); let marker_prefixes = versions .iter() - .map(|version| self.store.marker_version_path(*version)) + .map(|version| self.marker_version_path(*version)) .collect::>(); let marker_streams = marker_prefixes .into_iter() - .map(|prefix| self.store.object_store.list(Some(prefix))); + .map(|prefix| self.object_store.list(Some(prefix))); let mut paths = stream::iter(marker_streams) .flatten() .map_ok(|metadata| metadata.location) .try_collect::>() .await?; - for metadata in self.store.lease_metadata().await? { + for metadata in self.lease_metadata().await? { let (version, _) = parse_lease_metadata(&metadata)?; if versions.contains(&version) { paths.push(metadata.location); } } - self.store.delete_paths(paths).await?; - self.fences.clear(); - Ok(()) + self.delete_paths(paths).await } -} -impl VersionLeaseStore { async fn delete_paths(&self, paths: Vec) -> Result<()> { stream::iter(paths) .map(Ok) @@ -624,18 +949,49 @@ fn ttl_micros(ttl: Duration) -> Result { } fn expiration_from_ttl(observed_at: DateTime, ttl: Duration) -> Result> { - let ttl = TimeDelta::from_std(ttl).map_err(|error| { + let conservative_ttl = ttl + .checked_add(STORAGE_TIMESTAMP_PRECISION) + .ok_or_else(|| { + Error::invalid_input(format!( + "version lease TTL {ttl:?} overflows the storage timestamp precision interval" + )) + })?; + let ttl_delta = TimeDelta::from_std(conservative_ttl).map_err(|error| { Error::invalid_input(format!( "version lease TTL {ttl:?} is out of range: {error}" )) })?; - observed_at.checked_add_signed(ttl).ok_or_else(|| { + observed_at.checked_add_signed(ttl_delta).ok_or_else(|| { Error::invalid_input(format!( "version lease TTL {ttl:?} overflows its expiration" )) }) } +fn draining_owner_is_active( + draining_observed_at: DateTime, + current_observed_at: DateTime, +) -> Result { + Ok( + expiration_from_ttl(draining_observed_at, DRAINING_OWNERSHIP_TIMEOUT)? + > current_observed_at, + ) +} + +fn retirement_marker_payload(manifest_paths: &[Path]) -> Result { + if manifest_paths.is_empty() { + return Err(Error::internal( + "cannot create a retirement marker without a manifest identity", + )); + } + let marker = RetirementMarker { + manifest_paths: manifest_paths.iter().map(ToString::to_string).collect(), + }; + serde_json::to_vec(&marker) + .map(Bytes::from) + .map_err(|error| Error::internal(format!("failed to serialize retirement marker: {error}"))) +} + fn parse_lease_metadata(metadata: &ObjectMeta) -> Result<(u64, Duration)> { let file_name = metadata.location.filename().ok_or_else(|| { Error::corrupt_file( @@ -707,6 +1063,23 @@ mod tests { } } + fn manifest_paths(version: u64) -> HashMap> { + HashMap::from([( + version, + vec![Path::from(format!("manifests/{version}.manifest"))], + )]) + } + + fn assert_send(_: T) {} + + #[test] + fn retirement_futures_are_send() { + let store = memory_store(); + let manifests = manifest_paths(42); + assert_send(store.clone().recover_retirements()); + assert_send(store.fence_versions(&manifests)); + } + #[test] fn parses_lease_file_name() { let lease_id = Uuid::nil().simple(); @@ -720,7 +1093,7 @@ mod tests { async fn draining_lease_remains_renewable() { let store = memory_store(); let mut lease = store.acquire(42, Duration::from_secs(60)).await.unwrap(); - let mut guard = store.fence_versions(&HashSet::from([42])).await.unwrap(); + let mut guard = store.fence_versions(&manifest_paths(42)).await.unwrap(); assert!( store @@ -742,7 +1115,7 @@ mod tests { // Moving the cleanup host clock ahead does not affect the storage // timestamps used for lease liveness. MockClock::set_system_time(Duration::from_secs(160)); - let mut guard = store.fence_versions(&HashSet::from([42])).await.unwrap(); + let mut guard = store.fence_versions(&manifest_paths(42)).await.unwrap(); assert!( store .active_versions_at(&guard.observed_at(), false) @@ -757,7 +1130,7 @@ mod tests { async fn sealed_version_rejects_renewal() { let store = memory_store(); let mut lease = store.acquire(42, Duration::from_secs(60)).await.unwrap(); - let mut guard = store.fence_versions(&HashSet::from([42])).await.unwrap(); + let mut guard = store.fence_versions(&manifest_paths(42)).await.unwrap(); guard.seal_versions(&HashSet::from([42])).await.unwrap(); let error = lease.renew(Duration::from_secs(60)).await.unwrap_err(); @@ -766,6 +1139,99 @@ mod tests { guard.cancel_all().await.unwrap(); } + #[test] + fn lease_ttl_survives_coarse_storage_timestamps() { + let storage_second = DateTime::from_timestamp(100, 0).unwrap(); + let acquired_at = storage_second + TimeDelta::try_milliseconds(900).unwrap(); + let cleanup_started_at = storage_second + TimeDelta::try_milliseconds(1_001).unwrap(); + let ttl = Duration::from_millis(900); + + assert!( + cleanup_started_at < acquired_at + TimeDelta::from_std(ttl).unwrap(), + "the requested TTL is still active" + ); + let marker_last_modified = storage_second + TimeDelta::try_seconds(1).unwrap(); + assert!( + expiration_from_ttl(storage_second, ttl).unwrap() > marker_last_modified, + "coarse Last-Modified timestamps must not expire the lease early" + ); + } + + #[tokio::test] + async fn abandoned_drain_does_not_block_future_acquire() { + let store = memory_store(); + let guard = store.fence_versions(&manifest_paths(42)).await.unwrap(); + let draining_path = guard.fences[&42].draining_path.clone().unwrap(); + + // Dropping the owner proves this in-process drain cannot proceed to deletion. + drop(guard); + assert!(LOCALLY_ABANDONED_DRAINS.contains(&draining_path)); + + store.acquire(42, Duration::from_secs(60)).await.unwrap(); + } + + #[test] + fn draining_ownership_is_bounded() { + let started_at = DateTime::from_timestamp(100, 0).unwrap(); + let ownership_expired_at = + expiration_from_ttl(started_at, DRAINING_OWNERSHIP_TIMEOUT).unwrap(); + + assert!(!draining_owner_is_active(started_at, ownership_expired_at).unwrap()); + } + + #[tokio::test] + async fn sealed_retirement_is_recovered() { + let store = memory_store(); + let manifest_paths = manifest_paths(42); + let manifest_path = manifest_paths[&42][0].clone(); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + let mut guard = store.fence_versions(&manifest_paths).await.unwrap(); + guard.seal_versions(&HashSet::from([42])).await.unwrap(); + drop(guard); + + let versions_to_resume = store.clone().recover_retirements().await.unwrap(); + assert_eq!(versions_to_resume, HashSet::from([42])); + + store.object_store.delete(&manifest_path).await.unwrap(); + assert!( + store + .clone() + .recover_retirements() + .await + .unwrap() + .is_empty() + ); + assert!(store.version_marker_metadata(42).await.unwrap().is_empty()); + } + + #[tokio::test] + async fn sealed_recovery_waits_for_active_lease() { + let store = memory_store(); + let manifest_paths = manifest_paths(42); + let manifest_path = manifest_paths[&42][0].clone(); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + let lease = store.acquire(42, Duration::from_secs(60)).await.unwrap(); + let mut guard = store.fence_versions(&manifest_paths).await.unwrap(); + guard.seal_versions(&HashSet::from([42])).await.unwrap(); + drop(guard); + + assert!( + store + .clone() + .recover_retirements() + .await + .unwrap() + .is_empty() + ); + assert!(store.object_store.exists(&manifest_path).await.unwrap()); + + lease.release().await.unwrap(); + assert_eq!( + store.clone().recover_retirements().await.unwrap(), + HashSet::from([42]) + ); + } + #[tokio::test] async fn branch_deletion_removes_incarnation_state() { let store = memory_store(); From 724e0e0e7c51c7477dc4ff8d97e866f0a4e06976 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 18:51:21 +0000 Subject: [PATCH 04/13] fix: make version retirement commit durable --- rust/lance/src/dataset/cleanup.rs | 150 ++++++++--- rust/lance/src/dataset/refs.rs | 68 ++++- rust/lance/src/dataset/version_lease.rs | 319 ++++++++++++++++++++++-- 3 files changed, 470 insertions(+), 67 deletions(-) diff --git a/rust/lance/src/dataset/cleanup.rs b/rust/lance/src/dataset/cleanup.rs index 340936a65da..93f9046367e 100644 --- a/rust/lance/src/dataset/cleanup.rs +++ b/rust/lance/src/dataset/cleanup.rs @@ -511,21 +511,12 @@ impl<'a> CleanupTask<'a> { .cloned() .collect(); inspection = self - .retain_branch_lineage_files( - inspection, - &referenced_branches, - &ignored_manifests, - &forced_retirement_versions, - ) + .retain_branch_lineage_files(inspection, &referenced_branches, &ignored_manifests) .await? }; let (inspection, mut retirement) = self - .fence_old_versions_and_retain_new_leases( - inspection, - &version_lease_store, - &forced_retirement_versions, - ) + .fence_old_versions_and_retain_new_leases(inspection, &version_lease_store) .await?; let cleanup_result = self.delete_unreferenced_files(inspection).await?; @@ -608,13 +599,17 @@ impl<'a> CleanupTask<'a> { let is_tagged = tagged_versions.contains(&manifest.version); let is_leased = leased_versions.contains(&manifest.version); let is_forced_retirement = forced_retirement_versions.contains(&manifest.version); - let in_working_set = !is_forced_retirement - && (is_latest || !self.policy.should_clean(&manifest) || is_tagged || is_leased); + // Recovery may override a later retention policy, but never a durable + // reference or lease admitted before retirement began. + let in_working_set = is_latest + || is_tagged + || is_leased + || (!is_forced_retirement && !self.policy.should_clean(&manifest)); let mut inspection = inspection.lock().unwrap(); // Track tagged old versions in case we want to return a `CleanupError` later. // Only track tagged when it is old. - if is_tagged && !is_forced_retirement && !is_latest && self.policy.should_clean(&manifest) { + if is_tagged && !is_latest && self.policy.should_clean(&manifest) { inspection.tagged_old_versions.insert(manifest.version); } @@ -633,7 +628,6 @@ impl<'a> CleanupTask<'a> { &self, mut inspection: CleanupInspection, version_lease_store: &VersionLeaseStore, - forced_retirement_versions: &HashSet, ) -> Result<(CleanupInspection, Option)> { if !self.action.deletes_files() { return Ok((inspection, None)); @@ -663,7 +657,6 @@ impl<'a> CleanupTask<'a> { .await?; let retained_versions: HashSet<_> = versions_to_delete .intersection(&leased_versions) - .filter(|version| !forced_retirement_versions.contains(version)) .copied() .collect(); for version in &retained_versions { @@ -683,7 +676,6 @@ impl<'a> CleanupTask<'a> { .await?; let retained_versions: HashSet<_> = versions_to_seal .intersection(&leased_versions) - .filter(|version| !forced_retirement_versions.contains(version)) .copied() .collect(); for version in &retained_versions { @@ -691,6 +683,42 @@ impl<'a> CleanupTask<'a> { } guard.cancel_versions(&retained_versions).await?; + // A tag or child branch published before the draining marker must + // win. Publications racing the marker perform their own second + // admission check and roll back if retirement won instead. + let tags = self.dataset.tags().list().await?; + let current_branch = &self.dataset.manifest.branch; + let mut referenced_versions = tags + .values() + .filter(|tag| match (tag.branch.as_ref(), current_branch.as_ref()) { + (Some(tag_branch), Some(current_branch)) => tag_branch == current_branch, + (None, None) => true, + _ => false, + }) + .map(|tag| tag.version) + .collect::>(); + referenced_versions.extend( + self.find_referenced_branches() + .await? + .into_iter() + .map(|(_, version)| version), + ); + let retained_versions = versions_to_seal + .intersection(&referenced_versions) + .copied() + .collect::>(); + for version in &retained_versions { + self.retain_version(&mut inspection, *version).await?; + } + guard.cancel_versions(&retained_versions).await?; + + let versions_to_commit = inspection + .old_manifests + .values() + .copied() + .collect::>(); + guard.commit_versions(&versions_to_commit).await?; + let mut manifest_paths = HashMap::>::new(); for (path, version) in &inspection.old_manifests { manifest_paths @@ -1344,18 +1372,15 @@ impl<'a> CleanupTask<'a> { inspection: CleanupInspection, referenced_branches: &[(String, u64)], removed_branch_manifests: &HashSet, - forced_retirement_versions: &HashSet, ) -> Result { let inspection = Arc::new(Mutex::new(inspection)); let removed_branch_manifests = Arc::new(removed_branch_manifests.clone()); - let forced_retirement_versions = Arc::new(forced_retirement_versions.clone()); for (branch, referenced_version) in referenced_branches.iter().cloned() { // Use find_branch to get the branch path directly without checkout. // This avoids creating a dataset instance and prevents manifest deletion // during the retain operation. let branch_location = self.dataset.branch_location().find_branch(Some(&branch))?; let removed_branch_manifests = Arc::clone(&removed_branch_manifests); - let forced_retirement_versions = Arc::clone(&forced_retirement_versions); let branch_inspection = Arc::clone(&inspection); self.dataset .commit_handler @@ -1364,14 +1389,12 @@ impl<'a> CleanupTask<'a> { future::ready(!removed_branch_manifests.contains(&location.path)) }) .try_for_each_concurrent(self.dataset.object_store.io_parallelism(), |location| { - let forced_retirement_versions = Arc::clone(&forced_retirement_versions); let branch_inspection = Arc::clone(&branch_inspection); async move { self.process_branch_referenced_manifests( location, referenced_version, branch_inspection.as_ref(), - forced_retirement_versions.as_ref(), ) .await } @@ -1388,7 +1411,6 @@ impl<'a> CleanupTask<'a> { location: ManifestLocation, referenced_version: u64, inspection: &Mutex, - forced_retirement_versions: &HashSet, ) -> Result<()> { let manifest = read_manifest(&self.dataset.object_store, &location.path, location.size).await?; @@ -1458,7 +1480,7 @@ impl<'a> CleanupTask<'a> { } } } - if is_referenced && !forced_retirement_versions.contains(&referenced_version) { + if is_referenced { inspection .old_manifests .retain(|_path, version_number| *version_number != referenced_version); @@ -2488,6 +2510,69 @@ mod tests { ); } + #[tokio::test] + async fn forced_retirement_still_respects_active_lease() { + let fixture = MockDatasetFixture::try_new().unwrap(); + fixture.create_some_data().await.unwrap(); + fixture.overwrite_some_data().await.unwrap(); + let dataset = fixture.load().await.unwrap(); + let old_manifest = dataset + .commit_handler + .list_manifest_locations(&dataset.base, &dataset.object_store, false) + .try_filter(|location| future::ready(location.version == 1)) + .try_next() + .await + .unwrap() + .unwrap(); + let inspection = Mutex::new(CleanupInspection::default()); + let task = CleanupTask::new( + &dataset, + CleanupPolicyBuilder::default().build(), + CleanupAction::Execute, + ); + + task.process_manifest_file( + old_manifest, + &inspection, + &HashSet::new(), + &HashSet::from([1]), + &HashSet::from([1]), + ) + .await + .unwrap(); + + assert!(inspection.into_inner().unwrap().old_manifests.is_empty()); + } + + #[tokio::test] + async fn sealed_retirement_rejects_new_tags_and_branches() { + let fixture = MockDatasetFixture::try_new().unwrap(); + fixture.create_some_data().await.unwrap(); + let historical = fixture.load().await.unwrap(); + fixture.overwrite_some_data().await.unwrap(); + let dataset = fixture.load().await.unwrap(); + let store = VersionLeaseStore::for_dataset(&historical).await.unwrap(); + let manifest_paths = HashMap::from([( + historical.version().version, + vec![historical.manifest_location.path.clone()], + )]); + let mut guard = store.fence_versions(&manifest_paths).await.unwrap(); + guard.seal_versions(&HashSet::from([1])).await.unwrap(); + + let tag_error = dataset.tags().create("after-seal", 1).await.unwrap_err(); + let branch_error = dataset + .branches() + .create("after-seal", 1, None) + .await + .unwrap_err(); + + assert!(matches!(tag_error, Error::RefConflict { .. })); + assert!(matches!(branch_error, Error::RefConflict { .. })); + assert!(dataset.tags().get("after-seal").await.is_err()); + assert!(dataset.branches().get("after-seal").await.is_err()); + guard.cancel_all().await.unwrap(); + } + #[tokio::test] async fn explain_cleanup_does_not_delete_files() { let fixture = MockDatasetFixture::try_new().unwrap(); @@ -3170,12 +3255,7 @@ mod tests { .unwrap(); let referenced_branches = task.find_referenced_branches().await.unwrap(); let inspection = task - .retain_branch_lineage_files( - inspection, - &referenced_branches, - &HashSet::new(), - &HashSet::new(), - ) + .retain_branch_lineage_files(inspection, &referenced_branches, &HashSet::new()) .await .unwrap(); @@ -3745,10 +3825,14 @@ mod tests { let after_count = fixture.count_files().await.unwrap(); assert_eq!(removed.old_versions, 1); - assert_eq!( - removed.bytes_removed, - mid_count.num_bytes - after_count.num_bytes + let total_bytes_removed = mid_count.num_bytes - after_count.num_bytes; + assert!( + removed.bytes_removed > 0 && removed.bytes_removed <= total_bytes_removed, + "cleanup reported {} bytes removed out of {total_bytes_removed} total bytes removed", + removed.bytes_removed ); + // RemovalStats covers dataset files. Internal retirement markers are + // finalized in the same retry and may add to the object-store delta. assert_eq!(after_count.num_data_files, 1); assert_eq!(after_count.num_manifest_files, 1); diff --git a/rust/lance/src/dataset/refs.rs b/rust/lance/src/dataset/refs.rs index 0f9e9ba4465..4ec4ce4c336 100644 --- a/rust/lance/src/dataset/refs.rs +++ b/rust/lance/src/dataset/refs.rs @@ -14,7 +14,7 @@ use std::sync::Arc; use crate::dataset::branch_location::BranchLocation; use crate::dataset::refs::Ref::{Tag, Version, VersionNumber}; -use crate::dataset::version_lease::remove_branch_state; +use crate::dataset::version_lease::{ensure_reference_admissible, remove_branch_state}; use crate::utils::temporal::utc_now; use crate::{Error, Result}; use serde::de::DeserializeOwned; @@ -227,14 +227,30 @@ impl Tags<'_> { let tag_contents = self .build_tag_content_by_ref(reference, Some(now), Some(now)) .await?; + ensure_reference_admissible( + self.refs, + tag_contents.branch.as_deref(), + tag_contents.version, + ) + .await?; self.object_store() .put( &tag_file, serde_json::to_string_pretty(&tag_contents)?.as_bytes(), ) - .await - .map(|_| ()) + .await?; + if let Err(error) = ensure_reference_admissible( + self.refs, + tag_contents.branch.as_deref(), + tag_contents.version, + ) + .await + { + self.object_store().delete(&tag_file).await?; + return Err(error); + } + Ok(()) } pub async fn delete(&self, tag: &str) -> Result<()> { @@ -262,10 +278,17 @@ impl Tags<'_> { message: format!("tag {} does not exist", tag), }); } - let mut tag_contents = TagContents::from_path(&tag_file, self.object_store()).await?; + let original_tag_contents = TagContents::from_path(&tag_file, self.object_store()).await?; let updated_reference = self - .build_tag_content_by_ref(reference, tag_contents.created_at, Some(utc_now())) + .build_tag_content_by_ref(reference, original_tag_contents.created_at, Some(utc_now())) .await?; + ensure_reference_admissible( + self.refs, + updated_reference.branch.as_deref(), + updated_reference.version, + ) + .await?; + let mut tag_contents = original_tag_contents.clone(); tag_contents.branch = updated_reference.branch; tag_contents.version = updated_reference.version; tag_contents.created_at = updated_reference.created_at; @@ -277,8 +300,23 @@ impl Tags<'_> { &tag_file, serde_json::to_string_pretty(&tag_contents)?.as_bytes(), ) - .await - .map(|_| ()) + .await?; + if let Err(error) = ensure_reference_admissible( + self.refs, + tag_contents.branch.as_deref(), + tag_contents.version, + ) + .await + { + self.object_store() + .put( + &tag_file, + serde_json::to_string_pretty(&original_tag_contents)?.as_bytes(), + ) + .await?; + return Err(error); + } + Ok(()) } pub async fn replace_metadata( @@ -354,7 +392,6 @@ impl Tags<'_> { } else { self.object_store().size(&manifest_file.path).await? as usize }; - let tag_contents = TagContents { branch, version: manifest_file.version, @@ -479,6 +516,7 @@ impl Branches<'_> { message: format!("Manifest file {} does not exist", manifest_file.path), }); }; + ensure_reference_admissible(self.refs, source_branch.as_deref(), version_number).await?; let parent_branch_id = if let Some(ref parent_branch) = source_branch { let parent_file = branch_contents_path(&root_location.path, parent_branch); @@ -513,8 +551,18 @@ impl Branches<'_> { &branch_file, serde_json::to_string_pretty(&branch_contents)?.as_bytes(), ) - .await - .map(|_| ()) + .await?; + if let Err(error) = ensure_reference_admissible( + self.refs, + branch_contents.parent_branch.as_deref(), + version_number, + ) + .await + { + self.object_store().delete(&branch_file).await?; + return Err(error); + } + Ok(()) } pub async fn replace_metadata( diff --git a/rust/lance/src/dataset/version_lease.rs b/rust/lance/src/dataset/version_lease.rs index a8aca681562..c03913be0b3 100644 --- a/rust/lance/src/dataset/version_lease.rs +++ b/rust/lance/src/dataset/version_lease.rs @@ -18,7 +18,10 @@ use object_store::{ObjectMeta, ObjectStoreExt, PutMode, PutOptions, path::Path}; use serde::{Deserialize, Serialize}; use uuid::Uuid; -use super::{Dataset, refs::MAIN_BRANCH}; +use super::{ + Dataset, + refs::{MAIN_BRANCH, Refs}, +}; use crate::{Error, Result}; const LEASES_DIR: &str = "_refs/version_leases"; @@ -26,6 +29,7 @@ const LEASE_GC_MARKERS_DIR: &str = "_refs/version_lease_gc"; const LEASE_FILE_SUFFIX: &str = ".lease"; const DRAINING_MARKER_SUFFIX: &str = ".draining"; const SEALED_MARKER_SUFFIX: &str = ".sealed"; +const COMMITTED_MARKER_SUFFIX: &str = ".committed"; const STORAGE_CLOCK_PATH: &str = "_clock"; /// HTTP `Last-Modified`, used by supported cloud stores, has whole-second precision. @@ -156,6 +160,7 @@ struct LeaseFile { struct RetirementFence { draining_path: Option, sealed_path: Option, + committed_path: Option, observed_at: DateTime, manifest_paths: Vec, } @@ -165,20 +170,29 @@ struct RetirementMarker { manifest_paths: Vec, } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum RetirementState { + Draining, + Sealed, + Committed, +} + #[derive(Debug)] struct RetirementMarkerMetadata { version: u64, operation_id: String, - is_sealed: bool, + state: RetirementState, metadata: ObjectMeta, } /// Per-cleanup retirement markers for versions selected for deletion. /// /// Draining markers reject new leases while allowing already-admitted leases -/// to renew. Sealed markers reject both acquisition and renewal immediately -/// before deletion. Unique operation markers keep cancellation and finalization -/// safe when multiple cleanups overlap. +/// to renew. Sealed markers reject both acquisition and renewal while cleanup +/// performs its final lease and reference checks. Committed markers are the +/// durable, irreversible deletion boundary shared with reference admission. +/// Unique operation markers keep cancellation and finalization safe when +/// multiple cleanups overlap. #[derive(Debug)] pub(super) struct RetirementGuard { store: VersionLeaseStore, @@ -187,9 +201,22 @@ pub(super) struct RetirementGuard { impl VersionLeaseStore { pub(super) async fn for_dataset(dataset: &Dataset) -> Result { - let root = dataset.refs.root()?; let branch = dataset.manifest.branch.as_deref(); - let branch_identifier = dataset.branches().get_identifier(branch).await?; + Self::for_refs( + &dataset.refs, + branch, + Some(dataset.manifest_location.path.clone()), + ) + .await + } + + async fn for_refs( + refs: &Refs, + branch: Option<&str>, + manifest_path: Option, + ) -> Result { + let root = refs.root()?; + let branch_identifier = refs.branches().get_identifier(branch).await?; let namespace = if branch.is_none() { MAIN_BRANCH } else { @@ -206,10 +233,10 @@ impl VersionLeaseStore { }; Ok(Self { - object_store: Arc::clone(&dataset.object_store), + object_store: Arc::clone(&refs.object_store), leases_path: root.path.clone().join(LEASES_DIR).join(namespace), markers_path: root.path.join(LEASE_GC_MARKERS_DIR).join(namespace), - manifest_path: Some(dataset.manifest_location.path.clone()), + manifest_path, }) } @@ -342,6 +369,7 @@ impl VersionLeaseStore { RetirementFence { draining_path: Some(path), sealed_path: None, + committed_path: None, observed_at, manifest_paths, }, @@ -389,12 +417,15 @@ impl VersionLeaseStore { .version_marker_metadata(version) .await? .iter() - .any(|marker| marker.is_sealed)) + .any(|marker| marker.state != RetirementState::Draining)) } async fn has_active_retirement_marker(&self, version: u64) -> Result { let markers = self.version_marker_metadata(version).await?; - if markers.iter().any(|marker| marker.is_sealed) { + if markers + .iter() + .any(|marker| marker.state != RetirementState::Draining) + { return Ok(true); } @@ -489,11 +520,13 @@ impl VersionLeaseStore { ) })?; let file_name = &relative_parts[1]; - let (operation_id, is_sealed) = + let (operation_id, state) = if let Some(operation_id) = file_name.strip_suffix(DRAINING_MARKER_SUFFIX) { - (operation_id, false) + (operation_id, RetirementState::Draining) } else if let Some(operation_id) = file_name.strip_suffix(SEALED_MARKER_SUFFIX) { - (operation_id, true) + (operation_id, RetirementState::Sealed) + } else if let Some(operation_id) = file_name.strip_suffix(COMMITTED_MARKER_SUFFIX) { + (operation_id, RetirementState::Committed) } else { return Err(Error::corrupt_file( metadata.location, @@ -509,7 +542,7 @@ impl VersionLeaseStore { Ok(RetirementMarkerMetadata { version, operation_id: operation_id.to_string(), - is_sealed, + state, metadata, }) } @@ -535,20 +568,20 @@ impl VersionLeaseStore { let sealed_operations: HashSet<_> = markers .iter() - .filter(|marker| marker.is_sealed) + .filter(|marker| marker.state != RetirementState::Draining) .map(|marker| (marker.version, marker.operation_id.clone())) .collect(); let observed_at = self.storage_observed_at().await?; let sealed_observation = markers .iter() - .filter(|marker| marker.is_sealed) + .filter(|marker| marker.state != RetirementState::Draining) .map(|marker| (marker.version, observed_at)) .collect::>(); let actively_leased_versions = self.active_versions_at(&sealed_observation, true).await?; let mut stale_drains = Vec::new(); let mut sealed_manifest_paths = HashMap::>::new(); for marker in markers { - if marker.is_sealed { + if marker.state != RetirementState::Draining { let manifest_paths = self.read_retirement_marker(&marker.metadata).await?; sealed_manifest_paths .entry(marker.version) @@ -741,8 +774,100 @@ impl RetirementGuard { Ok(()) } + /// Publish the irreversible retirement boundary after the final lease and + /// durable-reference scans. Reference admission checks this marker before + /// and after publishing a tag or branch. + pub(super) async fn commit_versions(&mut self, versions: &HashSet) -> Result<()> { + let mut marker_paths = Vec::with_capacity(versions.len()); + for version in versions { + let fence = self.fences.get(version).ok_or_else(|| { + Error::internal(format!("missing sealed fence for version {version}")) + })?; + let sealed_path = fence.sealed_path.as_ref().ok_or_else(|| { + Error::internal(format!("version {version} has no sealed marker")) + })?; + let file_name = sealed_path.filename().ok_or_else(|| { + Error::internal(format!("sealed marker {sealed_path} has no filename")) + })?; + let operation_id = file_name + .strip_suffix(SEALED_MARKER_SUFFIX) + .ok_or_else(|| { + Error::internal(format!("sealed marker {sealed_path} has an invalid suffix")) + })?; + marker_paths.push(( + *version, + self.store + .marker_version_path(*version) + .join(format!("{operation_id}{COMMITTED_MARKER_SUFFIX}")), + retirement_marker_payload(&fence.manifest_paths)?, + )); + } + + let store = self.store.clone(); + let results = stream::iter(marker_paths) + .map(move |(version, path, payload)| { + let store = store.clone(); + async move { + let metadata = store.create_marker(&path, payload).await?; + Ok::<_, Error>((version, path, metadata.last_modified)) + } + }) + .buffer_unordered(self.store.object_store.io_parallelism()) + .collect::>() + .await; + + let mut first_error = None; + for result in results { + match result { + Ok((version, committed_path, observed_at)) => { + let fence = self.fences.get_mut(&version).ok_or_else(|| { + Error::internal(format!("missing sealed fence for version {version}")) + })?; + fence.committed_path = Some(committed_path); + fence.observed_at = observed_at; + } + Err(error) if first_error.is_none() => first_error = Some(error), + Err(_) => {} + } + } + if let Some(error) = first_error { + return Err(error); + } + + let sealed_paths = versions + .iter() + .filter_map(|version| { + self.fences + .get(version) + .and_then(|fence| fence.sealed_path.clone()) + }) + .collect::>(); + if let Err(error) = self.store.delete_paths(sealed_paths).await { + tracing::warn!( + error = %error, + "Failed to remove superseded sealed retirement markers" + ); + } else { + for version in versions { + if let Some(fence) = self.fences.get_mut(version) { + fence.sealed_path = None; + } + } + } + Ok(()) + } + pub(super) async fn cancel_versions(&mut self, versions: &HashSet) -> Result<()> { - let paths = versions + let cancellable_versions = versions + .iter() + .filter(|version| { + self.fences + .get(version) + .is_some_and(|fence| fence.committed_path.is_none()) + }) + .copied() + .collect::>(); + let paths = cancellable_versions .iter() .filter_map(|version| self.fences.get(version)) .flat_map(|fence| { @@ -752,7 +877,8 @@ impl RetirementGuard { }) .collect::>(); self.store.delete_paths(paths).await?; - self.fences.retain(|version, _| !versions.contains(version)); + self.fences + .retain(|version, _| !cancellable_versions.contains(version)); Ok(()) } @@ -849,18 +975,57 @@ impl VersionLeaseStore { let marker_streams = marker_prefixes .into_iter() .map(|prefix| self.object_store.list(Some(prefix))); - let mut paths = stream::iter(marker_streams) + let marker_metadata = stream::iter(marker_streams) .flatten() - .map_ok(|metadata| metadata.location) .try_collect::>() .await?; + let parsed_markers = marker_metadata + .into_iter() + .map(|metadata| self.parse_retirement_marker_metadata(metadata)) + .collect::>>()?; + + let mut dependent_marker_paths = Vec::new(); + let mut anchor_paths = Vec::new(); + for version in &versions { + let version_markers = parsed_markers + .iter() + .filter(|marker| marker.version == *version) + .collect::>(); + let has_committed = version_markers + .iter() + .any(|marker| marker.state == RetirementState::Committed); + if !has_committed + && !version_markers + .iter() + .any(|marker| marker.state == RetirementState::Sealed) + { + return Err(Error::internal(format!( + "cannot finalize retirement for version {version} without a durable marker" + ))); + } + for marker in version_markers { + let is_anchor = marker.state == RetirementState::Committed + || (!has_committed && marker.state == RetirementState::Sealed); + if is_anchor { + anchor_paths.push(marker.metadata.location.clone()); + } else { + dependent_marker_paths.push(marker.metadata.location.clone()); + } + } + } + + let mut lease_paths = Vec::new(); for metadata in self.lease_metadata().await? { let (version, _) = parse_lease_metadata(&metadata)?; if versions.contains(&version) { - paths.push(metadata.location); + lease_paths.push(metadata.location); } } - self.delete_paths(paths).await + // Leases and superseded marker states are dependent metadata. Keep at + // least one terminal marker as the retry anchor until they are gone. + self.delete_paths(lease_paths).await?; + self.delete_paths(dependent_marker_paths).await?; + self.delete_paths(anchor_paths).await } async fn delete_paths(&self, paths: Vec) -> Result<()> { @@ -877,6 +1042,44 @@ impl VersionLeaseStore { } } +/// Verify that a durable tag or branch can be admitted for a version. +/// +/// Callers check both before and after publishing their reference. Cleanup +/// publishes a draining marker before its final reference scan, so either the +/// reference is visible to that scan or the second admission check rejects and +/// rolls it back. +pub(super) async fn ensure_reference_admissible( + refs: &Refs, + branch: Option<&str>, + version: u64, +) -> Result<()> { + let branch_location = refs.base_location.find_branch(branch)?; + let manifest = refs + .commit_handler + .resolve_version_location(&branch_location.path, version, &refs.object_store.inner) + .await?; + if !refs.object_store.exists(&manifest.path).await? { + return Err(Error::VersionNotFound { + message: format!("version {version} no longer exists and cannot be referenced"), + }); + } + + let store = VersionLeaseStore::for_refs(refs, branch, None).await?; + if store.has_active_retirement_marker(version).await? { + return Err(Error::RefConflict { + message: format!( + "version {version} is retiring and cannot accept a new durable reference" + ), + }); + } + if !refs.object_store.exists(&manifest.path).await? { + return Err(Error::VersionNotFound { + message: format!("version {version} no longer exists and cannot be referenced"), + }); + } + Ok(()) +} + pub(super) async fn remove_branch_state( object_store: &ObjectStore, root_path: &Path, @@ -1050,6 +1253,8 @@ fn retiring_version_error(version: u64) -> Error { #[cfg(test)] mod tests { + use crate::utils::test::FailingProxyStore; + use lance_io::object_store::WrappingObjectStore; use mock_instant::thread_local::MockClock; use super::*; @@ -1139,6 +1344,72 @@ mod tests { guard.cancel_all().await.unwrap(); } + #[tokio::test] + async fn committed_retirement_cannot_be_cancelled() { + let store = memory_store(); + let mut guard = store.fence_versions(&manifest_paths(42)).await.unwrap(); + guard.seal_versions(&HashSet::from([42])).await.unwrap(); + guard.commit_versions(&HashSet::from([42])).await.unwrap(); + let committed_path = guard.fences[&42].committed_path.clone().unwrap(); + + guard.cancel_all().await.unwrap(); + + assert!(!guard.is_empty()); + assert!(store.object_store.exists(&committed_path).await.unwrap()); + let error = store + .acquire(42, Duration::from_secs(60)) + .await + .unwrap_err(); + assert!(error.to_string().contains("retiring"), "{error}"); + } + + #[tokio::test] + async fn terminal_marker_survives_lease_deletion_failure() { + let failing_store = Arc::new(FailingProxyStore::new()); + let mut object_store = ObjectStore::memory(); + object_store.inner = failing_store.wrap("memory", Arc::clone(&object_store.inner)); + let store = VersionLeaseStore { + object_store: Arc::new(object_store), + leases_path: Path::from("leases"), + markers_path: Path::from("markers"), + manifest_path: None, + }; + let lease = store.acquire(42, Duration::from_secs(60)).await.unwrap(); + let manifests = manifest_paths(42); + let mut guard = store.fence_versions(&manifests).await.unwrap(); + guard.seal_versions(&HashSet::from([42])).await.unwrap(); + guard.commit_versions(&HashSet::from([42])).await.unwrap(); + let committed_path = guard.fences[&42].committed_path.clone().unwrap(); + failing_store.fail_when( + "delete", + LEASE_FILE_SUFFIX, + "injected lease deletion failure", + ); + + let error = guard.finalize(&manifests).await.unwrap_err(); + + assert!( + error + .to_string() + .contains("injected lease deletion failure"), + "{error}" + ); + assert!(store.object_store.exists(&committed_path).await.unwrap()); + assert!(store.object_store.exists(&lease.path).await.unwrap()); + + failing_store.clear_fail_when("delete", LEASE_FILE_SUFFIX); + assert!( + store + .clone() + .recover_retirements() + .await + .unwrap() + .is_empty() + ); + assert!(!store.object_store.exists(&committed_path).await.unwrap()); + assert!(!store.object_store.exists(&lease.path).await.unwrap()); + } + #[test] fn lease_ttl_survives_coarse_storage_timestamps() { let storage_second = DateTime::from_timestamp(100, 0).unwrap(); From fb524c49b92ecf0d9ebf5cd50465ecb4b12bf4a4 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 20:26:41 +0000 Subject: [PATCH 05/13] fix: make durable reference admission race-safe --- rust/lance/src/dataset/cleanup.rs | 129 ++++++++--- rust/lance/src/dataset/refs.rs | 249 ++++++++++++++------ rust/lance/src/dataset/version_lease.rs | 293 ++++++++++++++++++++++-- 3 files changed, 553 insertions(+), 118 deletions(-) diff --git a/rust/lance/src/dataset/cleanup.rs b/rust/lance/src/dataset/cleanup.rs index 93f9046367e..4f7429d0bbc 100644 --- a/rust/lance/src/dataset/cleanup.rs +++ b/rust/lance/src/dataset/cleanup.rs @@ -504,19 +504,27 @@ impl<'a> CleanupTask<'a> { )); } + let ignored_branch_manifests: HashSet<_> = final_result + .removed_manifests + .union(&self.ignored_manifests) + .cloned() + .collect(); if !referenced_branches.is_empty() { - let ignored_manifests: HashSet<_> = final_result - .removed_manifests - .union(&self.ignored_manifests) - .cloned() - .collect(); inspection = self - .retain_branch_lineage_files(inspection, &referenced_branches, &ignored_manifests) + .retain_branch_lineage_files( + inspection, + &referenced_branches, + &ignored_branch_manifests, + ) .await? }; let (inspection, mut retirement) = self - .fence_old_versions_and_retain_new_leases(inspection, &version_lease_store) + .fence_old_versions_and_retain_new_leases( + inspection, + &version_lease_store, + &ignored_branch_manifests, + ) .await?; let cleanup_result = self.delete_unreferenced_files(inspection).await?; @@ -628,6 +636,7 @@ impl<'a> CleanupTask<'a> { &self, mut inspection: CleanupInspection, version_lease_store: &VersionLeaseStore, + ignored_branch_manifests: &HashSet, ) -> Result<(CleanupInspection, Option)> { if !self.action.deletes_files() { return Ok((inspection, None)); @@ -683,9 +692,10 @@ impl<'a> CleanupTask<'a> { } guard.cancel_versions(&retained_versions).await?; - // A tag or child branch published before the draining marker must - // win. Publications racing the marker perform their own second - // admission check and roll back if retirement won instead. + // A canonical tag, descendant manifest, or durable publication + // intent admitted before the draining marker must win. Descendant + // retention is based on actual lineage file references rather than + // cleanup policy or branch metadata alone. let tags = self.dataset.tags().list().await?; let current_branch = &self.dataset.manifest.branch; let mut referenced_versions = tags @@ -697,19 +707,25 @@ impl<'a> CleanupTask<'a> { }) .map(|tag| tag.version) .collect::>(); - referenced_versions.extend( - self.find_referenced_branches() - .await? - .into_iter() - .map(|(_, version)| version), - ); - let retained_versions = versions_to_seal + referenced_versions.extend(version_lease_store.active_reference_versions().await?); + let mut retained_versions = versions_to_seal .intersection(&referenced_versions) .copied() .collect::>(); for version in &retained_versions { self.retain_version(&mut inspection, *version).await?; } + + let all_descendants = self.all_referenced_branches().await?; + inspection = self + .retain_branch_lineage_files(inspection, &all_descendants, ignored_branch_manifests) + .await?; + let remaining_versions = inspection + .old_manifests + .values() + .copied() + .collect::>(); + retained_versions.extend(versions_to_seal.difference(&remaining_versions).copied()); guard.cancel_versions(&retained_versions).await?; let versions_to_commit = inspection @@ -717,7 +733,11 @@ impl<'a> CleanupTask<'a> { .values() .copied() .collect::>(); - guard.commit_versions(&versions_to_commit).await?; + let late_reference_versions = guard.commit_versions(&versions_to_commit).await?; + for version in &late_reference_versions { + self.retain_version(&mut inspection, *version).await?; + } + guard.cancel_versions(&late_reference_versions).await?; let mut manifest_paths = HashMap::>::new(); for (path, version) in &inspection.old_manifests { @@ -1260,9 +1280,7 @@ impl<'a> CleanupTask<'a> { } async fn find_referenced_branches(&self) -> Result> { - let current_branch_id = self.dataset.branch_identifier().await?; - let all_branches = self.dataset.branches().list().await?; - let children = current_branch_id.collect_referenced_versions(&all_branches); + let children = self.all_referenced_branches().await?; // Use a concurrent set to identify branches eligible for cleanup. // The filter below preserves the original (branch_name, version) tuples. @@ -1284,16 +1302,23 @@ impl<'a> CleanupTask<'a> { ) .await?; + // A descendant that has already detached from this + // lineage can outlive deletion of its original parent + // manifest. Confirm durable absence explicitly; once an + // object is present, every read error is safety-critical + // and must be propagated. + if !dataset.object_store.exists(&manifest_location.path).await? { + return Ok(()); + } + let manifest = read_manifest( &dataset.object_store, &manifest_location.path, manifest_location.size, ) - .await; + .await?; - if let Ok(manifest) = manifest - && policy.should_clean(&manifest) - { + if policy.should_clean(&manifest) { referenced_branches.insert(branch_name.clone()); } Ok::<(), Error>(()) @@ -1305,12 +1330,16 @@ impl<'a> CleanupTask<'a> { // Filter children to only include branches that should be cleaned. // The DashSet contains branch names found eligible during concurrent scan. - let referenced_branches = children - .iter() + Ok(children + .into_iter() .filter(|(branch_name, _)| referenced_branches.contains(branch_name)) - .cloned() - .collect(); - Ok(referenced_branches) + .collect()) + } + + async fn all_referenced_branches(&self) -> Result> { + let current_branch_id = self.dataset.branch_identifier().await?; + let all_branches = self.dataset.branches().list().await?; + Ok(current_branch_id.collect_referenced_versions(&all_branches)) } async fn clean_referenced_branches( @@ -2510,6 +2539,46 @@ mod tests { ); } + #[tokio::test] + async fn recovered_retirement_still_respects_descendant_branch() { + let fixture = MockDatasetFixture::try_new().unwrap(); + fixture.create_some_data().await.unwrap(); + let historical = fixture.load().await.unwrap(); + fixture.overwrite_some_data().await.unwrap(); + let mut dataset = fixture.load().await.unwrap(); + fixture + .create_branch_and_load(&mut dataset, "child", (None, Some(1))) + .await + .unwrap(); + + let store = VersionLeaseStore::for_dataset(&historical).await.unwrap(); + let manifest_paths = HashMap::from([( + historical.version().version, + vec![historical.manifest_location.path.clone()], + )]); + let mut guard = store.fence_versions(&manifest_paths).await.unwrap(); + guard.seal_versions(&HashSet::from([1])).await.unwrap(); + drop(guard); + + let removed = fixture + .run_cleanup_with_policy(CleanupPolicy { + before_version: Some(1), + ..Default::default() + }) + .await + .unwrap(); + + assert_eq!(removed.old_versions, 0); + assert!( + historical + .object_store + .exists(&manifest_paths[&1][0]) + .await + .unwrap() + ); + dataset.checkout_version(("child", None)).await.unwrap(); + } + #[tokio::test] async fn forced_retirement_still_respects_active_lease() { let fixture = MockDatasetFixture::try_new().unwrap(); diff --git a/rust/lance/src/dataset/refs.rs b/rust/lance/src/dataset/refs.rs index 4ec4ce4c336..dc6a0c1c82f 100644 --- a/rust/lance/src/dataset/refs.rs +++ b/rust/lance/src/dataset/refs.rs @@ -3,18 +3,21 @@ use std::ops::Range; +use bytes::Bytes; use chrono::{DateTime, Utc}; 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::{ObjectStoreExt, PutMode, PutOptions, UpdateVersion, path::Path}; use serde::{Deserialize, Serialize}; use std::sync::Arc; use crate::dataset::branch_location::BranchLocation; use crate::dataset::refs::Ref::{Tag, Version, VersionNumber}; -use crate::dataset::version_lease::{ensure_reference_admissible, remove_branch_state}; +use crate::dataset::version_lease::{ + ReferenceAdmission, begin_reference_admission, remove_branch_state, +}; use crate::utils::temporal::utc_now; use crate::{Error, Result}; use serde::de::DeserializeOwned; @@ -145,6 +148,117 @@ impl Branches<'_> { } } +async fn publish_new_reference( + object_store: &ObjectStore, + path: &Path, + payload: Vec, + admission: ReferenceAdmission, + reference_kind: &str, + reference_name: &str, +) -> Result<()> { + if let Err(error) = admission.ensure_owned().await { + admission.cancel_before_publish().await; + return Err(error); + } + match object_store + .inner + .put_opts( + path, + Bytes::from(payload).into(), + PutOptions { + mode: PutMode::Create, + ..Default::default() + }, + ) + .await + { + Ok(_) => { + admission.complete().await; + Ok(()) + } + Err( + object_store::Error::AlreadyExists { .. } | object_store::Error::Precondition { .. }, + ) => { + admission.cancel_before_publish().await; + Err(Error::RefConflict { + message: format!("{reference_kind} {reference_name} already exists"), + }) + } + // A generic transport failure can be ambiguous after the server has + // accepted the conditional write. Keep the intent until bounded + // recovery instead of risking deletion beneath a published reference. + Err(error) => Err(error.into()), + } +} + +async fn publish_updated_reference( + object_store: &ObjectStore, + path: &Path, + expected_payload: &[u8], + expected_version: UpdateVersion, + payload: Vec, + admission: ReferenceAdmission, + reference_label: String, +) -> Result<()> { + if let Err(error) = admission.ensure_owned().await { + admission.cancel_before_publish().await; + return Err(error); + } + let conditional_result = object_store + .inner + .put_opts( + path, + Bytes::from(payload.clone()).into(), + PutOptions { + mode: PutMode::Update(expected_version), + ..Default::default() + }, + ) + .await; + let result = match conditional_result { + Err( + object_store::Error::NotSupported { .. } | object_store::Error::NotImplemented { .. }, + ) => { + // The local object store does not implement UpdateVersion. Preserve + // local tag updates, but verify the exact object immediately before + // the atomic overwrite. Cloud stores retain native CAS semantics. + let current = object_store.inner.get(path).await?.bytes().await?; + if current.as_ref() != expected_payload { + admission.cancel_before_publish().await; + return Err(Error::RefConflict { + message: format!("{reference_label} changed during conditional update"), + }); + } + if let Err(error) = admission.ensure_owned().await { + admission.cancel_before_publish().await; + return Err(error); + } + object_store + .inner + .put(path, Bytes::from(payload).into()) + .await + } + result => result, + }; + match result { + Ok(_) => { + admission.complete().await; + Ok(()) + } + Err( + object_store::Error::AlreadyExists { .. } + | object_store::Error::Precondition { .. } + | object_store::Error::NotFound { .. }, + ) => { + admission.cancel_before_publish().await; + Err(Error::RefConflict { + message: format!("{reference_label} changed during conditional update"), + }) + } + Err(error) => Err(error.into()), + } +} + impl Tags<'_> { pub async fn fetch_tags(&self) -> Result> { let root_location = self.refs.root()?; @@ -227,30 +341,22 @@ impl Tags<'_> { let tag_contents = self .build_tag_content_by_ref(reference, Some(now), Some(now)) .await?; - ensure_reference_admissible( + let payload = serde_json::to_vec_pretty(&tag_contents)?; + let admission = begin_reference_admission( self.refs, tag_contents.branch.as_deref(), tag_contents.version, ) .await?; - - self.object_store() - .put( - &tag_file, - serde_json::to_string_pretty(&tag_contents)?.as_bytes(), - ) - .await?; - if let Err(error) = ensure_reference_admissible( - self.refs, - tag_contents.branch.as_deref(), - tag_contents.version, + publish_new_reference( + self.object_store(), + &tag_file, + payload, + admission, + "tag", + tag, ) .await - { - self.object_store().delete(&tag_file).await?; - return Err(error); - } - Ok(()) } pub async fn delete(&self, tag: &str) -> Result<()> { @@ -278,45 +384,39 @@ impl Tags<'_> { message: format!("tag {} does not exist", tag), }); } - let original_tag_contents = TagContents::from_path(&tag_file, self.object_store()).await?; + let original_get = self.object_store().inner.get(&tag_file).await?; + let expected_version = UpdateVersion { + e_tag: original_get.meta.e_tag.clone(), + version: original_get.meta.version.clone(), + }; + let original_payload = original_get.bytes().await?; + let original_tag_contents: TagContents = serde_json::from_slice(&original_payload)?; let updated_reference = self .build_tag_content_by_ref(reference, original_tag_contents.created_at, Some(utc_now())) .await?; - ensure_reference_admissible( - self.refs, - updated_reference.branch.as_deref(), - updated_reference.version, - ) - .await?; - let mut tag_contents = original_tag_contents.clone(); + let mut tag_contents = original_tag_contents; tag_contents.branch = updated_reference.branch; tag_contents.version = updated_reference.version; tag_contents.created_at = updated_reference.created_at; tag_contents.updated_at = updated_reference.updated_at; tag_contents.manifest_size = updated_reference.manifest_size; - - self.object_store() - .put( - &tag_file, - serde_json::to_string_pretty(&tag_contents)?.as_bytes(), - ) - .await?; - if let Err(error) = ensure_reference_admissible( + let payload = serde_json::to_vec_pretty(&tag_contents)?; + let admission = begin_reference_admission( self.refs, tag_contents.branch.as_deref(), tag_contents.version, ) + .await?; + publish_updated_reference( + self.object_store(), + &tag_file, + &original_payload, + expected_version, + payload, + admission, + format!("tag {tag}"), + ) .await - { - self.object_store() - .put( - &tag_file, - serde_json::to_string_pretty(&original_tag_contents)?.as_bytes(), - ) - .await?; - return Err(error); - } - Ok(()) } pub async fn replace_metadata( @@ -334,16 +434,31 @@ impl Tags<'_> { }); } - let mut tag_contents = TagContents::from_path(&tag_file, self.object_store()).await?; + let original_get = self.object_store().inner.get(&tag_file).await?; + let expected_version = UpdateVersion { + e_tag: original_get.meta.e_tag.clone(), + version: original_get.meta.version.clone(), + }; + let original_payload = original_get.bytes().await?; + let mut tag_contents: TagContents = serde_json::from_slice(&original_payload)?; tag_contents.metadata = metadata; - - self.object_store() - .put( - &tag_file, - serde_json::to_string_pretty(&tag_contents)?.as_bytes(), - ) - .await - .map(|_| ()) + let payload = serde_json::to_vec_pretty(&tag_contents)?; + let admission = begin_reference_admission( + self.refs, + tag_contents.branch.as_deref(), + tag_contents.version, + ) + .await?; + publish_updated_reference( + self.object_store(), + &tag_file, + &original_payload, + expected_version, + payload, + admission, + format!("tag {tag}"), + ) + .await } async fn build_tag_content_by_ref( @@ -516,8 +631,6 @@ impl Branches<'_> { message: format!("Manifest file {} does not exist", manifest_file.path), }); }; - ensure_reference_admissible(self.refs, source_branch.as_deref(), version_number).await?; - let parent_branch_id = if let Some(ref parent_branch) = source_branch { let parent_file = branch_contents_path(&root_location.path, parent_branch); if self.object_store().exists(&parent_file).await? { @@ -545,24 +658,22 @@ impl Branches<'_> { }, metadata: HashMap::new(), }; - - self.object_store() - .put( - &branch_file, - serde_json::to_string_pretty(&branch_contents)?.as_bytes(), - ) - .await?; - if let Err(error) = ensure_reference_admissible( + let payload = serde_json::to_vec_pretty(&branch_contents)?; + let admission = begin_reference_admission( self.refs, branch_contents.parent_branch.as_deref(), version_number, ) + .await?; + publish_new_reference( + self.object_store(), + &branch_file, + payload, + admission, + "branch", + branch_name, + ) .await - { - self.object_store().delete(&branch_file).await?; - return Err(error); - } - Ok(()) } pub async fn replace_metadata( diff --git a/rust/lance/src/dataset/version_lease.rs b/rust/lance/src/dataset/version_lease.rs index c03913be0b3..998e5ffbd3c 100644 --- a/rust/lance/src/dataset/version_lease.rs +++ b/rust/lance/src/dataset/version_lease.rs @@ -26,7 +26,9 @@ use crate::{Error, Result}; const LEASES_DIR: &str = "_refs/version_leases"; const LEASE_GC_MARKERS_DIR: &str = "_refs/version_lease_gc"; +const REFERENCE_INTENTS_DIR: &str = "_refs/version_reference_intents"; const LEASE_FILE_SUFFIX: &str = ".lease"; +const REFERENCE_INTENT_SUFFIX: &str = ".intent"; const DRAINING_MARKER_SUFFIX: &str = ".draining"; const SEALED_MARKER_SUFFIX: &str = ".sealed"; const COMMITTED_MARKER_SUFFIX: &str = ".committed"; @@ -41,6 +43,11 @@ const STORAGE_TIMESTAMP_PRECISION: Duration = Duration::from_secs(1); /// ignore and remove the abandoned drain without racing a live deletion. const DRAINING_OWNERSHIP_TIMEOUT: Duration = Duration::from_secs(15 * 60); +/// Reference publication is expected to be short-lived. An abandoned intent +/// protects its target for this ownership window and is then recoverable by +/// cleanup, preventing a crashed writer from retaining a version forever. +const REFERENCE_ADMISSION_TIMEOUT: Duration = Duration::from_secs(15 * 60); + /// Drains whose in-process owner was dropped before sealing. The operation UUID /// makes each path unique across stores, while process-wide scope lets a fresh /// dataset handle immediately disregard a drain that can no longer delete. @@ -146,9 +153,24 @@ pub(super) struct VersionLeaseStore { object_store: Arc, leases_path: Path, markers_path: Path, + reference_intents_path: Path, manifest_path: Option, } +/// Durable ownership of one in-flight tag or branch publication. +/// +/// This type intentionally has no `Drop` cleanup. Cancelling a future while a +/// remote conditional write is in flight must leave the intent durable until +/// cleanup can safely expire it. +#[derive(Debug)] +pub(super) struct ReferenceAdmission { + store: VersionLeaseStore, + path: Path, + manifest_path: Path, + version: u64, + created_at: DateTime, +} + #[derive(Debug)] struct LeaseFile { path: Path, @@ -185,6 +207,12 @@ struct RetirementMarkerMetadata { metadata: ObjectMeta, } +#[derive(Debug)] +struct ReferenceIntentMetadata { + version: u64, + metadata: ObjectMeta, +} + /// Per-cleanup retirement markers for versions selected for deletion. /// /// Draining markers reject new leases while allowing already-admitted leases @@ -235,7 +263,8 @@ impl VersionLeaseStore { Ok(Self { object_store: Arc::clone(&refs.object_store), leases_path: root.path.clone().join(LEASES_DIR).join(namespace), - markers_path: root.path.join(LEASE_GC_MARKERS_DIR).join(namespace), + markers_path: root.path.clone().join(LEASE_GC_MARKERS_DIR).join(namespace), + reference_intents_path: root.path.join(REFERENCE_INTENTS_DIR).join(namespace), manifest_path, }) } @@ -299,6 +328,103 @@ impl VersionLeaseStore { .collect() } + pub(super) async fn active_reference_versions(&self) -> Result> { + let observed_at = self.storage_observed_at().await?; + let mut active_versions = HashSet::new(); + let mut stale_paths = Vec::new(); + for intent in self.reference_intent_metadata().await? { + if reference_owner_is_active(intent.metadata.last_modified, observed_at)? { + active_versions.insert(intent.version); + } else { + stale_paths.push(intent.metadata.location); + } + } + self.delete_paths(stale_paths).await?; + Ok(active_versions) + } + + async fn create_reference_intent(&self, version: u64) -> Result<(Path, DateTime)> { + let path = self + .reference_intents_path + .clone() + .join(version.to_string()) + .join(format!( + "{}{}", + Uuid::new_v4().simple(), + REFERENCE_INTENT_SUFFIX + )); + self.object_store + .inner + .put_opts( + &path, + Bytes::new().into(), + PutOptions { + mode: PutMode::Create, + ..Default::default() + }, + ) + .await?; + // A failed HEAD is ambiguous after a successful create. Leave the + // intent for bounded recovery instead of deleting a write that may be + // protecting an in-flight canonical publication. + let metadata = self.object_store.inner.head(&path).await?; + Ok((path, metadata.last_modified)) + } + + async fn reference_intent_metadata(&self) -> Result> { + self.object_store + .list(Some(self.reference_intents_path.clone())) + .try_collect::>() + .await? + .into_iter() + .map(|metadata| self.parse_reference_intent_metadata(metadata)) + .collect() + } + + fn parse_reference_intent_metadata( + &self, + metadata: ObjectMeta, + ) -> Result { + let relative_parts = metadata + .location + .prefix_match(&self.reference_intents_path) + .ok_or_else(|| { + Error::corrupt_file( + metadata.location.clone(), + "reference intent is outside its namespace", + ) + })? + .map(|part| part.as_ref().to_string()) + .collect::>(); + if relative_parts.len() != 2 { + return Err(Error::corrupt_file( + metadata.location, + "reference intent path must contain a version and operation filename", + )); + } + let version = relative_parts[0].parse::().map_err(|error| { + Error::corrupt_file( + metadata.location.clone(), + format!("reference intent has invalid version: {error}"), + ) + })?; + let operation_id = relative_parts[1] + .strip_suffix(REFERENCE_INTENT_SUFFIX) + .ok_or_else(|| { + Error::corrupt_file( + metadata.location.clone(), + "reference intent filename has an invalid suffix", + ) + })?; + Uuid::parse_str(operation_id).map_err(|error| { + Error::corrupt_file( + metadata.location.clone(), + format!("reference intent has invalid operation id: {error}"), + ) + })?; + Ok(ReferenceIntentMetadata { version, metadata }) + } + pub(super) async fn active_versions_at( &self, observed_at: &HashMap>, @@ -578,6 +704,7 @@ impl VersionLeaseStore { .map(|marker| (marker.version, observed_at)) .collect::>(); let actively_leased_versions = self.active_versions_at(&sealed_observation, true).await?; + let actively_referenced_versions = self.active_reference_versions().await?; let mut stale_drains = Vec::new(); let mut sealed_manifest_paths = HashMap::>::new(); for marker in markers { @@ -611,7 +738,9 @@ impl VersionLeaseStore { // A seal is the renewal cutoff. If its owner disappeared before // confirming the final lease scan, wait out every lease that was // still entitled to its published TTL before resuming deletion. - if !actively_leased_versions.contains(&version) { + if !actively_leased_versions.contains(&version) + && !actively_referenced_versions.contains(&version) + { versions_to_resume.insert(version); } } else { @@ -775,11 +904,26 @@ impl RetirementGuard { } /// Publish the irreversible retirement boundary after the final lease and - /// durable-reference scans. Reference admission checks this marker before - /// and after publishing a tag or branch. - pub(super) async fn commit_versions(&mut self, versions: &HashSet) -> Result<()> { + /// durable-reference scans. A final intent scan keeps any publication that + /// crossed the caller's census on the cancellable side of the boundary. + pub(super) async fn commit_versions( + &mut self, + versions: &HashSet, + ) -> Result> { + // This scan runs inside the commit operation, after the caller's final + // reference census. An admission that won that narrow race remains + // cancellable and cannot cross the irreversible committed boundary. + let active_reference_versions = self.store.active_reference_versions().await?; + let retained_versions = versions + .intersection(&active_reference_versions) + .copied() + .collect::>(); + let versions = versions + .difference(&retained_versions) + .copied() + .collect::>(); let mut marker_paths = Vec::with_capacity(versions.len()); - for version in versions { + for version in &versions { let fence = self.fences.get(version).ok_or_else(|| { Error::internal(format!("missing sealed fence for version {version}")) })?; @@ -848,13 +992,13 @@ impl RetirementGuard { "Failed to remove superseded sealed retirement markers" ); } else { - for version in versions { + for version in &versions { if let Some(fence) = self.fences.get_mut(version) { fence.sealed_path = None; } } } - Ok(()) + Ok(retained_versions) } pub(super) async fn cancel_versions(&mut self, versions: &HashSet) -> Result<()> { @@ -1021,9 +1165,16 @@ impl VersionLeaseStore { lease_paths.push(metadata.location); } } + let mut intent_paths = Vec::new(); + for intent in self.reference_intent_metadata().await? { + if versions.contains(&intent.version) { + intent_paths.push(intent.metadata.location); + } + } // Leases and superseded marker states are dependent metadata. Keep at // least one terminal marker as the retry anchor until they are gone. self.delete_paths(lease_paths).await?; + self.delete_paths(intent_paths).await?; self.delete_paths(dependent_marker_paths).await?; self.delete_paths(anchor_paths).await } @@ -1042,17 +1193,12 @@ impl VersionLeaseStore { } } -/// Verify that a durable tag or branch can be admitted for a version. -/// -/// Callers check both before and after publishing their reference. Cleanup -/// publishes a draining marker before its final reference scan, so either the -/// reference is visible to that scan or the second admission check rejects and -/// rolls it back. -pub(super) async fn ensure_reference_admissible( +/// Begin a durable tag or branch admission before canonical publication. +pub(super) async fn begin_reference_admission( refs: &Refs, branch: Option<&str>, version: u64, -) -> Result<()> { +) -> Result { let branch_location = refs.base_location.find_branch(branch)?; let manifest = refs .commit_handler @@ -1065,6 +1211,74 @@ pub(super) async fn ensure_reference_admissible( } let store = VersionLeaseStore::for_refs(refs, branch, None).await?; + ensure_reference_available(&store, &manifest.path, version).await?; + let (path, created_at) = store.create_reference_intent(version).await?; + let admission = ReferenceAdmission { + store, + path, + manifest_path: manifest.path, + version, + created_at, + }; + if let Err(error) = admission.ensure_owned().await { + admission.cancel_before_publish().await; + return Err(error); + } + Ok(admission) +} + +impl ReferenceAdmission { + /// Recheck ownership immediately before the canonical conditional write. + pub(super) async fn ensure_owned(&self) -> Result<()> { + let metadata = self.store.object_store.inner.head(&self.path).await?; + let observed_at = self.store.storage_observed_at().await?; + if metadata.last_modified != self.created_at + || !reference_owner_is_active(metadata.last_modified, observed_at)? + { + return Err(Error::RefConflict { + message: format!( + "reference admission for version {} no longer owns its publication intent", + self.version + ), + }); + } + ensure_reference_available(&self.store, &self.manifest_path, self.version).await + } + + /// Remove this operation's uniquely owned intent after a definitive result + /// that did not publish the canonical reference. + pub(super) async fn cancel_before_publish(&self) { + if let Err(error) = self.store.object_store.delete(&self.path).await + && !error.is_not_found() + { + tracing::warn!( + path = %self.path, + error = %error, + "Failed to remove cancelled reference admission intent" + ); + } + } + + /// Canonical publication is durable. Intent removal is best-effort because + /// the canonical tag or branch now protects the version independently. + pub(super) async fn complete(&self) { + if let Err(error) = self.store.object_store.delete(&self.path).await + && !error.is_not_found() + { + tracing::warn!( + path = %self.path, + error = %error, + "Failed to remove completed reference admission intent" + ); + } + } +} + +async fn ensure_reference_available( + store: &VersionLeaseStore, + manifest_path: &Path, + version: u64, +) -> Result<()> { if store.has_active_retirement_marker(version).await? { return Err(Error::RefConflict { message: format!( @@ -1072,7 +1286,7 @@ pub(super) async fn ensure_reference_admissible( ), }); } - if !refs.object_store.exists(&manifest.path).await? { + if !store.object_store.exists(manifest_path).await? { return Err(Error::VersionNotFound { message: format!("version {version} no longer exists and cannot be referenced"), }); @@ -1094,6 +1308,10 @@ pub(super) async fn remove_branch_state( .clone() .join(LEASE_GC_MARKERS_DIR) .join(namespace.to_string()), + root_path + .clone() + .join(REFERENCE_INTENTS_DIR) + .join(namespace.to_string()), ] { match object_store.remove_dir_all(path).await { Ok(()) => {} @@ -1181,6 +1399,13 @@ fn draining_owner_is_active( ) } +fn reference_owner_is_active( + intent_observed_at: DateTime, + current_observed_at: DateTime, +) -> Result { + Ok(expiration_from_ttl(intent_observed_at, REFERENCE_ADMISSION_TIMEOUT)? > current_observed_at) +} + fn retirement_marker_payload(manifest_paths: &[Path]) -> Result { if manifest_paths.is_empty() { return Err(Error::internal( @@ -1264,6 +1489,7 @@ mod tests { object_store: Arc::new(ObjectStore::memory()), leases_path: Path::from("leases"), markers_path: Path::from("markers"), + reference_intents_path: Path::from("reference_intents"), manifest_path: None, } } @@ -1349,7 +1575,13 @@ mod tests { let store = memory_store(); let mut guard = store.fence_versions(&manifest_paths(42)).await.unwrap(); guard.seal_versions(&HashSet::from([42])).await.unwrap(); - guard.commit_versions(&HashSet::from([42])).await.unwrap(); + assert!( + guard + .commit_versions(&HashSet::from([42])) + .await + .unwrap() + .is_empty() + ); let committed_path = guard.fences[&42].committed_path.clone().unwrap(); guard.cancel_all().await.unwrap(); @@ -1363,6 +1595,22 @@ mod tests { assert!(error.to_string().contains("retiring"), "{error}"); } + #[tokio::test] + async fn reference_intent_blocks_retirement_commit() { + let store = memory_store(); + let (intent_path, _) = store.create_reference_intent(42).await.unwrap(); + let mut guard = store.fence_versions(&manifest_paths(42)).await.unwrap(); + guard.seal_versions(&HashSet::from([42])).await.unwrap(); + + let retained_versions = guard.commit_versions(&HashSet::from([42])).await.unwrap(); + + assert_eq!(retained_versions, HashSet::from([42])); + assert!(guard.fences[&42].committed_path.is_none()); + assert!(guard.fences[&42].sealed_path.is_some()); + store.object_store.delete(&intent_path).await.unwrap(); + guard.cancel_all().await.unwrap(); + } + #[tokio::test] async fn terminal_marker_survives_lease_deletion_failure() { let failing_store = Arc::new(FailingProxyStore::new()); @@ -1372,13 +1620,20 @@ mod tests { object_store: Arc::new(object_store), leases_path: Path::from("leases"), markers_path: Path::from("markers"), + reference_intents_path: Path::from("reference_intents"), manifest_path: None, }; let lease = store.acquire(42, Duration::from_secs(60)).await.unwrap(); let manifests = manifest_paths(42); let mut guard = store.fence_versions(&manifests).await.unwrap(); guard.seal_versions(&HashSet::from([42])).await.unwrap(); - guard.commit_versions(&HashSet::from([42])).await.unwrap(); + assert!( + guard + .commit_versions(&HashSet::from([42])) + .await + .unwrap() + .is_empty() + ); let committed_path = guard.fences[&42].committed_path.clone().unwrap(); failing_store.fail_when( "delete", From 95294fb55f00f4a5fef4a199cb5960be134c043a Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 22:01:51 +0000 Subject: [PATCH 06/13] fix: close reference retirement handoff races --- rust/lance/src/dataset/refs.rs | 240 +++++--------- rust/lance/src/dataset/version_lease.rs | 396 ++++++++++++++++++++++-- 2 files changed, 435 insertions(+), 201 deletions(-) diff --git a/rust/lance/src/dataset/refs.rs b/rust/lance/src/dataset/refs.rs index dc6a0c1c82f..87a15e7d083 100644 --- a/rust/lance/src/dataset/refs.rs +++ b/rust/lance/src/dataset/refs.rs @@ -3,20 +3,19 @@ use std::ops::Range; -use bytes::Bytes; use chrono::{DateTime, Utc}; use futures::stream::{StreamExt, TryStreamExt}; use itertools::Itertools; use lance_io::object_store::ObjectStore; use lance_table::io::commit::CommitHandler; -use object_store::{ObjectStoreExt, PutMode, PutOptions, UpdateVersion, path::Path}; +use object_store::{ObjectStoreExt, path::Path}; use serde::{Deserialize, Serialize}; use std::sync::Arc; use crate::dataset::branch_location::BranchLocation; use crate::dataset::refs::Ref::{Tag, Version, VersionNumber}; use crate::dataset::version_lease::{ - ReferenceAdmission, begin_reference_admission, remove_branch_state, + ReferenceMutation, begin_reference_admission, remove_branch_state, }; use crate::utils::temporal::utc_now; use crate::{Error, Result}; @@ -148,117 +147,6 @@ impl Branches<'_> { } } -async fn publish_new_reference( - object_store: &ObjectStore, - path: &Path, - payload: Vec, - admission: ReferenceAdmission, - reference_kind: &str, - reference_name: &str, -) -> Result<()> { - if let Err(error) = admission.ensure_owned().await { - admission.cancel_before_publish().await; - return Err(error); - } - match object_store - .inner - .put_opts( - path, - Bytes::from(payload).into(), - PutOptions { - mode: PutMode::Create, - ..Default::default() - }, - ) - .await - { - Ok(_) => { - admission.complete().await; - Ok(()) - } - Err( - object_store::Error::AlreadyExists { .. } | object_store::Error::Precondition { .. }, - ) => { - admission.cancel_before_publish().await; - Err(Error::RefConflict { - message: format!("{reference_kind} {reference_name} already exists"), - }) - } - // A generic transport failure can be ambiguous after the server has - // accepted the conditional write. Keep the intent until bounded - // recovery instead of risking deletion beneath a published reference. - Err(error) => Err(error.into()), - } -} - -async fn publish_updated_reference( - object_store: &ObjectStore, - path: &Path, - expected_payload: &[u8], - expected_version: UpdateVersion, - payload: Vec, - admission: ReferenceAdmission, - reference_label: String, -) -> Result<()> { - if let Err(error) = admission.ensure_owned().await { - admission.cancel_before_publish().await; - return Err(error); - } - let conditional_result = object_store - .inner - .put_opts( - path, - Bytes::from(payload.clone()).into(), - PutOptions { - mode: PutMode::Update(expected_version), - ..Default::default() - }, - ) - .await; - let result = match conditional_result { - Err( - object_store::Error::NotSupported { .. } | object_store::Error::NotImplemented { .. }, - ) => { - // The local object store does not implement UpdateVersion. Preserve - // local tag updates, but verify the exact object immediately before - // the atomic overwrite. Cloud stores retain native CAS semantics. - let current = object_store.inner.get(path).await?.bytes().await?; - if current.as_ref() != expected_payload { - admission.cancel_before_publish().await; - return Err(Error::RefConflict { - message: format!("{reference_label} changed during conditional update"), - }); - } - if let Err(error) = admission.ensure_owned().await { - admission.cancel_before_publish().await; - return Err(error); - } - object_store - .inner - .put(path, Bytes::from(payload).into()) - .await - } - result => result, - }; - match result { - Ok(_) => { - admission.complete().await; - Ok(()) - } - Err( - object_store::Error::AlreadyExists { .. } - | object_store::Error::Precondition { .. } - | object_store::Error::NotFound { .. }, - ) => { - admission.cancel_before_publish().await; - Err(Error::RefConflict { - message: format!("{reference_label} changed during conditional update"), - }) - } - Err(error) => Err(error.into()), - } -} - impl Tags<'_> { pub async fn fetch_tags(&self) -> Result> { let root_location = self.refs.root()?; @@ -342,21 +230,18 @@ impl Tags<'_> { .build_tag_content_by_ref(reference, Some(now), Some(now)) .await?; let payload = serde_json::to_vec_pretty(&tag_contents)?; + let mutation = ReferenceMutation::Create { + path: tag_file.to_string(), + payload, + }; let admission = begin_reference_admission( self.refs, tag_contents.branch.as_deref(), tag_contents.version, + mutation, ) .await?; - publish_new_reference( - self.object_store(), - &tag_file, - payload, - admission, - "tag", - tag, - ) - .await + admission.publish(format!("tag {tag} already exists")).await } pub async fn delete(&self, tag: &str) -> Result<()> { @@ -385,10 +270,8 @@ impl Tags<'_> { }); } let original_get = self.object_store().inner.get(&tag_file).await?; - let expected_version = UpdateVersion { - e_tag: original_get.meta.e_tag.clone(), - version: original_get.meta.version.clone(), - }; + let expected_etag = original_get.meta.e_tag.clone(); + let expected_version = original_get.meta.version.clone(); let original_payload = original_get.bytes().await?; let original_tag_contents: TagContents = serde_json::from_slice(&original_payload)?; let updated_reference = self @@ -401,22 +284,23 @@ impl Tags<'_> { tag_contents.updated_at = updated_reference.updated_at; tag_contents.manifest_size = updated_reference.manifest_size; let payload = serde_json::to_vec_pretty(&tag_contents)?; + let mutation = ReferenceMutation::Update { + path: tag_file.to_string(), + expected_payload: original_payload.to_vec(), + expected_etag, + expected_version, + payload, + }; let admission = begin_reference_admission( self.refs, tag_contents.branch.as_deref(), tag_contents.version, + mutation, ) .await?; - publish_updated_reference( - self.object_store(), - &tag_file, - &original_payload, - expected_version, - payload, - admission, - format!("tag {tag}"), - ) - .await + admission + .publish(format!("tag {tag} changed during conditional update")) + .await } pub async fn replace_metadata( @@ -435,30 +319,29 @@ impl Tags<'_> { } let original_get = self.object_store().inner.get(&tag_file).await?; - let expected_version = UpdateVersion { - e_tag: original_get.meta.e_tag.clone(), - version: original_get.meta.version.clone(), - }; + let expected_etag = original_get.meta.e_tag.clone(); + let expected_version = original_get.meta.version.clone(); let original_payload = original_get.bytes().await?; let mut tag_contents: TagContents = serde_json::from_slice(&original_payload)?; tag_contents.metadata = metadata; let payload = serde_json::to_vec_pretty(&tag_contents)?; + let mutation = ReferenceMutation::Update { + path: tag_file.to_string(), + expected_payload: original_payload.to_vec(), + expected_etag, + expected_version, + payload, + }; let admission = begin_reference_admission( self.refs, tag_contents.branch.as_deref(), tag_contents.version, + mutation, ) .await?; - publish_updated_reference( - self.object_store(), - &tag_file, - &original_payload, - expected_version, - payload, - admission, - format!("tag {tag}"), - ) - .await + admission + .publish(format!("tag {tag} changed during conditional update")) + .await } async fn build_tag_content_by_ref( @@ -659,21 +542,20 @@ impl Branches<'_> { metadata: HashMap::new(), }; let payload = serde_json::to_vec_pretty(&branch_contents)?; + let mutation = ReferenceMutation::Create { + path: branch_file.to_string(), + payload, + }; let admission = begin_reference_admission( self.refs, branch_contents.parent_branch.as_deref(), version_number, + mutation, ) .await?; - publish_new_reference( - self.object_store(), - &branch_file, - payload, - admission, - "branch", - branch_name, - ) - .await + admission + .publish(format!("branch {branch_name} already exists")) + .await } pub async fn replace_metadata( @@ -691,17 +573,37 @@ impl Branches<'_> { }); } - let mut branch_contents = - BranchContents::from_path(&branch_file, self.object_store(), branch).await?; + let original_get = self.object_store().inner.get(&branch_file).await?; + let expected_etag = original_get.meta.e_tag.clone(); + let expected_version = original_get.meta.version.clone(); + let original_payload = original_get.bytes().await?; + let mut branch_contents: BranchContents = serde_json::from_slice(&original_payload)?; + if branch_contents.identifier == BranchIdentifier::missing_identifier_sentinel() { + branch_contents.identifier = BranchIdentifier::synthetic_identifier( + branch, + branch_contents.parent_branch.as_deref(), + branch_contents.parent_version, + branch_contents.create_at, + ); + } branch_contents.metadata = metadata; - - self.object_store() - .put( - &branch_file, - serde_json::to_string_pretty(&branch_contents)?.as_bytes(), - ) + let mutation = ReferenceMutation::Update { + path: branch_file.to_string(), + expected_payload: original_payload.to_vec(), + expected_etag, + expected_version, + payload: serde_json::to_vec_pretty(&branch_contents)?, + }; + let admission = begin_reference_admission( + self.refs, + branch_contents.parent_branch.as_deref(), + branch_contents.parent_version, + mutation, + ) + .await?; + admission + .publish(format!("branch {branch} changed during conditional update")) .await - .map(|_| ()) } /// Delete a branch diff --git a/rust/lance/src/dataset/version_lease.rs b/rust/lance/src/dataset/version_lease.rs index 998e5ffbd3c..9c2dadcf493 100644 --- a/rust/lance/src/dataset/version_lease.rs +++ b/rust/lance/src/dataset/version_lease.rs @@ -14,7 +14,7 @@ use chrono::{DateTime, TimeDelta, Utc}; use dashmap::DashSet; use futures::{StreamExt, TryStreamExt, stream}; use lance_io::object_store::ObjectStore; -use object_store::{ObjectMeta, ObjectStoreExt, PutMode, PutOptions, path::Path}; +use object_store::{ObjectMeta, ObjectStoreExt, PutMode, PutOptions, UpdateVersion, path::Path}; use serde::{Deserialize, Serialize}; use uuid::Uuid; @@ -169,6 +169,36 @@ pub(super) struct ReferenceAdmission { manifest_path: Path, version: u64, created_at: DateTime, + mutation: ReferenceMutation, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase", tag = "kind")] +pub(super) enum ReferenceMutation { + Create { + path: String, + payload: Vec, + }, + Update { + path: String, + expected_payload: Vec, + expected_etag: Option, + expected_version: Option, + payload: Vec, + }, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +struct ReferenceIntent { + manifest_path: String, + mutation: ReferenceMutation, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum ReferenceMutationOutcome { + Published, + Conflict, } #[derive(Debug)] @@ -211,6 +241,7 @@ struct RetirementMarkerMetadata { struct ReferenceIntentMetadata { version: u64, metadata: ObjectMeta, + intent: ReferenceIntent, } /// Per-cleanup retirement markers for versions selected for deletion. @@ -336,6 +367,21 @@ impl VersionLeaseStore { if reference_owner_is_active(intent.metadata.last_modified, observed_at)? { active_versions.insert(intent.version); } else { + // The intent contains the exact conditional mutation, so an + // expired owner cannot resume a different or unconditional + // publication after recovery removes its intent. Applying the + // same mutation either publishes the reference or proves that + // its original precondition can no longer succeed. + let manifest_path = Path::parse(&intent.intent.manifest_path)?; + if self.object_store.exists(&manifest_path).await? + && !self.has_committed_marker(intent.version).await? + && self + .apply_reference_mutation(&intent.intent.mutation) + .await? + == ReferenceMutationOutcome::Published + { + active_versions.insert(intent.version); + } stale_paths.push(intent.metadata.location); } } @@ -343,7 +389,11 @@ impl VersionLeaseStore { Ok(active_versions) } - async fn create_reference_intent(&self, version: u64) -> Result<(Path, DateTime)> { + async fn create_reference_intent( + &self, + version: u64, + intent: &ReferenceIntent, + ) -> Result<(Path, DateTime)> { let path = self .reference_intents_path .clone() @@ -357,7 +407,7 @@ impl VersionLeaseStore { .inner .put_opts( &path, - Bytes::new().into(), + Bytes::from(serde_json::to_vec(intent)?).into(), PutOptions { mode: PutMode::Create, ..Default::default() @@ -371,19 +421,136 @@ impl VersionLeaseStore { Ok((path, metadata.last_modified)) } + async fn apply_reference_mutation( + &self, + mutation: &ReferenceMutation, + ) -> Result { + let (path, payload, result) = match mutation { + ReferenceMutation::Create { path, payload } => { + let path = Path::parse(path)?; + let result = self + .object_store + .inner + .put_opts( + &path, + Bytes::from(payload.clone()).into(), + PutOptions { + mode: PutMode::Create, + ..Default::default() + }, + ) + .await; + (path, payload, result) + } + ReferenceMutation::Update { + path, + expected_payload, + expected_etag, + expected_version, + payload, + } => { + let path = Path::parse(path)?; + let result = self + .object_store + .inner + .put_opts( + &path, + Bytes::from(payload.clone()).into(), + PutOptions { + mode: PutMode::Update(UpdateVersion { + e_tag: expected_etag.clone(), + version: expected_version.clone(), + }), + ..Default::default() + }, + ) + .await; + let result = match result { + Err( + object_store::Error::NotSupported { .. } + | object_store::Error::NotImplemented { .. }, + ) => { + // Local object stores do not implement UpdateVersion. + // Preserve their existing update support while checking + // the exact prior contents immediately before overwrite. + let current = self.object_store.inner.get(&path).await?.bytes().await?; + if current.as_ref() != expected_payload { + return Ok(ReferenceMutationOutcome::Conflict); + } + self.object_store + .inner + .put(&path, Bytes::from(payload.clone()).into()) + .await + } + result => result, + }; + (path, payload, result) + } + }; + + match result { + Ok(_) => Ok(ReferenceMutationOutcome::Published), + Err( + object_store::Error::AlreadyExists { .. } + | object_store::Error::Precondition { .. }, + ) => { + self.reference_mutation_conflict_outcome(&path, payload) + .await + } + Err(object_store::Error::NotFound { .. }) => Ok(ReferenceMutationOutcome::Conflict), + Err(error) => Err(error.into()), + } + } + + async fn reference_mutation_conflict_outcome( + &self, + path: &Path, + intended_payload: &[u8], + ) -> Result { + match self.object_store.inner.get(path).await { + Ok(result) => { + let current = result.bytes().await?; + Ok(if current.as_ref() == intended_payload { + ReferenceMutationOutcome::Published + } else { + ReferenceMutationOutcome::Conflict + }) + } + Err(object_store::Error::NotFound { .. }) => Ok(ReferenceMutationOutcome::Conflict), + Err(error) => Err(error.into()), + } + } + async fn reference_intent_metadata(&self) -> Result> { - self.object_store + let metadata = self + .object_store .list(Some(self.reference_intents_path.clone())) .try_collect::>() - .await? - .into_iter() - .map(|metadata| self.parse_reference_intent_metadata(metadata)) - .collect() + .await?; + stream::iter(metadata) + .map(|metadata| async move { + let result = match self.object_store.inner.get(&metadata.location).await { + Ok(result) => result, + Err(object_store::Error::NotFound { .. }) => return Ok(None), + Err(error) => return Err(Error::from(error)), + }; + let bytes = result.bytes().await?; + let intent = serde_json::from_slice(&bytes).map_err(|error| { + Error::corrupt_file(metadata.location.clone(), error.to_string()) + })?; + self.parse_reference_intent_metadata(metadata, intent) + .map(Some) + }) + .buffer_unordered(self.object_store.io_parallelism()) + .try_collect::>() + .await + .map(|intents| intents.into_iter().flatten().collect()) } fn parse_reference_intent_metadata( &self, metadata: ObjectMeta, + intent: ReferenceIntent, ) -> Result { let relative_parts = metadata .location @@ -422,7 +589,11 @@ impl VersionLeaseStore { format!("reference intent has invalid operation id: {error}"), ) })?; - Ok(ReferenceIntentMetadata { version, metadata }) + Ok(ReferenceIntentMetadata { + version, + metadata, + intent, + }) } pub(super) async fn active_versions_at( @@ -546,6 +717,14 @@ impl VersionLeaseStore { .any(|marker| marker.state != RetirementState::Draining)) } + async fn has_committed_marker(&self, version: u64) -> Result { + Ok(self + .version_marker_metadata(version) + .await? + .iter() + .any(|marker| marker.state == RetirementState::Committed)) + } + async fn has_active_retirement_marker(&self, version: u64) -> Result { let markers = self.version_marker_metadata(version).await?; if markers @@ -707,6 +886,8 @@ impl VersionLeaseStore { let actively_referenced_versions = self.active_reference_versions().await?; let mut stale_drains = Vec::new(); let mut sealed_manifest_paths = HashMap::>::new(); + let mut sealed_marker_paths = HashMap::>::new(); + let mut committed_versions = HashSet::new(); for marker in markers { if marker.state != RetirementState::Draining { let manifest_paths = self.read_retirement_marker(&marker.metadata).await?; @@ -714,6 +895,14 @@ impl VersionLeaseStore { .entry(marker.version) .or_default() .extend(manifest_paths); + if marker.state == RetirementState::Committed { + committed_versions.insert(marker.version); + } else { + sealed_marker_paths + .entry(marker.version) + .or_default() + .push(marker.metadata.location); + } } else if sealed_operations.contains(&(marker.version, marker.operation_id)) || LOCALLY_ABANDONED_DRAINS.contains(&marker.metadata.location) || !draining_owner_is_active(marker.metadata.last_modified, observed_at)? @@ -728,6 +917,7 @@ impl VersionLeaseStore { let mut terminal_manifests = HashMap::new(); let mut versions_to_resume = HashSet::new(); + let mut cancelled_seals = Vec::new(); for (version, manifest_paths) in sealed_manifest_paths { let manifest_paths = manifest_paths.into_iter().collect::>(); let mut has_existing_manifest = false; @@ -735,18 +925,24 @@ impl VersionLeaseStore { has_existing_manifest |= self.object_store.exists(path).await?; } if has_existing_manifest { - // A seal is the renewal cutoff. If its owner disappeared before - // confirming the final lease scan, wait out every lease that was - // still entitled to its published TTL before resuming deletion. - if !actively_leased_versions.contains(&version) - && !actively_referenced_versions.contains(&version) - { + let is_committed = committed_versions.contains(&version); + let is_retained = actively_leased_versions.contains(&version) + || actively_referenced_versions.contains(&version); + if is_committed || !is_retained { versions_to_resume.insert(version); + } else { + // Recovery may cancel a pre-commit seal when a lease or + // durable reference won the final census. Removing the + // seal restores renewal; committed retirement remains + // irreversible and is always resumed above. + cancelled_seals + .extend(sealed_marker_paths.remove(&version).unwrap_or_default()); } } else { terminal_manifests.insert(version, manifest_paths); } } + self.delete_paths(cancelled_seals).await?; self.finalize_versions(&terminal_manifests).await?; Ok(versions_to_resume) } @@ -1198,6 +1394,7 @@ pub(super) async fn begin_reference_admission( refs: &Refs, branch: Option<&str>, version: u64, + mutation: ReferenceMutation, ) -> Result { let branch_location = refs.base_location.find_branch(branch)?; let manifest = refs @@ -1212,13 +1409,18 @@ pub(super) async fn begin_reference_admission( let store = VersionLeaseStore::for_refs(refs, branch, None).await?; ensure_reference_available(&store, &manifest.path, version).await?; - let (path, created_at) = store.create_reference_intent(version).await?; + let intent = ReferenceIntent { + manifest_path: manifest.path.to_string(), + mutation: mutation.clone(), + }; + let (path, created_at) = store.create_reference_intent(version, &intent).await?; let admission = ReferenceAdmission { store, path, manifest_path: manifest.path, version, created_at, + mutation, }; if let Err(error) = admission.ensure_owned().await { admission.cancel_before_publish().await; @@ -1245,6 +1447,29 @@ impl ReferenceAdmission { ensure_reference_available(&self.store, &self.manifest_path, self.version).await } + pub(super) async fn publish(self, conflict_message: String) -> Result<()> { + if let Err(error) = self.ensure_owned().await { + self.cancel_before_publish().await; + return Err(error); + } + match self.store.apply_reference_mutation(&self.mutation).await { + Ok(ReferenceMutationOutcome::Published) => { + self.complete(); + Ok(()) + } + Ok(ReferenceMutationOutcome::Conflict) => { + self.cancel_before_publish().await; + Err(Error::RefConflict { + message: conflict_message, + }) + } + // A transport failure can be ambiguous after the server accepted + // the conditional mutation. Recovery retains and replays the exact + // operation recorded in the intent. + Err(error) => Err(error), + } + } + /// Remove this operation's uniquely owned intent after a definitive result /// that did not publish the canonical reference. pub(super) async fn cancel_before_publish(&self) { @@ -1259,19 +1484,10 @@ impl ReferenceAdmission { } } - /// Canonical publication is durable. Intent removal is best-effort because - /// the canonical tag or branch now protects the version independently. - pub(super) async fn complete(&self) { - if let Err(error) = self.store.object_store.delete(&self.path).await - && !error.is_not_found() - { - tracing::warn!( - path = %self.path, - error = %error, - "Failed to remove completed reference admission intent" - ); - } - } + /// Leave the intent as the durable handoff to the canonical reference. + /// Cleanup later observes the intent and, after its ownership timeout, + /// verifies or replays the exact conditional mutation before removing it. + fn complete(&self) {} } async fn ensure_reference_available( @@ -1501,6 +1717,16 @@ mod tests { )]) } + fn reference_intent(version: u64) -> ReferenceIntent { + ReferenceIntent { + manifest_path: format!("manifests/{version}.manifest"), + mutation: ReferenceMutation::Create { + path: format!("tags/version-{version}.json"), + payload: version.to_string().into_bytes(), + }, + } + } + fn assert_send(_: T) {} #[test] @@ -1598,7 +1824,10 @@ mod tests { #[tokio::test] async fn reference_intent_blocks_retirement_commit() { let store = memory_store(); - let (intent_path, _) = store.create_reference_intent(42).await.unwrap(); + let (intent_path, _) = store + .create_reference_intent(42, &reference_intent(42)) + .await + .unwrap(); let mut guard = store.fence_versions(&manifest_paths(42)).await.unwrap(); guard.seal_versions(&HashSet::from([42])).await.unwrap(); @@ -1736,7 +1965,7 @@ mod tests { let manifest_paths = manifest_paths(42); let manifest_path = manifest_paths[&42][0].clone(); store.object_store.put(&manifest_path, &[]).await.unwrap(); - let lease = store.acquire(42, Duration::from_secs(60)).await.unwrap(); + let mut lease = store.acquire(42, Duration::from_secs(60)).await.unwrap(); let mut guard = store.fence_versions(&manifest_paths).await.unwrap(); guard.seal_versions(&HashSet::from([42])).await.unwrap(); drop(guard); @@ -1750,12 +1979,115 @@ mod tests { .is_empty() ); assert!(store.object_store.exists(&manifest_path).await.unwrap()); + assert!(store.version_marker_metadata(42).await.unwrap().is_empty()); + lease.renew(Duration::from_secs(60)).await.unwrap(); lease.release().await.unwrap(); + assert!( + store + .clone() + .recover_retirements() + .await + .unwrap() + .is_empty() + ); + } + + #[tokio::test] + async fn sealed_recovery_cancels_for_active_reference() { + let store = memory_store(); + let manifest_paths = manifest_paths(42); + let manifest_path = manifest_paths[&42][0].clone(); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + let (intent_path, _) = store + .create_reference_intent(42, &reference_intent(42)) + .await + .unwrap(); + let mut guard = store.fence_versions(&manifest_paths).await.unwrap(); + guard.seal_versions(&HashSet::from([42])).await.unwrap(); + drop(guard); + + assert!( + store + .clone() + .recover_retirements() + .await + .unwrap() + .is_empty() + ); + assert!(store.object_store.exists(&manifest_path).await.unwrap()); + assert!(store.version_marker_metadata(42).await.unwrap().is_empty()); + + store.object_store.delete(&intent_path).await.unwrap(); + } + + #[tokio::test] + async fn canonical_reference_handoff_blocks_retirement_commit() { + let store = memory_store(); + let manifest_path = Path::from("manifests/42.manifest"); + let canonical_path = Path::from("tags/racing.json"); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + let intent = ReferenceIntent { + manifest_path: manifest_path.to_string(), + mutation: ReferenceMutation::Create { + path: canonical_path.to_string(), + payload: b"reference".to_vec(), + }, + }; + let (intent_path, created_at) = store.create_reference_intent(42, &intent).await.unwrap(); + let admission = ReferenceAdmission { + store: store.clone(), + path: intent_path.clone(), + manifest_path: manifest_path.clone(), + version: 42, + created_at, + mutation: intent.mutation, + }; + admission.ensure_owned().await.unwrap(); + + let manifests = HashMap::from([(42, vec![manifest_path])]); + let mut guard = store.fence_versions(&manifests).await.unwrap(); + guard.seal_versions(&HashSet::from([42])).await.unwrap(); + store + .object_store + .put(&canonical_path, b"reference") + .await + .unwrap(); + admission.complete(); + assert_eq!( - store.clone().recover_retirements().await.unwrap(), + guard.commit_versions(&HashSet::from([42])).await.unwrap(), HashSet::from([42]) ); + assert!(store.object_store.exists(&intent_path).await.unwrap()); + guard.cancel_all().await.unwrap(); + store.object_store.delete(&intent_path).await.unwrap(); + } + + #[tokio::test] + async fn conditional_reference_update_does_not_resurrect_deleted_reference() { + let store = memory_store(); + let canonical_path = Path::from("branches/child.json"); + let original = store + .object_store + .inner + .put(&canonical_path, Bytes::from_static(b"original").into()) + .await + .unwrap(); + let mutation = ReferenceMutation::Update { + path: canonical_path.to_string(), + expected_payload: b"original".to_vec(), + expected_etag: original.e_tag, + expected_version: original.version, + payload: b"updated".to_vec(), + }; + store.object_store.delete(&canonical_path).await.unwrap(); + + assert_eq!( + store.apply_reference_mutation(&mutation).await.unwrap(), + ReferenceMutationOutcome::Conflict + ); + assert!(!store.object_store.exists(&canonical_path).await.unwrap()); } #[tokio::test] From 77d91ebae405ad90445284b3223d91c13ffc1c08 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 23:04:48 +0000 Subject: [PATCH 07/13] fix: retire completed reference intents after handoff --- rust/lance/src/dataset/cleanup.rs | 5 +- rust/lance/src/dataset/version_lease.rs | 206 +++++++++++++++++++++--- 2 files changed, 186 insertions(+), 25 deletions(-) diff --git a/rust/lance/src/dataset/cleanup.rs b/rust/lance/src/dataset/cleanup.rs index 4f7429d0bbc..44a1f2f9e57 100644 --- a/rust/lance/src/dataset/cleanup.rs +++ b/rust/lance/src/dataset/cleanup.rs @@ -696,6 +696,9 @@ impl<'a> CleanupTask<'a> { // intent admitted before the draining marker must win. Descendant // retention is based on actual lineage file references rather than // cleanup policy or branch metadata alone. + let intent_versions = version_lease_store + .reference_versions_before_canonical_census() + .await?; let tags = self.dataset.tags().list().await?; let current_branch = &self.dataset.manifest.branch; let mut referenced_versions = tags @@ -707,7 +710,7 @@ impl<'a> CleanupTask<'a> { }) .map(|tag| tag.version) .collect::>(); - referenced_versions.extend(version_lease_store.active_reference_versions().await?); + referenced_versions.extend(intent_versions); let mut retained_versions = versions_to_seal .intersection(&referenced_versions) .copied() diff --git a/rust/lance/src/dataset/version_lease.rs b/rust/lance/src/dataset/version_lease.rs index 9c2dadcf493..98e0474cdba 100644 --- a/rust/lance/src/dataset/version_lease.rs +++ b/rust/lance/src/dataset/version_lease.rs @@ -193,6 +193,22 @@ pub(super) enum ReferenceMutation { struct ReferenceIntent { manifest_path: String, mutation: ReferenceMutation, + #[serde(default)] + state: ReferenceIntentState, +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +enum ReferenceIntentState { + #[default] + Pending, + Completed, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum CompletedReferenceIntentHandling { + DeferToCanonicalCensus, + RetainForCurrentScan, } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -360,29 +376,74 @@ impl VersionLeaseStore { } pub(super) async fn active_reference_versions(&self) -> Result> { + self.reference_versions(CompletedReferenceIntentHandling::RetainForCurrentScan) + .await + } + + /// Fence in-flight publication before the caller scans canonical tags and + /// branches. Completed intents can be removed without retaining their + /// target because the following canonical census observes their result. + pub(super) async fn reference_versions_before_canonical_census(&self) -> Result> { + self.reference_versions(CompletedReferenceIntentHandling::DeferToCanonicalCensus) + .await + } + + async fn reference_versions( + &self, + completed_handling: CompletedReferenceIntentHandling, + ) -> Result> { let observed_at = self.storage_observed_at().await?; let mut active_versions = HashSet::new(); let mut stale_paths = Vec::new(); for intent in self.reference_intent_metadata().await? { - if reference_owner_is_active(intent.metadata.last_modified, observed_at)? { - active_versions.insert(intent.version); - } else { - // The intent contains the exact conditional mutation, so an - // expired owner cannot resume a different or unconditional - // publication after recovery removes its intent. Applying the - // same mutation either publishes the reference or proves that - // its original precondition can no longer succeed. - let manifest_path = Path::parse(&intent.intent.manifest_path)?; - if self.object_store.exists(&manifest_path).await? - && !self.has_committed_marker(intent.version).await? - && self - .apply_reference_mutation(&intent.intent.mutation) - .await? - == ReferenceMutationOutcome::Published + match intent.intent.state { + ReferenceIntentState::Completed => { + // The completed intent is the durable handoff from an + // in-flight publication to its canonical object. A caller + // that already completed its canonical census retains a + // matching payload for this scan; a caller about to census + // canonical state can defer to that newer observation. + if completed_handling == CompletedReferenceIntentHandling::RetainForCurrentScan + { + let (path, payload) = match &intent.intent.mutation { + ReferenceMutation::Create { path, payload } + | ReferenceMutation::Update { path, payload, .. } => { + (Path::parse(path)?, payload) + } + }; + if self + .reference_mutation_conflict_outcome(&path, payload) + .await? + == ReferenceMutationOutcome::Published + { + active_versions.insert(intent.version); + } + } + stale_paths.push(intent.metadata.location); + } + ReferenceIntentState::Pending + if reference_owner_is_active(intent.metadata.last_modified, observed_at)? => { active_versions.insert(intent.version); } - stale_paths.push(intent.metadata.location); + ReferenceIntentState::Pending => { + // The intent contains the exact conditional mutation, so an + // expired owner cannot resume a different or unconditional + // publication after recovery removes its intent. Applying the + // same mutation either publishes the reference or proves that + // its original precondition can no longer succeed. + let manifest_path = Path::parse(&intent.intent.manifest_path)?; + if self.object_store.exists(&manifest_path).await? + && !self.has_committed_marker(intent.version).await? + && self + .apply_reference_mutation(&intent.intent.mutation) + .await? + == ReferenceMutationOutcome::Published + { + active_versions.insert(intent.version); + } + stale_paths.push(intent.metadata.location); + } } } self.delete_paths(stale_paths).await?; @@ -1412,6 +1473,7 @@ pub(super) async fn begin_reference_admission( let intent = ReferenceIntent { manifest_path: manifest.path.to_string(), mutation: mutation.clone(), + state: ReferenceIntentState::Pending, }; let (path, created_at) = store.create_reference_intent(version, &intent).await?; let admission = ReferenceAdmission { @@ -1454,7 +1516,7 @@ impl ReferenceAdmission { } match self.store.apply_reference_mutation(&self.mutation).await { Ok(ReferenceMutationOutcome::Published) => { - self.complete(); + self.complete().await; Ok(()) } Ok(ReferenceMutationOutcome::Conflict) => { @@ -1484,10 +1546,29 @@ impl ReferenceAdmission { } } - /// Leave the intent as the durable handoff to the canonical reference. - /// Cleanup later observes the intent and, after its ownership timeout, - /// verifies or replays the exact conditional mutation before removing it. - fn complete(&self) {} + /// Mark the intent as the durable handoff to the canonical reference. + /// Cleanup verifies the canonical payload, retains it for the current scan, + /// and removes the completed intent without waiting for the ownership timeout. + async fn complete(&self) { + let intent = ReferenceIntent { + manifest_path: self.manifest_path.to_string(), + mutation: self.mutation.clone(), + state: ReferenceIntentState::Completed, + }; + let result = match serde_json::to_vec(&intent) { + Ok(payload) => self.store.object_store.put(&self.path, &payload).await, + Err(error) => Err(error.into()), + }; + if let Err(error) = result { + // The pending intent remains a safe bounded fallback if completion + // cannot be recorded after the canonical write succeeded. + tracing::warn!( + path = %self.path, + error = %error, + "Failed to complete reference admission intent" + ); + } + } } async fn ensure_reference_available( @@ -1724,6 +1805,7 @@ mod tests { path: format!("tags/version-{version}.json"), payload: version.to_string().into_bytes(), }, + state: ReferenceIntentState::Pending, } } @@ -2033,6 +2115,7 @@ mod tests { path: canonical_path.to_string(), payload: b"reference".to_vec(), }, + state: ReferenceIntentState::Pending, }; let (intent_path, created_at) = store.create_reference_intent(42, &intent).await.unwrap(); let admission = ReferenceAdmission { @@ -2053,15 +2136,90 @@ mod tests { .put(&canonical_path, b"reference") .await .unwrap(); - admission.complete(); + admission.complete().await; assert_eq!( guard.commit_versions(&HashSet::from([42])).await.unwrap(), HashSet::from([42]) ); - assert!(store.object_store.exists(&intent_path).await.unwrap()); + assert!(!store.object_store.exists(&intent_path).await.unwrap()); guard.cancel_all().await.unwrap(); - store.object_store.delete(&intent_path).await.unwrap(); + } + + #[tokio::test] + async fn completed_intent_does_not_retain_deleted_reference() { + let store = memory_store(); + let manifest_path = Path::from("manifests/42.manifest"); + let canonical_path = Path::from("tags/removed.json"); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + let intent = ReferenceIntent { + manifest_path: manifest_path.to_string(), + mutation: ReferenceMutation::Create { + path: canonical_path.to_string(), + payload: b"reference".to_vec(), + }, + state: ReferenceIntentState::Pending, + }; + let (intent_path, created_at) = store.create_reference_intent(42, &intent).await.unwrap(); + let admission = ReferenceAdmission { + store: store.clone(), + path: intent_path.clone(), + manifest_path, + version: 42, + created_at, + mutation: intent.mutation, + }; + store + .object_store + .put(&canonical_path, b"reference") + .await + .unwrap(); + admission.complete().await; + store.object_store.delete(&canonical_path).await.unwrap(); + + assert!(store.active_reference_versions().await.unwrap().is_empty()); + assert!(!store.object_store.exists(&intent_path).await.unwrap()); + } + + #[tokio::test] + async fn completed_intent_defers_to_following_canonical_census() { + let store = memory_store(); + let manifest_path = Path::from("manifests/42.manifest"); + let canonical_path = Path::from("branches/child.json"); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + let intent = ReferenceIntent { + manifest_path: manifest_path.to_string(), + mutation: ReferenceMutation::Create { + path: canonical_path.to_string(), + payload: b"reference".to_vec(), + }, + state: ReferenceIntentState::Pending, + }; + let (intent_path, created_at) = store.create_reference_intent(42, &intent).await.unwrap(); + let admission = ReferenceAdmission { + store: store.clone(), + path: intent_path.clone(), + manifest_path, + version: 42, + created_at, + mutation: intent.mutation, + }; + store + .object_store + .put(&canonical_path, b"reference") + .await + .unwrap(); + admission.complete().await; + + assert!( + store + .reference_versions_before_canonical_census() + .await + .unwrap() + .is_empty() + ); + assert!(store.object_store.exists(&canonical_path).await.unwrap()); + assert!(!store.object_store.exists(&intent_path).await.unwrap()); } #[tokio::test] From ba3e9e91ddd83195a52ea37b2236d5b9a6ff3325 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Sat, 8 Aug 2026 00:24:10 +0000 Subject: [PATCH 08/13] fix: fence canonical reference lifecycle --- rust/lance/src/dataset/cleanup.rs | 11 +- rust/lance/src/dataset/refs.rs | 285 ++++++++++++++---- rust/lance/src/dataset/version_lease.rs | 384 +++++++++++++++++++++--- 3 files changed, 577 insertions(+), 103 deletions(-) diff --git a/rust/lance/src/dataset/cleanup.rs b/rust/lance/src/dataset/cleanup.rs index 44a1f2f9e57..06e903de971 100644 --- a/rust/lance/src/dataset/cleanup.rs +++ b/rust/lance/src/dataset/cleanup.rs @@ -696,7 +696,7 @@ impl<'a> CleanupTask<'a> { // intent admitted before the draining marker must win. Descendant // retention is based on actual lineage file references rather than // cleanup policy or branch metadata alone. - let intent_versions = version_lease_store + let reference_census = version_lease_store .reference_versions_before_canonical_census() .await?; let tags = self.dataset.tags().list().await?; @@ -710,7 +710,7 @@ impl<'a> CleanupTask<'a> { }) .map(|tag| tag.version) .collect::>(); - referenced_versions.extend(intent_versions); + referenced_versions.extend(reference_census.versions.iter().copied()); let mut retained_versions = versions_to_seal .intersection(&referenced_versions) .copied() @@ -736,7 +736,12 @@ impl<'a> CleanupTask<'a> { .values() .copied() .collect::>(); - let late_reference_versions = guard.commit_versions(&versions_to_commit).await?; + let late_reference_versions = guard + .commit_versions( + &versions_to_commit, + &reference_census.completed_intent_paths, + ) + .await?; for version in &late_reference_versions { self.retain_version(&mut inspection, *version).await?; } diff --git a/rust/lance/src/dataset/refs.rs b/rust/lance/src/dataset/refs.rs index 87a15e7d083..1a7595dc03a 100644 --- a/rust/lance/src/dataset/refs.rs +++ b/rust/lance/src/dataset/refs.rs @@ -1,21 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use std::ops::Range; - +use bytes::Bytes; use chrono::{DateTime, Utc}; use futures::stream::{StreamExt, TryStreamExt}; use itertools::Itertools; use lance_io::object_store::ObjectStore; use lance_table::io::commit::CommitHandler; -use object_store::{ObjectStoreExt, path::Path}; +use object_store::{ObjectMeta, ObjectStoreExt, PutMode, PutOptions, UpdateVersion, path::Path}; use serde::{Deserialize, Serialize}; use std::sync::Arc; use crate::dataset::branch_location::BranchLocation; use crate::dataset::refs::Ref::{Tag, Version, VersionNumber}; use crate::dataset::version_lease::{ - ReferenceMutation, begin_reference_admission, remove_branch_state, + REFERENCE_TOMBSTONE_PAYLOAD, ReferenceMutation, begin_reference_admission, + lock_local_reference, remove_branch_state, }; use crate::utils::temporal::utc_now; use crate::{Error, Result}; @@ -135,6 +135,91 @@ pub struct Branches<'a> { refs: &'a Refs, } +struct ReferenceSnapshot { + metadata: ObjectMeta, + payload: Bytes, +} + +impl ReferenceSnapshot { + fn is_tombstone(&self) -> bool { + self.payload.as_ref() == REFERENCE_TOMBSTONE_PAYLOAD + } +} + +async fn reference_snapshot( + object_store: &ObjectStore, + path: &Path, +) -> Result> { + let result = match object_store.inner.get(path).await { + Ok(result) => result, + Err(object_store::Error::NotFound { .. }) => return Ok(None), + Err(error) => return Err(error.into()), + }; + let metadata = result.meta.clone(); + let payload = result.bytes().await?; + Ok(Some(ReferenceSnapshot { metadata, payload })) +} + +async fn tombstone_reference( + object_store: &ObjectStore, + path: &Path, + snapshot: &ReferenceSnapshot, +) -> Result<()> { + let local_lock = lock_local_reference(object_store, path).await?; + let Some(current) = reference_snapshot(object_store, path).await? else { + return Err(Error::RefConflict { + message: format!("reference {path} changed during conditional deletion"), + }); + }; + if current.payload != snapshot.payload { + return Err(Error::RefConflict { + message: format!("reference {path} changed during conditional deletion"), + }); + } + let result = object_store + .inner + .put_opts( + path, + Bytes::from_static(REFERENCE_TOMBSTONE_PAYLOAD).into(), + PutOptions { + mode: PutMode::Update(UpdateVersion { + e_tag: snapshot.metadata.e_tag.clone(), + version: snapshot.metadata.version.clone(), + }), + ..Default::default() + }, + ) + .await; + match result { + Ok(_) => Ok(()), + Err( + object_store::Error::NotSupported { .. } | object_store::Error::NotImplemented { .. }, + ) if local_lock.is_some() => { + // Local reference mutations hold the same OS advisory lock, so an + // atomic overwrite cannot race another supported local mutation. + // The durable tombstone prevents later creates from reusing this + // path without a locked conditional mutation. + object_store + .inner + .put(path, Bytes::from_static(REFERENCE_TOMBSTONE_PAYLOAD).into()) + .await?; + Ok(()) + } + Err( + object_store::Error::NotSupported { .. } | object_store::Error::NotImplemented { .. }, + ) => Err(Error::not_supported(format!( + "object store {} does not support atomic conditional reference deletion for {path}", + object_store.scheme() + ))), + Err(object_store::Error::NotFound { .. } | object_store::Error::Precondition { .. }) => { + Err(Error::RefConflict { + message: format!("reference {path} changed during conditional deletion"), + }) + } + Err(error) => Err(error.into()), + } +} + impl Tags<'_> { fn object_store(&self) -> &ObjectStore { &self.refs.object_store @@ -162,14 +247,20 @@ impl Tags<'_> { let root_path = &root_location.path; futures::stream::iter(tag_names) .map(|tag_name| async move { - let contents = - TagContents::from_path(&tag_path(root_path, &tag_name), self.object_store()) - .await?; - Ok((tag_name, contents)) + let path = tag_path(root_path, &tag_name); + let Some(snapshot) = reference_snapshot(self.object_store(), &path).await? else { + return Ok(None); + }; + if snapshot.is_tombstone() { + return Ok(None); + } + let contents = serde_json::from_slice(&snapshot.payload)?; + Ok(Some((tag_name, contents))) }) .buffer_unordered(10) - .try_collect() + .try_collect::>() .await + .map(|tags| tags.into_iter().flatten().collect()) } pub async fn list(&self) -> Result> { @@ -205,13 +296,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? { + let Some(snapshot) = reference_snapshot(self.object_store(), &tag_file).await? else { + return Err(Error::RefNotFound { + message: format!("tag {} does not exist", tag), + }); + }; + if snapshot.is_tombstone() { return Err(Error::RefNotFound { message: format!("tag {} does not exist", tag), }); } - let tag_contents = TagContents::from_path(&tag_file, self.object_store()).await?; + let tag_contents = serde_json::from_slice(&snapshot.payload)?; Ok(tag_contents) } @@ -220,7 +316,11 @@ 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? { + let previous = reference_snapshot(self.object_store(), &tag_file).await?; + if previous + .as_ref() + .is_some_and(|snapshot| !snapshot.is_tombstone()) + { return Err(Error::RefConflict { message: format!("tag {} already exists", tag), }); @@ -230,9 +330,18 @@ impl Tags<'_> { .build_tag_content_by_ref(reference, Some(now), Some(now)) .await?; let payload = serde_json::to_vec_pretty(&tag_contents)?; - let mutation = ReferenceMutation::Create { - path: tag_file.to_string(), - payload, + let mutation = match previous { + Some(snapshot) => ReferenceMutation::Update { + path: tag_file.to_string(), + expected_payload: snapshot.payload.to_vec(), + expected_etag: snapshot.metadata.e_tag, + expected_version: snapshot.metadata.version, + payload, + }, + None => ReferenceMutation::Create { + path: tag_file.to_string(), + payload, + }, }; let admission = begin_reference_admission( self.refs, @@ -250,13 +359,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? { + let Some(snapshot) = reference_snapshot(self.object_store(), &tag_file).await? else { + return Err(Error::RefNotFound { + message: format!("tag {} does not exist", tag), + }); + }; + if snapshot.is_tombstone() { return Err(Error::RefNotFound { message: format!("tag {} does not exist", tag), }); } - self.object_store().delete(&tag_file).await + tombstone_reference(self.object_store(), &tag_file, &snapshot).await } pub async fn update(&self, tag: &str, reference: impl Into) -> Result<()> { @@ -264,15 +378,19 @@ 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? { + let Some(snapshot) = reference_snapshot(self.object_store(), &tag_file).await? else { + return Err(Error::RefNotFound { + message: format!("tag {} does not exist", tag), + }); + }; + if snapshot.is_tombstone() { return Err(Error::RefNotFound { message: format!("tag {} does not exist", tag), }); } - let original_get = self.object_store().inner.get(&tag_file).await?; - let expected_etag = original_get.meta.e_tag.clone(); - let expected_version = original_get.meta.version.clone(); - let original_payload = original_get.bytes().await?; + let expected_etag = snapshot.metadata.e_tag.clone(); + let expected_version = snapshot.metadata.version.clone(); + let original_payload = snapshot.payload; let original_tag_contents: TagContents = serde_json::from_slice(&original_payload)?; let updated_reference = self .build_tag_content_by_ref(reference, original_tag_contents.created_at, Some(utc_now())) @@ -312,16 +430,20 @@ 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? { + let Some(snapshot) = reference_snapshot(self.object_store(), &tag_file).await? else { + return Err(Error::RefNotFound { + message: format!("tag {} does not exist", tag), + }); + }; + if snapshot.is_tombstone() { return Err(Error::RefNotFound { message: format!("tag {} does not exist", tag), }); } - let original_get = self.object_store().inner.get(&tag_file).await?; - let expected_etag = original_get.meta.e_tag.clone(); - let expected_version = original_get.meta.version.clone(); - let original_payload = original_get.bytes().await?; + let expected_etag = snapshot.metadata.e_tag.clone(); + let expected_version = snapshot.metadata.version.clone(); + let original_payload = snapshot.payload; let mut tag_contents: TagContents = serde_json::from_slice(&original_payload)?; tag_contents.metadata = metadata; let payload = serde_json::to_vec_pretty(&tag_contents)?; @@ -430,17 +552,21 @@ impl Branches<'_> { let branch_path = &root_location.path; futures::stream::iter(branch_names) .map(|name| async move { - let contents = BranchContents::from_path( - &branch_contents_path(branch_path, &name), - self.object_store(), - &name, - ) - .await?; - Ok((name, contents)) + let path = branch_contents_path(branch_path, &name); + let Some(snapshot) = reference_snapshot(self.object_store(), &path).await? else { + return Ok(None); + }; + if snapshot.is_tombstone() { + return Ok(None); + } + let mut contents: BranchContents = serde_json::from_slice(&snapshot.payload)?; + contents.ensure_identifier(&name); + Ok(Some((name, contents))) }) .buffer_unordered(10) - .try_collect() + .try_collect::>() .await + .map(|branches| branches.into_iter().flatten().collect()) } pub async fn list(&self) -> Result> { @@ -455,14 +581,19 @@ impl Branches<'_> { let root_location = self.refs.root()?; let branch_file = branch_contents_path(&root_location.path, branch); - if !self.object_store().exists(&branch_file).await? { + let Some(snapshot) = reference_snapshot(self.object_store(), &branch_file).await? else { + return Err(Error::RefNotFound { + message: format!("branch {} does not exist", branch), + }); + }; + if snapshot.is_tombstone() { return Err(Error::RefNotFound { message: format!("branch {} does not exist", branch), }); } - let branch_contents = - BranchContents::from_path(&branch_file, self.object_store(), branch).await?; + let mut branch_contents: BranchContents = serde_json::from_slice(&snapshot.payload)?; + branch_contents.ensure_identifier(branch); Ok(branch_contents) } @@ -488,7 +619,11 @@ 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? { + let previous = reference_snapshot(self.object_store(), &branch_file).await?; + if previous + .as_ref() + .is_some_and(|snapshot| !snapshot.is_tombstone()) + { return Err(Error::RefConflict { message: format!("branch {} already exists", branch_name), }); @@ -542,9 +677,18 @@ impl Branches<'_> { metadata: HashMap::new(), }; let payload = serde_json::to_vec_pretty(&branch_contents)?; - let mutation = ReferenceMutation::Create { - path: branch_file.to_string(), - payload, + let mutation = match previous { + Some(snapshot) => ReferenceMutation::Update { + path: branch_file.to_string(), + expected_payload: snapshot.payload.to_vec(), + expected_etag: snapshot.metadata.e_tag, + expected_version: snapshot.metadata.version, + payload, + }, + None => ReferenceMutation::Create { + path: branch_file.to_string(), + payload, + }, }; let admission = begin_reference_admission( self.refs, @@ -567,16 +711,20 @@ impl Branches<'_> { let root_location = self.refs.root()?; let branch_file = branch_contents_path(&root_location.path, branch); - if !self.object_store().exists(&branch_file).await? { + let Some(snapshot) = reference_snapshot(self.object_store(), &branch_file).await? else { + return Err(Error::RefNotFound { + message: format!("branch {} does not exist", branch), + }); + }; + if snapshot.is_tombstone() { return Err(Error::RefNotFound { message: format!("branch {} does not exist", branch), }); } - let original_get = self.object_store().inner.get(&branch_file).await?; - let expected_etag = original_get.meta.e_tag.clone(); - let expected_version = original_get.meta.version.clone(); - let original_payload = original_get.bytes().await?; + let expected_etag = snapshot.metadata.e_tag.clone(); + let expected_version = snapshot.metadata.version.clone(); + let original_payload = snapshot.payload; let mut branch_contents: BranchContents = serde_json::from_slice(&original_payload)?; if branch_contents.identifier == BranchIdentifier::missing_identifier_sentinel() { branch_contents.identifier = BranchIdentifier::synthetic_identifier( @@ -637,8 +785,10 @@ impl Branches<'_> { let root_location = self.refs.root()?; let branch_file = branch_contents_path(&root_location.path, branch); - if self.object_store().exists(&branch_file).await? { - self.object_store().delete(&branch_file).await?; + if let Some(snapshot) = reference_snapshot(self.object_store(), &branch_file).await? + && !snapshot.is_tombstone() + { + tombstone_reference(self.object_store(), &branch_file, &snapshot).await?; } // Clean up branch directories and operational state scoped to this @@ -987,16 +1137,17 @@ async fn from_path(path: &Path, object_store: &ObjectStore) -> Result where T: DeserializeOwned, { - let tag_reader = object_store.open(path).await?; - let tag_bytes = tag_reader - .get_range(Range { - start: 0, - end: tag_reader.size().await?, - }) - .await?; - let json_str = String::from_utf8(tag_bytes.to_vec()) - .map_err(|e| Error::corrupt_file(path.clone(), e.to_string()))?; - Ok(serde_json::from_str(&json_str)?) + let Some(snapshot) = reference_snapshot(object_store, path).await? else { + return Err(Error::RefNotFound { + message: format!("reference {path} does not exist"), + }); + }; + if snapshot.is_tombstone() { + return Err(Error::RefNotFound { + message: format!("reference {path} does not exist"), + }); + } + Ok(serde_json::from_slice(&snapshot.payload)?) } impl TagContents { @@ -1012,18 +1163,22 @@ impl BranchContents { branch_name: &str, ) -> Result { let mut contents: Self = from_path(path, object_store).await?; - if contents.identifier == BranchIdentifier::missing_identifier_sentinel() { + contents.ensure_identifier(branch_name); + Ok(contents) + } + + fn ensure_identifier(&mut self, branch_name: &str) { + if self.identifier == BranchIdentifier::missing_identifier_sentinel() { // Legacy branch files do not store an identifier. Derive a deterministic fallback // from stable branch metadata so repeated reads expose the same public // branch_identifier. - contents.identifier = BranchIdentifier::synthetic_identifier( + self.identifier = BranchIdentifier::synthetic_identifier( branch_name, - contents.parent_branch.as_deref(), - contents.parent_version, - contents.create_at, + self.parent_branch.as_deref(), + self.parent_version, + self.create_at, ); } - Ok(contents) } } diff --git a/rust/lance/src/dataset/version_lease.rs b/rust/lance/src/dataset/version_lease.rs index 98e0474cdba..10ebab6b858 100644 --- a/rust/lance/src/dataset/version_lease.rs +++ b/rust/lance/src/dataset/version_lease.rs @@ -5,6 +5,8 @@ use std::{ collections::{HashMap, HashSet}, + fs::{File, OpenOptions}, + path::PathBuf, sync::{Arc, LazyLock}, time::Duration, }; @@ -34,6 +36,11 @@ const SEALED_MARKER_SUFFIX: &str = ".sealed"; const COMMITTED_MARKER_SUFFIX: &str = ".committed"; const STORAGE_CLOCK_PATH: &str = "_clock"; +/// A deleted canonical reference remains as a tombstone so an abandoned +/// `PutMode::Create` publication cannot recreate the name after deletion. +pub(super) const REFERENCE_TOMBSTONE_PAYLOAD: &[u8] = + br#"{"_lanceReferenceState":"deleted","formatVersion":1}"#; + /// HTTP `Last-Modified`, used by supported cloud stores, has whole-second precision. /// Adding one interval guarantees a lease is never shortened by timestamp truncation. const STORAGE_TIMESTAMP_PRECISION: Duration = Duration::from_secs(1); @@ -172,6 +179,16 @@ pub(super) struct ReferenceAdmission { mutation: ReferenceMutation, } +/// An operating-system advisory lock for one local canonical reference. +/// +/// The lock file is intentionally retained after unlock. Reusing one inode +/// avoids a delete/recreate race between local processes, while the OS releases +/// the held lock automatically if its owner exits. +#[derive(Debug)] +pub(super) struct LocalReferenceLock { + _file: File, +} + #[derive(Clone, Debug, Deserialize, Serialize)] #[serde(rename_all = "camelCase", tag = "kind")] pub(super) enum ReferenceMutation { @@ -260,6 +277,12 @@ struct ReferenceIntentMetadata { intent: ReferenceIntent, } +#[derive(Debug)] +pub(super) struct CanonicalReferenceCensus { + pub(super) versions: HashSet, + pub(super) completed_intent_paths: HashSet, +} + /// Per-cleanup retirement markers for versions selected for deletion. /// /// Draining markers reject new leases while allowing already-admitted leases @@ -376,28 +399,45 @@ impl VersionLeaseStore { } pub(super) async fn active_reference_versions(&self) -> Result> { - self.reference_versions(CompletedReferenceIntentHandling::RetainForCurrentScan) - .await + Ok(self + .reference_versions( + CompletedReferenceIntentHandling::RetainForCurrentScan, + &HashSet::new(), + ) + .await? + .versions) } /// Fence in-flight publication before the caller scans canonical tags and /// branches. Completed intents can be removed without retaining their /// target because the following canonical census observes their result. - pub(super) async fn reference_versions_before_canonical_census(&self) -> Result> { - self.reference_versions(CompletedReferenceIntentHandling::DeferToCanonicalCensus) - .await + pub(super) async fn reference_versions_before_canonical_census( + &self, + ) -> Result { + self.reference_versions( + CompletedReferenceIntentHandling::DeferToCanonicalCensus, + &HashSet::new(), + ) + .await } async fn reference_versions( &self, completed_handling: CompletedReferenceIntentHandling, - ) -> Result> { + completed_intents_observed_before_census: &HashSet, + ) -> Result { let observed_at = self.storage_observed_at().await?; let mut active_versions = HashSet::new(); + let mut completed_intent_paths = HashSet::new(); let mut stale_paths = Vec::new(); for intent in self.reference_intent_metadata().await? { match intent.intent.state { ReferenceIntentState::Completed => { + if completed_intents_observed_before_census.contains(&intent.metadata.location) + { + stale_paths.push(intent.metadata.location); + continue; + } // The completed intent is the durable handoff from an // in-flight publication to its canonical object. A caller // that already completed its canonical census retains a @@ -419,7 +459,19 @@ impl VersionLeaseStore { active_versions.insert(intent.version); } } - stale_paths.push(intent.metadata.location); + // The pre-census scan must not consume the only durable + // handoff while a pre-commit seal exists. If cleanup is + // interrupted before its canonical census cancels that + // seal, recovery still needs the completed intent to + // recognize that the canonical reference won. + if completed_handling + == CompletedReferenceIntentHandling::DeferToCanonicalCensus + && self.has_precommit_sealed_marker(intent.version).await? + { + completed_intent_paths.insert(intent.metadata.location); + } else { + stale_paths.push(intent.metadata.location); + } } ReferenceIntentState::Pending if reference_owner_is_active(intent.metadata.last_modified, observed_at)? => @@ -427,16 +479,16 @@ impl VersionLeaseStore { active_versions.insert(intent.version); } ReferenceIntentState::Pending => { - // The intent contains the exact conditional mutation, so an - // expired owner cannot resume a different or unconditional - // publication after recovery removes its intent. Applying the - // same mutation either publishes the reference or proves that - // its original precondition can no longer succeed. + // Never replay an expired publication: an absent canonical + // path cannot distinguish an operation that never published + // from one that published and was subsequently deleted. + // An abandoned create instead installs a durable tombstone, + // atomically fencing any owner that resumes after expiry. let manifest_path = Path::parse(&intent.intent.manifest_path)?; if self.object_store.exists(&manifest_path).await? && !self.has_committed_marker(intent.version).await? && self - .apply_reference_mutation(&intent.intent.mutation) + .expired_reference_mutation_outcome(&intent.intent.mutation) .await? == ReferenceMutationOutcome::Published { @@ -447,7 +499,51 @@ impl VersionLeaseStore { } } self.delete_paths(stale_paths).await?; - Ok(active_versions) + Ok(CanonicalReferenceCensus { + versions: active_versions, + completed_intent_paths, + }) + } + + async fn expired_reference_mutation_outcome( + &self, + mutation: &ReferenceMutation, + ) -> Result { + match mutation { + ReferenceMutation::Create { path, payload } => { + let path = Path::parse(path)?; + let _local_lock = lock_local_reference(&self.object_store, &path).await?; + match self + .object_store + .inner + .put_opts( + &path, + Bytes::from_static(REFERENCE_TOMBSTONE_PAYLOAD).into(), + PutOptions { + mode: PutMode::Create, + ..Default::default() + }, + ) + .await + { + // The tombstone won the create race, so this abandoned + // operation can never publish the deleted name. + Ok(_) => Ok(ReferenceMutationOutcome::Conflict), + Err( + object_store::Error::AlreadyExists { .. } + | object_store::Error::Precondition { .. }, + ) => { + self.reference_mutation_conflict_outcome(&path, payload) + .await + } + Err(error) => Err(error.into()), + } + } + ReferenceMutation::Update { path, payload, .. } => { + self.reference_mutation_conflict_outcome(&Path::parse(path)?, payload) + .await + } + } } async fn create_reference_intent( @@ -482,9 +578,10 @@ impl VersionLeaseStore { Ok((path, metadata.last_modified)) } - async fn apply_reference_mutation( + async fn apply_reference_mutation_inner( &self, mutation: &ReferenceMutation, + has_local_lock: bool, ) -> Result { let (path, payload, result) = match mutation { ReferenceMutation::Create { path, payload } => { @@ -530,10 +627,7 @@ impl VersionLeaseStore { Err( object_store::Error::NotSupported { .. } | object_store::Error::NotImplemented { .. }, - ) => { - // Local object stores do not implement UpdateVersion. - // Preserve their existing update support while checking - // the exact prior contents immediately before overwrite. + ) if has_local_lock => { let current = self.object_store.inner.get(&path).await?.bytes().await?; if current.as_ref() != expected_payload { return Ok(ReferenceMutationOutcome::Conflict); @@ -543,6 +637,15 @@ impl VersionLeaseStore { .put(&path, Bytes::from(payload.clone()).into()) .await } + Err( + object_store::Error::NotSupported { .. } + | object_store::Error::NotImplemented { .. }, + ) => { + return Err(Error::not_supported(format!( + "object store {} does not support atomic conditional reference updates for {path}", + self.object_store.scheme() + ))); + } result => result, }; (path, payload, result) @@ -778,6 +881,14 @@ impl VersionLeaseStore { .any(|marker| marker.state != RetirementState::Draining)) } + async fn has_precommit_sealed_marker(&self, version: u64) -> Result { + Ok(self + .version_marker_metadata(version) + .await? + .iter() + .any(|marker| marker.state == RetirementState::Sealed)) + } + async fn has_committed_marker(&self, version: u64) -> Result { Ok(self .version_marker_metadata(version) @@ -1166,11 +1277,19 @@ impl RetirementGuard { pub(super) async fn commit_versions( &mut self, versions: &HashSet, + completed_intents_observed_before_census: &HashSet, ) -> Result> { // This scan runs inside the commit operation, after the caller's final // reference census. An admission that won that narrow race remains // cancellable and cannot cross the irreversible committed boundary. - let active_reference_versions = self.store.active_reference_versions().await?; + let active_reference_versions = self + .store + .reference_versions( + CompletedReferenceIntentHandling::RetainForCurrentScan, + completed_intents_observed_before_census, + ) + .await? + .versions; let retained_versions = versions .intersection(&active_reference_versions) .copied() @@ -1510,11 +1629,21 @@ impl ReferenceAdmission { } pub(super) async fn publish(self, conflict_message: String) -> Result<()> { + let canonical_path = match &self.mutation { + ReferenceMutation::Create { path, .. } | ReferenceMutation::Update { path, .. } => { + Path::parse(path)? + } + }; + let local_lock = lock_local_reference(&self.store.object_store, &canonical_path).await?; if let Err(error) = self.ensure_owned().await { self.cancel_before_publish().await; return Err(error); } - match self.store.apply_reference_mutation(&self.mutation).await { + match self + .store + .apply_reference_mutation_inner(&self.mutation, local_lock.is_some()) + .await + { Ok(ReferenceMutationOutcome::Published) => { self.complete().await; Ok(()) @@ -1526,8 +1655,8 @@ impl ReferenceAdmission { }) } // A transport failure can be ambiguous after the server accepted - // the conditional mutation. Recovery retains and replays the exact - // operation recorded in the intent. + // the conditional mutation. Recovery retains the intent and + // verifies canonical state without replaying the operation. Err(error) => Err(error), } } @@ -1571,6 +1700,39 @@ impl ReferenceAdmission { } } +pub(super) async fn lock_local_reference( + object_store: &ObjectStore, + path: &Path, +) -> Result> { + if !object_store.is_local() && object_store.scheme() != "file-object-store" { + return Ok(None); + } + + let path = PathBuf::from(path.as_ref()); + let path = if path.is_absolute() { + path + } else { + PathBuf::from(std::path::MAIN_SEPARATOR_STR).join(path) + }; + let lock_path = PathBuf::from(format!("{}.lance-reference.lock", path.display())); + tokio::task::spawn_blocking(move || -> Result { + if let Some(parent) = lock_path.parent() { + std::fs::create_dir_all(parent)?; + } + let file = OpenOptions::new() + .create(true) + .truncate(false) + .read(true) + .write(true) + .open(lock_path)?; + file.lock()?; + Ok(LocalReferenceLock { _file: file }) + }) + .await + .map_err(|error| Error::internal(format!("failed to acquire local reference lock: {error}")))? + .map(Some) +} + async fn ensure_reference_available( store: &VersionLeaseStore, manifest_path: &Path, @@ -1885,7 +2047,7 @@ mod tests { guard.seal_versions(&HashSet::from([42])).await.unwrap(); assert!( guard - .commit_versions(&HashSet::from([42])) + .commit_versions(&HashSet::from([42]), &HashSet::new()) .await .unwrap() .is_empty() @@ -1913,7 +2075,10 @@ mod tests { let mut guard = store.fence_versions(&manifest_paths(42)).await.unwrap(); guard.seal_versions(&HashSet::from([42])).await.unwrap(); - let retained_versions = guard.commit_versions(&HashSet::from([42])).await.unwrap(); + let retained_versions = guard + .commit_versions(&HashSet::from([42]), &HashSet::new()) + .await + .unwrap(); assert_eq!(retained_versions, HashSet::from([42])); assert!(guard.fences[&42].committed_path.is_none()); @@ -1940,7 +2105,7 @@ mod tests { guard.seal_versions(&HashSet::from([42])).await.unwrap(); assert!( guard - .commit_versions(&HashSet::from([42])) + .commit_versions(&HashSet::from([42]), &HashSet::new()) .await .unwrap() .is_empty() @@ -2139,7 +2304,10 @@ mod tests { admission.complete().await; assert_eq!( - guard.commit_versions(&HashSet::from([42])).await.unwrap(), + guard + .commit_versions(&HashSet::from([42]), &HashSet::new()) + .await + .unwrap(), HashSet::from([42]) ); assert!(!store.object_store.exists(&intent_path).await.unwrap()); @@ -2211,13 +2379,12 @@ mod tests { .unwrap(); admission.complete().await; - assert!( - store - .reference_versions_before_canonical_census() - .await - .unwrap() - .is_empty() - ); + let census = store + .reference_versions_before_canonical_census() + .await + .unwrap(); + assert!(census.versions.is_empty()); + assert!(census.completed_intent_paths.is_empty()); assert!(store.object_store.exists(&canonical_path).await.unwrap()); assert!(!store.object_store.exists(&intent_path).await.unwrap()); } @@ -2242,12 +2409,159 @@ mod tests { store.object_store.delete(&canonical_path).await.unwrap(); assert_eq!( - store.apply_reference_mutation(&mutation).await.unwrap(), + store + .apply_reference_mutation_inner(&mutation, false) + .await + .unwrap(), ReferenceMutationOutcome::Conflict ); assert!(!store.object_store.exists(&canonical_path).await.unwrap()); } + #[tokio::test] + async fn expired_pending_create_intent_does_not_resurrect_deleted_reference() { + let store = memory_store(); + let canonical_path = Path::from("tags/deleted.json"); + let mutation = ReferenceMutation::Create { + path: canonical_path.to_string(), + payload: b"reference".to_vec(), + }; + store + .object_store + .put(&canonical_path, b"reference") + .await + .unwrap(); + store.object_store.delete(&canonical_path).await.unwrap(); + + assert_eq!( + store + .expired_reference_mutation_outcome(&mutation) + .await + .unwrap(), + ReferenceMutationOutcome::Conflict + ); + assert_eq!( + store + .object_store + .inner + .get(&canonical_path) + .await + .unwrap() + .bytes() + .await + .unwrap() + .as_ref(), + REFERENCE_TOMBSTONE_PAYLOAD + ); + } + + #[tokio::test] + async fn unsupported_conditional_reference_update_is_rejected() { + let temp_dir = tempfile::tempdir().unwrap(); + let uri = format!("file://{}", temp_dir.path().display()); + let (object_store, base_path) = ObjectStore::from_uri(&uri).await.unwrap(); + let store = VersionLeaseStore { + object_store, + leases_path: base_path.clone().join("leases"), + markers_path: base_path.clone().join("markers"), + reference_intents_path: base_path.clone().join("reference_intents"), + manifest_path: None, + }; + let canonical_path = base_path.join("branches/child.json"); + let original = store + .object_store + .inner + .put(&canonical_path, Bytes::from_static(b"original").into()) + .await + .unwrap(); + let mutation = ReferenceMutation::Update { + path: canonical_path.to_string(), + expected_payload: b"original".to_vec(), + expected_etag: original.e_tag, + expected_version: original.version, + payload: b"updated".to_vec(), + }; + + let error = store + .apply_reference_mutation_inner(&mutation, false) + .await + .unwrap_err(); + + assert!(matches!(error, Error::NotSupported { .. })); + assert!(error.to_string().contains("atomic conditional"), "{error}"); + assert_eq!( + store + .object_store + .inner + .get(&canonical_path) + .await + .unwrap() + .bytes() + .await + .unwrap() + .as_ref(), + b"original" + ); + } + + #[tokio::test] + async fn canonical_only_recovery_cancels_seal_after_completed_handoff() { + let store = memory_store(); + let manifest_path = Path::from("manifests/42.manifest"); + let canonical_path = Path::from("tags/recovery.json"); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + let intent = ReferenceIntent { + manifest_path: manifest_path.to_string(), + mutation: ReferenceMutation::Create { + path: canonical_path.to_string(), + payload: b"reference".to_vec(), + }, + state: ReferenceIntentState::Pending, + }; + let (intent_path, created_at) = store.create_reference_intent(42, &intent).await.unwrap(); + let admission = ReferenceAdmission { + store: store.clone(), + path: intent_path.clone(), + manifest_path: manifest_path.clone(), + version: 42, + created_at, + mutation: intent.mutation, + }; + store + .object_store + .put(&canonical_path, b"reference") + .await + .unwrap(); + admission.complete().await; + + let manifests = HashMap::from([(42, vec![manifest_path])]); + let mut guard = store.fence_versions(&manifests).await.unwrap(); + guard.seal_versions(&HashSet::from([42])).await.unwrap(); + let census = store + .reference_versions_before_canonical_census() + .await + .unwrap(); + assert!(census.versions.is_empty()); + assert_eq!( + census.completed_intent_paths, + HashSet::from([intent_path.clone()]) + ); + assert!(store.object_store.exists(&intent_path).await.unwrap()); + drop(guard); + + assert!( + store + .clone() + .recover_retirements() + .await + .unwrap() + .is_empty() + ); + assert!(store.version_marker_metadata(42).await.unwrap().is_empty()); + assert!(!store.object_store.exists(&intent_path).await.unwrap()); + store.acquire(42, Duration::from_secs(60)).await.unwrap(); + } + #[tokio::test] async fn branch_deletion_removes_incarnation_state() { let store = memory_store(); From a3a30bfc837e603749cc0a0039094f88de72b9b8 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Sat, 8 Aug 2026 01:27:03 +0000 Subject: [PATCH 09/13] fix: add reference lifecycle generations --- rust/lance/src/dataset/cleanup.rs | 1 + rust/lance/src/dataset/refs.rs | 269 ++-- rust/lance/src/dataset/version_lease.rs | 1527 +++++++++++++++++++---- 3 files changed, 1411 insertions(+), 386 deletions(-) diff --git a/rust/lance/src/dataset/cleanup.rs b/rust/lance/src/dataset/cleanup.rs index 06e903de971..4f5828d32d1 100644 --- a/rust/lance/src/dataset/cleanup.rs +++ b/rust/lance/src/dataset/cleanup.rs @@ -740,6 +740,7 @@ impl<'a> CleanupTask<'a> { .commit_versions( &versions_to_commit, &reference_census.completed_intent_paths, + &reference_census.lifecycle_generations, ) .await?; for version in &late_reference_versions { diff --git a/rust/lance/src/dataset/refs.rs b/rust/lance/src/dataset/refs.rs index 1a7595dc03a..06d64a663eb 100644 --- a/rust/lance/src/dataset/refs.rs +++ b/rust/lance/src/dataset/refs.rs @@ -7,15 +7,16 @@ use futures::stream::{StreamExt, TryStreamExt}; use itertools::Itertools; use lance_io::object_store::ObjectStore; use lance_table::io::commit::CommitHandler; -use object_store::{ObjectMeta, ObjectStoreExt, PutMode, PutOptions, UpdateVersion, path::Path}; +use object_store::{ObjectMeta, ObjectStoreExt, path::Path}; use serde::{Deserialize, Serialize}; use std::sync::Arc; use crate::dataset::branch_location::BranchLocation; use crate::dataset::refs::Ref::{Tag, Version, VersionNumber}; use crate::dataset::version_lease::{ - REFERENCE_TOMBSTONE_PAYLOAD, ReferenceMutation, begin_reference_admission, - lock_local_reference, remove_branch_state, + ReferenceMutation, begin_branch_state_removal, begin_reference_admission, + cancel_branch_state_removal, canonical_reference_is_visible, delete_canonical_reference, + remove_branch_state, }; use crate::utils::temporal::utc_now; use crate::{Error, Result}; @@ -140,14 +141,9 @@ struct ReferenceSnapshot { payload: Bytes, } -impl ReferenceSnapshot { - fn is_tombstone(&self) -> bool { - self.payload.as_ref() == REFERENCE_TOMBSTONE_PAYLOAD - } -} - async fn reference_snapshot( object_store: &ObjectStore, + root_path: &Path, path: &Path, ) -> Result> { let result = match object_store.inner.get(path).await { @@ -157,67 +153,12 @@ async fn reference_snapshot( }; let metadata = result.meta.clone(); let payload = result.bytes().await?; - Ok(Some(ReferenceSnapshot { metadata, payload })) -} - -async fn tombstone_reference( - object_store: &ObjectStore, - path: &Path, - snapshot: &ReferenceSnapshot, -) -> Result<()> { - let local_lock = lock_local_reference(object_store, path).await?; - let Some(current) = reference_snapshot(object_store, path).await? else { - return Err(Error::RefConflict { - message: format!("reference {path} changed during conditional deletion"), - }); - }; - if current.payload != snapshot.payload { - return Err(Error::RefConflict { - message: format!("reference {path} changed during conditional deletion"), - }); - } - let result = object_store - .inner - .put_opts( - path, - Bytes::from_static(REFERENCE_TOMBSTONE_PAYLOAD).into(), - PutOptions { - mode: PutMode::Update(UpdateVersion { - e_tag: snapshot.metadata.e_tag.clone(), - version: snapshot.metadata.version.clone(), - }), - ..Default::default() - }, - ) - .await; - match result { - Ok(_) => Ok(()), - Err( - object_store::Error::NotSupported { .. } | object_store::Error::NotImplemented { .. }, - ) if local_lock.is_some() => { - // Local reference mutations hold the same OS advisory lock, so an - // atomic overwrite cannot race another supported local mutation. - // The durable tombstone prevents later creates from reusing this - // path without a locked conditional mutation. - object_store - .inner - .put(path, Bytes::from_static(REFERENCE_TOMBSTONE_PAYLOAD).into()) - .await?; - Ok(()) - } - Err( - object_store::Error::NotSupported { .. } | object_store::Error::NotImplemented { .. }, - ) => Err(Error::not_supported(format!( - "object store {} does not support atomic conditional reference deletion for {path}", - object_store.scheme() - ))), - Err(object_store::Error::NotFound { .. } | object_store::Error::Precondition { .. }) => { - Err(Error::RefConflict { - message: format!("reference {path} changed during conditional deletion"), - }) - } - Err(error) => Err(error.into()), + if !canonical_reference_is_visible(Arc::new(object_store.clone()), root_path, path, &payload) + .await? + { + return Ok(None); } + Ok(Some(ReferenceSnapshot { metadata, payload })) } impl Tags<'_> { @@ -248,12 +189,11 @@ impl Tags<'_> { futures::stream::iter(tag_names) .map(|tag_name| async move { let path = tag_path(root_path, &tag_name); - let Some(snapshot) = reference_snapshot(self.object_store(), &path).await? else { + let Some(snapshot) = + reference_snapshot(self.object_store(), root_path, &path).await? + else { return Ok(None); }; - if snapshot.is_tombstone() { - return Ok(None); - } let contents = serde_json::from_slice(&snapshot.payload)?; Ok(Some((tag_name, contents))) }) @@ -296,16 +236,13 @@ impl Tags<'_> { let root_location = self.refs.root()?; let tag_file = tag_path(&root_location.path, tag); - let Some(snapshot) = reference_snapshot(self.object_store(), &tag_file).await? else { + let Some(snapshot) = + reference_snapshot(self.object_store(), &root_location.path, &tag_file).await? + else { return Err(Error::RefNotFound { message: format!("tag {} does not exist", tag), }); }; - if snapshot.is_tombstone() { - return Err(Error::RefNotFound { - message: format!("tag {} does not exist", tag), - }); - } let tag_contents = serde_json::from_slice(&snapshot.payload)?; Ok(tag_contents) @@ -316,11 +253,9 @@ impl Tags<'_> { let root_location = self.refs.root()?; let tag_file = tag_path(&root_location.path, tag); - let previous = reference_snapshot(self.object_store(), &tag_file).await?; - if previous - .as_ref() - .is_some_and(|snapshot| !snapshot.is_tombstone()) - { + let previous = + reference_snapshot(self.object_store(), &root_location.path, &tag_file).await?; + if previous.is_some() { return Err(Error::RefConflict { message: format!("tag {} already exists", tag), }); @@ -359,18 +294,21 @@ impl Tags<'_> { let root_location = self.refs.root()?; let tag_file = tag_path(&root_location.path, tag); - let Some(snapshot) = reference_snapshot(self.object_store(), &tag_file).await? else { + let Some(snapshot) = + reference_snapshot(self.object_store(), &root_location.path, &tag_file).await? + else { return Err(Error::RefNotFound { message: format!("tag {} does not exist", tag), }); }; - if snapshot.is_tombstone() { - return Err(Error::RefNotFound { - message: format!("tag {} does not exist", tag), - }); - } - - tombstone_reference(self.object_store(), &tag_file, &snapshot).await + delete_canonical_reference( + Arc::clone(&self.refs.object_store), + &root_location.path, + &tag_file, + &snapshot.metadata, + &snapshot.payload, + ) + .await } pub async fn update(&self, tag: &str, reference: impl Into) -> Result<()> { @@ -378,16 +316,13 @@ impl Tags<'_> { let root_location = self.refs.root()?; let tag_file = tag_path(&root_location.path, tag); - let Some(snapshot) = reference_snapshot(self.object_store(), &tag_file).await? else { + let Some(snapshot) = + reference_snapshot(self.object_store(), &root_location.path, &tag_file).await? + else { return Err(Error::RefNotFound { message: format!("tag {} does not exist", tag), }); }; - if snapshot.is_tombstone() { - return Err(Error::RefNotFound { - message: format!("tag {} does not exist", tag), - }); - } let expected_etag = snapshot.metadata.e_tag.clone(); let expected_version = snapshot.metadata.version.clone(); let original_payload = snapshot.payload; @@ -430,16 +365,13 @@ impl Tags<'_> { let root_location = self.refs.root()?; let tag_file = tag_path(&root_location.path, tag); - let Some(snapshot) = reference_snapshot(self.object_store(), &tag_file).await? else { + let Some(snapshot) = + reference_snapshot(self.object_store(), &root_location.path, &tag_file).await? + else { return Err(Error::RefNotFound { message: format!("tag {} does not exist", tag), }); }; - if snapshot.is_tombstone() { - return Err(Error::RefNotFound { - message: format!("tag {} does not exist", tag), - }); - } let expected_etag = snapshot.metadata.e_tag.clone(); let expected_version = snapshot.metadata.version.clone(); @@ -553,12 +485,11 @@ impl Branches<'_> { futures::stream::iter(branch_names) .map(|name| async move { let path = branch_contents_path(branch_path, &name); - let Some(snapshot) = reference_snapshot(self.object_store(), &path).await? else { + let Some(snapshot) = + reference_snapshot(self.object_store(), branch_path, &path).await? + else { return Ok(None); }; - if snapshot.is_tombstone() { - return Ok(None); - } let mut contents: BranchContents = serde_json::from_slice(&snapshot.payload)?; contents.ensure_identifier(&name); Ok(Some((name, contents))) @@ -581,16 +512,13 @@ impl Branches<'_> { let root_location = self.refs.root()?; let branch_file = branch_contents_path(&root_location.path, branch); - let Some(snapshot) = reference_snapshot(self.object_store(), &branch_file).await? else { + let Some(snapshot) = + reference_snapshot(self.object_store(), &root_location.path, &branch_file).await? + else { return Err(Error::RefNotFound { message: format!("branch {} does not exist", branch), }); }; - if snapshot.is_tombstone() { - return Err(Error::RefNotFound { - message: format!("branch {} does not exist", branch), - }); - } let mut branch_contents: BranchContents = serde_json::from_slice(&snapshot.payload)?; branch_contents.ensure_identifier(branch); @@ -619,11 +547,9 @@ 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); - let previous = reference_snapshot(self.object_store(), &branch_file).await?; - if previous - .as_ref() - .is_some_and(|snapshot| !snapshot.is_tombstone()) - { + let previous = + reference_snapshot(self.object_store(), &root_location.path, &branch_file).await?; + if previous.is_some() { return Err(Error::RefConflict { message: format!("branch {} already exists", branch_name), }); @@ -650,16 +576,7 @@ impl Branches<'_> { }); }; let parent_branch_id = if let Some(ref parent_branch) = source_branch { - let parent_file = branch_contents_path(&root_location.path, parent_branch); - if self.object_store().exists(&parent_file).await? { - BranchContents::from_path(&parent_file, self.object_store(), parent_branch) - .await? - .identifier - } else { - return Err(Error::RefNotFound { - message: format!("Parent branch {} does not exist", branch_name), - }); - } + self.get(parent_branch).await?.identifier } else { BranchIdentifier::main() }; @@ -711,16 +628,13 @@ impl Branches<'_> { let root_location = self.refs.root()?; let branch_file = branch_contents_path(&root_location.path, branch); - let Some(snapshot) = reference_snapshot(self.object_store(), &branch_file).await? else { + let Some(snapshot) = + reference_snapshot(self.object_store(), &root_location.path, &branch_file).await? + else { return Err(Error::RefNotFound { message: format!("branch {} does not exist", branch), }); }; - if snapshot.is_tombstone() { - return Err(Error::RefNotFound { - message: format!("branch {} does not exist", branch), - }); - } let expected_etag = snapshot.metadata.e_tag.clone(); let expected_version = snapshot.metadata.version.clone(); @@ -761,13 +675,47 @@ impl Branches<'_> { pub async fn delete(&self, branch: &str, force: bool) -> Result<()> { check_valid_branch(branch)?; - let all_branches = self.list().await?; - let branch_id = all_branches - .get(branch) + let root_location = self.refs.root()?; + let branch_contents = match self.get(branch).await { + Ok(contents) => Some(contents), + Err(Error::RefNotFound { .. }) if force => None, + Err(error) => return Err(error), + }; + let branch_id = branch_contents + .as_ref() .map(|contents| contents.identifier.clone()); + let namespace = branch_id + .as_ref() + .and_then(|identifier| identifier.version_mapping.last().map(|(_, id)| id.as_str())); + if let Some(namespace) = namespace { + begin_branch_state_removal(self.object_store(), &root_location.path, namespace).await?; + } + + let all_branches = match self.list().await { + Ok(branches) => branches, + Err(error) => { + if let Some(namespace) = namespace { + cancel_branch_state_removal( + self.object_store(), + &root_location.path, + namespace, + ) + .await?; + } + return Err(error); + } + }; if let Some(branch_id) = branch_id.as_ref() { let referenced_versions = branch_id.collect_referenced_versions(&all_branches); if !referenced_versions.is_empty() { + if let Some(namespace) = namespace { + cancel_branch_state_removal( + self.object_store(), + &root_location.path, + namespace, + ) + .await?; + } return Err(Error::RefConflict { message: format!( "Branch {} is referenced by {:?} versions, can not delete", @@ -775,32 +723,36 @@ impl Branches<'_> { ), }); } - } else if !force { - return Err(Error::RefNotFound { - message: format!("Branch {} does not exist", branch), - }); } else { log::warn!("BranchContents of {} does not exist", branch); } - let root_location = self.refs.root()?; let branch_file = branch_contents_path(&root_location.path, branch); - if let Some(snapshot) = reference_snapshot(self.object_store(), &branch_file).await? - && !snapshot.is_tombstone() + if let Some(snapshot) = + reference_snapshot(self.object_store(), &root_location.path, &branch_file).await? + && let Err(error) = delete_canonical_reference( + Arc::clone(&self.refs.object_store), + &root_location.path, + &branch_file, + &snapshot.metadata, + &snapshot.payload, + ) + .await { - tombstone_reference(self.object_store(), &branch_file, &snapshot).await?; + if let Some(namespace) = namespace { + cancel_branch_state_removal(self.object_store(), &root_location.path, namespace) + .await?; + } + return Err(error); } + if let Some(namespace) = namespace { + remove_branch_state(self.object_store(), &root_location.path, namespace).await?; + } // Clean up branch directories and operational state scoped to this // branch incarnation. The identifier prevents a recreated branch with // the same name from inheriting leases or retirement markers. self.cleanup_branch_directories(branch).await?; - if let Some(namespace) = branch_id - .as_ref() - .and_then(|identifier| identifier.version_mapping.last().map(|(_, id)| id.as_str())) - { - remove_branch_state(self.object_store(), &root_location.path, namespace).await?; - } Ok(()) } @@ -1137,17 +1089,8 @@ async fn from_path(path: &Path, object_store: &ObjectStore) -> Result where T: DeserializeOwned, { - let Some(snapshot) = reference_snapshot(object_store, path).await? else { - return Err(Error::RefNotFound { - message: format!("reference {path} does not exist"), - }); - }; - if snapshot.is_tombstone() { - return Err(Error::RefNotFound { - message: format!("reference {path} does not exist"), - }); - } - Ok(serde_json::from_slice(&snapshot.payload)?) + let result = object_store.inner.get(path).await?; + Ok(serde_json::from_slice(&result.bytes().await?)?) } impl TagContents { diff --git a/rust/lance/src/dataset/version_lease.rs b/rust/lance/src/dataset/version_lease.rs index 10ebab6b858..ad5209e05ca 100644 --- a/rust/lance/src/dataset/version_lease.rs +++ b/rust/lance/src/dataset/version_lease.rs @@ -22,13 +22,15 @@ use uuid::Uuid; use super::{ Dataset, - refs::{MAIN_BRANCH, Refs}, + refs::{BranchIdentifier, MAIN_BRANCH, Refs}, }; use crate::{Error, Result}; const LEASES_DIR: &str = "_refs/version_leases"; const LEASE_GC_MARKERS_DIR: &str = "_refs/version_lease_gc"; const REFERENCE_INTENTS_DIR: &str = "_refs/version_reference_intents"; +const REFERENCE_STATES_DIR: &str = "_refs/version_reference_states"; +const BRANCH_TERMINATIONS_DIR: &str = "_refs/version_branch_terminations"; const LEASE_FILE_SUFFIX: &str = ".lease"; const REFERENCE_INTENT_SUFFIX: &str = ".intent"; const DRAINING_MARKER_SUFFIX: &str = ".draining"; @@ -36,10 +38,7 @@ const SEALED_MARKER_SUFFIX: &str = ".sealed"; const COMMITTED_MARKER_SUFFIX: &str = ".committed"; const STORAGE_CLOCK_PATH: &str = "_clock"; -/// A deleted canonical reference remains as a tombstone so an abandoned -/// `PutMode::Create` publication cannot recreate the name after deletion. -pub(super) const REFERENCE_TOMBSTONE_PAYLOAD: &[u8] = - br#"{"_lanceReferenceState":"deleted","formatVersion":1}"#; +const REFERENCE_GENERATION_FIELD: &str = "_lanceReferenceGeneration"; /// HTTP `Last-Modified`, used by supported cloud stores, has whole-second precision. /// Adding one interval guarantees a lease is never shortened by timestamp truncation. @@ -158,10 +157,20 @@ impl Drop for VersionLease { #[derive(Clone, Debug)] pub(super) struct VersionLeaseStore { object_store: Arc, + root_path: Path, + namespace: String, leases_path: Path, markers_path: Path, reference_intents_path: Path, manifest_path: Option, + canonical_references: Option, +} + +#[derive(Clone, Debug)] +struct CanonicalReferenceContext { + refs: Refs, + branch: Option, + branch_identifier: BranchIdentifier, } /// Durable ownership of one in-flight tag or branch publication. @@ -176,6 +185,7 @@ pub(super) struct ReferenceAdmission { manifest_path: Path, version: u64, created_at: DateTime, + operation_id: String, mutation: ReferenceMutation, } @@ -211,9 +221,52 @@ struct ReferenceIntent { manifest_path: String, mutation: ReferenceMutation, #[serde(default)] + operation_id: String, + #[serde(default)] state: ReferenceIntentState, } +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +struct ReferenceTarget { + namespace: String, + version: u64, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase", tag = "state")] +enum ReferenceLifecycleState { + Pending { + canonical_path: String, + operation_id: String, + target: ReferenceTarget, + previous: Option, + previous_was_legacy: bool, + }, + Live { + canonical_path: String, + live: ReferenceLiveState, + }, + Deleting { + canonical_path: String, + operation_id: String, + }, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +struct ReferenceLiveState { + generation: String, + target: ReferenceTarget, +} + +#[derive(Debug)] +struct ReferenceLifecycleSnapshot { + metadata: ObjectMeta, + payload: Bytes, + state: ReferenceLifecycleState, +} + #[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] #[serde(rename_all = "camelCase")] enum ReferenceIntentState { @@ -281,6 +334,7 @@ struct ReferenceIntentMetadata { pub(super) struct CanonicalReferenceCensus { pub(super) versions: HashSet, pub(super) completed_intent_paths: HashSet, + pub(super) lifecycle_generations: HashMap, } /// Per-cleanup retirement markers for versions selected for deletion. @@ -332,13 +386,294 @@ impl VersionLeaseStore { Ok(Self { object_store: Arc::clone(&refs.object_store), + root_path: root.path.clone(), + namespace: namespace.to_string(), leases_path: root.path.clone().join(LEASES_DIR).join(namespace), markers_path: root.path.clone().join(LEASE_GC_MARKERS_DIR).join(namespace), - reference_intents_path: root.path.join(REFERENCE_INTENTS_DIR).join(namespace), + reference_intents_path: root + .path + .clone() + .join(REFERENCE_INTENTS_DIR) + .join(namespace), manifest_path, + canonical_references: Some(CanonicalReferenceContext { + refs: refs.clone(), + branch: branch.map(ToOwned::to_owned), + branch_identifier, + }), }) } + fn lifecycle_only(object_store: Arc, root_path: Path) -> Self { + Self { + object_store, + root_path: root_path.clone(), + namespace: String::new(), + leases_path: root_path.clone().join(LEASES_DIR), + markers_path: root_path.clone().join(LEASE_GC_MARKERS_DIR), + reference_intents_path: root_path.join(REFERENCE_INTENTS_DIR), + manifest_path: None, + canonical_references: None, + } + } + + fn reference_state_path(&self, canonical_path: &Path) -> Path { + self.root_path + .clone() + .join(REFERENCE_STATES_DIR) + .join(format!( + "{}.state", + stable_reference_path_id(canonical_path.as_ref()) + )) + } + + async fn reference_lifecycle_snapshot( + &self, + canonical_path: &Path, + ) -> Result> { + let state_path = self.reference_state_path(canonical_path); + let result = match self.object_store.inner.get(&state_path).await { + Ok(result) => result, + Err(object_store::Error::NotFound { .. }) => return Ok(None), + Err(error) => return Err(error.into()), + }; + let metadata = result.meta.clone(); + let payload = result.bytes().await?; + let state: ReferenceLifecycleState = serde_json::from_slice(&payload) + .map_err(|error| Error::corrupt_file(state_path, error.to_string()))?; + let recorded_path = match &state { + ReferenceLifecycleState::Pending { canonical_path, .. } + | ReferenceLifecycleState::Live { canonical_path, .. } + | ReferenceLifecycleState::Deleting { canonical_path, .. } => canonical_path, + }; + if recorded_path != canonical_path.as_ref() { + return Err(Error::corrupt_file( + metadata.location.clone(), + format!( + "reference lifecycle state belongs to {recorded_path}, not {canonical_path}" + ), + )); + } + Ok(Some(ReferenceLifecycleSnapshot { + metadata, + payload, + state, + })) + } + + async fn put_reference_lifecycle( + &self, + canonical_path: &Path, + expected: Option<&ReferenceLifecycleSnapshot>, + state: &ReferenceLifecycleState, + has_local_lock: bool, + ) -> Result> { + let state_path = self.reference_state_path(canonical_path); + let payload = Bytes::from(serde_json::to_vec(state)?); + let mode = expected.map_or(PutMode::Create, |snapshot| { + PutMode::Update(UpdateVersion { + e_tag: snapshot.metadata.e_tag.clone(), + version: snapshot.metadata.version.clone(), + }) + }); + let result = self + .object_store + .inner + .put_opts( + &state_path, + payload.clone().into(), + PutOptions { + mode, + ..Default::default() + }, + ) + .await; + let result = match result { + Err( + object_store::Error::NotSupported { .. } + | object_store::Error::NotImplemented { .. }, + ) if has_local_lock => { + let current = match self.object_store.inner.get(&state_path).await { + Ok(result) => Some(result.bytes().await?), + Err(object_store::Error::NotFound { .. }) => None, + Err(error) => return Err(error.into()), + }; + let matches_expected = match (expected, current.as_deref()) { + (None, None) => true, + (Some(expected), Some(current)) => current == expected.payload.as_ref(), + _ => false, + }; + if !matches_expected { + return Ok(None); + } + self.object_store + .inner + .put(&state_path, payload.into()) + .await + } + Err( + object_store::Error::NotSupported { .. } + | object_store::Error::NotImplemented { .. }, + ) => { + return Err(Error::not_supported(format!( + "object store {} does not support atomic reference lifecycle updates for {canonical_path}", + self.object_store.scheme() + ))); + } + result => result, + }; + match result { + Ok(_) => self.reference_lifecycle_snapshot(canonical_path).await, + Err( + object_store::Error::AlreadyExists { .. } + | object_store::Error::Precondition { .. } + | object_store::Error::NotFound { .. }, + ) => Ok(None), + Err(error) => Err(error.into()), + } + } + + async fn claim_reference_lifecycle( + &self, + canonical_path: &Path, + operation_id: &str, + version: u64, + previous_was_legacy: bool, + ) -> Result<()> { + let local_lock = lock_local_reference(&self.object_store, canonical_path).await?; + let current = self.reference_lifecycle_snapshot(canonical_path).await?; + let previous = match current.as_ref().map(|snapshot| &snapshot.state) { + None => None, + Some(ReferenceLifecycleState::Live { live, .. }) => Some(live.clone()), + Some(ReferenceLifecycleState::Pending { .. }) => { + return Err(Error::RefConflict { + message: format!("reference {canonical_path} has an in-flight mutation"), + }); + } + Some(ReferenceLifecycleState::Deleting { .. }) => { + return Err(Error::RefConflict { + message: format!("reference {canonical_path} is being deleted"), + }); + } + }; + let pending = ReferenceLifecycleState::Pending { + canonical_path: canonical_path.to_string(), + operation_id: operation_id.to_string(), + target: ReferenceTarget { + namespace: self.namespace.clone(), + version, + }, + previous, + previous_was_legacy, + }; + if self + .put_reference_lifecycle( + canonical_path, + current.as_ref(), + &pending, + local_lock.is_some(), + ) + .await? + .is_none() + { + return Err(Error::RefConflict { + message: format!("reference {canonical_path} lifecycle changed during admission"), + }); + } + Ok(()) + } + + async fn complete_reference_lifecycle( + &self, + canonical_path: &Path, + operation_id: &str, + version: u64, + has_local_lock: bool, + ) -> Result { + let Some(current) = self.reference_lifecycle_snapshot(canonical_path).await? else { + return Ok(false); + }; + if !matches!( + ¤t.state, + ReferenceLifecycleState::Pending { + operation_id: current_operation, + .. + } if current_operation == operation_id + ) { + return Ok(false); + } + let live = ReferenceLifecycleState::Live { + canonical_path: canonical_path.to_string(), + live: ReferenceLiveState { + generation: operation_id.to_string(), + target: ReferenceTarget { + namespace: self.namespace.clone(), + version, + }, + }, + }; + Ok(self + .put_reference_lifecycle(canonical_path, Some(¤t), &live, has_local_lock) + .await? + .is_some()) + } + + async fn cancel_reference_lifecycle( + &self, + canonical_path: &Path, + operation_id: &str, + has_local_lock: bool, + ) -> Result { + let Some(current) = self.reference_lifecycle_snapshot(canonical_path).await? else { + return Ok(false); + }; + let previous = match ¤t.state { + ReferenceLifecycleState::Pending { + operation_id: current_operation, + previous, + .. + } if current_operation == operation_id => previous.clone(), + _ => return Ok(false), + }; + if let Some(previous) = previous { + let live = ReferenceLifecycleState::Live { + canonical_path: canonical_path.to_string(), + live: previous, + }; + return Ok(self + .put_reference_lifecycle(canonical_path, Some(¤t), &live, has_local_lock) + .await? + .is_some()); + } + let state_path = self.reference_state_path(canonical_path); + match self.object_store.delete(&state_path).await { + Ok(()) => Ok(true), + Err(error) if error.is_not_found() => Ok(false), + Err(error) => Err(error), + } + } + + fn branch_termination_path(&self) -> Path { + self.root_path + .clone() + .join(BRANCH_TERMINATIONS_DIR) + .join(format!("{}.deleted", self.namespace)) + } + + async fn ensure_branch_incarnation_active(&self) -> Result<()> { + if self.namespace != MAIN_BRANCH + && self + .object_store + .exists(&self.branch_termination_path()) + .await? + { + return Err(Error::RefConflict { + message: "parent branch incarnation is being deleted".to_string(), + }); + } + Ok(()) + } + async fn acquire(&self, version: u64, ttl: Duration) -> Result { self.ensure_version_available(version).await?; let lease_file = self.create_lease_file(version, ttl).await?; @@ -399,13 +734,150 @@ impl VersionLeaseStore { } pub(super) async fn active_reference_versions(&self) -> Result> { - Ok(self + let mut versions = self .reference_versions( CompletedReferenceIntentHandling::RetainForCurrentScan, &HashSet::new(), ) .await? - .versions) + .versions; + versions.extend(self.lifecycle_reference_versions().await?); + versions.extend(self.canonical_reference_versions().await?); + Ok(versions) + } + + async fn lifecycle_reference_versions(&self) -> Result> { + let state_path = self.root_path.clone().join(REFERENCE_STATES_DIR); + let metadata = self + .object_store + .list(Some(state_path)) + .try_collect::>() + .await?; + stream::iter(metadata) + .map(|metadata| async move { + let payload = self + .object_store + .inner + .get(&metadata.location) + .await? + .bytes() + .await?; + let state: ReferenceLifecycleState = + serde_json::from_slice(&payload).map_err(|error| { + Error::corrupt_file(metadata.location.clone(), error.to_string()) + })?; + let target = match state { + ReferenceLifecycleState::Pending { target, .. } => Some(target), + ReferenceLifecycleState::Live { live, .. } => Some(live.target), + ReferenceLifecycleState::Deleting { .. } => None, + }; + Ok::<_, Error>(target.and_then(|target| { + (target.namespace == self.namespace).then_some(target.version) + })) + }) + .buffer_unordered(self.object_store.io_parallelism()) + .try_collect::>() + .await + .map(|versions| versions.into_iter().flatten().collect()) + } + + async fn lifecycle_generation_snapshot(&self) -> Result> { + let state_path = self.root_path.clone().join(REFERENCE_STATES_DIR); + let metadata = self + .object_store + .list(Some(state_path)) + .try_collect::>() + .await?; + stream::iter(metadata) + .map(|metadata| async move { + let payload = self + .object_store + .inner + .get(&metadata.location) + .await? + .bytes() + .await?; + let state: ReferenceLifecycleState = + serde_json::from_slice(&payload).map_err(|error| { + Error::corrupt_file(metadata.location.clone(), error.to_string()) + })?; + let generation = match state { + ReferenceLifecycleState::Pending { operation_id, .. } => Some(operation_id), + ReferenceLifecycleState::Live { live, .. } => Some(live.generation), + ReferenceLifecycleState::Deleting { .. } => None, + }; + Ok::<_, Error>(generation.map(|generation| (metadata.location, generation))) + }) + .buffer_unordered(self.object_store.io_parallelism()) + .try_collect::>() + .await + .map(|generations| generations.into_iter().flatten().collect()) + } + + async fn lifecycle_reference_versions_since( + &self, + observed_generations: &HashMap, + ) -> Result> { + let state_path = self.root_path.clone().join(REFERENCE_STATES_DIR); + let metadata = self + .object_store + .list(Some(state_path)) + .try_collect::>() + .await?; + stream::iter(metadata) + .map(|metadata| async move { + let payload = self + .object_store + .inner + .get(&metadata.location) + .await? + .bytes() + .await?; + let state: ReferenceLifecycleState = + serde_json::from_slice(&payload).map_err(|error| { + Error::corrupt_file(metadata.location.clone(), error.to_string()) + })?; + let (generation, target) = match state { + ReferenceLifecycleState::Pending { + operation_id, + target, + .. + } => (operation_id, Some(target)), + ReferenceLifecycleState::Live { live, .. } => { + (live.generation, Some(live.target)) + } + ReferenceLifecycleState::Deleting { operation_id, .. } => (operation_id, None), + }; + let changed = observed_generations.get(&metadata.location) != Some(&generation); + Ok::<_, Error>(target.and_then(|target| { + (changed && target.namespace == self.namespace).then_some(target.version) + })) + }) + .buffer_unordered(self.object_store.io_parallelism()) + .try_collect::>() + .await + .map(|versions| versions.into_iter().flatten().collect()) + } + + async fn canonical_reference_versions(&self) -> Result> { + let Some(context) = &self.canonical_references else { + return Ok(HashSet::new()); + }; + let tags = context.refs.tags().list().await?; + let branches = context.refs.branches().list().await?; + let mut versions = tags + .values() + .filter(|tag| tag.branch == context.branch) + .map(|tag| tag.version) + .collect::>(); + versions.extend( + context + .branch_identifier + .collect_referenced_versions(&branches) + .into_iter() + .map(|(_, version)| version), + ); + Ok(versions) } /// Fence in-flight publication before the caller scans canonical tags and @@ -414,11 +886,14 @@ impl VersionLeaseStore { pub(super) async fn reference_versions_before_canonical_census( &self, ) -> Result { - self.reference_versions( - CompletedReferenceIntentHandling::DeferToCanonicalCensus, - &HashSet::new(), - ) - .await + let mut census = self + .reference_versions( + CompletedReferenceIntentHandling::DeferToCanonicalCensus, + &HashSet::new(), + ) + .await?; + census.lifecycle_generations = self.lifecycle_generation_snapshot().await?; + Ok(census) } async fn reference_versions( @@ -479,18 +954,18 @@ impl VersionLeaseStore { active_versions.insert(intent.version); } ReferenceIntentState::Pending => { - // Never replay an expired publication: an absent canonical - // path cannot distinguish an operation that never published - // from one that published and was subsequently deleted. - // An abandoned create instead installs a durable tombstone, - // atomically fencing any owner that resumes after expiry. let manifest_path = Path::parse(&intent.intent.manifest_path)?; - if self.object_store.exists(&manifest_path).await? - && !self.has_committed_marker(intent.version).await? - && self - .expired_reference_mutation_outcome(&intent.intent.mutation) - .await? - == ReferenceMutationOutcome::Published + let allow_publish = self.object_store.exists(&manifest_path).await? + && !self.has_committed_marker(intent.version).await?; + if self + .expired_reference_mutation_outcome( + &intent.intent.mutation, + &intent.intent.operation_id, + intent.version, + allow_publish, + ) + .await? + == ReferenceMutationOutcome::Published { active_versions.insert(intent.version); } @@ -502,47 +977,151 @@ impl VersionLeaseStore { Ok(CanonicalReferenceCensus { versions: active_versions, completed_intent_paths, + lifecycle_generations: HashMap::new(), }) } async fn expired_reference_mutation_outcome( &self, mutation: &ReferenceMutation, + operation_id: &str, + version: u64, + allow_publish: bool, ) -> Result { - match mutation { - ReferenceMutation::Create { path, payload } => { - let path = Path::parse(path)?; - let _local_lock = lock_local_reference(&self.object_store, &path).await?; - match self - .object_store - .inner - .put_opts( - &path, - Bytes::from_static(REFERENCE_TOMBSTONE_PAYLOAD).into(), - PutOptions { - mode: PutMode::Create, - ..Default::default() - }, - ) - .await - { - // The tombstone won the create race, so this abandoned - // operation can never publish the deleted name. - Ok(_) => Ok(ReferenceMutationOutcome::Conflict), - Err( - object_store::Error::AlreadyExists { .. } - | object_store::Error::Precondition { .. }, - ) => { - self.reference_mutation_conflict_outcome(&path, payload) - .await - } - Err(error) => Err(error.into()), - } + let (path, intended_payload) = match mutation { + ReferenceMutation::Create { path, payload } + | ReferenceMutation::Update { path, payload, .. } => { + (Path::parse(path)?, payload.as_slice()) } - ReferenceMutation::Update { path, payload, .. } => { - self.reference_mutation_conflict_outcome(&Path::parse(path)?, payload) - .await + }; + let local_lock = lock_local_reference(&self.object_store, &path).await?; + if operation_id.is_empty() { + return self + .reference_mutation_conflict_outcome(&path, intended_payload) + .await; + } + + let lifecycle = self.reference_lifecycle_snapshot(&path).await?; + match lifecycle.as_ref().map(|snapshot| &snapshot.state) { + Some(ReferenceLifecycleState::Live { live, .. }) if live.generation == operation_id => { + return Ok(ReferenceMutationOutcome::Published); + } + Some(ReferenceLifecycleState::Pending { + operation_id: current_operation, + .. + }) if current_operation == operation_id => {} + _ => return Ok(ReferenceMutationOutcome::Conflict), + } + + let canonical = match self.object_store.inner.get(&path).await { + Ok(result) => { + let metadata = result.meta.clone(); + Some((metadata, result.bytes().await?)) + } + Err(object_store::Error::NotFound { .. }) => None, + Err(error) => return Err(error.into()), + }; + let canonical_is_intended = canonical + .as_ref() + .is_some_and(|(_, payload)| payload.as_ref() == intended_payload); + if canonical_is_intended && allow_publish { + if self + .complete_reference_lifecycle(&path, operation_id, version, local_lock.is_some()) + .await? + { + return Ok(ReferenceMutationOutcome::Published); + } + return Ok(ReferenceMutationOutcome::Conflict); + } + + match (mutation, canonical) { + ( + ReferenceMutation::Update { + expected_payload, .. + }, + Some((metadata, current_payload)), + ) if current_payload.as_ref() == expected_payload => { + self.rewrite_reference_payload( + &path, + &metadata, + ¤t_payload, + ¤t_payload, + local_lock.is_some(), + ) + .await?; + } + ( + ReferenceMutation::Update { + expected_payload, .. + }, + Some((metadata, current_payload)), + ) if canonical_is_intended => { + self.rewrite_reference_payload( + &path, + &metadata, + ¤t_payload, + expected_payload, + local_lock.is_some(), + ) + .await?; + } + (ReferenceMutation::Create { .. }, Some(_)) if canonical_is_intended => { + self.object_store.delete(&path).await?; + } + _ => {} + } + self.cancel_reference_lifecycle(&path, operation_id, local_lock.is_some()) + .await?; + Ok(ReferenceMutationOutcome::Conflict) + } + + async fn rewrite_reference_payload( + &self, + path: &Path, + metadata: &ObjectMeta, + expected_payload: &[u8], + replacement_payload: &[u8], + has_local_lock: bool, + ) -> Result<()> { + let result = self + .object_store + .inner + .put_opts( + path, + Bytes::copy_from_slice(replacement_payload).into(), + PutOptions { + mode: PutMode::Update(UpdateVersion { + e_tag: metadata.e_tag.clone(), + version: metadata.version.clone(), + }), + ..Default::default() + }, + ) + .await; + match result { + Ok(_) => Ok(()), + Err( + object_store::Error::NotSupported { .. } + | object_store::Error::NotImplemented { .. }, + ) if has_local_lock => { + let current = self.object_store.inner.get(path).await?.bytes().await?; + if current.as_ref() != expected_payload { + return Err(Error::RefConflict { + message: format!("reference {path} changed while fencing an expired owner"), + }); + } + self.object_store + .inner + .put(path, Bytes::copy_from_slice(replacement_payload).into()) + .await?; + Ok(()) } + Err( + object_store::Error::Precondition { .. } | object_store::Error::NotFound { .. }, + ) => Err(Error::RefConflict { + message: format!("reference {path} changed while fencing an expired owner"), + }), + Err(error) => Err(error.into()), } } @@ -582,7 +1161,26 @@ impl VersionLeaseStore { &self, mutation: &ReferenceMutation, has_local_lock: bool, + expected_operation_id: Option<&str>, ) -> Result { + if let Some(expected_operation_id) = expected_operation_id { + self.ensure_branch_incarnation_active().await?; + let canonical_path = match mutation { + ReferenceMutation::Create { path, .. } | ReferenceMutation::Update { path, .. } => { + Path::parse(path)? + } + }; + let lifecycle = self.reference_lifecycle_snapshot(&canonical_path).await?; + if !matches!( + lifecycle.as_ref().map(|snapshot| &snapshot.state), + Some(ReferenceLifecycleState::Pending { + operation_id, + .. + }) if operation_id == expected_operation_id + ) { + return Ok(ReferenceMutationOutcome::Conflict); + } + } let (path, payload, result) = match mutation { ReferenceMutation::Create { path, payload } => { let path = Path::parse(path)?; @@ -1278,11 +1876,12 @@ impl RetirementGuard { &mut self, versions: &HashSet, completed_intents_observed_before_census: &HashSet, + lifecycle_generations_observed_before_census: &HashMap, ) -> Result> { // This scan runs inside the commit operation, after the caller's final // reference census. An admission that won that narrow race remains // cancellable and cannot cross the irreversible committed boundary. - let active_reference_versions = self + let mut active_reference_versions = self .store .reference_versions( CompletedReferenceIntentHandling::RetainForCurrentScan, @@ -1290,6 +1889,11 @@ impl RetirementGuard { ) .await? .versions; + active_reference_versions.extend( + self.store + .lifecycle_reference_versions_since(lifecycle_generations_observed_before_census) + .await?, + ); let retained_versions = versions .intersection(&active_reference_versions) .copied() @@ -1574,7 +2178,7 @@ pub(super) async fn begin_reference_admission( refs: &Refs, branch: Option<&str>, version: u64, - mutation: ReferenceMutation, + mut mutation: ReferenceMutation, ) -> Result { let branch_location = refs.base_location.find_branch(branch)?; let manifest = refs @@ -1589,18 +2193,45 @@ pub(super) async fn begin_reference_admission( let store = VersionLeaseStore::for_refs(refs, branch, None).await?; ensure_reference_available(&store, &manifest.path, version).await?; + store.ensure_branch_incarnation_active().await?; + let operation_id = Uuid::new_v4().simple().to_string(); + let (canonical_path, previous_was_legacy) = match &mut mutation { + ReferenceMutation::Create { path, payload } => { + *payload = set_payload_reference_generation(payload, &operation_id)?; + (Path::parse(path)?, false) + } + ReferenceMutation::Update { + path, + expected_payload, + payload, + .. + } => { + let previous_was_legacy = payload_reference_generation(expected_payload)?.is_none(); + *payload = set_payload_reference_generation(payload, &operation_id)?; + (Path::parse(path)?, previous_was_legacy) + } + }; let intent = ReferenceIntent { manifest_path: manifest.path.to_string(), mutation: mutation.clone(), + operation_id: operation_id.clone(), state: ReferenceIntentState::Pending, }; let (path, created_at) = store.create_reference_intent(version, &intent).await?; + if let Err(error) = store + .claim_reference_lifecycle(&canonical_path, &operation_id, version, previous_was_legacy) + .await + { + let _ = store.object_store.delete(&path).await; + return Err(error); + } let admission = ReferenceAdmission { store, path, manifest_path: manifest.path, version, created_at, + operation_id, mutation, }; if let Err(error) = admission.ensure_owned().await { @@ -1625,31 +2256,76 @@ impl ReferenceAdmission { ), }); } - ensure_reference_available(&self.store, &self.manifest_path, self.version).await + ensure_reference_available(&self.store, &self.manifest_path, self.version).await?; + self.store.ensure_branch_incarnation_active().await?; + let canonical_path = self.canonical_path()?; + let lifecycle = self + .store + .reference_lifecycle_snapshot(&canonical_path) + .await?; + if !matches!( + lifecycle.as_ref().map(|snapshot| &snapshot.state), + Some(ReferenceLifecycleState::Pending { + operation_id, + .. + }) if operation_id == &self.operation_id + ) { + return Err(Error::RefConflict { + message: format!( + "reference admission for version {} lost its lifecycle generation", + self.version + ), + }); + } + Ok(()) } - pub(super) async fn publish(self, conflict_message: String) -> Result<()> { - let canonical_path = match &self.mutation { + fn canonical_path(&self) -> Result { + match &self.mutation { ReferenceMutation::Create { path, .. } | ReferenceMutation::Update { path, .. } => { - Path::parse(path)? + Ok(Path::parse(path)?) } - }; + } + } + + pub(super) async fn publish(self, conflict_message: String) -> Result<()> { + let canonical_path = self.canonical_path()?; let local_lock = lock_local_reference(&self.store.object_store, &canonical_path).await?; if let Err(error) = self.ensure_owned().await { - self.cancel_before_publish().await; + self.cancel_before_publish_inner(local_lock.is_some()).await; return Err(error); } match self .store - .apply_reference_mutation_inner(&self.mutation, local_lock.is_some()) + .apply_reference_mutation_inner( + &self.mutation, + local_lock.is_some(), + Some(&self.operation_id), + ) .await { Ok(ReferenceMutationOutcome::Published) => { - self.complete().await; - Ok(()) + if self + .store + .complete_reference_lifecycle( + &canonical_path, + &self.operation_id, + self.version, + local_lock.is_some(), + ) + .await? + { + self.complete().await; + Ok(()) + } else { + self.cancel_before_publish_inner(local_lock.is_some()).await; + Err(Error::RefConflict { + message: conflict_message, + }) + } } Ok(ReferenceMutationOutcome::Conflict) => { - self.cancel_before_publish().await; + self.cancel_before_publish_inner(local_lock.is_some()).await; Err(Error::RefConflict { message: conflict_message, }) @@ -1664,6 +2340,38 @@ impl ReferenceAdmission { /// Remove this operation's uniquely owned intent after a definitive result /// that did not publish the canonical reference. pub(super) async fn cancel_before_publish(&self) { + let local_lock = match self.canonical_path() { + Ok(canonical_path) => { + match lock_local_reference(&self.store.object_store, &canonical_path).await { + Ok(local_lock) => local_lock, + Err(error) => { + tracing::warn!( + path = %canonical_path, + error = %error, + "Failed to lock cancelled reference admission" + ); + None + } + } + } + Err(_) => None, + }; + self.cancel_before_publish_inner(local_lock.is_some()).await; + } + + async fn cancel_before_publish_inner(&self, has_local_lock: bool) { + if let Ok(canonical_path) = self.canonical_path() + && let Err(error) = self + .store + .cancel_reference_lifecycle(&canonical_path, &self.operation_id, has_local_lock) + .await + { + tracing::warn!( + path = %canonical_path, + error = %error, + "Failed to cancel reference lifecycle admission" + ); + } if let Err(error) = self.store.object_store.delete(&self.path).await && !error.is_not_found() { @@ -1679,58 +2387,201 @@ impl ReferenceAdmission { /// Cleanup verifies the canonical payload, retains it for the current scan, /// and removes the completed intent without waiting for the ownership timeout. async fn complete(&self) { - let intent = ReferenceIntent { - manifest_path: self.manifest_path.to_string(), - mutation: self.mutation.clone(), - state: ReferenceIntentState::Completed, - }; - let result = match serde_json::to_vec(&intent) { - Ok(payload) => self.store.object_store.put(&self.path, &payload).await, - Err(error) => Err(error.into()), - }; - if let Err(error) = result { - // The pending intent remains a safe bounded fallback if completion - // cannot be recorded after the canonical write succeeded. + if let Err(error) = self.store.object_store.delete(&self.path).await + && !error.is_not_found() + { + // The live lifecycle generation is already the durable handoff. tracing::warn!( path = %self.path, error = %error, - "Failed to complete reference admission intent" + "Failed to remove completed reference admission intent" ); } } -} +} + +pub(super) async fn lock_local_reference( + object_store: &ObjectStore, + path: &Path, +) -> Result> { + if !object_store.is_local() && object_store.scheme() != "file-object-store" { + return Ok(None); + } + + let path = PathBuf::from(path.as_ref()); + let path = if path.is_absolute() { + path + } else { + PathBuf::from(std::path::MAIN_SEPARATOR_STR).join(path) + }; + let lock_path = PathBuf::from(format!("{}.lance-reference.lock", path.display())); + tokio::task::spawn_blocking(move || -> Result { + if let Some(parent) = lock_path.parent() { + std::fs::create_dir_all(parent)?; + } + let file = OpenOptions::new() + .create(true) + .truncate(false) + .read(true) + .write(true) + .open(lock_path)?; + file.lock()?; + Ok(LocalReferenceLock { _file: file }) + }) + .await + .map_err(|error| Error::internal(format!("failed to acquire local reference lock: {error}")))? + .map(Some) +} + +fn stable_reference_path_id(path: &str) -> String { + const FNV_OFFSET: u64 = 0xcbf29ce484222325; + const FNV_PRIME: u64 = 0x100000001b3; + + fn hash(path: &[u8], seed: u64) -> u64 { + path.iter().fold(seed, |hash, byte| { + (hash ^ u64::from(*byte)).wrapping_mul(FNV_PRIME) + }) + } + + format!( + "{:016x}{:016x}", + hash(path.as_bytes(), FNV_OFFSET), + hash(path.as_bytes(), FNV_OFFSET ^ 0x9e3779b97f4a7c15) + ) +} + +fn payload_reference_generation(payload: &[u8]) -> Result> { + let value: serde_json::Value = serde_json::from_slice(payload)?; + Ok(value + .as_object() + .and_then(|object| object.get(REFERENCE_GENERATION_FIELD)) + .and_then(serde_json::Value::as_str) + .map(ToOwned::to_owned)) +} + +fn set_payload_reference_generation(payload: &[u8], generation: &str) -> Result> { + let mut value: serde_json::Value = serde_json::from_slice(payload)?; + let object = value.as_object_mut().ok_or_else(|| { + Error::internal("canonical reference payload must be a JSON object".to_string()) + })?; + object.insert( + REFERENCE_GENERATION_FIELD.to_string(), + serde_json::Value::String(generation.to_string()), + ); + Ok(serde_json::to_vec_pretty(&value)?) +} + +pub(super) async fn canonical_reference_is_visible( + object_store: Arc, + root_path: &Path, + canonical_path: &Path, + payload: &[u8], +) -> Result { + let generation = payload_reference_generation(payload)?; + let store = VersionLeaseStore::lifecycle_only(object_store, root_path.clone()); + let Some(snapshot) = store.reference_lifecycle_snapshot(canonical_path).await? else { + return Ok(generation.is_none()); + }; + Ok(match snapshot.state { + ReferenceLifecycleState::Live { live, .. } => { + generation.as_deref() == Some(live.generation.as_str()) + } + ReferenceLifecycleState::Pending { + previous, + previous_was_legacy, + .. + } => match previous { + Some(previous) => generation.as_deref() == Some(previous.generation.as_str()), + None => previous_was_legacy && generation.is_none(), + }, + ReferenceLifecycleState::Deleting { .. } => false, + }) +} + +pub(super) async fn delete_canonical_reference( + object_store: Arc, + root_path: &Path, + canonical_path: &Path, + expected_metadata: &ObjectMeta, + expected_payload: &[u8], +) -> Result<()> { + let store = VersionLeaseStore::lifecycle_only(Arc::clone(&object_store), root_path.clone()); + let local_lock = lock_local_reference(&object_store, canonical_path).await?; + let current = object_store + .inner + .get(canonical_path) + .await + .map_err(|error| { + if matches!(error, object_store::Error::NotFound { .. }) { + Error::RefConflict { + message: format!( + "reference {canonical_path} changed during conditional deletion" + ), + } + } else { + error.into() + } + })?; + let current_metadata = current.meta.clone(); + let current_payload = current.bytes().await?; + if current_payload.as_ref() != expected_payload + || current_metadata.e_tag != expected_metadata.e_tag + || current_metadata.version != expected_metadata.version + { + return Err(Error::RefConflict { + message: format!("reference {canonical_path} changed during conditional deletion"), + }); + } -pub(super) async fn lock_local_reference( - object_store: &ObjectStore, - path: &Path, -) -> Result> { - if !object_store.is_local() && object_store.scheme() != "file-object-store" { - return Ok(None); + let lifecycle = store.reference_lifecycle_snapshot(canonical_path).await?; + match lifecycle.as_ref().map(|snapshot| &snapshot.state) { + Some(ReferenceLifecycleState::Pending { .. }) => { + return Err(Error::RefConflict { + message: format!("reference {canonical_path} has an in-flight mutation"), + }); + } + Some(ReferenceLifecycleState::Deleting { .. }) => { + return Err(Error::RefConflict { + message: format!("reference {canonical_path} is already being deleted"), + }); + } + Some(ReferenceLifecycleState::Live { live, .. }) + if payload_reference_generation(¤t_payload)?.as_deref() + != Some(live.generation.as_str()) => + { + return Err(Error::RefConflict { + message: format!("reference {canonical_path} lifecycle changed during deletion"), + }); + } + _ => {} } - let path = PathBuf::from(path.as_ref()); - let path = if path.is_absolute() { - path - } else { - PathBuf::from(std::path::MAIN_SEPARATOR_STR).join(path) + let deleting = ReferenceLifecycleState::Deleting { + canonical_path: canonical_path.to_string(), + operation_id: Uuid::new_v4().simple().to_string(), }; - let lock_path = PathBuf::from(format!("{}.lance-reference.lock", path.display())); - tokio::task::spawn_blocking(move || -> Result { - if let Some(parent) = lock_path.parent() { - std::fs::create_dir_all(parent)?; - } - let file = OpenOptions::new() - .create(true) - .truncate(false) - .read(true) - .write(true) - .open(lock_path)?; - file.lock()?; - Ok(LocalReferenceLock { _file: file }) - }) - .await - .map_err(|error| Error::internal(format!("failed to acquire local reference lock: {error}")))? - .map(Some) + if store + .put_reference_lifecycle( + canonical_path, + lifecycle.as_ref(), + &deleting, + local_lock.is_some(), + ) + .await? + .is_none() + { + return Err(Error::RefConflict { + message: format!("reference {canonical_path} lifecycle changed during deletion"), + }); + } + + object_store.delete(canonical_path).await?; + let state_path = store.reference_state_path(canonical_path); + match object_store.delete(&state_path).await { + Ok(()) => Ok(()), + Err(error) if error.is_not_found() => Ok(()), + Err(error) => Err(error), + } } async fn ensure_reference_available( @@ -1758,6 +2609,38 @@ pub(super) async fn remove_branch_state( root_path: &Path, namespace: &str, ) -> Result<()> { + begin_branch_state_removal(object_store, root_path, namespace).await?; + let lifecycle_store = + VersionLeaseStore::lifecycle_only(Arc::new(object_store.clone()), root_path.clone()); + let intent_prefix = root_path + .clone() + .join(REFERENCE_INTENTS_DIR) + .join(namespace.to_string()); + let intent_metadata = object_store + .list(Some(intent_prefix)) + .try_collect::>() + .await?; + for metadata in &intent_metadata { + let result = match object_store.inner.get(&metadata.location).await { + Ok(result) => result, + Err(object_store::Error::NotFound { .. }) => continue, + Err(error) => return Err(error.into()), + }; + let intent: ReferenceIntent = serde_json::from_slice(&result.bytes().await?) + .map_err(|error| Error::corrupt_file(metadata.location.clone(), error.to_string()))?; + if intent.operation_id.is_empty() { + continue; + } + let canonical_path = match &intent.mutation { + ReferenceMutation::Create { path, .. } | ReferenceMutation::Update { path, .. } => { + Path::parse(path)? + } + }; + let local_lock = lock_local_reference(object_store, &canonical_path).await?; + lifecycle_store + .cancel_reference_lifecycle(&canonical_path, &intent.operation_id, local_lock.is_some()) + .await?; + } for path in [ root_path .clone() @@ -1778,7 +2661,49 @@ pub(super) async fn remove_branch_state( Err(error) => return Err(error), } } - Ok(()) + cancel_branch_state_removal(object_store, root_path, namespace).await +} + +pub(super) async fn begin_branch_state_removal( + object_store: &ObjectStore, + root_path: &Path, + namespace: &str, +) -> Result<()> { + let path = root_path + .clone() + .join(BRANCH_TERMINATIONS_DIR) + .join(format!("{namespace}.deleted")); + match object_store + .inner + .put_opts( + &path, + Bytes::new().into(), + PutOptions { + mode: PutMode::Create, + ..Default::default() + }, + ) + .await + { + Ok(_) | Err(object_store::Error::AlreadyExists { .. }) => Ok(()), + Err(error) => Err(error.into()), + } +} + +pub(super) async fn cancel_branch_state_removal( + object_store: &ObjectStore, + root_path: &Path, + namespace: &str, +) -> Result<()> { + let path = root_path + .clone() + .join(BRANCH_TERMINATIONS_DIR) + .join(format!("{namespace}.deleted")); + match object_store.delete(&path).await { + Ok(()) => Ok(()), + Err(error) if error.is_not_found() => Ok(()), + Err(error) => Err(error), + } } impl Dataset { @@ -1937,6 +2862,7 @@ fn retiring_version_error(version: u64) -> Error { #[cfg(test)] mod tests { + use crate::dataset::refs::{BranchContents, TagContents}; use crate::utils::test::FailingProxyStore; use lance_io::object_store::WrappingObjectStore; use mock_instant::thread_local::MockClock; @@ -1946,10 +2872,13 @@ mod tests { fn memory_store() -> VersionLeaseStore { VersionLeaseStore { object_store: Arc::new(ObjectStore::memory()), + root_path: Path::from(""), + namespace: MAIN_BRANCH.to_string(), leases_path: Path::from("leases"), markers_path: Path::from("markers"), reference_intents_path: Path::from("reference_intents"), manifest_path: None, + canonical_references: None, } } @@ -1967,7 +2896,45 @@ mod tests { path: format!("tags/version-{version}.json"), payload: version.to_string().into_bytes(), }, + operation_id: String::new(), + state: ReferenceIntentState::Pending, + } + } + + async fn create_test_reference_admission( + store: &VersionLeaseStore, + version: u64, + manifest_path: &Path, + canonical_path: &Path, + payload: &[u8], + ) -> ReferenceAdmission { + let operation_id = Uuid::new_v4().simple().to_string(); + let mutation = ReferenceMutation::Create { + path: canonical_path.to_string(), + payload: set_payload_reference_generation(payload, &operation_id).unwrap(), + }; + let intent = ReferenceIntent { + manifest_path: manifest_path.to_string(), + mutation: mutation.clone(), + operation_id: operation_id.clone(), state: ReferenceIntentState::Pending, + }; + let (path, created_at) = store + .create_reference_intent(version, &intent) + .await + .unwrap(); + store + .claim_reference_lifecycle(canonical_path, &operation_id, version, false) + .await + .unwrap(); + ReferenceAdmission { + store: store.clone(), + path, + manifest_path: manifest_path.clone(), + version, + created_at, + operation_id, + mutation, } } @@ -1981,6 +2948,23 @@ mod tests { assert_send(store.fence_versions(&manifests)); } + #[test] + fn generated_reference_payloads_remain_readable_by_released_clients() { + let tag = set_payload_reference_generation( + br#"{"branch":null,"version":42,"manifestSize":0,"metadata":{}}"#, + "generation", + ) + .unwrap(); + let branch = set_payload_reference_generation( + br#"{"parentBranch":null,"identifier":{"version_mapping":[]},"parentVersion":42,"createAt":0,"manifestSize":0,"metadata":{}}"#, + "generation", + ) + .unwrap(); + + assert!(serde_json::from_slice::(&tag).is_ok()); + assert!(serde_json::from_slice::(&branch).is_ok()); + } + #[test] fn parses_lease_file_name() { let lease_id = Uuid::nil().simple(); @@ -2047,7 +3031,7 @@ mod tests { guard.seal_versions(&HashSet::from([42])).await.unwrap(); assert!( guard - .commit_versions(&HashSet::from([42]), &HashSet::new()) + .commit_versions(&HashSet::from([42]), &HashSet::new(), &HashMap::new()) .await .unwrap() .is_empty() @@ -2076,7 +3060,7 @@ mod tests { guard.seal_versions(&HashSet::from([42])).await.unwrap(); let retained_versions = guard - .commit_versions(&HashSet::from([42]), &HashSet::new()) + .commit_versions(&HashSet::from([42]), &HashSet::new(), &HashMap::new()) .await .unwrap(); @@ -2094,10 +3078,13 @@ mod tests { object_store.inner = failing_store.wrap("memory", Arc::clone(&object_store.inner)); let store = VersionLeaseStore { object_store: Arc::new(object_store), + root_path: Path::from(""), + namespace: MAIN_BRANCH.to_string(), leases_path: Path::from("leases"), markers_path: Path::from("markers"), reference_intents_path: Path::from("reference_intents"), manifest_path: None, + canonical_references: None, }; let lease = store.acquire(42, Duration::from_secs(60)).await.unwrap(); let manifests = manifest_paths(42); @@ -2105,7 +3092,7 @@ mod tests { guard.seal_versions(&HashSet::from([42])).await.unwrap(); assert!( guard - .commit_versions(&HashSet::from([42]), &HashSet::new()) + .commit_versions(&HashSet::from([42]), &HashSet::new(), &HashMap::new()) .await .unwrap() .is_empty() @@ -2274,38 +3261,47 @@ mod tests { let manifest_path = Path::from("manifests/42.manifest"); let canonical_path = Path::from("tags/racing.json"); store.object_store.put(&manifest_path, &[]).await.unwrap(); - let intent = ReferenceIntent { - manifest_path: manifest_path.to_string(), - mutation: ReferenceMutation::Create { - path: canonical_path.to_string(), - payload: b"reference".to_vec(), - }, - state: ReferenceIntentState::Pending, - }; - let (intent_path, created_at) = store.create_reference_intent(42, &intent).await.unwrap(); - let admission = ReferenceAdmission { - store: store.clone(), - path: intent_path.clone(), - manifest_path: manifest_path.clone(), - version: 42, - created_at, - mutation: intent.mutation, - }; + let admission = create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + br#"{"version":42}"#, + ) + .await; + let intent_path = admission.path.clone(); admission.ensure_owned().await.unwrap(); let manifests = HashMap::from([(42, vec![manifest_path])]); let mut guard = store.fence_versions(&manifests).await.unwrap(); guard.seal_versions(&HashSet::from([42])).await.unwrap(); - store + assert_eq!( + store + .apply_reference_mutation_inner( + &admission.mutation, + false, + Some(&admission.operation_id), + ) + .await + .unwrap(), + ReferenceMutationOutcome::Published + ); + assert!( + store + .complete_reference_lifecycle(&canonical_path, &admission.operation_id, 42, false,) + .await + .unwrap() + ); + admission + .store .object_store - .put(&canonical_path, b"reference") + .delete(&admission.path) .await .unwrap(); - admission.complete().await; assert_eq!( guard - .commit_versions(&HashSet::from([42]), &HashSet::new()) + .commit_versions(&HashSet::from([42]), &HashSet::new(), &HashMap::new()) .await .unwrap(), HashSet::from([42]) @@ -2320,30 +3316,28 @@ mod tests { let manifest_path = Path::from("manifests/42.manifest"); let canonical_path = Path::from("tags/removed.json"); store.object_store.put(&manifest_path, &[]).await.unwrap(); - let intent = ReferenceIntent { - manifest_path: manifest_path.to_string(), - mutation: ReferenceMutation::Create { - path: canonical_path.to_string(), - payload: b"reference".to_vec(), - }, - state: ReferenceIntentState::Pending, - }; - let (intent_path, created_at) = store.create_reference_intent(42, &intent).await.unwrap(); - let admission = ReferenceAdmission { - store: store.clone(), - path: intent_path.clone(), - manifest_path, - version: 42, - created_at, - mutation: intent.mutation, - }; - store - .object_store - .put(&canonical_path, b"reference") - .await - .unwrap(); - admission.complete().await; - store.object_store.delete(&canonical_path).await.unwrap(); + let admission = create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + br#"{"version":42}"#, + ) + .await; + let intent_path = admission.path.clone(); + admission.publish("conflict".to_string()).await.unwrap(); + let result = store.object_store.inner.get(&canonical_path).await.unwrap(); + let metadata = result.meta.clone(); + let payload = result.bytes().await.unwrap(); + delete_canonical_reference( + Arc::clone(&store.object_store), + &store.root_path, + &canonical_path, + &metadata, + &payload, + ) + .await + .unwrap(); assert!(store.active_reference_versions().await.unwrap().is_empty()); assert!(!store.object_store.exists(&intent_path).await.unwrap()); @@ -2355,29 +3349,16 @@ mod tests { let manifest_path = Path::from("manifests/42.manifest"); let canonical_path = Path::from("branches/child.json"); store.object_store.put(&manifest_path, &[]).await.unwrap(); - let intent = ReferenceIntent { - manifest_path: manifest_path.to_string(), - mutation: ReferenceMutation::Create { - path: canonical_path.to_string(), - payload: b"reference".to_vec(), - }, - state: ReferenceIntentState::Pending, - }; - let (intent_path, created_at) = store.create_reference_intent(42, &intent).await.unwrap(); - let admission = ReferenceAdmission { - store: store.clone(), - path: intent_path.clone(), - manifest_path, - version: 42, - created_at, - mutation: intent.mutation, - }; - store - .object_store - .put(&canonical_path, b"reference") - .await - .unwrap(); - admission.complete().await; + let admission = create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + br#"{"parentVersion":42}"#, + ) + .await; + let intent_path = admission.path.clone(); + admission.publish("conflict".to_string()).await.unwrap(); let census = store .reference_versions_before_canonical_census() @@ -2410,7 +3391,7 @@ mod tests { assert_eq!( store - .apply_reference_mutation_inner(&mutation, false) + .apply_reference_mutation_inner(&mutation, false, None) .await .unwrap(), ReferenceMutationOutcome::Conflict @@ -2421,21 +3402,90 @@ mod tests { #[tokio::test] async fn expired_pending_create_intent_does_not_resurrect_deleted_reference() { let store = memory_store(); + let manifest_path = Path::from("manifests/42.manifest"); let canonical_path = Path::from("tags/deleted.json"); - let mutation = ReferenceMutation::Create { + store.object_store.put(&manifest_path, &[]).await.unwrap(); + let admission = create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + br#"{"version":42}"#, + ) + .await; + assert_eq!( + store + .apply_reference_mutation_inner( + &admission.mutation, + false, + Some(&admission.operation_id), + ) + .await + .unwrap(), + ReferenceMutationOutcome::Published + ); + store.object_store.delete(&canonical_path).await.unwrap(); + + assert_eq!( + store + .expired_reference_mutation_outcome( + &admission.mutation, + &admission.operation_id, + 42, + false, + ) + .await + .unwrap(), + ReferenceMutationOutcome::Conflict + ); + assert_eq!( + store + .apply_reference_mutation_inner( + &admission.mutation, + false, + Some(&admission.operation_id), + ) + .await + .unwrap(), + ReferenceMutationOutcome::Conflict + ); + assert!(!store.object_store.exists(&canonical_path).await.unwrap()); + } + + #[tokio::test] + async fn expired_update_owner_cannot_publish_after_retirement_commit() { + let store = memory_store(); + let canonical_path = Path::from("tags/updated.json"); + let original_payload = br#"{"version":1}"#; + let original = store + .object_store + .inner + .put(&canonical_path, Bytes::from_static(original_payload).into()) + .await + .unwrap(); + let operation_id = Uuid::new_v4().simple().to_string(); + let mutation = ReferenceMutation::Update { path: canonical_path.to_string(), - payload: b"reference".to_vec(), + expected_payload: original_payload.to_vec(), + expected_etag: original.e_tag, + expected_version: original.version, + payload: set_payload_reference_generation(br#"{"version":42}"#, &operation_id).unwrap(), }; store - .object_store - .put(&canonical_path, b"reference") + .claim_reference_lifecycle(&canonical_path, &operation_id, 42, true) .await .unwrap(); - store.object_store.delete(&canonical_path).await.unwrap(); assert_eq!( store - .expired_reference_mutation_outcome(&mutation) + .expired_reference_mutation_outcome(&mutation, &operation_id, 42, false) + .await + .unwrap(), + ReferenceMutationOutcome::Conflict + ); + assert_eq!( + store + .apply_reference_mutation_inner(&mutation, false, Some(&operation_id)) .await .unwrap(), ReferenceMutationOutcome::Conflict @@ -2451,7 +3501,7 @@ mod tests { .await .unwrap() .as_ref(), - REFERENCE_TOMBSTONE_PAYLOAD + original_payload ); } @@ -2462,10 +3512,13 @@ mod tests { let (object_store, base_path) = ObjectStore::from_uri(&uri).await.unwrap(); let store = VersionLeaseStore { object_store, + root_path: base_path.clone(), + namespace: MAIN_BRANCH.to_string(), leases_path: base_path.clone().join("leases"), markers_path: base_path.clone().join("markers"), reference_intents_path: base_path.clone().join("reference_intents"), manifest_path: None, + canonical_references: None, }; let canonical_path = base_path.join("branches/child.json"); let original = store @@ -2483,7 +3536,7 @@ mod tests { }; let error = store - .apply_reference_mutation_inner(&mutation, false) + .apply_reference_mutation_inner(&mutation, false, None) .await .unwrap_err(); @@ -2510,43 +3563,27 @@ mod tests { let manifest_path = Path::from("manifests/42.manifest"); let canonical_path = Path::from("tags/recovery.json"); store.object_store.put(&manifest_path, &[]).await.unwrap(); - let intent = ReferenceIntent { - manifest_path: manifest_path.to_string(), - mutation: ReferenceMutation::Create { - path: canonical_path.to_string(), - payload: b"reference".to_vec(), - }, - state: ReferenceIntentState::Pending, - }; - let (intent_path, created_at) = store.create_reference_intent(42, &intent).await.unwrap(); - let admission = ReferenceAdmission { - store: store.clone(), - path: intent_path.clone(), - manifest_path: manifest_path.clone(), - version: 42, - created_at, - mutation: intent.mutation, - }; - store - .object_store - .put(&canonical_path, b"reference") + let admission = create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + br#"{"version":42}"#, + ) + .await; + let intent_path = admission.path.clone(); + admission.publish("conflict".to_string()).await.unwrap(); + assert!(!store.object_store.exists(&intent_path).await.unwrap()); + let pre_seal_census = store + .reference_versions_before_canonical_census() .await .unwrap(); - admission.complete().await; + assert!(pre_seal_census.versions.is_empty()); + assert!(pre_seal_census.completed_intent_paths.is_empty()); let manifests = HashMap::from([(42, vec![manifest_path])]); let mut guard = store.fence_versions(&manifests).await.unwrap(); guard.seal_versions(&HashSet::from([42])).await.unwrap(); - let census = store - .reference_versions_before_canonical_census() - .await - .unwrap(); - assert!(census.versions.is_empty()); - assert_eq!( - census.completed_intent_paths, - HashSet::from([intent_path.clone()]) - ); - assert!(store.object_store.exists(&intent_path).await.unwrap()); drop(guard); assert!( @@ -2558,7 +3595,6 @@ mod tests { .is_empty() ); assert!(store.version_marker_metadata(42).await.unwrap().is_empty()); - assert!(!store.object_store.exists(&intent_path).await.unwrap()); store.acquire(42, Duration::from_secs(60)).await.unwrap(); } @@ -2583,4 +3619,49 @@ mod tests { assert!(!store.object_store.exists(&lease_path).await.unwrap()); assert!(!store.object_store.exists(&marker_path).await.unwrap()); } + + #[tokio::test] + async fn deleting_parent_branch_state_fences_admitted_child_publication() { + let root = Path::from("dataset"); + let namespace = "branch-id"; + let store = VersionLeaseStore { + object_store: Arc::new(ObjectStore::memory()), + root_path: root.clone(), + namespace: namespace.to_string(), + leases_path: root.clone().join(LEASES_DIR).join(namespace), + markers_path: root.clone().join(LEASE_GC_MARKERS_DIR).join(namespace), + reference_intents_path: root.clone().join(REFERENCE_INTENTS_DIR).join(namespace), + manifest_path: None, + canonical_references: None, + }; + let manifest_path = Path::from("dataset/branches/parent/versions/42.manifest"); + let canonical_path = Path::from("dataset/_refs/branches/child.json"); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + let admission = create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + br#"{"parentVersion":42}"#, + ) + .await; + admission.ensure_owned().await.unwrap(); + + remove_branch_state(&store.object_store, &root, namespace) + .await + .unwrap(); + + assert_eq!( + store + .apply_reference_mutation_inner( + &admission.mutation, + false, + Some(&admission.operation_id), + ) + .await + .unwrap(), + ReferenceMutationOutcome::Conflict + ); + assert!(!store.object_store.exists(&canonical_path).await.unwrap()); + } } From 9097168fab1653db0b8e59fa8b62ec94355a19e2 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Sat, 8 Aug 2026 01:56:43 +0000 Subject: [PATCH 10/13] fix: reconcile partial reference transitions --- rust/lance/src/dataset/version_lease.rs | 613 ++++++++++++++++++++---- 1 file changed, 525 insertions(+), 88 deletions(-) diff --git a/rust/lance/src/dataset/version_lease.rs b/rust/lance/src/dataset/version_lease.rs index ad5209e05ca..7d237c9b760 100644 --- a/rust/lance/src/dataset/version_lease.rs +++ b/rust/lance/src/dataset/version_lease.rs @@ -247,9 +247,14 @@ enum ReferenceLifecycleState { canonical_path: String, live: ReferenceLiveState, }, + Legacy { + canonical_path: String, + }, Deleting { canonical_path: String, operation_id: String, + previous: Option, + previous_was_legacy: bool, }, } @@ -444,6 +449,7 @@ impl VersionLeaseStore { let recorded_path = match &state { ReferenceLifecycleState::Pending { canonical_path, .. } | ReferenceLifecycleState::Live { canonical_path, .. } + | ReferenceLifecycleState::Legacy { canonical_path } | ReferenceLifecycleState::Deleting { canonical_path, .. } => canonical_path, }; if recorded_path != canonical_path.as_ref() { @@ -542,14 +548,24 @@ impl VersionLeaseStore { ) -> Result<()> { let local_lock = lock_local_reference(&self.object_store, canonical_path).await?; let current = self.reference_lifecycle_snapshot(canonical_path).await?; - let previous = match current.as_ref().map(|snapshot| &snapshot.state) { - None => None, - Some(ReferenceLifecycleState::Live { live, .. }) => Some(live.clone()), + let (previous, previous_was_legacy) = match current.as_ref().map(|snapshot| &snapshot.state) + { + None => (None, previous_was_legacy), + Some(ReferenceLifecycleState::Live { live, .. }) => (Some(live.clone()), false), + Some(ReferenceLifecycleState::Legacy { .. }) => (None, true), Some(ReferenceLifecycleState::Pending { .. }) => { return Err(Error::RefConflict { message: format!("reference {canonical_path} has an in-flight mutation"), }); } + Some(ReferenceLifecycleState::Deleting { .. }) + if !self.object_store.exists(canonical_path).await? => + { + // A crash after canonical deletion can leave the sidecar in + // Deleting. Reusing that generation directly avoids requiring + // an unsafe delete/recreate race on the lifecycle object. + (None, false) + } Some(ReferenceLifecycleState::Deleting { .. }) => { return Err(Error::RefConflict { message: format!("reference {canonical_path} is being deleted"), @@ -645,6 +661,21 @@ impl VersionLeaseStore { .await? .is_some()); } + if matches!( + ¤t.state, + ReferenceLifecycleState::Pending { + previous_was_legacy: true, + .. + } + ) { + let legacy = ReferenceLifecycleState::Legacy { + canonical_path: canonical_path.to_string(), + }; + return Ok(self + .put_reference_lifecycle(canonical_path, Some(¤t), &legacy, has_local_lock) + .await? + .is_some()); + } let state_path = self.reference_state_path(canonical_path); match self.object_store.delete(&state_path).await { Ok(()) => Ok(true), @@ -769,7 +800,10 @@ impl VersionLeaseStore { let target = match state { ReferenceLifecycleState::Pending { target, .. } => Some(target), ReferenceLifecycleState::Live { live, .. } => Some(live.target), - ReferenceLifecycleState::Deleting { .. } => None, + ReferenceLifecycleState::Deleting { previous, .. } => { + previous.map(|live| live.target) + } + ReferenceLifecycleState::Legacy { .. } => None, }; Ok::<_, Error>(target.and_then(|target| { (target.namespace == self.namespace).then_some(target.version) @@ -804,7 +838,8 @@ impl VersionLeaseStore { let generation = match state { ReferenceLifecycleState::Pending { operation_id, .. } => Some(operation_id), ReferenceLifecycleState::Live { live, .. } => Some(live.generation), - ReferenceLifecycleState::Deleting { .. } => None, + ReferenceLifecycleState::Deleting { operation_id, .. } => Some(operation_id), + ReferenceLifecycleState::Legacy { .. } => Some("legacy".to_string()), }; Ok::<_, Error>(generation.map(|generation| (metadata.location, generation))) }) @@ -846,7 +881,12 @@ impl VersionLeaseStore { ReferenceLifecycleState::Live { live, .. } => { (live.generation, Some(live.target)) } - ReferenceLifecycleState::Deleting { operation_id, .. } => (operation_id, None), + ReferenceLifecycleState::Deleting { + operation_id, + previous, + .. + } => (operation_id, previous.map(|live| live.target)), + ReferenceLifecycleState::Legacy { .. } => ("legacy".to_string(), None), }; let changed = observed_generations.get(&metadata.location) != Some(&generation); Ok::<_, Error>(target.and_then(|target| { @@ -1041,11 +1081,12 @@ impl VersionLeaseStore { }, Some((metadata, current_payload)), ) if current_payload.as_ref() == expected_payload => { + let fenced_payload = fenced_reference_payload(¤t_payload); self.rewrite_reference_payload( &path, &metadata, ¤t_payload, - ¤t_payload, + &fenced_payload, local_lock.is_some(), ) .await?; @@ -1163,7 +1204,7 @@ impl VersionLeaseStore { has_local_lock: bool, expected_operation_id: Option<&str>, ) -> Result { - if let Some(expected_operation_id) = expected_operation_id { + let pending_version = if let Some(expected_operation_id) = expected_operation_id { self.ensure_branch_incarnation_active().await?; let canonical_path = match mutation { ReferenceMutation::Create { path, .. } | ReferenceMutation::Update { path, .. } => { @@ -1171,16 +1212,17 @@ impl VersionLeaseStore { } }; let lifecycle = self.reference_lifecycle_snapshot(&canonical_path).await?; - if !matches!( - lifecycle.as_ref().map(|snapshot| &snapshot.state), + match lifecycle.as_ref().map(|snapshot| &snapshot.state) { Some(ReferenceLifecycleState::Pending { operation_id, + target, .. - }) if operation_id == expected_operation_id - ) { - return Ok(ReferenceMutationOutcome::Conflict); + }) if operation_id == expected_operation_id => Some(target.version), + _ => return Ok(ReferenceMutationOutcome::Conflict), } - } + } else { + None + }; let (path, payload, result) = match mutation { ReferenceMutation::Create { path, payload } => { let path = Path::parse(path)?; @@ -1250,18 +1292,117 @@ impl VersionLeaseStore { } }; - match result { - Ok(_) => Ok(ReferenceMutationOutcome::Published), + let outcome = match result { + Ok(_) => ReferenceMutationOutcome::Published, Err( object_store::Error::AlreadyExists { .. } | object_store::Error::Precondition { .. }, ) => { self.reference_mutation_conflict_outcome(&path, payload) - .await + .await? } - Err(object_store::Error::NotFound { .. }) => Ok(ReferenceMutationOutcome::Conflict), - Err(error) => Err(error.into()), + Err(object_store::Error::NotFound { .. }) => ReferenceMutationOutcome::Conflict, + Err(error) => return Err(error.into()), + }; + if outcome == ReferenceMutationOutcome::Published + && let (Some(operation_id), Some(version)) = (expected_operation_id, pending_version) + { + self.finish_reference_mutation(mutation, operation_id, version, has_local_lock) + .await + } else { + Ok(outcome) + } + } + + async fn finish_reference_mutation( + &self, + mutation: &ReferenceMutation, + operation_id: &str, + version: u64, + has_local_lock: bool, + ) -> Result { + let canonical_path = match mutation { + ReferenceMutation::Create { path, .. } | ReferenceMutation::Update { path, .. } => { + Path::parse(path)? + } + }; + let branch_is_active = match self.ensure_branch_incarnation_active().await { + Ok(()) => true, + Err(Error::RefConflict { .. }) => false, + Err(error) => return Err(error), + }; + let lifecycle_completed = branch_is_active + && self + .complete_reference_lifecycle( + &canonical_path, + operation_id, + version, + has_local_lock, + ) + .await?; + if lifecycle_completed { + return Ok(ReferenceMutationOutcome::Published); } + + self.rollback_lost_reference_mutation(mutation, has_local_lock) + .await?; + self.cancel_reference_lifecycle(&canonical_path, operation_id, has_local_lock) + .await?; + Ok(ReferenceMutationOutcome::Conflict) + } + + async fn rollback_lost_reference_mutation( + &self, + mutation: &ReferenceMutation, + has_local_lock: bool, + ) -> Result<()> { + let (path, intended_payload) = match mutation { + ReferenceMutation::Create { path, payload } + | ReferenceMutation::Update { path, payload, .. } => { + (Path::parse(path)?, payload.as_slice()) + } + }; + let result = match self.object_store.inner.get(&path).await { + Ok(result) => result, + Err(object_store::Error::NotFound { .. }) => return Ok(()), + Err(error) => return Err(error.into()), + }; + let metadata = result.meta.clone(); + let current_payload = result.bytes().await?; + if current_payload.as_ref() != intended_payload { + return Ok(()); + } + + let replacement = match mutation { + ReferenceMutation::Create { payload, .. } => { + set_payload_reference_generation(payload, &Uuid::new_v4().simple().to_string())? + } + ReferenceMutation::Update { + expected_payload, .. + } => expected_payload.clone(), + }; + match self + .rewrite_reference_payload( + &path, + &metadata, + ¤t_payload, + &replacement, + has_local_lock, + ) + .await + { + Ok(()) => {} + Err(Error::RefConflict { .. }) => return Ok(()), + Err(error) => return Err(error), + } + if matches!(mutation, ReferenceMutation::Create { .. }) { + match self.object_store.delete(&path).await { + Ok(()) => {} + Err(error) if error.is_not_found() => {} + Err(error) => return Err(error), + } + } + Ok(()) } async fn reference_mutation_conflict_outcome( @@ -2305,24 +2446,8 @@ impl ReferenceAdmission { .await { Ok(ReferenceMutationOutcome::Published) => { - if self - .store - .complete_reference_lifecycle( - &canonical_path, - &self.operation_id, - self.version, - local_lock.is_some(), - ) - .await? - { - self.complete().await; - Ok(()) - } else { - self.cancel_before_publish_inner(local_lock.is_some()).await; - Err(Error::RefConflict { - message: conflict_message, - }) - } + self.complete().await; + Ok(()) } Ok(ReferenceMutationOutcome::Conflict) => { self.cancel_before_publish_inner(local_lock.is_some()).await; @@ -2471,6 +2596,13 @@ fn set_payload_reference_generation(payload: &[u8], generation: &str) -> Result< Ok(serde_json::to_vec_pretty(&value)?) } +fn fenced_reference_payload(payload: &[u8]) -> Vec { + let mut fenced = Vec::with_capacity(payload.len() + 1); + fenced.extend_from_slice(payload); + fenced.push(b'\n'); + fenced +} + pub(super) async fn canonical_reference_is_visible( object_store: Arc, root_path: &Path, @@ -2478,23 +2610,27 @@ pub(super) async fn canonical_reference_is_visible( payload: &[u8], ) -> Result { let generation = payload_reference_generation(payload)?; + // Released writers do not preserve fields they do not know about. Their + // rewrites remain authoritative canonical references and must stay visible + // while the sidecar is reconciled by a later current-client mutation. + if generation.is_none() { + return Ok(true); + } let store = VersionLeaseStore::lifecycle_only(object_store, root_path.clone()); let Some(snapshot) = store.reference_lifecycle_snapshot(canonical_path).await? else { - return Ok(generation.is_none()); + return Ok(false); }; Ok(match snapshot.state { ReferenceLifecycleState::Live { live, .. } => { generation.as_deref() == Some(live.generation.as_str()) } - ReferenceLifecycleState::Pending { - previous, - previous_was_legacy, - .. - } => match previous { + ReferenceLifecycleState::Legacy { .. } => false, + ReferenceLifecycleState::Pending { previous, .. } => match previous { Some(previous) => generation.as_deref() == Some(previous.generation.as_str()), - None => previous_was_legacy && generation.is_none(), + None => false, }, - ReferenceLifecycleState::Deleting { .. } => false, + ReferenceLifecycleState::Deleting { previous, .. } => previous + .is_some_and(|previous| generation.as_deref() == Some(previous.generation.as_str())), }) } @@ -2534,21 +2670,22 @@ pub(super) async fn delete_canonical_reference( } let lifecycle = store.reference_lifecycle_snapshot(canonical_path).await?; + let canonical_generation = payload_reference_generation(¤t_payload)?; match lifecycle.as_ref().map(|snapshot| &snapshot.state) { Some(ReferenceLifecycleState::Pending { .. }) => { return Err(Error::RefConflict { message: format!("reference {canonical_path} has an in-flight mutation"), }); } - Some(ReferenceLifecycleState::Deleting { .. }) => { + Some(ReferenceLifecycleState::Live { live, .. }) + if canonical_generation.as_deref().is_some() + && canonical_generation.as_deref() != Some(live.generation.as_str()) => + { return Err(Error::RefConflict { - message: format!("reference {canonical_path} is already being deleted"), + message: format!("reference {canonical_path} lifecycle changed during deletion"), }); } - Some(ReferenceLifecycleState::Live { live, .. }) - if payload_reference_generation(¤t_payload)?.as_deref() - != Some(live.generation.as_str()) => - { + Some(ReferenceLifecycleState::Legacy { .. }) if canonical_generation.is_some() => { return Err(Error::RefConflict { message: format!("reference {canonical_path} lifecycle changed during deletion"), }); @@ -2556,32 +2693,114 @@ pub(super) async fn delete_canonical_reference( _ => {} } - let deleting = ReferenceLifecycleState::Deleting { - canonical_path: canonical_path.to_string(), - operation_id: Uuid::new_v4().simple().to_string(), + let is_retry = lifecycle + .as_ref() + .is_some_and(|snapshot| matches!(snapshot.state, ReferenceLifecycleState::Deleting { .. })); + let deleting_snapshot = if is_retry { + let Some(snapshot) = lifecycle else { + return Err(Error::internal(format!( + "reference {canonical_path} lost its deleting lifecycle state" + ))); + }; + snapshot + } else { + let previous = lifecycle + .as_ref() + .and_then(|snapshot| match &snapshot.state { + ReferenceLifecycleState::Live { live, .. } => Some(live.clone()), + _ => None, + }); + let deleting = ReferenceLifecycleState::Deleting { + canonical_path: canonical_path.to_string(), + operation_id: Uuid::new_v4().simple().to_string(), + previous, + previous_was_legacy: canonical_generation.is_none(), + }; + let Some(snapshot) = store + .put_reference_lifecycle( + canonical_path, + lifecycle.as_ref(), + &deleting, + local_lock.is_some(), + ) + .await? + else { + return Err(Error::RefConflict { + message: format!("reference {canonical_path} lifecycle changed during deletion"), + }); + }; + snapshot + }; + + if let Err(delete_error) = object_store.delete(canonical_path).await { + match object_store.inner.get(canonical_path).await { + Ok(_) => { + restore_reference_lifecycle( + &store, + canonical_path, + &deleting_snapshot, + local_lock.is_some(), + ) + .await?; + return Err(delete_error); + } + Err(object_store::Error::NotFound { .. }) => {} + Err(_) => return Err(delete_error), + } + } + let state_path = store.reference_state_path(canonical_path); + match object_store.delete(&state_path).await { + Ok(()) => Ok(()), + Err(error) if error.is_not_found() => Ok(()), + Err(error) => Err(error), + } +} + +async fn restore_reference_lifecycle( + store: &VersionLeaseStore, + canonical_path: &Path, + deleting_snapshot: &ReferenceLifecycleSnapshot, + has_local_lock: bool, +) -> Result<()> { + let ReferenceLifecycleState::Deleting { + previous, + previous_was_legacy, + .. + } = &deleting_snapshot.state + else { + return Err(Error::internal(format!( + "reference {canonical_path} has non-deleting rollback state" + ))); + }; + let restored = if let Some(previous) = previous { + ReferenceLifecycleState::Live { + canonical_path: canonical_path.to_string(), + live: previous.clone(), + } + } else if *previous_was_legacy { + ReferenceLifecycleState::Legacy { + canonical_path: canonical_path.to_string(), + } + } else { + return Err(Error::internal(format!( + "reference {canonical_path} deletion has no state to restore" + ))); }; if store .put_reference_lifecycle( canonical_path, - lifecycle.as_ref(), - &deleting, - local_lock.is_some(), + Some(deleting_snapshot), + &restored, + has_local_lock, ) .await? .is_none() { return Err(Error::RefConflict { - message: format!("reference {canonical_path} lifecycle changed during deletion"), + message: format!("reference {canonical_path} lifecycle changed during rollback"), }); } - - object_store.delete(canonical_path).await?; - let state_path = store.reference_state_path(canonical_path); - match object_store.delete(&state_path).await { - Ok(()) => Ok(()), - Err(error) if error.is_not_found() => Ok(()), - Err(error) => Err(error), - } + Ok(()) } async fn ensure_reference_available( @@ -2965,6 +3184,105 @@ mod tests { assert!(serde_json::from_slice::(&branch).is_ok()); } + #[tokio::test] + async fn released_client_rewrite_keeps_generated_references_visible() { + let store = memory_store(); + let manifest_path = Path::from("manifests/42.manifest"); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + + let tag_path = Path::from("tags/released.json"); + create_test_reference_admission( + &store, + 42, + &manifest_path, + &tag_path, + br#"{"branch":null,"version":42,"manifestSize":0,"metadata":{}}"#, + ) + .await + .publish("conflict".to_string()) + .await + .unwrap(); + let generated_tag = store + .object_store + .inner + .get(&tag_path) + .await + .unwrap() + .bytes() + .await + .unwrap(); + let released_tag = serde_json::to_vec_pretty( + &serde_json::from_slice::(&generated_tag).unwrap(), + ) + .unwrap(); + assert!( + payload_reference_generation(&released_tag) + .unwrap() + .is_none() + ); + store + .object_store + .put(&tag_path, &released_tag) + .await + .unwrap(); + assert!( + canonical_reference_is_visible( + Arc::clone(&store.object_store), + &store.root_path, + &tag_path, + &released_tag, + ) + .await + .unwrap() + ); + + let branch_path = Path::from("branches/released.json"); + create_test_reference_admission( + &store, + 42, + &manifest_path, + &branch_path, + br#"{"parentBranch":null,"identifier":{"version_mapping":[]},"parentVersion":42,"createAt":0,"manifestSize":0,"metadata":{}}"#, + ) + .await + .publish("conflict".to_string()) + .await + .unwrap(); + let generated_branch = store + .object_store + .inner + .get(&branch_path) + .await + .unwrap() + .bytes() + .await + .unwrap(); + let released_branch = serde_json::to_vec_pretty( + &serde_json::from_slice::(&generated_branch).unwrap(), + ) + .unwrap(); + assert!( + payload_reference_generation(&released_branch) + .unwrap() + .is_none() + ); + store + .object_store + .put(&branch_path, &released_branch) + .await + .unwrap(); + assert!( + canonical_reference_is_visible( + Arc::clone(&store.object_store), + &store.root_path, + &branch_path, + &released_branch, + ) + .await + .unwrap() + ); + } + #[test] fn parses_lease_file_name() { let lease_id = Uuid::nil().simple(); @@ -3286,12 +3604,6 @@ mod tests { .unwrap(), ReferenceMutationOutcome::Published ); - assert!( - store - .complete_reference_lifecycle(&canonical_path, &admission.operation_id, 42, false,) - .await - .unwrap() - ); admission .store .object_store @@ -3415,11 +3727,7 @@ mod tests { .await; assert_eq!( store - .apply_reference_mutation_inner( - &admission.mutation, - false, - Some(&admission.operation_id), - ) + .apply_reference_mutation_inner(&admission.mutation, false, None,) .await .unwrap(), ReferenceMutationOutcome::Published @@ -3490,18 +3798,19 @@ mod tests { .unwrap(), ReferenceMutationOutcome::Conflict ); + let fenced_payload = store + .object_store + .inner + .get(&canonical_path) + .await + .unwrap() + .bytes() + .await + .unwrap(); + assert_ne!(fenced_payload.as_ref(), original_payload); assert_eq!( - store - .object_store - .inner - .get(&canonical_path) - .await - .unwrap() - .bytes() - .await - .unwrap() - .as_ref(), - original_payload + serde_json::from_slice::(&fenced_payload).unwrap(), + serde_json::from_slice::(original_payload).unwrap() ); } @@ -3664,4 +3973,132 @@ mod tests { ); assert!(!store.object_store.exists(&canonical_path).await.unwrap()); } + + #[tokio::test] + async fn checked_child_create_is_rolled_back_after_parent_state_removal() { + let root = Path::from("dataset"); + let namespace = "branch-id"; + let store = VersionLeaseStore { + object_store: Arc::new(ObjectStore::memory()), + root_path: root.clone(), + namespace: namespace.to_string(), + leases_path: root.clone().join(LEASES_DIR).join(namespace), + markers_path: root.clone().join(LEASE_GC_MARKERS_DIR).join(namespace), + reference_intents_path: root.clone().join(REFERENCE_INTENTS_DIR).join(namespace), + manifest_path: None, + canonical_references: None, + }; + let manifest_path = Path::from("dataset/branches/parent/versions/42.manifest"); + let canonical_path = Path::from("dataset/_refs/branches/child.json"); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + let admission = create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + br#"{"parentVersion":42}"#, + ) + .await; + admission.ensure_owned().await.unwrap(); + + remove_branch_state(&store.object_store, &root, namespace) + .await + .unwrap(); + assert_eq!( + store + .apply_reference_mutation_inner(&admission.mutation, false, None) + .await + .unwrap(), + ReferenceMutationOutcome::Published + ); + assert_eq!( + store + .finish_reference_mutation(&admission.mutation, &admission.operation_id, 42, false,) + .await + .unwrap(), + ReferenceMutationOutcome::Conflict + ); + assert!(!store.object_store.exists(&canonical_path).await.unwrap()); + } + + #[tokio::test] + async fn deletion_failure_restores_reference_lifecycle() { + let failing_store = Arc::new(FailingProxyStore::new()); + let mut object_store = ObjectStore::memory(); + object_store.inner = failing_store.wrap("memory", Arc::clone(&object_store.inner)); + let store = VersionLeaseStore { + object_store: Arc::new(object_store), + root_path: Path::from(""), + namespace: MAIN_BRANCH.to_string(), + leases_path: Path::from("leases"), + markers_path: Path::from("markers"), + reference_intents_path: Path::from("reference_intents"), + manifest_path: None, + canonical_references: None, + }; + let manifest_path = Path::from("manifests/42.manifest"); + let canonical_path = Path::from("tags/deletion-failure.json"); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + br#"{"branch":null,"version":42,"manifestSize":0,"metadata":{}}"#, + ) + .await + .publish("conflict".to_string()) + .await + .unwrap(); + let result = store.object_store.inner.get(&canonical_path).await.unwrap(); + let metadata = result.meta.clone(); + let payload = result.bytes().await.unwrap(); + failing_store.fail_when( + "delete", + "deletion-failure.json", + "injected canonical deletion failure", + ); + + let error = delete_canonical_reference( + Arc::clone(&store.object_store), + &store.root_path, + &canonical_path, + &metadata, + &payload, + ) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("injected canonical deletion failure"), + "{error}" + ); + failing_store.clear_fail_when("delete", "deletion-failure.json"); + + let remaining = store.object_store.inner.get(&canonical_path).await.unwrap(); + let remaining_metadata = remaining.meta.clone(); + let remaining_payload = remaining.bytes().await.unwrap(); + assert!( + canonical_reference_is_visible( + Arc::clone(&store.object_store), + &store.root_path, + &canonical_path, + &remaining_payload, + ) + .await + .unwrap(), + "a failed delete must not hide the still-present reference" + ); + delete_canonical_reference( + Arc::clone(&store.object_store), + &store.root_path, + &canonical_path, + &remaining_metadata, + &remaining_payload, + ) + .await + .unwrap(); + assert!(!store.object_store.exists(&canonical_path).await.unwrap()); + } } From fcfa6fcc0fafefc8fa5ff3559cab7dc917cf9ecb Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Sat, 8 Aug 2026 07:19:23 +0000 Subject: [PATCH 11/13] fix: make reference reconciliation crash safe --- rust/lance/src/dataset/version_lease.rs | 1150 ++++++++++++++++++----- 1 file changed, 920 insertions(+), 230 deletions(-) diff --git a/rust/lance/src/dataset/version_lease.rs b/rust/lance/src/dataset/version_lease.rs index 7d237c9b760..82b8601a59d 100644 --- a/rust/lance/src/dataset/version_lease.rs +++ b/rust/lance/src/dataset/version_lease.rs @@ -39,6 +39,11 @@ const COMMITTED_MARKER_SUFFIX: &str = ".committed"; const STORAGE_CLOCK_PATH: &str = "_clock"; const REFERENCE_GENERATION_FIELD: &str = "_lanceReferenceGeneration"; +/// Marks a canonical payload as a logical deletion fence. The generic object +/// store API has conditional writes but no conditional delete, so retaining a +/// byte-changing tombstone is what prevents a later released-client write from +/// being erased by an unconditional physical delete. +const REFERENCE_DELETED_FIELD: &str = "_lanceReferenceDeleted"; /// HTTP `Last-Modified`, used by supported cloud stores, has whole-second precision. /// Adding one interval guarantees a lease is never shortened by timestamp truncation. @@ -250,11 +255,23 @@ enum ReferenceLifecycleState { Legacy { canonical_path: String, }, + Vacant { + canonical_path: String, + }, + Revoking { + canonical_path: String, + operation_id: String, + target: ReferenceTarget, + previous: Option, + previous_was_legacy: bool, + mutation: ReferenceMutation, + }, Deleting { canonical_path: String, operation_id: String, previous: Option, previous_was_legacy: bool, + expected_payload: Vec, }, } @@ -450,6 +467,8 @@ impl VersionLeaseStore { ReferenceLifecycleState::Pending { canonical_path, .. } | ReferenceLifecycleState::Live { canonical_path, .. } | ReferenceLifecycleState::Legacy { canonical_path } + | ReferenceLifecycleState::Vacant { canonical_path } + | ReferenceLifecycleState::Revoking { canonical_path, .. } | ReferenceLifecycleState::Deleting { canonical_path, .. } => canonical_path, }; if recorded_path != canonical_path.as_ref() { @@ -547,56 +566,75 @@ impl VersionLeaseStore { previous_was_legacy: bool, ) -> Result<()> { let local_lock = lock_local_reference(&self.object_store, canonical_path).await?; - let current = self.reference_lifecycle_snapshot(canonical_path).await?; - let (previous, previous_was_legacy) = match current.as_ref().map(|snapshot| &snapshot.state) - { - None => (None, previous_was_legacy), - Some(ReferenceLifecycleState::Live { live, .. }) => (Some(live.clone()), false), - Some(ReferenceLifecycleState::Legacy { .. }) => (None, true), - Some(ReferenceLifecycleState::Pending { .. }) => { - return Err(Error::RefConflict { - message: format!("reference {canonical_path} has an in-flight mutation"), - }); - } - Some(ReferenceLifecycleState::Deleting { .. }) - if !self.object_store.exists(canonical_path).await? => + loop { + let current = self.reference_lifecycle_snapshot(canonical_path).await?; + let (previous, previous_was_legacy) = match current + .as_ref() + .map(|snapshot| &snapshot.state) { - // A crash after canonical deletion can leave the sidecar in - // Deleting. Reusing that generation directly avoids requiring - // an unsafe delete/recreate race on the lifecycle object. - (None, false) - } - Some(ReferenceLifecycleState::Deleting { .. }) => { - return Err(Error::RefConflict { - message: format!("reference {canonical_path} is being deleted"), - }); + None | Some(ReferenceLifecycleState::Vacant { .. }) => (None, previous_was_legacy), + Some(ReferenceLifecycleState::Live { live, .. }) => (Some(live.clone()), false), + Some(ReferenceLifecycleState::Legacy { .. }) => (None, true), + Some(ReferenceLifecycleState::Pending { .. }) => { + return Err(Error::RefConflict { + message: format!("reference {canonical_path} has an in-flight mutation"), + }); + } + Some(ReferenceLifecycleState::Revoking { .. }) => { + let current = current.as_ref().ok_or_else(|| { + Error::internal(format!( + "reference {canonical_path} lost its revoking state" + )) + })?; + self.reconcile_revoking_reference( + canonical_path, + current, + local_lock.is_some(), + ) + .await?; + continue; + } + Some(ReferenceLifecycleState::Deleting { .. }) => { + let current = current.as_ref().ok_or_else(|| { + Error::internal(format!( + "reference {canonical_path} lost its deleting state" + )) + })?; + self.reconcile_deleting_reference( + canonical_path, + current, + local_lock.is_some(), + ) + .await?; + continue; + } + }; + let pending = ReferenceLifecycleState::Pending { + canonical_path: canonical_path.to_string(), + operation_id: operation_id.to_string(), + target: ReferenceTarget { + namespace: self.namespace.clone(), + version, + }, + previous, + previous_was_legacy, + }; + if self + .put_reference_lifecycle( + canonical_path, + current.as_ref(), + &pending, + local_lock.is_some(), + ) + .await? + .is_some() + { + return Ok(()); } - }; - let pending = ReferenceLifecycleState::Pending { - canonical_path: canonical_path.to_string(), - operation_id: operation_id.to_string(), - target: ReferenceTarget { - namespace: self.namespace.clone(), - version, - }, - previous, - previous_was_legacy, - }; - if self - .put_reference_lifecycle( - canonical_path, - current.as_ref(), - &pending, - local_lock.is_some(), - ) - .await? - .is_none() - { return Err(Error::RefConflict { message: format!("reference {canonical_path} lifecycle changed during admission"), }); } - Ok(()) } async fn complete_reference_lifecycle( @@ -618,6 +656,16 @@ impl VersionLeaseStore { ) { return Ok(false); } + let canonical = match self.object_store.inner.get(canonical_path).await { + Ok(result) => result.bytes().await?, + Err(object_store::Error::NotFound { .. }) => return Ok(false), + Err(error) => return Err(error.into()), + }; + if payload_reference_generation(&canonical)?.as_deref() != Some(operation_id) + || payload_reference_is_deleted(&canonical)? + { + return Ok(false); + } let live = ReferenceLifecycleState::Live { canonical_path: canonical_path.to_string(), live: ReferenceLiveState { @@ -676,12 +724,13 @@ impl VersionLeaseStore { .await? .is_some()); } - let state_path = self.reference_state_path(canonical_path); - match self.object_store.delete(&state_path).await { - Ok(()) => Ok(true), - Err(error) if error.is_not_found() => Ok(false), - Err(error) => Err(error), - } + let vacant = ReferenceLifecycleState::Vacant { + canonical_path: canonical_path.to_string(), + }; + Ok(self + .put_reference_lifecycle(canonical_path, Some(¤t), &vacant, has_local_lock) + .await? + .is_some()) } fn branch_termination_path(&self) -> Path { @@ -797,14 +846,7 @@ impl VersionLeaseStore { serde_json::from_slice(&payload).map_err(|error| { Error::corrupt_file(metadata.location.clone(), error.to_string()) })?; - let target = match state { - ReferenceLifecycleState::Pending { target, .. } => Some(target), - ReferenceLifecycleState::Live { live, .. } => Some(live.target), - ReferenceLifecycleState::Deleting { previous, .. } => { - previous.map(|live| live.target) - } - ReferenceLifecycleState::Legacy { .. } => None, - }; + let target = self.retained_lifecycle_target(&state).await?; Ok::<_, Error>(target.and_then(|target| { (target.namespace == self.namespace).then_some(target.version) })) @@ -838,8 +880,10 @@ impl VersionLeaseStore { let generation = match state { ReferenceLifecycleState::Pending { operation_id, .. } => Some(operation_id), ReferenceLifecycleState::Live { live, .. } => Some(live.generation), - ReferenceLifecycleState::Deleting { operation_id, .. } => Some(operation_id), + ReferenceLifecycleState::Revoking { operation_id, .. } + | ReferenceLifecycleState::Deleting { operation_id, .. } => Some(operation_id), ReferenceLifecycleState::Legacy { .. } => Some("legacy".to_string()), + ReferenceLifecycleState::Vacant { .. } => Some("vacant".to_string()), }; Ok::<_, Error>(generation.map(|generation| (metadata.location, generation))) }) @@ -872,22 +916,17 @@ impl VersionLeaseStore { serde_json::from_slice(&payload).map_err(|error| { Error::corrupt_file(metadata.location.clone(), error.to_string()) })?; - let (generation, target) = match state { - ReferenceLifecycleState::Pending { - operation_id, - target, - .. - } => (operation_id, Some(target)), - ReferenceLifecycleState::Live { live, .. } => { - (live.generation, Some(live.target)) + let generation = match &state { + ReferenceLifecycleState::Pending { operation_id, .. } => operation_id.clone(), + ReferenceLifecycleState::Live { live, .. } => live.generation.clone(), + ReferenceLifecycleState::Revoking { operation_id, .. } + | ReferenceLifecycleState::Deleting { operation_id, .. } => { + operation_id.clone() } - ReferenceLifecycleState::Deleting { - operation_id, - previous, - .. - } => (operation_id, previous.map(|live| live.target)), - ReferenceLifecycleState::Legacy { .. } => ("legacy".to_string(), None), + ReferenceLifecycleState::Legacy { .. } => "legacy".to_string(), + ReferenceLifecycleState::Vacant { .. } => "vacant".to_string(), }; + let target = self.retained_lifecycle_target(&state).await?; let changed = observed_generations.get(&metadata.location) != Some(&generation); Ok::<_, Error>(target.and_then(|target| { (changed && target.namespace == self.namespace).then_some(target.version) @@ -899,6 +938,45 @@ impl VersionLeaseStore { .map(|versions| versions.into_iter().flatten().collect()) } + async fn retained_lifecycle_target( + &self, + state: &ReferenceLifecycleState, + ) -> Result> { + let (canonical_path, expected) = match state { + ReferenceLifecycleState::Pending { target, .. } => return Ok(Some(target.clone())), + ReferenceLifecycleState::Live { + canonical_path, + live, + } => (canonical_path, Some(live)), + ReferenceLifecycleState::Revoking { + canonical_path, + previous, + .. + } + | ReferenceLifecycleState::Deleting { + canonical_path, + previous, + .. + } => (canonical_path, previous.as_ref()), + ReferenceLifecycleState::Legacy { .. } | ReferenceLifecycleState::Vacant { .. } => { + return Ok(None); + } + }; + let canonical_path = Path::parse(canonical_path)?; + let payload = match self.object_store.inner.get(&canonical_path).await { + Ok(result) => result.bytes().await?, + Err(object_store::Error::NotFound { .. }) => return Ok(None), + Err(error) => return Err(error.into()), + }; + if payload_reference_is_deleted(&payload)? { + return Ok(None); + } + let generation = payload_reference_generation(&payload)?; + Ok(expected.and_then(|live| { + (generation.as_deref() == Some(live.generation.as_str())).then(|| live.target.clone()) + })) + } + async fn canonical_reference_versions(&self) -> Result> { let Some(context) = &self.canonical_references else { return Ok(HashSet::new()); @@ -1050,6 +1128,19 @@ impl VersionLeaseStore { operation_id: current_operation, .. }) if current_operation == operation_id => {} + Some(ReferenceLifecycleState::Revoking { + operation_id: current_operation, + .. + }) if current_operation == operation_id => { + let lifecycle = lifecycle.as_ref().ok_or_else(|| { + Error::internal(format!( + "reference {path} lost its revoking lifecycle state" + )) + })?; + self.reconcile_revoking_reference(&path, lifecycle, local_lock.is_some()) + .await?; + return Ok(ReferenceMutationOutcome::Conflict); + } _ => return Ok(ReferenceMutationOutcome::Conflict), } @@ -1073,21 +1164,98 @@ impl VersionLeaseStore { } return Ok(ReferenceMutationOutcome::Conflict); } + let Some(revoking) = self + .revoke_reference_lifecycle(&path, operation_id, mutation, local_lock.is_some()) + .await? + else { + return Ok(ReferenceMutationOutcome::Conflict); + }; + self.reconcile_revoking_reference(&path, &revoking, local_lock.is_some()) + .await?; + Ok(ReferenceMutationOutcome::Conflict) + } + + async fn revoke_reference_lifecycle( + &self, + canonical_path: &Path, + operation_id: &str, + mutation: &ReferenceMutation, + has_local_lock: bool, + ) -> Result> { + let Some(current) = self.reference_lifecycle_snapshot(canonical_path).await? else { + return Ok(None); + }; + if let ReferenceLifecycleState::Revoking { + operation_id: current_operation, + .. + } = ¤t.state + && current_operation == operation_id + { + return Ok(Some(current)); + } + let ReferenceLifecycleState::Pending { + operation_id: current_operation, + target, + previous, + previous_was_legacy, + .. + } = ¤t.state + else { + return Ok(None); + }; + if current_operation != operation_id { + return Ok(None); + } + let revoking = ReferenceLifecycleState::Revoking { + canonical_path: canonical_path.to_string(), + operation_id: operation_id.to_string(), + target: target.clone(), + previous: previous.clone(), + previous_was_legacy: *previous_was_legacy, + mutation: mutation.clone(), + }; + self.put_reference_lifecycle(canonical_path, Some(¤t), &revoking, has_local_lock) + .await + } - match (mutation, canonical) { + async fn reconcile_revoking_reference( + &self, + canonical_path: &Path, + revoking: &ReferenceLifecycleSnapshot, + has_local_lock: bool, + ) -> Result { + let ReferenceLifecycleState::Revoking { + operation_id, + previous, + mutation, + .. + } = &revoking.state + else { + return Ok(false); + }; + let current = match self.object_store.inner.get(canonical_path).await { + Ok(result) => { + let metadata = result.meta.clone(); + Some((metadata, result.bytes().await?)) + } + Err(object_store::Error::NotFound { .. }) => None, + Err(error) => return Err(error.into()), + }; + match (mutation, current.as_ref()) { ( ReferenceMutation::Update { - expected_payload, .. + expected_payload, + payload, + .. }, Some((metadata, current_payload)), - ) if current_payload.as_ref() == expected_payload => { - let fenced_payload = fenced_reference_payload(¤t_payload); + ) if current_payload.as_ref() == payload => { self.rewrite_reference_payload( - &path, - &metadata, - ¤t_payload, - &fenced_payload, - local_lock.is_some(), + canonical_path, + metadata, + current_payload, + expected_payload, + has_local_lock, ) .await?; } @@ -1096,24 +1264,161 @@ impl VersionLeaseStore { expected_payload, .. }, Some((metadata, current_payload)), - ) if canonical_is_intended => { + ) if current_payload.as_ref() == expected_payload => { + let fenced_payload = fenced_reference_payload(current_payload); self.rewrite_reference_payload( - &path, - &metadata, - ¤t_payload, - expected_payload, - local_lock.is_some(), + canonical_path, + metadata, + current_payload, + &fenced_payload, + has_local_lock, ) .await?; } - (ReferenceMutation::Create { .. }, Some(_)) if canonical_is_intended => { - self.object_store.delete(&path).await?; + (ReferenceMutation::Create { payload, .. }, None) => { + let tombstone = set_payload_reference_deleted(payload, operation_id)?; + match self + .object_store + .inner + .put_opts( + canonical_path, + Bytes::from(tombstone).into(), + PutOptions { + mode: PutMode::Create, + ..Default::default() + }, + ) + .await + { + Ok(_) | Err(object_store::Error::AlreadyExists { .. }) => {} + Err(error) => return Err(error.into()), + } + } + (ReferenceMutation::Create { payload, .. }, Some((metadata, current_payload))) + if current_payload.as_ref() == payload => + { + let tombstone = set_payload_reference_deleted(payload, operation_id)?; + self.rewrite_reference_payload( + canonical_path, + metadata, + current_payload, + &tombstone, + has_local_lock, + ) + .await?; } _ => {} } - self.cancel_reference_lifecycle(&path, operation_id, local_lock.is_some()) + let terminal = self + .terminal_reference_lifecycle(canonical_path, previous.as_ref()) .await?; - Ok(ReferenceMutationOutcome::Conflict) + Ok(self + .put_reference_lifecycle(canonical_path, Some(revoking), &terminal, has_local_lock) + .await? + .is_some()) + } + + async fn reconcile_deleting_reference( + &self, + canonical_path: &Path, + deleting: &ReferenceLifecycleSnapshot, + has_local_lock: bool, + ) -> Result { + let ReferenceLifecycleState::Deleting { + operation_id, + previous, + expected_payload, + .. + } = &deleting.state + else { + return Ok(false); + }; + let current = match self.object_store.inner.get(canonical_path).await { + Ok(result) => Some(result), + Err(object_store::Error::NotFound { .. }) => None, + Err(error) => return Err(error.into()), + }; + if let Some(result) = current { + let metadata = result.meta.clone(); + let current_payload = result.bytes().await?; + if current_payload.as_ref() == expected_payload { + // This conditional update is the deletion linearization point. + // A released writer that lands first changes the incarnation + // and makes this update fail; one that lands later replaces the + // tombstone and becomes the visible canonical reference. + let tombstone = set_payload_reference_deleted(expected_payload, operation_id)?; + match self + .rewrite_reference_payload( + canonical_path, + &metadata, + ¤t_payload, + &tombstone, + has_local_lock, + ) + .await + { + Ok(()) | Err(Error::RefConflict { .. }) => {} + Err(error) => return Err(error), + } + } + } + let terminal = self + .terminal_reference_lifecycle(canonical_path, previous.as_ref()) + .await?; + let is_deleted = matches!(terminal, ReferenceLifecycleState::Vacant { .. }); + if self + .put_reference_lifecycle(canonical_path, Some(deleting), &terminal, has_local_lock) + .await? + .is_none() + { + return Err(Error::RefConflict { + message: format!( + "reference {canonical_path} lifecycle changed during reconciliation" + ), + }); + } + Ok(is_deleted) + } + + async fn terminal_reference_lifecycle( + &self, + canonical_path: &Path, + previous: Option<&ReferenceLiveState>, + ) -> Result { + let payload = match self.object_store.inner.get(canonical_path).await { + Ok(result) => Some(result.bytes().await?), + Err(object_store::Error::NotFound { .. }) => None, + Err(error) => return Err(error.into()), + }; + let Some(payload) = payload else { + return Ok(ReferenceLifecycleState::Vacant { + canonical_path: canonical_path.to_string(), + }); + }; + if payload_reference_is_deleted(&payload)? { + return Ok(ReferenceLifecycleState::Vacant { + canonical_path: canonical_path.to_string(), + }); + } + let generation = payload_reference_generation(&payload)?; + if let Some(previous) = previous + && generation.as_deref() == Some(previous.generation.as_str()) + { + return Ok(ReferenceLifecycleState::Live { + canonical_path: canonical_path.to_string(), + live: previous.clone(), + }); + } + if generation.is_none() { + return Ok(ReferenceLifecycleState::Legacy { + canonical_path: canonical_path.to_string(), + }); + } + Err(Error::RefConflict { + message: format!( + "reference {canonical_path} canonical generation is not reconciled with lifecycle state" + ), + }) } async fn rewrite_reference_payload( @@ -1343,69 +1648,18 @@ impl VersionLeaseStore { if lifecycle_completed { return Ok(ReferenceMutationOutcome::Published); } - - self.rollback_lost_reference_mutation(mutation, has_local_lock) - .await?; - self.cancel_reference_lifecycle(&canonical_path, operation_id, has_local_lock) + let Some(revoking) = self + .revoke_reference_lifecycle(&canonical_path, operation_id, mutation, has_local_lock) + .await? + else { + return Ok(ReferenceMutationOutcome::Conflict); + }; + self.reconcile_revoking_reference(&canonical_path, &revoking, has_local_lock) .await?; Ok(ReferenceMutationOutcome::Conflict) } - async fn rollback_lost_reference_mutation( - &self, - mutation: &ReferenceMutation, - has_local_lock: bool, - ) -> Result<()> { - let (path, intended_payload) = match mutation { - ReferenceMutation::Create { path, payload } - | ReferenceMutation::Update { path, payload, .. } => { - (Path::parse(path)?, payload.as_slice()) - } - }; - let result = match self.object_store.inner.get(&path).await { - Ok(result) => result, - Err(object_store::Error::NotFound { .. }) => return Ok(()), - Err(error) => return Err(error.into()), - }; - let metadata = result.meta.clone(); - let current_payload = result.bytes().await?; - if current_payload.as_ref() != intended_payload { - return Ok(()); - } - - let replacement = match mutation { - ReferenceMutation::Create { payload, .. } => { - set_payload_reference_generation(payload, &Uuid::new_v4().simple().to_string())? - } - ReferenceMutation::Update { - expected_payload, .. - } => expected_payload.clone(), - }; - match self - .rewrite_reference_payload( - &path, - &metadata, - ¤t_payload, - &replacement, - has_local_lock, - ) - .await - { - Ok(()) => {} - Err(Error::RefConflict { .. }) => return Ok(()), - Err(error) => return Err(error), - } - if matches!(mutation, ReferenceMutation::Create { .. }) { - match self.object_store.delete(&path).await { - Ok(()) => {} - Err(error) if error.is_not_found() => {} - Err(error) => return Err(error), - } - } - Ok(()) - } - - async fn reference_mutation_conflict_outcome( + async fn reference_mutation_conflict_outcome( &self, path: &Path, intended_payload: &[u8], @@ -2352,6 +2606,26 @@ pub(super) async fn begin_reference_admission( (Path::parse(path)?, previous_was_legacy) } }; + if let ReferenceMutation::Create { path, payload } = &mutation { + let existing = match store.object_store.inner.get(&canonical_path).await { + Ok(result) => Some(result), + Err(object_store::Error::NotFound { .. }) => None, + Err(error) => return Err(error.into()), + }; + if let Some(result) = existing { + let metadata = result.meta.clone(); + let existing_payload = result.bytes().await?; + if payload_reference_is_deleted(&existing_payload)? { + mutation = ReferenceMutation::Update { + path: path.clone(), + expected_payload: existing_payload.to_vec(), + expected_etag: metadata.e_tag, + expected_version: metadata.version, + payload: payload.clone(), + }; + } + } + } let intent = ReferenceIntent { manifest_path: manifest.path.to_string(), mutation: mutation.clone(), @@ -2363,7 +2637,26 @@ pub(super) async fn begin_reference_admission( .claim_reference_lifecycle(&canonical_path, &operation_id, version, previous_was_legacy) .await { - let _ = store.object_store.delete(&path).await; + let claim_is_definitively_absent = + match store.reference_lifecycle_snapshot(&canonical_path).await { + Ok(Some(ReferenceLifecycleSnapshot { + state: + ReferenceLifecycleState::Pending { + operation_id: current_operation, + .. + } + | ReferenceLifecycleState::Revoking { + operation_id: current_operation, + .. + }, + .. + })) => current_operation != operation_id, + Ok(_) => true, + Err(_) => false, + }; + if claim_is_definitively_absent { + let _ = store.object_store.delete(&path).await; + } return Err(error); } let admission = ReferenceAdmission { @@ -2596,6 +2889,31 @@ fn set_payload_reference_generation(payload: &[u8], generation: &str) -> Result< Ok(serde_json::to_vec_pretty(&value)?) } +fn set_payload_reference_deleted(payload: &[u8], generation: &str) -> Result> { + let mut value: serde_json::Value = serde_json::from_slice(payload)?; + let object = value.as_object_mut().ok_or_else(|| { + Error::internal("canonical reference payload must be a JSON object".to_string()) + })?; + object.insert( + REFERENCE_GENERATION_FIELD.to_string(), + serde_json::Value::String(generation.to_string()), + ); + object.insert( + REFERENCE_DELETED_FIELD.to_string(), + serde_json::Value::Bool(true), + ); + Ok(serde_json::to_vec_pretty(&value)?) +} + +fn payload_reference_is_deleted(payload: &[u8]) -> Result { + let value: serde_json::Value = serde_json::from_slice(payload)?; + Ok(value + .as_object() + .and_then(|object| object.get(REFERENCE_DELETED_FIELD)) + .and_then(serde_json::Value::as_bool) + .unwrap_or(false)) +} + fn fenced_reference_payload(payload: &[u8]) -> Vec { let mut fenced = Vec::with_capacity(payload.len() + 1); fenced.extend_from_slice(payload); @@ -2609,6 +2927,9 @@ pub(super) async fn canonical_reference_is_visible( canonical_path: &Path, payload: &[u8], ) -> Result { + if payload_reference_is_deleted(payload)? { + return Ok(false); + } let generation = payload_reference_generation(payload)?; // Released writers do not preserve fields they do not know about. Their // rewrites remain authoritative canonical references and must stay visible @@ -2625,11 +2946,13 @@ pub(super) async fn canonical_reference_is_visible( generation.as_deref() == Some(live.generation.as_str()) } ReferenceLifecycleState::Legacy { .. } => false, + ReferenceLifecycleState::Vacant { .. } => false, ReferenceLifecycleState::Pending { previous, .. } => match previous { Some(previous) => generation.as_deref() == Some(previous.generation.as_str()), None => false, }, - ReferenceLifecycleState::Deleting { previous, .. } => previous + ReferenceLifecycleState::Revoking { previous, .. } + | ReferenceLifecycleState::Deleting { previous, .. } => previous .is_some_and(|previous| generation.as_deref() == Some(previous.generation.as_str())), }) } @@ -2672,7 +2995,9 @@ pub(super) async fn delete_canonical_reference( let lifecycle = store.reference_lifecycle_snapshot(canonical_path).await?; let canonical_generation = payload_reference_generation(¤t_payload)?; match lifecycle.as_ref().map(|snapshot| &snapshot.state) { - Some(ReferenceLifecycleState::Pending { .. }) => { + Some( + ReferenceLifecycleState::Pending { .. } | ReferenceLifecycleState::Revoking { .. }, + ) => { return Err(Error::RefConflict { message: format!("reference {canonical_path} has an in-flight mutation"), }); @@ -2690,6 +3015,11 @@ pub(super) async fn delete_canonical_reference( message: format!("reference {canonical_path} lifecycle changed during deletion"), }); } + Some(ReferenceLifecycleState::Vacant { .. }) if canonical_generation.is_some() => { + return Err(Error::RefConflict { + message: format!("reference {canonical_path} lifecycle changed during deletion"), + }); + } _ => {} } @@ -2715,6 +3045,7 @@ pub(super) async fn delete_canonical_reference( operation_id: Uuid::new_v4().simple().to_string(), previous, previous_was_legacy: canonical_generation.is_none(), + expected_payload: current_payload.to_vec(), }; let Some(snapshot) = store .put_reference_lifecycle( @@ -2731,73 +3062,12 @@ pub(super) async fn delete_canonical_reference( }; snapshot }; - - if let Err(delete_error) = object_store.delete(canonical_path).await { - match object_store.inner.get(canonical_path).await { - Ok(_) => { - restore_reference_lifecycle( - &store, - canonical_path, - &deleting_snapshot, - local_lock.is_some(), - ) - .await?; - return Err(delete_error); - } - Err(object_store::Error::NotFound { .. }) => {} - Err(_) => return Err(delete_error), - } - } - let state_path = store.reference_state_path(canonical_path); - match object_store.delete(&state_path).await { - Ok(()) => Ok(()), - Err(error) if error.is_not_found() => Ok(()), - Err(error) => Err(error), - } -} - -async fn restore_reference_lifecycle( - store: &VersionLeaseStore, - canonical_path: &Path, - deleting_snapshot: &ReferenceLifecycleSnapshot, - has_local_lock: bool, -) -> Result<()> { - let ReferenceLifecycleState::Deleting { - previous, - previous_was_legacy, - .. - } = &deleting_snapshot.state - else { - return Err(Error::internal(format!( - "reference {canonical_path} has non-deleting rollback state" - ))); - }; - let restored = if let Some(previous) = previous { - ReferenceLifecycleState::Live { - canonical_path: canonical_path.to_string(), - live: previous.clone(), - } - } else if *previous_was_legacy { - ReferenceLifecycleState::Legacy { - canonical_path: canonical_path.to_string(), - } - } else { - return Err(Error::internal(format!( - "reference {canonical_path} deletion has no state to restore" - ))); - }; - if store - .put_reference_lifecycle( - canonical_path, - Some(deleting_snapshot), - &restored, - has_local_lock, - ) + if !store + .reconcile_deleting_reference(canonical_path, &deleting_snapshot, local_lock.is_some()) .await? - .is_none() { return Err(Error::RefConflict { - message: format!("reference {canonical_path} lifecycle changed during rollback"), + message: format!("reference {canonical_path} changed during conditional deletion"), }); } Ok(()) @@ -2856,9 +3126,19 @@ pub(super) async fn remove_branch_state( } }; let local_lock = lock_local_reference(object_store, &canonical_path).await?; - lifecycle_store - .cancel_reference_lifecycle(&canonical_path, &intent.operation_id, local_lock.is_some()) - .await?; + if let Some(revoking) = lifecycle_store + .revoke_reference_lifecycle( + &canonical_path, + &intent.operation_id, + &intent.mutation, + local_lock.is_some(), + ) + .await? + { + lifecycle_store + .reconcile_revoking_reference(&canonical_path, &revoking, local_lock.is_some()) + .await?; + } } for path in [ root_path @@ -3757,7 +4037,16 @@ mod tests { .unwrap(), ReferenceMutationOutcome::Conflict ); - assert!(!store.object_store.exists(&canonical_path).await.unwrap()); + let rollback = store + .object_store + .inner + .get(&canonical_path) + .await + .unwrap() + .bytes() + .await + .unwrap(); + assert!(payload_reference_is_deleted(&rollback).unwrap()); } #[tokio::test] @@ -3971,7 +4260,16 @@ mod tests { .unwrap(), ReferenceMutationOutcome::Conflict ); - assert!(!store.object_store.exists(&canonical_path).await.unwrap()); + let rollback = store + .object_store + .inner + .get(&canonical_path) + .await + .unwrap() + .bytes() + .await + .unwrap(); + assert!(payload_reference_is_deleted(&rollback).unwrap()); } #[tokio::test] @@ -4009,16 +4307,399 @@ mod tests { .apply_reference_mutation_inner(&admission.mutation, false, None) .await .unwrap(), + ReferenceMutationOutcome::Conflict + ); + let rollback = store + .object_store + .inner + .get(&canonical_path) + .await + .unwrap() + .bytes() + .await + .unwrap(); + assert!(payload_reference_is_deleted(&rollback).unwrap()); + } + + #[tokio::test] + async fn released_client_delete_does_not_leave_stale_lifecycle_retention() { + let store = memory_store(); + let manifest_path = Path::from("manifests/42.manifest"); + let canonical_path = Path::from("tags/released-delete.json"); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + br#"{"version":42}"#, + ) + .await + .publish("conflict".to_string()) + .await + .unwrap(); + assert_eq!( + store.active_reference_versions().await.unwrap(), + HashSet::from([42]) + ); + + store.object_store.delete(&canonical_path).await.unwrap(); + + assert!(store.active_reference_versions().await.unwrap().is_empty()); + } + + #[tokio::test] + async fn expired_update_completion_cannot_win_after_canonical_rollback() { + let store = memory_store(); + let canonical_path = Path::from("tags/expired-update.json"); + let original_payload = br#"{"version":1}"#; + let original = store + .object_store + .inner + .put(&canonical_path, Bytes::from_static(original_payload).into()) + .await + .unwrap(); + let operation_id = Uuid::new_v4().simple().to_string(); + let mutation = ReferenceMutation::Update { + path: canonical_path.to_string(), + expected_payload: original_payload.to_vec(), + expected_etag: original.e_tag, + expected_version: original.version, + payload: set_payload_reference_generation(br#"{"version":42}"#, &operation_id).unwrap(), + }; + store + .claim_reference_lifecycle(&canonical_path, &operation_id, 42, true) + .await + .unwrap(); + assert_eq!( + store + .apply_reference_mutation_inner(&mutation, false, None) + .await + .unwrap(), ReferenceMutationOutcome::Published ); + let stale_pending = store + .reference_lifecycle_snapshot(&canonical_path) + .await + .unwrap() + .unwrap(); + assert_eq!( store - .finish_reference_mutation(&admission.mutation, &admission.operation_id, 42, false,) + .expired_reference_mutation_outcome(&mutation, &operation_id, 42, false) .await .unwrap(), ReferenceMutationOutcome::Conflict ); - assert!(!store.object_store.exists(&canonical_path).await.unwrap()); + let stale_live = ReferenceLifecycleState::Live { + canonical_path: canonical_path.to_string(), + live: ReferenceLiveState { + generation: operation_id, + target: ReferenceTarget { + namespace: MAIN_BRANCH.to_string(), + version: 42, + }, + }, + }; + assert!( + store + .put_reference_lifecycle(&canonical_path, Some(&stale_pending), &stale_live, false,) + .await + .unwrap() + .is_none() + ); + } + + #[tokio::test] + async fn rollback_create_failure_keeps_recoverable_revocation() { + let failing_store = Arc::new(FailingProxyStore::new()); + let mut object_store = ObjectStore::memory(); + object_store.inner = failing_store.wrap("memory", Arc::clone(&object_store.inner)); + let root = Path::from("dataset"); + let namespace = "branch-id"; + let store = VersionLeaseStore { + object_store: Arc::new(object_store), + root_path: root.clone(), + namespace: namespace.to_string(), + leases_path: root.clone().join(LEASES_DIR).join(namespace), + markers_path: root.clone().join(LEASE_GC_MARKERS_DIR).join(namespace), + reference_intents_path: root.clone().join(REFERENCE_INTENTS_DIR).join(namespace), + manifest_path: None, + canonical_references: None, + }; + let manifest_path = Path::from("dataset/branches/parent/versions/42.manifest"); + let canonical_path = Path::from("dataset/_refs/branches/rollback-crash.json"); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + let admission = create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + br#"{"parentVersion":42}"#, + ) + .await; + assert_eq!( + store + .apply_reference_mutation_inner(&admission.mutation, false, None) + .await + .unwrap(), + ReferenceMutationOutcome::Published + ); + failing_store.fail_when( + "put", + "rollback-crash.json", + "injected rollback tombstone failure", + ); + + store + .expired_reference_mutation_outcome( + &admission.mutation, + &admission.operation_id, + 42, + false, + ) + .await + .unwrap_err(); + assert!(matches!( + store + .reference_lifecycle_snapshot(&canonical_path) + .await + .unwrap() + .map(|snapshot| snapshot.state), + Some(ReferenceLifecycleState::Revoking { .. }) + )); + assert!(store.object_store.exists(&admission.path).await.unwrap()); + + failing_store.clear_fail_when("put", "rollback-crash.json"); + assert_eq!( + store + .expired_reference_mutation_outcome( + &admission.mutation, + &admission.operation_id, + 42, + false, + ) + .await + .unwrap(), + ReferenceMutationOutcome::Conflict + ); + let rollback = store + .object_store + .inner + .get(&canonical_path) + .await + .unwrap() + .bytes() + .await + .unwrap(); + assert!(payload_reference_is_deleted(&rollback).unwrap()); + } + + #[tokio::test] + async fn ambiguous_lifecycle_readback_keeps_recovery_intent() { + let failing_store = Arc::new(FailingProxyStore::new()); + let mut object_store = ObjectStore::memory(); + object_store.inner = failing_store.wrap("memory", Arc::clone(&object_store.inner)); + let store = VersionLeaseStore { + object_store: Arc::new(object_store), + root_path: Path::from(""), + namespace: MAIN_BRANCH.to_string(), + leases_path: Path::from("leases"), + markers_path: Path::from("markers"), + reference_intents_path: Path::from("reference_intents"), + manifest_path: None, + canonical_references: None, + }; + let canonical_path = Path::from("tags/ambiguous-state.json"); + let operation_id = Uuid::new_v4().simple().to_string(); + let intent = ReferenceIntent { + manifest_path: "manifests/42.manifest".to_string(), + mutation: ReferenceMutation::Create { + path: canonical_path.to_string(), + payload: br#"{"version":42}"#.to_vec(), + }, + operation_id: operation_id.clone(), + state: ReferenceIntentState::Pending, + }; + let (intent_path, _) = store.create_reference_intent(42, &intent).await.unwrap(); + failing_store.fail_after_n( + "get_opts", + "version_reference_states", + 1, + "injected lifecycle readback failure", + ); + + store + .claim_reference_lifecycle(&canonical_path, &operation_id, 42, false) + .await + .unwrap_err(); + + assert!(store.object_store.exists(&intent_path).await.unwrap()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn released_writer_update_during_delete_is_not_erased() { + use lance_core::utils::testing::{ProxyObjectStore, ProxyObjectStorePolicy}; + + let mut object_store = ObjectStore::memory(); + let underlying = Arc::clone(&object_store.inner); + let policy = Arc::new(std::sync::Mutex::new(ProxyObjectStorePolicy::new())); + object_store.inner = Arc::new(ProxyObjectStore::new( + Arc::clone(&underlying), + Arc::clone(&policy), + )); + let store = VersionLeaseStore { + object_store: Arc::new(object_store), + root_path: Path::from(""), + namespace: MAIN_BRANCH.to_string(), + leases_path: Path::from("leases"), + markers_path: Path::from("markers"), + reference_intents_path: Path::from("reference_intents"), + manifest_path: None, + canonical_references: None, + }; + let manifest_path = Path::from("manifests/42.manifest"); + let canonical_path = Path::from("tags/mixed-client-delete.json"); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + br#"{"branch":null,"version":42,"manifestSize":0,"metadata":{}}"#, + ) + .await + .publish("conflict".to_string()) + .await + .unwrap(); + let result = store.object_store.inner.get(&canonical_path).await.unwrap(); + let metadata = result.meta.clone(); + let payload = result.bytes().await.unwrap(); + + let (put_entered_tx, put_entered_rx) = tokio::sync::oneshot::channel(); + let put_entered_tx = Arc::new(std::sync::Mutex::new(Some(put_entered_tx))); + let (resume_put_tx, resume_put_rx) = std::sync::mpsc::channel(); + let resume_put_rx = Arc::new(std::sync::Mutex::new(resume_put_rx)); + let canonical_path_string = canonical_path.to_string(); + policy.lock().unwrap().set_before_policy( + "pause_canonical_tombstone", + Arc::new(move |method, path| { + if method == "put" + && path.as_ref() == canonical_path_string + && let Some(sender) = put_entered_tx.lock().unwrap().take() + { + sender.send(()).unwrap(); + resume_put_rx.lock().unwrap().recv().unwrap(); + } + Ok(()) + }), + ); + + let delete_store = Arc::clone(&store.object_store); + let delete_root = store.root_path.clone(); + let delete_path = canonical_path.clone(); + let delete_task = tokio::spawn(async move { + delete_canonical_reference( + delete_store, + &delete_root, + &delete_path, + &metadata, + &payload, + ) + .await + }); + put_entered_rx.await.unwrap(); + + let released_payload = + Bytes::from_static(br#"{"branch":null,"version":99,"manifestSize":0,"metadata":{}}"#); + underlying + .put(&canonical_path, released_payload.clone().into()) + .await + .unwrap(); + resume_put_tx.send(()).unwrap(); + assert!(matches!( + delete_task.await.unwrap(), + Err(Error::RefConflict { .. }) + )); + + let remaining = underlying + .get(&canonical_path) + .await + .unwrap() + .bytes() + .await + .unwrap(); + assert_eq!(remaining, released_payload); + } + + #[tokio::test] + async fn state_completion_failure_does_not_retain_deleted_reference_version() { + let failing_store = Arc::new(FailingProxyStore::new()); + let mut object_store = ObjectStore::memory(); + object_store.inner = failing_store.wrap("memory", Arc::clone(&object_store.inner)); + let store = VersionLeaseStore { + object_store: Arc::new(object_store), + root_path: Path::from(""), + namespace: MAIN_BRANCH.to_string(), + leases_path: Path::from("leases"), + markers_path: Path::from("markers"), + reference_intents_path: Path::from("reference_intents"), + manifest_path: None, + canonical_references: None, + }; + let manifest_path = Path::from("manifests/42.manifest"); + let canonical_path = Path::from("tags/state-completion-failure.json"); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + br#"{"branch":null,"version":42,"manifestSize":0,"metadata":{}}"#, + ) + .await + .publish("conflict".to_string()) + .await + .unwrap(); + let result = store.object_store.inner.get(&canonical_path).await.unwrap(); + let metadata = result.meta.clone(); + let payload = result.bytes().await.unwrap(); + failing_store.fail_after_n( + "put", + "version_reference_states", + 1, + "injected lifecycle completion failure", + ); + + delete_canonical_reference( + Arc::clone(&store.object_store), + &store.root_path, + &canonical_path, + &metadata, + &payload, + ) + .await + .unwrap_err(); + + assert!(store.active_reference_versions().await.unwrap().is_empty()); + let tombstone = store + .object_store + .inner + .get(&canonical_path) + .await + .unwrap() + .bytes() + .await + .unwrap(); + assert!(payload_reference_is_deleted(&tombstone).unwrap()); + assert!(matches!( + store + .reference_lifecycle_snapshot(&canonical_path) + .await + .unwrap() + .map(|snapshot| snapshot.state), + Some(ReferenceLifecycleState::Deleting { .. }) + )); } #[tokio::test] @@ -4054,9 +4735,9 @@ mod tests { let metadata = result.meta.clone(); let payload = result.bytes().await.unwrap(); failing_store.fail_when( - "delete", + "put", "deletion-failure.json", - "injected canonical deletion failure", + "injected canonical tombstone failure", ); let error = delete_canonical_reference( @@ -4071,10 +4752,10 @@ mod tests { assert!( error .to_string() - .contains("injected canonical deletion failure"), + .contains("injected canonical tombstone failure"), "{error}" ); - failing_store.clear_fail_when("delete", "deletion-failure.json"); + failing_store.clear_fail_when("put", "deletion-failure.json"); let remaining = store.object_store.inner.get(&canonical_path).await.unwrap(); let remaining_metadata = remaining.meta.clone(); @@ -4099,6 +4780,15 @@ mod tests { ) .await .unwrap(); - assert!(!store.object_store.exists(&canonical_path).await.unwrap()); + let tombstone = store + .object_store + .inner + .get(&canonical_path) + .await + .unwrap() + .bytes() + .await + .unwrap(); + assert!(payload_reference_is_deleted(&tombstone).unwrap()); } } From 692343f72107529ed6d5d0bc9fb021873a7e3029 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Sat, 8 Aug 2026 08:54:03 +0000 Subject: [PATCH 12/13] fix: make reference deletion incarnation-safe --- Cargo.lock | 1 + java/lance-jni/Cargo.lock | 1 + python/Cargo.lock | 1 + python/python/tests/test_dataset.py | 82 ++- rust/lance-io/Cargo.toml | 1 + rust/lance-io/src/object_store.rs | 480 ++++++++++++++- rust/lance-io/src/object_store/providers.rs | 2 + .../src/object_store/providers/aws.rs | 40 +- .../src/object_store/providers/azure.rs | 38 +- .../src/object_store/providers/gcp.rs | 38 +- .../src/object_store/providers/goosefs.rs | 1 + .../src/object_store/providers/huggingface.rs | 1 + .../src/object_store/providers/local.rs | 8 + .../src/object_store/providers/memory.rs | 1 + .../src/object_store/providers/oss.rs | 1 + .../src/object_store/providers/tencent.rs | 1 + .../src/object_store/providers/tos.rs | 1 + rust/lance/src/dataset/version_lease.rs | 571 +++++++++++------- 18 files changed, 963 insertions(+), 306 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 48b60c7b514..187b36984c6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4934,6 +4934,7 @@ dependencies = [ "pin-project", "prost", "rand 0.9.5", + "reqwest 0.12.28", "rstest", "serde", "tempfile", diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index 1d8c42337a7..0d3eb9a31a1 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -4103,6 +4103,7 @@ dependencies = [ "pin-project", "prost", "rand 0.9.5", + "reqwest 0.12.28", "serde", "tempfile", "tokio", diff --git a/python/Cargo.lock b/python/Cargo.lock index 10dc6efc545..eaac432cd24 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -4431,6 +4431,7 @@ dependencies = [ "pin-project", "prost", "rand 0.9.5", + "reqwest 0.12.28", "serde", "tempfile", "tokio", diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index db7feb9d7a7..39c28b596e3 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -650,10 +650,10 @@ def test_v2_manifest_paths_migration(tmp_path: Path): def test_tag(tmp_path: Path): table = pa.Table.from_pydict({"colA": [1, 2, 3], "colB": [4, 5, 6]}) - base_dir = tmp_path / "test" + base_dir = "shared-memory://test-tag/dataset" - lance.write_dataset(table, base_dir) - ds = lance.write_dataset(table, base_dir, mode="append") + ds = lance.write_dataset(table, base_dir) + ds = lance.write_dataset(table, ds, mode="append") assert len(ds.tags.list()) == 0 @@ -689,11 +689,11 @@ def test_tag(tmp_path: Path): assert ds.checkout_version("tag1").version == 1 - ds = lance.dataset(base_dir, "tag1") + ds = ds.checkout_version("tag1") assert ds.version == 1 - with pytest.raises(ValueError): - lance.dataset(base_dir, "missing-tag") + with pytest.raises(OSError): + ds.checkout_version("missing-tag") # test tag update with pytest.raises( @@ -712,7 +712,7 @@ def test_tag(tmp_path: Path): assert tag1_created_at is not None assert tag1_updated_at is not None ds.tags.replace_metadata("tag1", {"description": "updated tag"}) - ds = lance.dataset(base_dir, "tag1") + ds = ds.checkout_version("tag1") assert ds.version == 1 replaced_tag1_meta = ds.tags.list()["tag1"] assert replaced_tag1_meta["metadata"] == {"description": "updated tag"} @@ -728,17 +728,17 @@ def test_tag(tmp_path: Path): assert updated_tag1_meta["created_at"] == tag1_created_at assert updated_tag1_meta["updated_at"] is not None assert updated_tag1_meta["updated_at"] >= tag1_updated_at - ds = lance.dataset(base_dir, "tag1") + ds = ds.checkout_version("tag1") assert ds.version == 2 assert ds.tags.list()["tag1"]["metadata"] == {"owner": "ml-team"} ds.tags.replace_metadata("tag1", {}) - ds = lance.dataset(base_dir, "tag1") + ds = ds.checkout_version("tag1") assert ds.version == 2 assert ds.tags.list()["tag1"]["metadata"] == {} ds.tags.update("tag1", 1) - ds = lance.dataset(base_dir, "tag1") + ds = ds.checkout_version("tag1") assert ds.version == 1 assert ds.tags.list()["tag1"]["metadata"] == {} @@ -775,11 +775,11 @@ def test_tag(tmp_path: Path): def test_tag_order(tmp_path: Path): table = pa.Table.from_pydict({"colA": [1, 2, 3], "colB": [4, 5, 6]}) - base_dir = tmp_path / "test" + base_dir = "shared-memory://test-tag-order/dataset" - for i in range(3): - mode = "append" if i > 0 else "create" - ds = lance.write_dataset(table, base_dir, mode=mode) + ds = lance.write_dataset(table, base_dir, mode="create") + for _ in range(2): + ds = lance.write_dataset(table, ds, mode="append") expected_tags = {"tag3": 3, "tag2": 2, "tag1": 1} for name, version in expected_tags.items(): @@ -1653,14 +1653,12 @@ def test_explain_cleanup_old_versions(tmp_path): def test_cleanup_error_when_tagged_old_versions(tmp_path): table = pa.Table.from_pydict({"a": range(100), "b": range(100)}) - base_dir = tmp_path / "test" - lance.write_dataset(table, base_dir) - lance.write_dataset(table, base_dir, mode="overwrite") + base_dir = "shared-memory://cleanup-error-tagged/dataset" + dataset = lance.write_dataset(table, base_dir) + dataset = lance.write_dataset(table, dataset, mode="overwrite") time.sleep(0.1) moment = datetime.now() - lance.write_dataset(table, base_dir, mode="overwrite") - - dataset = lance.dataset(base_dir) + dataset = lance.write_dataset(table, dataset, mode="overwrite") dataset.tags.create("old-tag", 1) dataset.tags.create("another-old-tag", 2) @@ -1682,14 +1680,12 @@ def test_cleanup_error_when_tagged_old_versions(tmp_path): def test_cleanup_around_tagged_old_versions(tmp_path): table = pa.Table.from_pydict({"a": range(100), "b": range(100)}) - base_dir = tmp_path / "test" - lance.write_dataset(table, base_dir) - lance.write_dataset(table, base_dir, mode="overwrite") + base_dir = "shared-memory://cleanup-around-tagged/dataset" + dataset = lance.write_dataset(table, base_dir) + dataset = lance.write_dataset(table, dataset, mode="overwrite") time.sleep(0.1) moment = datetime.now() - lance.write_dataset(table, base_dir, mode="overwrite") - - dataset = lance.dataset(base_dir) + dataset = lance.write_dataset(table, dataset, mode="overwrite") dataset.tags.create("old-tag", 1) dataset.tags.create("another-old-tag", 2) dataset.tags.create("tag-latest", 3) @@ -6243,10 +6239,10 @@ def test_update_config_transaction(tmp_path: Path): def test_shallow_clone(tmp_path: Path): - """Shallow clone a filesystem dataset by version number and by tag. + """Shallow clone a filesystem dataset by version and released tag. Arrange: - - Create a source dataset at a filesystem path with two versions + - Create a filesystem source dataset with two versions (create v1, then overwrite to v2). - Create a tag "v1" pointing to version 1. Act: @@ -6256,19 +6252,27 @@ def test_shallow_clone(tmp_path: Path): - Re-open cloned datasets and verify their tables equal the source version they were cloned from (schema and record count). - This test uses pathlib paths and tmp_path for cross-platform compatibility - and should not skip on Windows. + The tag fixture uses the stable released-client JSON representation. """ # Prepare source dataset with two versions src_dir = tmp_path / "shallow_src" table_v1 = pa.table({"a": [1, 2, 3], "b": [10, 20, 30]}) - lance.write_dataset(table_v1, src_dir, mode="create") + ds = lance.write_dataset( + table_v1, src_dir, mode="create", enable_v2_manifest_paths=False + ) table_v2 = pa.table({"a": [4, 5, 6], "b": [40, 50, 60]}) - ds = lance.write_dataset(table_v2, src_dir, mode="overwrite") + ds = lance.write_dataset(table_v2, ds, mode="overwrite") - # Create a tag pointing to version 1 - ds.tags.create("v1", 1) + # Emulate the stable canonical payload emitted by a released client. Local + # stores cannot safely execute current conditional reference deletion. + tag_path = src_dir / "_refs" / "tags" / "v1.json" + tag_path.parent.mkdir(parents=True) + manifest_size = (src_dir / "_versions" / "1.manifest").stat().st_size + tag_path.write_text( + f'{{"branch":null,"version":1,"manifestSize":{manifest_size},"metadata":{{}}}}', + encoding="utf-8", + ) # Clone by numeric version (v2) and assert equality clone_v2_dir = tmp_path / "clone_v2" @@ -6282,18 +6286,10 @@ def test_shallow_clone(tmp_path: Path): assert ds_clone_v1_tag.to_table() == table_v1 assert lance.dataset(clone_v1_tag_dir).to_table() == table_v1 - table_v3 = pa.table({"a": [7, 8, 9], "b": [40, 50, 60]}) - branch = ds.create_branch("branch", 2) - lance.write_dataset(table_v3, branch.uri, mode="overwrite") - clone_branch_v3 = tmp_path / "clone_branch_v3" - cloned_by_branch = branch.shallow_clone(clone_branch_v3, 3) - assert cloned_by_branch.to_table() == table_v3 - assert lance.dataset(clone_branch_v3).to_table() == table_v3 - def test_branches(tmp_path: Path): # Step 1: create branch1 from main → append to branch1 → create branch2 from tag - base_dir = tmp_path / "test_branches" + base_dir = "shared-memory://branches/dataset" main_table = pa.Table.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]}) ds_main = lance.write_dataset(main_table, base_dir) diff --git a/rust/lance-io/Cargo.toml b/rust/lance-io/Cargo.toml index 03c6780263a..34a314fdcd2 100644 --- a/rust/lance-io/Cargo.toml +++ b/rust/lance-io/Cargo.toml @@ -41,6 +41,7 @@ tracing.workspace = true url.workspace = true path_abs.workspace = true rand.workspace = true +reqwest = { version = "0.12", default-features = false, features = ["rustls-tls"] } tempfile.workspace = true [target.'cfg(target_os = "linux")'.dependencies] diff --git a/rust/lance-io/src/object_store.rs b/rust/lance-io/src/object_store.rs index a8230a578ca..eb0029af4a7 100644 --- a/rust/lance-io/src/object_store.rs +++ b/rust/lance-io/src/object_store.rs @@ -8,7 +8,7 @@ use std::collections::{HashMap, HashSet}; use std::ops::Range; use std::pin::Pin; use std::str::FromStr; -use std::sync::Arc; +use std::sync::{Arc, LazyLock, Mutex}; use std::time::Duration; use async_trait::async_trait; @@ -24,9 +24,14 @@ use object_store::DynObjectStore; use object_store::ObjectStoreExt as OSObjectStoreExt; #[cfg(feature = "aws")] use object_store::aws::AwsCredentialProvider; +use object_store::signer::Signer; #[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::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore as OSObjectStore, PutMode, PutMultipartOptions, PutOptions, PutPayload, PutResult, + RenameOptions, UpdateVersion, path::Path, +}; use providers::local::FileStoreProvider; use providers::memory::MemoryStoreProvider; use tokio::io::AsyncWriteExt; @@ -156,6 +161,363 @@ pub struct ObjectStore { /// which usually cannot be found in the URL such as Azure account name. The prefix plus the /// path uniquely identifies any object inside the store. pub store_prefix: String, + /// Whether this store can atomically delete one observed object incarnation. + conditional_delete: Option, +} + +#[derive(Debug, Clone)] +enum ConditionalDeleteConfig { + Serialized, + SignedUrl(Arc), +} + +/// Result of deleting one object only if its storage identity still matches. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ConditionalDeleteResult { + /// The observed object incarnation was deleted. + Deleted, + /// The object was already absent. + NotFound, + /// A different object incarnation now occupies the path. + IdentityMismatch, +} + +/// Internal marker carried through ordinary object-store wrappers so failure +/// injection, tracing, and accounting still observe the conditional operation. +#[derive(Debug, Clone)] +struct ConditionalDeleteRequest; + +static CONDITIONAL_DELETE_LOCKS: LazyLock>>>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +/// Adds an atomic compare-and-delete operation to stores whose complete writer +/// population is in this process, such as the in-memory test stores. +/// +/// The upstream object-store trait has conditional puts but no conditional +/// delete. Encoding the operation as an implementation-specific put extension +/// lets it pass through all outer wrappers before this layer performs the +/// comparison and delete under the same lock used by ordinary mutations. +#[derive(Debug)] +struct ConditionalDeleteStore { + target: Arc, + mutation_lock: Arc>, +} + +/// Performs a provider-native conditional DELETE through a short-lived signed +/// URL. S3, GCS, and Azure all enforce `If-Match` at the deletion linearization +/// point, including against writers using older Lance clients. +#[derive(Debug)] +struct SignedConditionalDeleteStore { + target: Arc, + signer: Arc, + client: reqwest::Client, +} + +impl SignedConditionalDeleteStore { + fn new(target: Arc, signer: Arc) -> Self { + Self { + target, + signer, + client: reqwest::Client::new(), + } + } +} + +impl std::fmt::Display for SignedConditionalDeleteStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "SignedConditionalDeleteStore({})", self.target) + } +} + +#[async_trait] +impl OSObjectStore for SignedConditionalDeleteStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> object_store::Result { + if opts.extensions.get::().is_some() { + let PutMode::Update(expected) = &opts.mode else { + return Err(ConditionalDeleteStore::precondition_error( + location, + "conditional delete requires an observed object identity", + )); + }; + let Some(e_tag) = expected.e_tag.as_ref() else { + return Err(object_store::Error::NotSupported { + source: std::io::Error::other( + "signed conditional delete requires an object ETag", + ) + .into(), + }); + }; + let url = self + .signer + .signed_url(reqwest::Method::DELETE, location, Duration::from_secs(300)) + .await?; + let response = self + .client + .delete(url) + .header(reqwest::header::IF_MATCH, e_tag) + .send() + .await + .map_err(|source| object_store::Error::Generic { + store: "signed conditional delete", + source: source.into(), + })?; + let status = response.status(); + if status.is_success() { + return Ok(PutResult { + e_tag: None, + version: None, + }); + } + if status == reqwest::StatusCode::NOT_FOUND { + return Err(object_store::Error::NotFound { + path: location.to_string(), + source: std::io::Error::new( + std::io::ErrorKind::NotFound, + "object is already absent", + ) + .into(), + }); + } + if status == reqwest::StatusCode::PRECONDITION_FAILED + || status == reqwest::StatusCode::CONFLICT + { + return Err(ConditionalDeleteStore::precondition_error( + location, + "object identity changed before conditional delete", + )); + } + Err(object_store::Error::Generic { + store: "signed conditional delete", + source: std::io::Error::other(format!( + "conditional delete for {location} failed with HTTP status {status}" + )) + .into(), + }) + } else { + self.target.put_opts(location, payload, opts).await + } + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> object_store::Result> { + self.target.put_multipart_opts(location, opts).await + } + + async fn get_opts( + &self, + location: &Path, + options: GetOptions, + ) -> object_store::Result { + self.target.get_opts(location, options).await + } + + async fn get_ranges( + &self, + location: &Path, + ranges: &[Range], + ) -> object_store::Result> { + self.target.get_ranges(location, ranges).await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, object_store::Result>, + ) -> BoxStream<'static, object_store::Result> { + self.target.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, object_store::Result> { + self.target.list(prefix) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> object_store::Result { + self.target.list_with_delimiter(prefix).await + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + opts: CopyOptions, + ) -> object_store::Result<()> { + self.target.copy_opts(from, to, opts).await + } + + async fn rename_opts( + &self, + from: &Path, + to: &Path, + opts: RenameOptions, + ) -> object_store::Result<()> { + self.target.rename_opts(from, to, opts).await + } +} + +impl ConditionalDeleteStore { + fn new(target: Arc, store_prefix: &str) -> Self { + let mutation_lock = CONDITIONAL_DELETE_LOCKS + .lock() + .expect("conditional delete lock registry poisoned") + .entry(store_prefix.to_string()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); + Self { + target, + mutation_lock, + } + } + + fn precondition_error(path: &Path, message: &'static str) -> object_store::Error { + object_store::Error::Precondition { + path: path.to_string(), + source: std::io::Error::other(message).into(), + } + } + + fn metadata_matches(metadata: &ObjectMeta, expected: &UpdateVersion) -> bool { + let has_identity = expected.e_tag.is_some() || expected.version.is_some(); + has_identity + && expected + .e_tag + .as_ref() + .is_none_or(|e_tag| metadata.e_tag.as_ref() == Some(e_tag)) + && expected + .version + .as_ref() + .is_none_or(|version| metadata.version.as_ref() == Some(version)) + } +} + +impl std::fmt::Display for ConditionalDeleteStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "ConditionalDeleteStore({})", self.target) + } +} + +#[async_trait] +impl OSObjectStore for ConditionalDeleteStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> object_store::Result { + let _guard = self.mutation_lock.lock().await; + if opts.extensions.get::().is_some() { + let PutMode::Update(expected) = &opts.mode else { + return Err(Self::precondition_error( + location, + "conditional delete requires an observed object identity", + )); + }; + let metadata = match self.target.head(location).await { + Ok(metadata) => metadata, + Err(object_store::Error::NotFound { .. }) => { + return Err(object_store::Error::NotFound { + path: location.to_string(), + source: std::io::Error::new( + std::io::ErrorKind::NotFound, + "object is already absent", + ) + .into(), + }); + } + Err(error) => return Err(error), + }; + if !Self::metadata_matches(&metadata, expected) { + return Err(Self::precondition_error( + location, + "object identity changed before conditional delete", + )); + } + self.target.delete(location).await?; + return Ok(PutResult { + e_tag: None, + version: None, + }); + } + self.target.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> object_store::Result> { + self.target.put_multipart_opts(location, opts).await + } + + async fn get_opts( + &self, + location: &Path, + options: GetOptions, + ) -> object_store::Result { + self.target.get_opts(location, options).await + } + + async fn get_ranges( + &self, + location: &Path, + ranges: &[Range], + ) -> object_store::Result> { + self.target.get_ranges(location, ranges).await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, object_store::Result>, + ) -> BoxStream<'static, object_store::Result> { + let target = Arc::clone(&self.target); + let mutation_lock = Arc::clone(&self.mutation_lock); + locations + .then(move |location| { + let target = Arc::clone(&target); + let mutation_lock = Arc::clone(&mutation_lock); + async move { + let location = location?; + let _guard = mutation_lock.lock().await; + target.delete(&location).await?; + Ok(location) + } + }) + .boxed() + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, object_store::Result> { + self.target.list(prefix) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> object_store::Result { + self.target.list_with_delimiter(prefix).await + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + opts: CopyOptions, + ) -> object_store::Result<()> { + let _guard = self.mutation_lock.lock().await; + self.target.copy_opts(from, to, opts).await + } + + async fn rename_opts( + &self, + from: &Path, + to: &Path, + opts: RenameOptions, + ) -> object_store::Result<()> { + let _guard = self.mutation_lock.lock().await; + self.target.rename_opts(from, to, opts).await + } } impl DeepSizeOf for ObjectStore { @@ -522,6 +884,7 @@ impl ObjectStore { download_retry_count: DEFAULT_DOWNLOAD_RETRY_COUNT, io_tracker, store_prefix, + conditional_delete: None, }; let path = Path::parse(path.path())?; return Ok((Arc::new(store), path)); @@ -581,11 +944,13 @@ impl ObjectStore { /// Create a in-memory object store directly for testing. pub fn memory() -> Self { let provider = MemoryStoreProvider; - provider + let mut store = provider .new_store(Url::parse("memory:///").unwrap(), &Default::default()) .now_or_never() .unwrap() - .unwrap() + .unwrap(); + store.install_conditional_delete_layer(); + store } /// Returns true if the object store pointed to a local file system. @@ -612,6 +977,66 @@ impl ObjectStore { &self.scheme } + fn install_conditional_delete_layer(&mut self) { + match self.conditional_delete.clone() { + Some(ConditionalDeleteConfig::Serialized) => { + self.inner = Arc::new(ConditionalDeleteStore::new( + Arc::clone(&self.inner), + &self.store_prefix, + )); + } + Some(ConditionalDeleteConfig::SignedUrl(signer)) => { + self.inner = Arc::new(SignedConditionalDeleteStore::new( + Arc::clone(&self.inner), + signer, + )); + } + None => {} + } + } + + /// Whether this store can atomically remove one observed object incarnation. + pub fn supports_conditional_delete(&self) -> bool { + self.conditional_delete.is_some() + } + + /// Delete `path` only if it still has the identity in `metadata`. + pub async fn delete_if_matches( + &self, + path: &Path, + expected: &UpdateVersion, + ) -> Result { + if self.conditional_delete.is_none() { + return Err(Error::not_supported(format!( + "object store {} does not support atomic conditional delete for {path}", + self.scheme + ))); + } + if expected.e_tag.is_none() && expected.version.is_none() { + return Err(Error::not_supported(format!( + "object store {} did not provide an identity for conditional delete of {path}", + self.scheme + ))); + } + let mut options = PutOptions { + mode: PutMode::Update(expected.clone()), + ..Default::default() + }; + options.extensions.insert(ConditionalDeleteRequest); + match self + .inner + .put_opts(path, Bytes::new().into(), options) + .await + { + Ok(_) => Ok(ConditionalDeleteResult::Deleted), + Err(object_store::Error::NotFound { .. }) => Ok(ConditionalDeleteResult::NotFound), + Err(object_store::Error::Precondition { .. }) => { + Ok(ConditionalDeleteResult::IdentityMismatch) + } + Err(error) => Err(error.into()), + } + } + pub fn block_size(&self) -> usize { self.block_size } @@ -1226,6 +1651,7 @@ impl ObjectStore { download_retry_count, io_tracker, store_prefix, + conditional_delete: None, } } } @@ -1326,6 +1752,52 @@ mod tests { ); } + #[tokio::test] + async fn test_conditional_delete_preserves_newer_incarnation() { + let store = ObjectStore::memory(); + let path = Path::from("conditional-delete"); + store.put(&path, b"first").await.unwrap(); + let first = store.inner.head(&path).await.unwrap(); + + store.put(&path, b"second").await.unwrap(); + let stale_identity = UpdateVersion { + e_tag: first.e_tag, + version: first.version, + }; + assert_eq!( + store + .delete_if_matches(&path, &stale_identity) + .await + .unwrap(), + ConditionalDeleteResult::IdentityMismatch + ); + assert_eq!( + store.read_one_all(&path).await.unwrap(), + b"second".as_slice() + ); + + let second = store.inner.head(&path).await.unwrap(); + let current_identity = UpdateVersion { + e_tag: second.e_tag, + version: second.version, + }; + assert_eq!( + store + .delete_if_matches(&path, ¤t_identity) + .await + .unwrap(), + ConditionalDeleteResult::Deleted + ); + assert!(!store.exists(&path).await.unwrap()); + assert_eq!( + store + .delete_if_matches(&path, ¤t_identity) + .await + .unwrap(), + ConditionalDeleteResult::NotFound + ); + } + #[tokio::test] async fn test_absolute_paths() { let tmp_path = TempStrDir::default(); diff --git a/rust/lance-io/src/object_store/providers.rs b/rust/lance-io/src/object_store/providers.rs index b84dc7362b8..dc22dd73eef 100644 --- a/rust/lance-io/src/object_store/providers.rs +++ b/rust/lance-io/src/object_store/providers.rs @@ -276,6 +276,8 @@ impl ObjectStoreRegistry { // Always wrap with IO tracking store.inner = store.io_tracker.wrap("", store.inner); + store.install_conditional_delete_layer(); + let store = Arc::new(store); { diff --git a/rust/lance-io/src/object_store/providers/aws.rs b/rust/lance-io/src/object_store/providers/aws.rs index ee70b76a042..452a875a8f9 100644 --- a/rust/lance-io/src/object_store/providers/aws.rs +++ b/rust/lance-io/src/object_store/providers/aws.rs @@ -23,7 +23,7 @@ use object_store::{ ClientOptions, CredentialProvider, Result as ObjectStoreResult, RetryConfig, StaticCredentialProvider, aws::{ - AmazonS3Builder, AmazonS3ConfigKey, AwsCredential as ObjectStoreAwsCredential, + AmazonS3, AmazonS3Builder, AmazonS3ConfigKey, AwsCredential as ObjectStoreAwsCredential, AwsCredentialProvider, }, }; @@ -31,8 +31,9 @@ use tokio::sync::RwLock; use url::Url; use crate::object_store::{ - DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, - ObjectStoreParams, ObjectStoreProvider, StorageOptions, StorageOptionsAccessor, + ConditionalDeleteConfig, DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, + DEFAULT_MAX_IOP_SIZE, ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions, + StorageOptionsAccessor, dynamic_credentials::{NamespaceCredentialsProvider, build_dynamic_credential_provider}, throttle::{AimdThrottleConfig, AimdThrottleState, AimdThrottledStore, cloud_http_connector}, }; @@ -49,7 +50,7 @@ impl AwsStoreProvider { storage_options: &StorageOptions, is_s3_express: bool, throttle_state: Option<&AimdThrottleState>, - ) -> Result> { + ) -> Result> { // Use a low retry count since the AIMD throttle layer handles // throttle recovery with its own retry loop. let retry_config = RetryConfig { @@ -104,7 +105,7 @@ impl AwsStoreProvider { builder = builder.with_http_connector(cloud_http_connector(throttle_state, store_prefix)); - Ok(Arc::new(builder.build()?) as Arc) + Ok(Arc::new(builder.build()?)) } async fn build_opendal_s3_store( @@ -179,20 +180,28 @@ impl ObjectStoreProvider for AwsStoreProvider { Some(AimdThrottleState::new(throttle_config)?) }; - let inner = if use_opendal { + let (inner, conditional_delete) = if use_opendal { // Use OpenDAL implementation - self.build_opendal_s3_store(&base_path, &storage_options) - .await? + ( + self.build_opendal_s3_store(&base_path, &storage_options) + .await?, + None, + ) } else { // Use default Amazon S3 implementation - self.build_amazon_s3_store( - &mut base_path, - params, - &storage_options, - is_s3_express, - throttle_state.as_ref(), + let native = self + .build_amazon_s3_store( + &mut base_path, + params, + &storage_options, + is_s3_express, + throttle_state.as_ref(), + ) + .await?; + ( + Arc::clone(&native) as Arc, + Some(ConditionalDeleteConfig::SignedUrl(native)), ) - .await? }; let inner = if let Some(throttle_state) = throttle_state { Arc::new(AimdThrottledStore::new_with_state( @@ -216,6 +225,7 @@ impl ObjectStoreProvider for AwsStoreProvider { io_tracker: Default::default(), store_prefix: self .calculate_object_store_prefix(&base_path, params.storage_options())?, + conditional_delete, }) } } diff --git a/rust/lance-io/src/object_store/providers/azure.rs b/rust/lance-io/src/object_store/providers/azure.rs index 2d407cd6df2..9b1a077c084 100644 --- a/rust/lance-io/src/object_store/providers/azure.rs +++ b/rust/lance-io/src/object_store/providers/azure.rs @@ -14,13 +14,14 @@ use opendal::{Operator, services::Azblob, services::Azdls}; use object_store::{ RetryConfig, - azure::{AzureConfigKey, AzureCredential, MicrosoftAzureBuilder}, + azure::{AzureConfigKey, AzureCredential, MicrosoftAzure, MicrosoftAzureBuilder}, }; use url::Url; use crate::object_store::{ - DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, - ObjectStoreParams, ObjectStoreProvider, StorageOptions, StorageOptionsAccessor, + ConditionalDeleteConfig, DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, + DEFAULT_MAX_IOP_SIZE, ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions, + StorageOptionsAccessor, dynamic_credentials::build_dynamic_credential_provider, throttle::{AimdThrottleConfig, AimdThrottleState, AimdThrottledStore, cloud_http_connector}, }; @@ -163,7 +164,7 @@ impl AzureBlobStoreProvider { storage_options: &StorageOptions, accessor: Option>, throttle_state: Option<&AimdThrottleState>, - ) -> Result> { + ) -> Result> { // Use a low retry count since the AIMD throttle layer handles // throttle recovery with its own retry loop. let retry_config = RetryConfig { @@ -190,7 +191,7 @@ impl AzureBlobStoreProvider { self.calculate_object_store_prefix(base_path, Some(&storage_options.0))?; builder = builder.with_http_connector(cloud_http_connector(throttle_state, store_prefix)); - Ok(Arc::new(builder.build()?) as Arc) + Ok(Arc::new(builder.build()?)) } fn calculate_object_store_prefix_with_env( @@ -262,19 +263,27 @@ impl ObjectStoreProvider for AzureBlobStoreProvider { Some(AimdThrottleState::new(throttle_config)?) }; - let inner: Arc = if use_opendal { + let (inner, conditional_delete) = if use_opendal { // OpenDAL Azure intentionally uses static/environment-backed configuration only. // Namespace-vended dynamic credentials are supported on the native object_store path. - self.build_opendal_azure_store(&base_path, &storage_options) - .await? + ( + self.build_opendal_azure_store(&base_path, &storage_options) + .await?, + None, + ) } else { - self.build_microsoft_azure_store( - &base_path, - &storage_options, - accessor, - throttle_state.as_ref(), + let native = self + .build_microsoft_azure_store( + &base_path, + &storage_options, + accessor, + throttle_state.as_ref(), + ) + .await?; + ( + Arc::clone(&native) as Arc, + Some(ConditionalDeleteConfig::SignedUrl(native)), ) - .await? }; let inner = if let Some(throttle_state) = throttle_state { Arc::new(AimdThrottledStore::new_with_state( @@ -298,6 +307,7 @@ impl ObjectStoreProvider for AzureBlobStoreProvider { io_tracker: Default::default(), store_prefix: self .calculate_object_store_prefix(&base_path, params.storage_options())?, + conditional_delete, }) } diff --git a/rust/lance-io/src/object_store/providers/gcp.rs b/rust/lance-io/src/object_store/providers/gcp.rs index 1a93c3ac9f0..8757d00a646 100644 --- a/rust/lance-io/src/object_store/providers/gcp.rs +++ b/rust/lance-io/src/object_store/providers/gcp.rs @@ -9,13 +9,14 @@ use opendal::{Operator, services::Gcs}; use object_store::{ RetryConfig, StaticCredentialProvider, - gcp::{GcpCredential, GoogleCloudStorageBuilder, GoogleConfigKey}, + gcp::{GcpCredential, GoogleCloudStorage, GoogleCloudStorageBuilder, GoogleConfigKey}, }; use url::Url; use crate::object_store::{ - DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, - ObjectStoreParams, ObjectStoreProvider, StorageOptions, StorageOptionsAccessor, + ConditionalDeleteConfig, DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, + DEFAULT_MAX_IOP_SIZE, ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions, + StorageOptionsAccessor, dynamic_credentials::build_dynamic_credential_provider, throttle::{AimdThrottleConfig, AimdThrottleState, AimdThrottledStore, cloud_http_connector}, }; @@ -60,7 +61,7 @@ impl GcsStoreProvider { storage_options: &StorageOptions, accessor: Option>, throttle_state: Option<&AimdThrottleState>, - ) -> Result> { + ) -> Result> { // Use a low retry count since the AIMD throttle layer handles // throttle recovery with its own retry loop. let retry_config = RetryConfig { @@ -93,7 +94,7 @@ impl GcsStoreProvider { self.calculate_object_store_prefix(base_path, Some(&storage_options.0))?; builder = builder.with_http_connector(cloud_http_connector(throttle_state, store_prefix)); - Ok(Arc::new(builder.build()?) as Arc) + Ok(Arc::new(builder.build()?)) } } @@ -121,19 +122,27 @@ impl ObjectStoreProvider for GcsStoreProvider { Some(AimdThrottleState::new(throttle_config)?) }; - let inner = if use_opendal { + let (inner, conditional_delete) = if use_opendal { // OpenDAL GCS intentionally uses static/environment-backed configuration only. // Namespace-vended dynamic credentials are supported on the native object_store path. - self.build_opendal_gcs_store(&base_path, &storage_options) - .await? + ( + self.build_opendal_gcs_store(&base_path, &storage_options) + .await?, + None, + ) } else { - self.build_google_cloud_store( - &base_path, - &storage_options, - accessor, - throttle_state.as_ref(), + let native = self + .build_google_cloud_store( + &base_path, + &storage_options, + accessor, + throttle_state.as_ref(), + ) + .await?; + ( + Arc::clone(&native) as Arc, + Some(ConditionalDeleteConfig::SignedUrl(native)), ) - .await? }; let inner = if let Some(throttle_state) = throttle_state { Arc::new(AimdThrottledStore::new_with_state( @@ -157,6 +166,7 @@ impl ObjectStoreProvider for GcsStoreProvider { io_tracker: Default::default(), store_prefix: self .calculate_object_store_prefix(&base_path, params.storage_options())?, + conditional_delete, }) } } diff --git a/rust/lance-io/src/object_store/providers/goosefs.rs b/rust/lance-io/src/object_store/providers/goosefs.rs index fe3002bccc0..3010c51c581 100644 --- a/rust/lance-io/src/object_store/providers/goosefs.rs +++ b/rust/lance-io/src/object_store/providers/goosefs.rs @@ -190,6 +190,7 @@ impl ObjectStoreProvider for GooseFsStoreProvider { io_tracker: Default::default(), store_prefix: self .calculate_object_store_prefix(&base_path, params.storage_options())?, + conditional_delete: None, }) } diff --git a/rust/lance-io/src/object_store/providers/huggingface.rs b/rust/lance-io/src/object_store/providers/huggingface.rs index cda56e36fbe..1af63c228b3 100644 --- a/rust/lance-io/src/object_store/providers/huggingface.rs +++ b/rust/lance-io/src/object_store/providers/huggingface.rs @@ -216,6 +216,7 @@ impl ObjectStoreProvider for HuggingfaceStoreProvider { io_tracker: Default::default(), store_prefix: self .calculate_object_store_prefix(&base_path, params.storage_options())?, + conditional_delete: None, }) } diff --git a/rust/lance-io/src/object_store/providers/local.rs b/rust/lance-io/src/object_store/providers/local.rs index 9f0762916f7..f06610dec1b 100644 --- a/rust/lance-io/src/object_store/providers/local.rs +++ b/rust/lance-io/src/object_store/providers/local.rs @@ -33,6 +33,14 @@ impl ObjectStoreProvider for FileStoreProvider { io_tracker: Default::default(), store_prefix: self .calculate_object_store_prefix(&base_path, params.storage_options())?, + // Local filesystems have no storage-atomic compare-and-unlink + // primitive, so production reference mutations fail closed. The + // test utility feature serializes the complete in-process writer + // population used by reference integration tests. + #[cfg(feature = "test-util")] + conditional_delete: Some(super::super::ConditionalDeleteConfig::Serialized), + #[cfg(not(feature = "test-util"))] + conditional_delete: None, }) } diff --git a/rust/lance-io/src/object_store/providers/memory.rs b/rust/lance-io/src/object_store/providers/memory.rs index dd72edc4627..f47199e695a 100644 --- a/rust/lance-io/src/object_store/providers/memory.rs +++ b/rust/lance-io/src/object_store/providers/memory.rs @@ -33,6 +33,7 @@ impl ObjectStoreProvider for MemoryStoreProvider { io_tracker: Default::default(), store_prefix: self .calculate_object_store_prefix(&base_path, params.storage_options())?, + conditional_delete: Some(super::super::ConditionalDeleteConfig::Serialized), }) } diff --git a/rust/lance-io/src/object_store/providers/oss.rs b/rust/lance-io/src/object_store/providers/oss.rs index 3d116e2e3cc..5186716c030 100644 --- a/rust/lance-io/src/object_store/providers/oss.rs +++ b/rust/lance-io/src/object_store/providers/oss.rs @@ -144,6 +144,7 @@ impl ObjectStoreProvider for OssStoreProvider { download_retry_count: storage_options.download_retry_count(), io_tracker: Default::default(), store_prefix: self.calculate_object_store_prefix(&url, params.storage_options())?, + conditional_delete: None, }) } } diff --git a/rust/lance-io/src/object_store/providers/tencent.rs b/rust/lance-io/src/object_store/providers/tencent.rs index d29d5a6ad62..ecbc96df434 100644 --- a/rust/lance-io/src/object_store/providers/tencent.rs +++ b/rust/lance-io/src/object_store/providers/tencent.rs @@ -100,6 +100,7 @@ impl ObjectStoreProvider for TencentStoreProvider { download_retry_count: storage_options.download_retry_count(), io_tracker: Default::default(), store_prefix: self.calculate_object_store_prefix(&url, params.storage_options())?, + conditional_delete: None, }) } } diff --git a/rust/lance-io/src/object_store/providers/tos.rs b/rust/lance-io/src/object_store/providers/tos.rs index 7dee659f5f9..92ecf62a9ee 100644 --- a/rust/lance-io/src/object_store/providers/tos.rs +++ b/rust/lance-io/src/object_store/providers/tos.rs @@ -150,6 +150,7 @@ impl ObjectStoreProvider for TosStoreProvider { download_retry_count: storage_options.download_retry_count(), io_tracker: Default::default(), store_prefix: self.calculate_object_store_prefix(&url, params.storage_options())?, + conditional_delete: None, }) } } diff --git a/rust/lance/src/dataset/version_lease.rs b/rust/lance/src/dataset/version_lease.rs index 82b8601a59d..1e27485c62d 100644 --- a/rust/lance/src/dataset/version_lease.rs +++ b/rust/lance/src/dataset/version_lease.rs @@ -15,7 +15,7 @@ use bytes::Bytes; use chrono::{DateTime, TimeDelta, Utc}; use dashmap::DashSet; use futures::{StreamExt, TryStreamExt, stream}; -use lance_io::object_store::ObjectStore; +use lance_io::object_store::{ConditionalDeleteResult, ObjectStore}; use object_store::{ObjectMeta, ObjectStoreExt, PutMode, PutOptions, UpdateVersion, path::Path}; use serde::{Deserialize, Serialize}; use uuid::Uuid; @@ -39,11 +39,6 @@ const COMMITTED_MARKER_SUFFIX: &str = ".committed"; const STORAGE_CLOCK_PATH: &str = "_clock"; const REFERENCE_GENERATION_FIELD: &str = "_lanceReferenceGeneration"; -/// Marks a canonical payload as a logical deletion fence. The generic object -/// store API has conditional writes but no conditional delete, so retaining a -/// byte-changing tombstone is what prevents a later released-client write from -/// being erased by an unconditional physical delete. -const REFERENCE_DELETED_FIELD: &str = "_lanceReferenceDeleted"; /// HTTP `Last-Modified`, used by supported cloud stores, has whole-second precision. /// Adding one interval guarantees a lease is never shortened by timestamp truncation. @@ -272,6 +267,8 @@ enum ReferenceLifecycleState { previous: Option, previous_was_legacy: bool, expected_payload: Vec, + expected_etag: Option, + expected_version: Option, }, } @@ -558,6 +555,39 @@ impl VersionLeaseStore { } } + async fn put_terminal_reference_lifecycle( + &self, + canonical_path: &Path, + expected: &ReferenceLifecycleSnapshot, + terminal: &ReferenceLifecycleState, + has_local_lock: bool, + ) -> Result { + let Some(snapshot) = self + .put_reference_lifecycle(canonical_path, Some(expected), terminal, has_local_lock) + .await? + else { + return Ok(false); + }; + if matches!(terminal, ReferenceLifecycleState::Vacant { .. }) + && matches!(snapshot.state, ReferenceLifecycleState::Vacant { .. }) + { + let expected = UpdateVersion { + e_tag: snapshot.metadata.e_tag.clone(), + version: snapshot.metadata.version.clone(), + }; + match self + .object_store + .delete_if_matches(&snapshot.metadata.location, &expected) + .await? + { + ConditionalDeleteResult::Deleted + | ConditionalDeleteResult::NotFound + | ConditionalDeleteResult::IdentityMismatch => {} + } + } + Ok(true) + } + async fn claim_reference_lifecycle( &self, canonical_path: &Path, @@ -586,12 +616,16 @@ impl VersionLeaseStore { "reference {canonical_path} lost its revoking state" )) })?; - self.reconcile_revoking_reference( - canonical_path, - current, - local_lock.is_some(), - ) - .await?; + if !self + .reconcile_revoking_reference(canonical_path, current, local_lock.is_some()) + .await? + { + return Err(Error::RefConflict { + message: format!( + "reference {canonical_path} is settling a revoked mutation" + ), + }); + } continue; } Some(ReferenceLifecycleState::Deleting { .. }) => { @@ -661,9 +695,7 @@ impl VersionLeaseStore { Err(object_store::Error::NotFound { .. }) => return Ok(false), Err(error) => return Err(error.into()), }; - if payload_reference_generation(&canonical)?.as_deref() != Some(operation_id) - || payload_reference_is_deleted(&canonical)? - { + if payload_reference_generation(&canonical)?.as_deref() != Some(operation_id) { return Ok(false); } let live = ReferenceLifecycleState::Live { @@ -727,10 +759,8 @@ impl VersionLeaseStore { let vacant = ReferenceLifecycleState::Vacant { canonical_path: canonical_path.to_string(), }; - Ok(self - .put_reference_lifecycle(canonical_path, Some(¤t), &vacant, has_local_lock) - .await? - .is_some()) + self.put_terminal_reference_lifecycle(canonical_path, ¤t, &vacant, has_local_lock) + .await } fn branch_termination_path(&self) -> Path { @@ -814,6 +844,7 @@ impl VersionLeaseStore { } pub(super) async fn active_reference_versions(&self) -> Result> { + self.reconcile_reference_lifecycles().await?; let mut versions = self .reference_versions( CompletedReferenceIntentHandling::RetainForCurrentScan, @@ -826,6 +857,72 @@ impl VersionLeaseStore { Ok(versions) } + async fn reconcile_reference_lifecycles(&self) -> Result<()> { + let state_path = self.root_path.clone().join(REFERENCE_STATES_DIR); + let metadata = self + .object_store + .list(Some(state_path)) + .try_collect::>() + .await?; + for listed in metadata { + let result = match self.object_store.inner.get(&listed.location).await { + Ok(result) => result, + Err(object_store::Error::NotFound { .. }) => continue, + Err(error) => return Err(error.into()), + }; + let current_metadata = result.meta.clone(); + let payload = result.bytes().await?; + let state: ReferenceLifecycleState = serde_json::from_slice(&payload) + .map_err(|error| Error::corrupt_file(listed.location.clone(), error.to_string()))?; + let canonical_path = match &state { + ReferenceLifecycleState::Revoking { canonical_path, .. } + | ReferenceLifecycleState::Deleting { canonical_path, .. } => { + Some(Path::parse(canonical_path)?) + } + ReferenceLifecycleState::Vacant { .. } => { + if self.object_store.supports_conditional_delete() { + let expected = UpdateVersion { + e_tag: current_metadata.e_tag, + version: current_metadata.version, + }; + self.object_store + .delete_if_matches(&listed.location, &expected) + .await?; + } + None + } + _ => None, + }; + let Some(canonical_path) = canonical_path else { + continue; + }; + let local_lock = lock_local_reference(&self.object_store, &canonical_path).await?; + let Some(snapshot) = self.reference_lifecycle_snapshot(&canonical_path).await? else { + continue; + }; + match &snapshot.state { + ReferenceLifecycleState::Revoking { .. } => { + self.reconcile_revoking_reference( + &canonical_path, + &snapshot, + local_lock.is_some(), + ) + .await?; + } + ReferenceLifecycleState::Deleting { .. } => { + self.reconcile_deleting_reference( + &canonical_path, + &snapshot, + local_lock.is_some(), + ) + .await?; + } + _ => {} + } + } + Ok(()) + } + async fn lifecycle_reference_versions(&self) -> Result> { let state_path = self.root_path.clone().join(REFERENCE_STATES_DIR); let metadata = self @@ -968,9 +1065,6 @@ impl VersionLeaseStore { Err(object_store::Error::NotFound { .. }) => return Ok(None), Err(error) => return Err(error.into()), }; - if payload_reference_is_deleted(&payload)? { - return Ok(None); - } let generation = payload_reference_generation(&payload)?; Ok(expected.and_then(|live| { (generation.as_deref() == Some(live.generation.as_str())).then(|| live.target.clone()) @@ -1004,6 +1098,7 @@ impl VersionLeaseStore { pub(super) async fn reference_versions_before_canonical_census( &self, ) -> Result { + self.reconcile_reference_lifecycles().await?; let mut census = self .reference_versions( CompletedReferenceIntentHandling::DeferToCanonicalCensus, @@ -1225,10 +1320,7 @@ impl VersionLeaseStore { has_local_lock: bool, ) -> Result { let ReferenceLifecycleState::Revoking { - operation_id, - previous, - mutation, - .. + previous, mutation, .. } = &revoking.state else { return Ok(false); @@ -1241,6 +1333,18 @@ impl VersionLeaseStore { Err(object_store::Error::NotFound { .. }) => None, Err(error) => return Err(error.into()), }; + if matches!(mutation, ReferenceMutation::Create { .. }) + && current.is_none() + && reference_owner_is_active( + revoking.metadata.last_modified, + self.storage_observed_at().await?, + )? + { + // A conditional create may already be in flight when its owner is + // revoked. Keep the generation until the request timeout window is + // closed so a late exact-incarnation write remains recoverable. + return Ok(false); + } match (mutation, current.as_ref()) { ( ReferenceMutation::Update { @@ -1275,47 +1379,25 @@ impl VersionLeaseStore { ) .await?; } - (ReferenceMutation::Create { payload, .. }, None) => { - let tombstone = set_payload_reference_deleted(payload, operation_id)?; - match self - .object_store - .inner - .put_opts( - canonical_path, - Bytes::from(tombstone).into(), - PutOptions { - mode: PutMode::Create, - ..Default::default() - }, - ) - .await - { - Ok(_) | Err(object_store::Error::AlreadyExists { .. }) => {} - Err(error) => return Err(error.into()), - } - } + (ReferenceMutation::Create { .. }, None) => {} (ReferenceMutation::Create { payload, .. }, Some((metadata, current_payload))) if current_payload.as_ref() == payload => { - let tombstone = set_payload_reference_deleted(payload, operation_id)?; - self.rewrite_reference_payload( - canonical_path, - metadata, - current_payload, - &tombstone, - has_local_lock, - ) - .await?; + let expected = UpdateVersion { + e_tag: metadata.e_tag.clone(), + version: metadata.version.clone(), + }; + self.object_store + .delete_if_matches(canonical_path, &expected) + .await?; } _ => {} } let terminal = self .terminal_reference_lifecycle(canonical_path, previous.as_ref()) .await?; - Ok(self - .put_reference_lifecycle(canonical_path, Some(revoking), &terminal, has_local_lock) - .await? - .is_some()) + self.put_terminal_reference_lifecycle(canonical_path, revoking, &terminal, has_local_lock) + .await } async fn reconcile_deleting_reference( @@ -1325,51 +1407,28 @@ impl VersionLeaseStore { has_local_lock: bool, ) -> Result { let ReferenceLifecycleState::Deleting { - operation_id, previous, - expected_payload, + expected_etag, + expected_version, .. } = &deleting.state else { return Ok(false); }; - let current = match self.object_store.inner.get(canonical_path).await { - Ok(result) => Some(result), - Err(object_store::Error::NotFound { .. }) => None, - Err(error) => return Err(error.into()), + let expected = UpdateVersion { + e_tag: expected_etag.clone(), + version: expected_version.clone(), }; - if let Some(result) = current { - let metadata = result.meta.clone(); - let current_payload = result.bytes().await?; - if current_payload.as_ref() == expected_payload { - // This conditional update is the deletion linearization point. - // A released writer that lands first changes the incarnation - // and makes this update fail; one that lands later replaces the - // tombstone and becomes the visible canonical reference. - let tombstone = set_payload_reference_deleted(expected_payload, operation_id)?; - match self - .rewrite_reference_payload( - canonical_path, - &metadata, - ¤t_payload, - &tombstone, - has_local_lock, - ) - .await - { - Ok(()) | Err(Error::RefConflict { .. }) => {} - Err(error) => return Err(error), - } - } - } + self.object_store + .delete_if_matches(canonical_path, &expected) + .await?; let terminal = self .terminal_reference_lifecycle(canonical_path, previous.as_ref()) .await?; let is_deleted = matches!(terminal, ReferenceLifecycleState::Vacant { .. }); - if self - .put_reference_lifecycle(canonical_path, Some(deleting), &terminal, has_local_lock) + if !self + .put_terminal_reference_lifecycle(canonical_path, deleting, &terminal, has_local_lock) .await? - .is_none() { return Err(Error::RefConflict { message: format!( @@ -1395,11 +1454,6 @@ impl VersionLeaseStore { canonical_path: canonical_path.to_string(), }); }; - if payload_reference_is_deleted(&payload)? { - return Ok(ReferenceLifecycleState::Vacant { - canonical_path: canonical_path.to_string(), - }); - } let generation = payload_reference_generation(&payload)?; if let Some(previous) = previous && generation.as_deref() == Some(previous.generation.as_str()) @@ -2592,6 +2646,12 @@ pub(super) async fn begin_reference_admission( let operation_id = Uuid::new_v4().simple().to_string(); let (canonical_path, previous_was_legacy) = match &mut mutation { ReferenceMutation::Create { path, payload } => { + if !store.object_store.supports_conditional_delete() { + return Err(Error::not_supported(format!( + "object store {} cannot safely roll back reference creation for {path} because atomic conditional delete is unavailable", + store.object_store.scheme() + ))); + } *payload = set_payload_reference_generation(payload, &operation_id)?; (Path::parse(path)?, false) } @@ -2606,26 +2666,6 @@ pub(super) async fn begin_reference_admission( (Path::parse(path)?, previous_was_legacy) } }; - if let ReferenceMutation::Create { path, payload } = &mutation { - let existing = match store.object_store.inner.get(&canonical_path).await { - Ok(result) => Some(result), - Err(object_store::Error::NotFound { .. }) => None, - Err(error) => return Err(error.into()), - }; - if let Some(result) = existing { - let metadata = result.meta.clone(); - let existing_payload = result.bytes().await?; - if payload_reference_is_deleted(&existing_payload)? { - mutation = ReferenceMutation::Update { - path: path.clone(), - expected_payload: existing_payload.to_vec(), - expected_etag: metadata.e_tag, - expected_version: metadata.version, - payload: payload.clone(), - }; - } - } - } let intent = ReferenceIntent { manifest_path: manifest.path.to_string(), mutation: mutation.clone(), @@ -2778,19 +2818,46 @@ impl ReferenceAdmission { } async fn cancel_before_publish_inner(&self, has_local_lock: bool) { - if let Ok(canonical_path) = self.canonical_path() - && let Err(error) = self - .store - .cancel_reference_lifecycle(&canonical_path, &self.operation_id, has_local_lock) - .await + let Ok(canonical_path) = self.canonical_path() else { + return; + }; + if let Err(error) = self + .store + .cancel_reference_lifecycle(&canonical_path, &self.operation_id, has_local_lock) + .await { tracing::warn!( path = %canonical_path, error = %error, "Failed to cancel reference lifecycle admission" ); + return; } - if let Err(error) = self.store.object_store.delete(&self.path).await + let matching_pending_is_absent = match self + .store + .reference_lifecycle_snapshot(&canonical_path) + .await + { + Ok(Some(ReferenceLifecycleSnapshot { + state: + ReferenceLifecycleState::Pending { + operation_id: current_operation, + .. + }, + .. + })) => current_operation != self.operation_id, + Ok(_) => true, + Err(error) => { + tracing::warn!( + path = %canonical_path, + error = %error, + "Failed to confirm cancelled reference lifecycle admission" + ); + false + } + }; + if matching_pending_is_absent + && let Err(error) = self.store.object_store.delete(&self.path).await && !error.is_not_found() { tracing::warn!( @@ -2889,31 +2956,6 @@ fn set_payload_reference_generation(payload: &[u8], generation: &str) -> Result< Ok(serde_json::to_vec_pretty(&value)?) } -fn set_payload_reference_deleted(payload: &[u8], generation: &str) -> Result> { - let mut value: serde_json::Value = serde_json::from_slice(payload)?; - let object = value.as_object_mut().ok_or_else(|| { - Error::internal("canonical reference payload must be a JSON object".to_string()) - })?; - object.insert( - REFERENCE_GENERATION_FIELD.to_string(), - serde_json::Value::String(generation.to_string()), - ); - object.insert( - REFERENCE_DELETED_FIELD.to_string(), - serde_json::Value::Bool(true), - ); - Ok(serde_json::to_vec_pretty(&value)?) -} - -fn payload_reference_is_deleted(payload: &[u8]) -> Result { - let value: serde_json::Value = serde_json::from_slice(payload)?; - Ok(value - .as_object() - .and_then(|object| object.get(REFERENCE_DELETED_FIELD)) - .and_then(serde_json::Value::as_bool) - .unwrap_or(false)) -} - fn fenced_reference_payload(payload: &[u8]) -> Vec { let mut fenced = Vec::with_capacity(payload.len() + 1); fenced.extend_from_slice(payload); @@ -2927,9 +2969,6 @@ pub(super) async fn canonical_reference_is_visible( canonical_path: &Path, payload: &[u8], ) -> Result { - if payload_reference_is_deleted(payload)? { - return Ok(false); - } let generation = payload_reference_generation(payload)?; // Released writers do not preserve fields they do not know about. Their // rewrites remain authoritative canonical references and must stay visible @@ -2964,6 +3003,14 @@ pub(super) async fn delete_canonical_reference( expected_metadata: &ObjectMeta, expected_payload: &[u8], ) -> Result<()> { + if !object_store.supports_conditional_delete() + || (expected_metadata.e_tag.is_none() && expected_metadata.version.is_none()) + { + return Err(Error::not_supported(format!( + "object store {} cannot safely delete reference {canonical_path} because atomic conditional delete is unavailable", + object_store.scheme() + ))); + } let store = VersionLeaseStore::lifecycle_only(Arc::clone(&object_store), root_path.clone()); let local_lock = lock_local_reference(&object_store, canonical_path).await?; let current = object_store @@ -3046,6 +3093,8 @@ pub(super) async fn delete_canonical_reference( previous, previous_was_legacy: canonical_generation.is_none(), expected_payload: current_payload.to_vec(), + expected_etag: current_metadata.e_tag, + expected_version: current_metadata.version, }; let Some(snapshot) = store .put_reference_lifecycle( @@ -4037,16 +4086,15 @@ mod tests { .unwrap(), ReferenceMutationOutcome::Conflict ); - let rollback = store - .object_store - .inner - .get(&canonical_path) - .await - .unwrap() - .bytes() - .await - .unwrap(); - assert!(payload_reference_is_deleted(&rollback).unwrap()); + assert!(!store.object_store.exists(&canonical_path).await.unwrap()); + assert!(matches!( + store + .reference_lifecycle_snapshot(&canonical_path) + .await + .unwrap() + .map(|snapshot| snapshot.state), + Some(ReferenceLifecycleState::Revoking { .. }) + )); } #[tokio::test] @@ -4260,16 +4308,7 @@ mod tests { .unwrap(), ReferenceMutationOutcome::Conflict ); - let rollback = store - .object_store - .inner - .get(&canonical_path) - .await - .unwrap() - .bytes() - .await - .unwrap(); - assert!(payload_reference_is_deleted(&rollback).unwrap()); + assert!(!store.object_store.exists(&canonical_path).await.unwrap()); } #[tokio::test] @@ -4307,18 +4346,16 @@ mod tests { .apply_reference_mutation_inner(&admission.mutation, false, None) .await .unwrap(), + ReferenceMutationOutcome::Published + ); + assert_eq!( + store + .finish_reference_mutation(&admission.mutation, &admission.operation_id, 42, false,) + .await + .unwrap(), ReferenceMutationOutcome::Conflict ); - let rollback = store - .object_store - .inner - .get(&canonical_path) - .await - .unwrap() - .bytes() - .await - .unwrap(); - assert!(payload_reference_is_deleted(&rollback).unwrap()); + assert!(!store.object_store.exists(&canonical_path).await.unwrap()); } #[tokio::test] @@ -4448,7 +4485,7 @@ mod tests { failing_store.fail_when( "put", "rollback-crash.json", - "injected rollback tombstone failure", + "injected conditional rollback delete failure", ); store @@ -4483,16 +4520,7 @@ mod tests { .unwrap(), ReferenceMutationOutcome::Conflict ); - let rollback = store - .object_store - .inner - .get(&canonical_path) - .await - .unwrap() - .bytes() - .await - .unwrap(); - assert!(payload_reference_is_deleted(&rollback).unwrap()); + assert!(!store.object_store.exists(&canonical_path).await.unwrap()); } #[tokio::test] @@ -4537,6 +4565,120 @@ mod tests { assert!(store.object_store.exists(&intent_path).await.unwrap()); } + #[tokio::test] + async fn cancelled_admission_lifecycle_failure_keeps_recovery_anchor() { + let failing_store = Arc::new(FailingProxyStore::new()); + let mut object_store = ObjectStore::memory(); + object_store.inner = failing_store.wrap("memory", Arc::clone(&object_store.inner)); + let store = VersionLeaseStore { + object_store: Arc::new(object_store), + root_path: Path::from(""), + namespace: MAIN_BRANCH.to_string(), + leases_path: Path::from("leases"), + markers_path: Path::from("markers"), + reference_intents_path: Path::from("reference_intents"), + manifest_path: None, + canonical_references: None, + }; + let manifest_path = Path::from("manifests/42.manifest"); + let canonical_path = Path::from("tags/cancel-failure.json"); + store.object_store.put(&manifest_path, &[]).await.unwrap(); + let admission = create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + br#"{"version":42}"#, + ) + .await; + failing_store.fail_when( + "put", + "version_reference_states", + "injected lifecycle cancellation failure", + ); + + admission.cancel_before_publish().await; + + assert!(store.object_store.exists(&admission.path).await.unwrap()); + assert!(matches!( + store + .reference_lifecycle_snapshot(&canonical_path) + .await + .unwrap() + .map(|snapshot| snapshot.state), + Some(ReferenceLifecycleState::Pending { operation_id, .. }) + if operation_id == admission.operation_id + )); + failing_store.clear_fail_when("put", "version_reference_states"); + admission.cancel_before_publish().await; + assert!(!store.object_store.exists(&admission.path).await.unwrap()); + assert!( + store + .reference_lifecycle_snapshot(&canonical_path) + .await + .unwrap() + .is_none() + ); + } + + #[tokio::test] + async fn conditional_delete_is_absent_to_released_tag_clients() { + let store = memory_store(); + let manifest_path = Path::from("manifests/42.manifest"); + let canonical_path = Path::from("tags/released-delete.json"); + let released_payload = br#"{"branch":null,"version":42,"manifestSize":0,"metadata":{}}"#; + store.object_store.put(&manifest_path, &[]).await.unwrap(); + create_test_reference_admission( + &store, + 42, + &manifest_path, + &canonical_path, + released_payload, + ) + .await + .publish("conflict".to_string()) + .await + .unwrap(); + let result = store.object_store.inner.get(&canonical_path).await.unwrap(); + let metadata = result.meta.clone(); + let payload = result.bytes().await.unwrap(); + serde_json::from_slice::(&payload).unwrap(); + + delete_canonical_reference( + Arc::clone(&store.object_store), + &store.root_path, + &canonical_path, + &metadata, + &payload, + ) + .await + .unwrap(); + + assert!(!store.object_store.exists(&canonical_path).await.unwrap()); + assert!( + !store + .object_store + .read_dir(Path::from("tags")) + .await + .unwrap() + .iter() + .any(|name| name == "released-delete.json") + ); + store + .object_store + .inner + .put_opts( + &canonical_path, + Bytes::from_static(released_payload).into(), + PutOptions { + mode: PutMode::Create, + ..Default::default() + }, + ) + .await + .unwrap(); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn released_writer_update_during_delete_is_not_erased() { use lance_core::utils::testing::{ProxyObjectStore, ProxyObjectStorePolicy}; @@ -4582,7 +4724,7 @@ mod tests { let resume_put_rx = Arc::new(std::sync::Mutex::new(resume_put_rx)); let canonical_path_string = canonical_path.to_string(); policy.lock().unwrap().set_before_policy( - "pause_canonical_tombstone", + "pause_conditional_delete", Arc::new(move |method, path| { if method == "put" && path.as_ref() == canonical_path_string @@ -4681,25 +4823,24 @@ mod tests { .await .unwrap_err(); - assert!(store.active_reference_versions().await.unwrap().is_empty()); - let tombstone = store - .object_store - .inner - .get(&canonical_path) + assert!(!store.object_store.exists(&canonical_path).await.unwrap()); + let snapshot = store + .reference_lifecycle_snapshot(&canonical_path) .await .unwrap() - .bytes() - .await .unwrap(); - assert!(payload_reference_is_deleted(&tombstone).unwrap()); assert!(matches!( + &snapshot.state, + ReferenceLifecycleState::Deleting { .. } + )); + assert!( store - .reference_lifecycle_snapshot(&canonical_path) + .retained_lifecycle_target(&snapshot.state) .await .unwrap() - .map(|snapshot| snapshot.state), - Some(ReferenceLifecycleState::Deleting { .. }) - )); + .is_none(), + "an absent canonical reference must not retain its deleted version" + ); } #[tokio::test] @@ -4737,7 +4878,7 @@ mod tests { failing_store.fail_when( "put", "deletion-failure.json", - "injected canonical tombstone failure", + "injected conditional delete failure", ); let error = delete_canonical_reference( @@ -4752,7 +4893,7 @@ mod tests { assert!( error .to_string() - .contains("injected canonical tombstone failure"), + .contains("injected conditional delete failure"), "{error}" ); failing_store.clear_fail_when("put", "deletion-failure.json"); @@ -4780,15 +4921,13 @@ mod tests { ) .await .unwrap(); - let tombstone = store - .object_store - .inner - .get(&canonical_path) - .await - .unwrap() - .bytes() - .await - .unwrap(); - assert!(payload_reference_is_deleted(&tombstone).unwrap()); + assert!(!store.object_store.exists(&canonical_path).await.unwrap()); + assert!( + store + .reference_lifecycle_snapshot(&canonical_path) + .await + .unwrap() + .is_none() + ); } } From a9d014b837c744710361d2094df952f8da4f7558 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Sat, 8 Aug 2026 09:49:29 +0000 Subject: [PATCH 13/13] fix: restore non-cloud reference test CI --- java/src/test/java/org/lance/CleanupTest.java | 2 +- java/src/test/java/org/lance/DatasetTest.java | 12 +++--- java/src/test/java/org/lance/TestUtils.java | 4 ++ python/python/tests/test_namespace_dir.py | 41 +++++++++++-------- rust/lance-io/Cargo.toml | 31 ++++++++++++-- rust/lance-io/src/object_store.rs | 7 ++++ 6 files changed, 70 insertions(+), 27 deletions(-) diff --git a/java/src/test/java/org/lance/CleanupTest.java b/java/src/test/java/org/lance/CleanupTest.java index 5fc8ceeaa3f..a6fa524a39f 100644 --- a/java/src/test/java/org/lance/CleanupTest.java +++ b/java/src/test/java/org/lance/CleanupTest.java @@ -114,7 +114,7 @@ public void testCleanupBeforeTimestamp(@TempDir Path tempDir) throws Exception { @Test public void testCleanupTaggedVersion(@TempDir Path tempDir) throws Exception { - String datasetPath = tempDir.resolve("test_dataset_for_cleanup").toString(); + String datasetPath = TestUtils.sharedMemoryUri(tempDir, "test_dataset_for_cleanup"); try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { TestUtils.SimpleTestDataset testDataset = new TestUtils.SimpleTestDataset(allocator, datasetPath); diff --git a/java/src/test/java/org/lance/DatasetTest.java b/java/src/test/java/org/lance/DatasetTest.java index 5999cbc0b19..024424920e5 100644 --- a/java/src/test/java/org/lance/DatasetTest.java +++ b/java/src/test/java/org/lance/DatasetTest.java @@ -308,7 +308,7 @@ void testDatasetCheckoutVersion(@TempDir Path tempDir) { @Test void testTags(@TempDir Path tempDir) { - String datasetPath = tempDir.resolve("dataset_tags").toString(); + String datasetPath = TestUtils.sharedMemoryUri(tempDir, "dataset_tags"); try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { TestUtils.SimpleTestDataset testDataset = new TestUtils.SimpleTestDataset(allocator, datasetPath); @@ -453,7 +453,7 @@ void testTags(@TempDir Path tempDir) { @Test void testDatasetRestore(@TempDir Path tempDir) { - String datasetPath = tempDir.resolve("dataset_restore").toString(); + String datasetPath = TestUtils.sharedMemoryUri(tempDir, "dataset_restore"); try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { TestUtils.SimpleTestDataset testDataset = new TestUtils.SimpleTestDataset(allocator, datasetPath); @@ -1785,9 +1785,9 @@ public ByteBuffer readManifest(Path filePath) throws IOException { @Test void testShallowClone(@TempDir Path tempDir) { - String srcPath = tempDir.resolve("shallow_clone_version_src").toString(); - String dstPathByVersion = tempDir.resolve("shallow_clone_version_dst").toString(); - String dstPathByTag = tempDir.resolve("shallow_clone_tag_dst").toString(); + String srcPath = TestUtils.sharedMemoryUri(tempDir, "shallow_clone_version_src"); + String dstPathByVersion = TestUtils.sharedMemoryUri(tempDir, "shallow_clone_version_dst"); + String dstPathByTag = TestUtils.sharedMemoryUri(tempDir, "shallow_clone_tag_dst"); try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { // Prepare a simple source dataset with some rows @@ -1841,7 +1841,7 @@ void testShallowClone(@TempDir Path tempDir) { @Test void testBranches(@TempDir Path tempDir) { - String datasetPath = tempDir.resolve("testBranches").toString(); + String datasetPath = TestUtils.sharedMemoryUri(tempDir, "testBranches"); try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { TestUtils.SimpleTestDataset suite = new TestUtils.SimpleTestDataset(allocator, datasetPath); diff --git a/java/src/test/java/org/lance/TestUtils.java b/java/src/test/java/org/lance/TestUtils.java index 1989a5243bc..3e044bf22a7 100644 --- a/java/src/test/java/org/lance/TestUtils.java +++ b/java/src/test/java/org/lance/TestUtils.java @@ -67,6 +67,10 @@ import static org.junit.jupiter.api.Assertions.assertTrue; public class TestUtils { + public static String sharedMemoryUri(Path tempDir, String datasetName) { + return String.format("shared-memory://%s/%s", tempDir.getFileName(), datasetName); + } + private abstract static class TestDataset { protected final BufferAllocator allocator; protected final String datasetPath; diff --git a/python/python/tests/test_namespace_dir.py b/python/python/tests/test_namespace_dir.py index fa1bc93b422..06377652ec7 100644 --- a/python/python/tests/test_namespace_dir.py +++ b/python/python/tests/test_namespace_dir.py @@ -259,6 +259,15 @@ def memory_ns_client(request): yield _wrap_if_custom(ns_client, use_custom) +@pytest.fixture(params=[False, True], ids=["DirectoryNamespace", "CustomNamespace"]) +def shared_memory_ns_client(request): + """Create a memory namespace shared across object-store instances.""" + use_custom = request.param + unique_id = uuid.uuid4().hex[:8] + ns_client = connect("dir", {"root": f"shared-memory://test-{unique_id}"}) + yield _wrap_if_custom(ns_client, use_custom) + + class TestCreateTable: """Tests for create_table operation - mirrors Rust test_create_table.""" @@ -591,17 +600,17 @@ class TestTableBranchOperations: """Branch CRUD through the python bindings - mirrors the Rust branch CRUD tests.""" - def test_branch_crud_round_trip(self, temp_ns_client): + def test_branch_crud_round_trip(self, shared_memory_ns_client): create_ns_req = CreateNamespaceRequest(id=["workspace"]) - temp_ns_client.create_namespace(create_ns_req) + shared_memory_ns_client.create_namespace(create_ns_req) ipc_data = table_to_ipc_bytes(create_test_data()) table_id = ["workspace", "branched_table"] - temp_ns_client.create_table(CreateTableRequest(id=table_id), ipc_data) + shared_memory_ns_client.create_table(CreateTableRequest(id=table_id), ipc_data) - temp_ns_client.create_table_branch( + shared_memory_ns_client.create_table_branch( CreateTableBranchRequest(id=table_id, name="dev") ) - listed = temp_ns_client.list_table_branches( + listed = shared_memory_ns_client.list_table_branches( ListTableBranchesRequest(id=table_id) ) assert "dev" in listed.branches @@ -609,41 +618,41 @@ def test_branch_crud_round_trip(self, temp_ns_client): # Duplicate creation and deleting a missing branch surface the typed # branch errors (codes 23 and 22), not InternalError. - temp_ns_client.create_table_branch( + shared_memory_ns_client.create_table_branch( CreateTableBranchRequest(id=table_id, name="dev2") ) with pytest.raises(TableBranchAlreadyExistsError): - temp_ns_client.create_table_branch( + shared_memory_ns_client.create_table_branch( CreateTableBranchRequest(id=table_id, name="dev2") ) - temp_ns_client.delete_table_branch( + shared_memory_ns_client.delete_table_branch( DeleteTableBranchRequest(id=table_id, name="dev") ) - listed = temp_ns_client.list_table_branches( + listed = shared_memory_ns_client.list_table_branches( ListTableBranchesRequest(id=table_id) ) assert "dev" not in listed.branches with pytest.raises(TableBranchNotFoundError): - temp_ns_client.delete_table_branch( + shared_memory_ns_client.delete_table_branch( DeleteTableBranchRequest(id=table_id, name="dev") ) - def test_create_branch_from_other_branch(self, temp_ns_client): + def test_create_branch_from_other_branch(self, shared_memory_ns_client): """Forking from a non-main source branch records the right parent.""" create_ns_req = CreateNamespaceRequest(id=["workspace"]) - temp_ns_client.create_namespace(create_ns_req) + shared_memory_ns_client.create_namespace(create_ns_req) ipc_data = table_to_ipc_bytes(create_test_data()) table_id = ["workspace", "fork_table"] - temp_ns_client.create_table(CreateTableRequest(id=table_id), ipc_data) + shared_memory_ns_client.create_table(CreateTableRequest(id=table_id), ipc_data) - temp_ns_client.create_table_branch( + shared_memory_ns_client.create_table_branch( CreateTableBranchRequest(id=table_id, name="dev") ) - temp_ns_client.create_table_branch( + shared_memory_ns_client.create_table_branch( CreateTableBranchRequest(id=table_id, name="child", from_branch="dev") ) - listed = temp_ns_client.list_table_branches( + listed = shared_memory_ns_client.list_table_branches( ListTableBranchesRequest(id=table_id) ) assert listed.branches["child"].parent_branch == "dev" diff --git a/rust/lance-io/Cargo.toml b/rust/lance-io/Cargo.toml index 34a314fdcd2..e7f2516932a 100644 --- a/rust/lance-io/Cargo.toml +++ b/rust/lance-io/Cargo.toml @@ -41,7 +41,9 @@ tracing.workspace = true url.workspace = true path_abs.workspace = true rand.workspace = true -reqwest = { version = "0.12", default-features = false, features = ["rustls-tls"] } +reqwest = { version = "0.12", default-features = false, features = [ + "rustls-tls", +], optional = true } tempfile.workspace = true [target.'cfg(target_os = "linux")'.dependencies] @@ -67,9 +69,30 @@ default = ["aws", "azure", "gcp"] metrics = ["dep:metrics"] gcs-test = [] goosefs-test = [] -gcp = ["object_store/gcp", "dep:opendal", "opendal/services-gcs", "dep:object_store_opendal"] -aws = ["object_store/aws", "dep:aws-config", "dep:aws-credential-types", "dep:opendal", "opendal/services-s3", "dep:object_store_opendal"] -azure = ["object_store/azure", "dep:opendal", "opendal/services-azblob", "opendal/services-azdls", "dep:object_store_opendal"] +gcp = [ + "object_store/gcp", + "dep:opendal", + "opendal/services-gcs", + "dep:object_store_opendal", + "dep:reqwest", +] +aws = [ + "object_store/aws", + "dep:aws-config", + "dep:aws-credential-types", + "dep:opendal", + "opendal/services-s3", + "dep:object_store_opendal", + "dep:reqwest", +] +azure = [ + "object_store/azure", + "dep:opendal", + "opendal/services-azblob", + "opendal/services-azdls", + "dep:object_store_opendal", + "dep:reqwest", +] oss = ["dep:opendal", "opendal/services-oss", "dep:object_store_opendal"] goosefs = ["dep:opendal", "opendal/services-goosefs", "dep:object_store_opendal"] tencent = ["dep:opendal", "opendal/services-cos", "dep:object_store_opendal"] diff --git a/rust/lance-io/src/object_store.rs b/rust/lance-io/src/object_store.rs index eb0029af4a7..4c1d413b4ea 100644 --- a/rust/lance-io/src/object_store.rs +++ b/rust/lance-io/src/object_store.rs @@ -24,6 +24,7 @@ use object_store::DynObjectStore; use object_store::ObjectStoreExt as OSObjectStoreExt; #[cfg(feature = "aws")] use object_store::aws::AwsCredentialProvider; +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] use object_store::signer::Signer; #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] use object_store::{ClientOptions, HeaderMap, HeaderValue}; @@ -168,6 +169,7 @@ pub struct ObjectStore { #[derive(Debug, Clone)] enum ConditionalDeleteConfig { Serialized, + #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] SignedUrl(Arc), } @@ -207,12 +209,14 @@ struct ConditionalDeleteStore { /// URL. S3, GCS, and Azure all enforce `If-Match` at the deletion linearization /// point, including against writers using older Lance clients. #[derive(Debug)] +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] struct SignedConditionalDeleteStore { target: Arc, signer: Arc, client: reqwest::Client, } +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] impl SignedConditionalDeleteStore { fn new(target: Arc, signer: Arc) -> Self { Self { @@ -223,6 +227,7 @@ impl SignedConditionalDeleteStore { } } +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] impl std::fmt::Display for SignedConditionalDeleteStore { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "SignedConditionalDeleteStore({})", self.target) @@ -230,6 +235,7 @@ impl std::fmt::Display for SignedConditionalDeleteStore { } #[async_trait] +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] impl OSObjectStore for SignedConditionalDeleteStore { async fn put_opts( &self, @@ -985,6 +991,7 @@ impl ObjectStore { &self.store_prefix, )); } + #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] Some(ConditionalDeleteConfig::SignedUrl(signer)) => { self.inner = Arc::new(SignedConditionalDeleteStore::new( Arc::clone(&self.inner),