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/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/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/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..065a129218c 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,38 @@ 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 + -------- + >>> 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)) + 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..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(): @@ -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" @@ -1633,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) @@ -1662,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) @@ -6223,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: @@ -6236,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" @@ -6262,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/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/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-io/Cargo.toml b/rust/lance-io/Cargo.toml index 03c6780263a..e7f2516932a 100644 --- a/rust/lance-io/Cargo.toml +++ b/rust/lance-io/Cargo.toml @@ -41,6 +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", +], optional = true } tempfile.workspace = true [target.'cfg(target_os = "linux")'.dependencies] @@ -66,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 a8230a578ca..4c1d413b4ea 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; @@ -25,8 +25,14 @@ 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}; -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 +162,368 @@ 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, + #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] + 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)] +#[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 { + target, + signer, + client: reqwest::Client::new(), + } + } +} + +#[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) + } +} + +#[async_trait] +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +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 +890,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 +950,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 +983,67 @@ 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, + )); + } + #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] + 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 +1658,7 @@ impl ObjectStore { download_retry_count, io_tracker, store_prefix, + conditional_delete: None, } } } @@ -1326,6 +1759,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.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..4f5828d32d1 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::{RetirementGuard, 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::{Arc, Mutex}, time::Duration, }; use tokio::time::{MissedTickBehavior, interval}; @@ -314,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; @@ -467,7 +474,28 @@ 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 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. + 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, + &forced_retirement_versions, + ) + .await?; if self.policy.error_if_tagged_old_versions && !inspection.tagged_old_versions.is_empty() { return Err(tagged_old_versions_cleanup_error( @@ -476,21 +504,37 @@ 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? }; - final_result.merge( - self.delete_unreferenced_files(inspection).await?, - candidate_file_limit, - ); + let (inspection, mut retirement) = self + .fence_old_versions_and_retain_new_leases( + inspection, + &version_lease_store, + &ignored_branch_manifests, + ) + .await?; + + 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) } @@ -498,6 +542,8 @@ impl<'a> CleanupTask<'a> { async fn process_manifests( &'a self, tagged_versions: &HashSet, + leased_versions: &HashSet, + forced_retirement_versions: &HashSet, ) -> Result { let inspection = Mutex::new(CleanupInspection::default()); self.dataset @@ -505,7 +551,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) + self.process_manifest_file( + location, + &inspection, + tagged_versions, + leased_versions, + forced_retirement_versions, + ) }) .await?; Ok(inspection.into_inner().unwrap()) @@ -516,6 +568,8 @@ impl<'a> CleanupTask<'a> { location: ManifestLocation, 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 @@ -551,7 +605,14 @@ 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 is_forced_retirement = forced_retirement_versions.contains(&manifest.version); + // 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. @@ -566,24 +627,209 @@ 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); + 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, + ignored_branch_manifests: &HashSet, + ) -> Result<(CleanupInspection, Option)> { + if !self.action.deletes_files() { + return Ok((inspection, None)); + } + + 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(&manifest_paths_to_delete) + .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?; + + // 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 reference_census = 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 + .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(reference_census.versions.iter().copied()); + 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 + .old_manifests + .values() + .copied() + .collect::>(); + let late_reference_versions = guard + .commit_versions( + &versions_to_commit, + &reference_census.completed_intent_paths, + &reference_census.lifecycle_generations, + ) + .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 { + 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" + ); } - } else { - inspection.earliest_retained_manifest_time = Some(commit_ts); + Err(error) } } + } + + 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); + } + } + fn process_manifest( &self, 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. @@ -1043,9 +1289,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. @@ -1067,16 +1311,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>(()) @@ -1088,12 +1339,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( @@ -1156,27 +1411,36 @@ impl<'a> CleanupTask<'a> { referenced_branches: &[(String, u64)], removed_branch_manifests: &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()); + 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 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 branch_inspection = Arc::clone(&branch_inspection); + async move { + self.process_branch_referenced_manifests( + location, + referenced_version, + branch_inspection.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()) } @@ -2136,11 +2400,257 @@ mod tests { old_manifest, &Mutex::new(CleanupInspection::default()), &HashSet::new(), + &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); + lease + .renew(Duration::from_secs(30 * 24 * 60 * 60)) + .await + .unwrap(); + 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); + 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() + .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_millis(10)) + .await + .unwrap(); + 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()); + + 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("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 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(); + 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(); @@ -2817,7 +3327,10 @@ mod tests { .build(), CleanupAction::Execute, ); - let inspection = task.process_manifests(&HashSet::new()).await.unwrap(); + let inspection = task + .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()) @@ -2881,7 +3394,10 @@ mod tests { .build(), CleanupAction::Execute, ); - let inspection = task.process_manifests(&HashSet::new()).await.unwrap(); + let inspection = task + .process_manifests(&HashSet::new(), &HashSet::new(), &HashSet::new()) + .await + .unwrap(); let kept: HashSet<&Path> = inspection .referenced_files .data_paths @@ -3387,10 +3903,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 0d3f65f7959..06d64a663eb 100644 --- a/rust/lance/src/dataset/refs.rs +++ b/rust/lance/src/dataset/refs.rs @@ -1,19 +1,23 @@ // 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::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::{ + 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}; use serde::de::DeserializeOwned; @@ -132,6 +136,31 @@ pub struct Branches<'a> { refs: &'a Refs, } +struct ReferenceSnapshot { + metadata: ObjectMeta, + payload: Bytes, +} + +async fn reference_snapshot( + object_store: &ObjectStore, + root_path: &Path, + 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?; + 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<'_> { fn object_store(&self) -> &ObjectStore { &self.refs.object_store @@ -159,14 +188,19 @@ 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(), root_path, &path).await? + else { + 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> { @@ -202,13 +236,15 @@ 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(), &root_location.path, &tag_file).await? + else { 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) } @@ -217,7 +253,9 @@ 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(), &root_location.path, &tag_file).await?; + if previous.is_some() { return Err(Error::RefConflict { message: format!("tag {} already exists", tag), }); @@ -226,14 +264,28 @@ impl Tags<'_> { let tag_contents = self .build_tag_content_by_ref(reference, Some(now), Some(now)) .await?; - - self.object_store() - .put( - &tag_file, - serde_json::to_string_pretty(&tag_contents)?.as_bytes(), - ) - .await - .map(|_| ()) + let payload = serde_json::to_vec_pretty(&tag_contents)?; + 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, + tag_contents.branch.as_deref(), + tag_contents.version, + mutation, + ) + .await?; + admission.publish(format!("tag {tag} already exists")).await } pub async fn delete(&self, tag: &str) -> Result<()> { @@ -242,13 +294,21 @@ 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(), &root_location.path, &tag_file).await? + else { return Err(Error::RefNotFound { message: format!("tag {} does not exist", tag), }); - } - - self.object_store().delete(&tag_file).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<()> { @@ -256,28 +316,44 @@ 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(), &root_location.path, &tag_file).await? + else { return Err(Error::RefNotFound { message: format!("tag {} does not exist", tag), }); - } - let mut tag_contents = TagContents::from_path(&tag_file, self.object_store()).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, tag_contents.created_at, Some(utc_now())) + .build_tag_content_by_ref(reference, original_tag_contents.created_at, Some(utc_now())) .await?; + 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(), - ) + 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?; + admission + .publish(format!("tag {tag} changed during conditional update")) .await - .map(|_| ()) } pub async fn replace_metadata( @@ -289,22 +365,37 @@ 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(), &root_location.path, &tag_file).await? + else { return Err(Error::RefNotFound { message: format!("tag {} does not exist", tag), }); - } + }; - let mut tag_contents = TagContents::from_path(&tag_file, self.object_store()).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; - - self.object_store() - .put( - &tag_file, - serde_json::to_string_pretty(&tag_contents)?.as_bytes(), - ) + 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?; + admission + .publish(format!("tag {tag} changed during conditional update")) .await - .map(|_| ()) } async fn build_tag_content_by_ref( @@ -353,7 +444,6 @@ impl Tags<'_> { } else { self.object_store().size(&manifest_file.path).await? as usize }; - let tag_contents = TagContents { branch, version: manifest_file.version, @@ -394,17 +484,20 @@ 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(), branch_path, &path).await? + else { + 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> { @@ -419,14 +512,16 @@ 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(), &root_location.path, &branch_file).await? + else { 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) } @@ -452,7 +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); - if self.object_store().exists(&branch_file).await? { + 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), }); @@ -478,18 +575,8 @@ impl Branches<'_> { message: format!("Manifest file {} does not exist", manifest_file.path), }); }; - 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() }; @@ -506,14 +593,30 @@ impl Branches<'_> { }, metadata: HashMap::new(), }; - - self.object_store() - .put( - &branch_file, - serde_json::to_string_pretty(&branch_contents)?.as_bytes(), - ) + let payload = serde_json::to_vec_pretty(&branch_contents)?; + 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, + branch_contents.parent_branch.as_deref(), + version_number, + mutation, + ) + .await?; + admission + .publish(format!("branch {branch_name} already exists")) .await - .map(|_| ()) } pub async fn replace_metadata( @@ -525,23 +628,44 @@ 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(), &root_location.path, &branch_file).await? + else { return Err(Error::RefNotFound { message: format!("branch {} does not exist", branch), }); - } + }; - let mut branch_contents = - BranchContents::from_path(&branch_file, self.object_store(), branch).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( + 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 @@ -551,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()); - if let Some(branch_id) = branch_id { + 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", @@ -565,22 +723,37 @@ 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 self.object_store().exists(&branch_file).await? { - self.object_store().delete(&branch_file).await?; + 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 + { + if let Some(namespace) = namespace { + cancel_branch_state_removal(self.object_store(), &root_location.path, namespace) + .await?; + } + return Err(error); } - // Clean up branch directories - self.cleanup_branch_directories(branch).await + 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?; + Ok(()) } pub async fn list_ordered( @@ -916,16 +1089,8 @@ 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 result = object_store.inner.get(path).await?; + Ok(serde_json::from_slice(&result.bytes().await?)?) } impl TagContents { @@ -941,18 +1106,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 new file mode 100644 index 00000000000..1e27485c62d --- /dev/null +++ b/rust/lance/src/dataset/version_lease.rs @@ -0,0 +1,4933 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Advisory leases that protect dataset versions from cleanup. + +use std::{ + collections::{HashMap, HashSet}, + fs::{File, OpenOptions}, + path::PathBuf, + 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::{ConditionalDeleteResult, ObjectStore}; +use object_store::{ObjectMeta, ObjectStoreExt, PutMode, PutOptions, UpdateVersion, path::Path}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use super::{ + Dataset, + 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"; +const SEALED_MARKER_SUFFIX: &str = ".sealed"; +const COMMITTED_MARKER_SUFFIX: &str = ".committed"; +const STORAGE_CLOCK_PATH: &str = "_clock"; + +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. +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); + +/// 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. +static LOCALLY_ABANDONED_DRAINS: LazyLock> = LazyLock::new(DashSet::new); + +/// 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. 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<()> { + if !self.store.object_store.exists(&self.path).await? { + return Err(expired_lease_error(self.version, self.expires_at)); + } + 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 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); + } + + 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 = lease_file.path; + self.expires_at = lease_file.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, + 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. +/// +/// 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, + operation_id: String, + 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 { + 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, + #[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, + }, + 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, + expected_etag: Option, + expected_version: Option, + }, +} + +#[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 { + #[default] + Pending, + Completed, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum CompletedReferenceIntentHandling { + DeferToCanonicalCensus, + RetainForCurrentScan, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum ReferenceMutationOutcome { + Published, + Conflict, +} + +#[derive(Debug)] +struct LeaseFile { + path: Path, + created_at: DateTime, + expires_at: DateTime, +} + +#[derive(Debug)] +struct RetirementFence { + draining_path: Option, + sealed_path: Option, + committed_path: Option, + observed_at: DateTime, + manifest_paths: Vec, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +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, + state: RetirementState, + metadata: ObjectMeta, +} + +#[derive(Debug)] +struct ReferenceIntentMetadata { + version: u64, + metadata: ObjectMeta, + intent: ReferenceIntent, +} + +#[derive(Debug)] +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. +/// +/// Draining markers reject new leases while allowing already-admitted leases +/// 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, + fences: HashMap, +} + +impl VersionLeaseStore { + pub(super) async fn for_dataset(dataset: &Dataset) -> Result { + let branch = dataset.manifest.branch.as_deref(); + 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 { + 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(&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 + .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::Legacy { canonical_path } + | ReferenceLifecycleState::Vacant { canonical_path } + | ReferenceLifecycleState::Revoking { 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 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, + operation_id: &str, + version: u64, + previous_was_legacy: bool, + ) -> Result<()> { + let local_lock = lock_local_reference(&self.object_store, 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) + { + 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" + )) + })?; + 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 { .. }) => { + 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(()); + } + return Err(Error::RefConflict { + message: format!("reference {canonical_path} lifecycle changed during admission"), + }); + } + } + + 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 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) { + 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()); + } + 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 vacant = ReferenceLifecycleState::Vacant { + canonical_path: canonical_path.to_string(), + }; + self.put_terminal_reference_lifecycle(canonical_path, ¤t, &vacant, has_local_lock) + .await + } + + 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?; + 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: lease_file.path, + version, + expires_at: lease_file.expires_at, + released: false, + }) + } + + 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, + ttl_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?; + 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 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_reference_versions(&self) -> Result> { + self.reconcile_reference_lifecycles().await?; + let mut versions = self + .reference_versions( + CompletedReferenceIntentHandling::RetainForCurrentScan, + &HashSet::new(), + ) + .await? + .versions; + versions.extend(self.lifecycle_reference_versions().await?); + versions.extend(self.canonical_reference_versions().await?); + 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 + .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 = self.retained_lifecycle_target(&state).await?; + 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::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))) + }) + .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 = 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::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) + })) + }) + .buffer_unordered(self.object_store.io_parallelism()) + .try_collect::>() + .await + .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()), + }; + 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()); + }; + 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 + /// 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.reconcile_reference_lifecycles().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( + &self, + completed_handling: CompletedReferenceIntentHandling, + 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 + // 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); + } + } + // 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)? => + { + active_versions.insert(intent.version); + } + ReferenceIntentState::Pending => { + let manifest_path = Path::parse(&intent.intent.manifest_path)?; + 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); + } + stale_paths.push(intent.metadata.location); + } + } + } + self.delete_paths(stale_paths).await?; + 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 { + let (path, intended_payload) = match mutation { + ReferenceMutation::Create { path, payload } + | ReferenceMutation::Update { path, payload, .. } => { + (Path::parse(path)?, payload.as_slice()) + } + }; + 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 => {} + 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), + } + + 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); + } + 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 + } + + async fn reconcile_revoking_reference( + &self, + canonical_path: &Path, + revoking: &ReferenceLifecycleSnapshot, + has_local_lock: bool, + ) -> Result { + let ReferenceLifecycleState::Revoking { + 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()), + }; + 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 { + expected_payload, + payload, + .. + }, + Some((metadata, current_payload)), + ) if current_payload.as_ref() == payload => { + self.rewrite_reference_payload( + canonical_path, + metadata, + current_payload, + expected_payload, + has_local_lock, + ) + .await?; + } + ( + ReferenceMutation::Update { + expected_payload, .. + }, + Some((metadata, current_payload)), + ) if current_payload.as_ref() == expected_payload => { + let fenced_payload = fenced_reference_payload(current_payload); + self.rewrite_reference_payload( + canonical_path, + metadata, + current_payload, + &fenced_payload, + has_local_lock, + ) + .await?; + } + (ReferenceMutation::Create { .. }, None) => {} + (ReferenceMutation::Create { payload, .. }, Some((metadata, current_payload))) + if current_payload.as_ref() == payload => + { + 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?; + self.put_terminal_reference_lifecycle(canonical_path, revoking, &terminal, has_local_lock) + .await + } + + async fn reconcile_deleting_reference( + &self, + canonical_path: &Path, + deleting: &ReferenceLifecycleSnapshot, + has_local_lock: bool, + ) -> Result { + let ReferenceLifecycleState::Deleting { + previous, + expected_etag, + expected_version, + .. + } = &deleting.state + else { + return Ok(false); + }; + let expected = UpdateVersion { + e_tag: expected_etag.clone(), + version: expected_version.clone(), + }; + 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_terminal_reference_lifecycle(canonical_path, deleting, &terminal, has_local_lock) + .await? + { + 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(), + }); + }; + 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( + &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()), + } + } + + async fn create_reference_intent( + &self, + version: u64, + intent: &ReferenceIntent, + ) -> 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::from(serde_json::to_vec(intent)?).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 apply_reference_mutation_inner( + &self, + mutation: &ReferenceMutation, + has_local_lock: bool, + expected_operation_id: Option<&str>, + ) -> Result { + 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, .. } => { + Path::parse(path)? + } + }; + let lifecycle = self.reference_lifecycle_snapshot(&canonical_path).await?; + match lifecycle.as_ref().map(|snapshot| &snapshot.state) { + Some(ReferenceLifecycleState::Pending { + operation_id, + target, + .. + }) 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)?; + 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 { .. }, + ) 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); + } + self.object_store + .inner + .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) + } + }; + + let outcome = match result { + 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 { .. }) => 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); + } + 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 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> { + let metadata = self + .object_store + .list(Some(self.reference_intents_path.clone())) + .try_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 + .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, + intent, + }) + } + + 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 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(metadata.location); + } + } + + 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, + manifest_paths: &HashMap>, + ) -> Result { + let operation_id = Uuid::new_v4().simple().to_string(); + 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 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()) + .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, manifest_paths)) => { + fences.insert( + version, + RetirementFence { + draining_path: Some(path), + sealed_path: None, + committed_path: None, + observed_at, + manifest_paths, + }, + ); + } + 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_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 {version} no longer exists and cannot be leased"), + }) + } else if self.has_active_retirement_marker(version).await? { + Err(retiring_version_error(version)) + } else { + Ok(()) + } + } + + async fn ensure_not_sealed(&self, version: u64) -> Result<()> { + if self.has_sealed_marker(version).await? { + Err(retiring_version_error(version)) + } else { + Ok(()) + } + } + + async fn has_sealed_marker(&self, version: u64) -> Result { + Ok(self + .version_marker_metadata(version) + .await? + .iter() + .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) + .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 + .iter() + .any(|marker| marker.state != RetirementState::Draining) + { + 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, state) = + if let Some(operation_id) = file_name.strip_suffix(DRAINING_MARKER_SUFFIX) { + (operation_id, RetirementState::Draining) + } else if let Some(operation_id) = file_name.strip_suffix(SEALED_MARKER_SUFFIX) { + (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, + "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(), + state, + 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.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.state != RetirementState::Draining) + .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(); + 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?; + sealed_manifest_paths + .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)? + { + 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(); + 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; + for path in &manifest_paths { + has_existing_manifest |= self.object_store.exists(path).await?; + } + if has_existing_manifest { + 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) + } + + 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())) + .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, payload: Bytes) -> Result { + self.object_store + .inner + .put_opts( + path, + payload.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 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, payload, draining_observed_at)| { + let store = store.clone(); + async move { + 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)) + } + }) + .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(()) + } + + /// Publish the irreversible retirement boundary after the final lease and + /// 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, + 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 mut active_reference_versions = self + .store + .reference_versions( + CompletedReferenceIntentHandling::RetainForCurrentScan, + completed_intents_observed_before_census, + ) + .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() + .collect::>(); + let versions = versions + .difference(&retained_versions) + .copied() + .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 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(retained_versions) + } + + pub(super) async fn cancel_versions(&mut self, versions: &HashSet) -> Result<()> { + 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| { + [fence.draining_path.clone(), fence.sealed_path.clone()] + .into_iter() + .flatten() + }) + .collect::>(); + self.store.delete_paths(paths).await?; + self.fences + .retain(|version, _| !cancellable_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<()> { + 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.object_store.exists(path).await? { + return Err(Error::internal(format!( + "cannot finalize retirement for version {version}: manifest {path} still exists" + ))); + } + } + } + + let versions: HashSet<_> = manifest_paths.keys().copied().collect(); + let marker_prefixes = versions + .iter() + .map(|version| self.marker_version_path(*version)) + .collect::>(); + let marker_streams = marker_prefixes + .into_iter() + .map(|prefix| self.object_store.list(Some(prefix))); + let marker_metadata = stream::iter(marker_streams) + .flatten() + .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) { + 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 + } + + 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 + } +} + +/// Begin a durable tag or branch admission before canonical publication. +pub(super) async fn begin_reference_admission( + refs: &Refs, + branch: Option<&str>, + version: u64, + mut mutation: ReferenceMutation, +) -> 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?; + 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 } => { + 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) + } + 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 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 { + store, + path, + manifest_path: manifest.path, + version, + created_at, + operation_id, + mutation, + }; + 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?; + 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(()) + } + + fn canonical_path(&self) -> Result { + match &self.mutation { + ReferenceMutation::Create { path, .. } | ReferenceMutation::Update { 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_inner(local_lock.is_some()).await; + return Err(error); + } + match self + .store + .apply_reference_mutation_inner( + &self.mutation, + local_lock.is_some(), + Some(&self.operation_id), + ) + .await + { + Ok(ReferenceMutationOutcome::Published) => { + self.complete().await; + Ok(()) + } + Ok(ReferenceMutationOutcome::Conflict) => { + self.cancel_before_publish_inner(local_lock.is_some()).await; + Err(Error::RefConflict { + message: conflict_message, + }) + } + // A transport failure can be ambiguous after the server accepted + // the conditional mutation. Recovery retains the intent and + // verifies canonical state without replaying the operation. + 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) { + 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) { + 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; + } + 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!( + path = %self.path, + error = %error, + "Failed to remove cancelled reference admission intent" + ); + } + } + + /// 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) { + 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 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)?) +} + +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, + canonical_path: &Path, + 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(false); + }; + Ok(match snapshot.state { + ReferenceLifecycleState::Live { live, .. } => { + 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::Revoking { previous, .. } + | ReferenceLifecycleState::Deleting { previous, .. } => previous + .is_some_and(|previous| generation.as_deref() == Some(previous.generation.as_str())), + }) +} + +pub(super) async fn delete_canonical_reference( + object_store: Arc, + root_path: &Path, + canonical_path: &Path, + 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 + .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"), + }); + } + + 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 { .. } | ReferenceLifecycleState::Revoking { .. }, + ) => { + return Err(Error::RefConflict { + message: format!("reference {canonical_path} has an in-flight mutation"), + }); + } + 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} lifecycle changed during deletion"), + }); + } + Some(ReferenceLifecycleState::Legacy { .. }) if canonical_generation.is_some() => { + return Err(Error::RefConflict { + 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"), + }); + } + _ => {} + } + + 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(), + expected_payload: current_payload.to_vec(), + expected_etag: current_metadata.e_tag, + expected_version: current_metadata.version, + }; + 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 !store + .reconcile_deleting_reference(canonical_path, &deleting_snapshot, local_lock.is_some()) + .await? + { + return Err(Error::RefConflict { + message: format!("reference {canonical_path} changed during conditional deletion"), + }); + } + Ok(()) +} + +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!( + "version {version} is retiring and cannot accept a new durable reference" + ), + }); + } + if !store.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, + 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?; + 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 + .clone() + .join(LEASES_DIR) + .join(namespace.to_string()), + root_path + .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(()) => {} + Err(error) if error.is_not_found() => {} + Err(error) => return Err(error), + } + } + 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 { + /// 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 ttl_micros(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}" + )) + })?; + ttl.num_microseconds().ok_or_else(|| { + Error::invalid_input(format!( + "version lease TTL {ttl:?} cannot be represented in microseconds" + )) + }) +} + +fn expiration_from_ttl(observed_at: DateTime, ttl: Duration) -> Result> { + 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_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 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( + "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( + 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}") + })?; + 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 ttl_micros = parts + .next() + .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 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, 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 crate::dataset::refs::{BranchContents, TagContents}; + use crate::utils::test::FailingProxyStore; + use lance_io::object_store::WrappingObjectStore; + use mock_instant::thread_local::MockClock; + + use super::*; + + 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, + } + } + + fn manifest_paths(version: u64) -> HashMap> { + HashMap::from([( + version, + vec![Path::from(format!("manifests/{version}.manifest"))], + )]) + } + + 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(), + }, + 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, + } + } + + 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 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()); + } + + #[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(); + assert_eq!( + parse_lease_file_name(&format!("42-123-{lease_id}.lease")).unwrap(), + (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(&manifest_paths(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(&manifest_paths(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(&manifest_paths(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 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(); + assert!( + guard + .commit_versions(&HashSet::from([42]), &HashSet::new(), &HashMap::new()) + .await + .unwrap() + .is_empty() + ); + 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 reference_intent_blocks_retirement_commit() { + let store = memory_store(); + 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(); + + let retained_versions = guard + .commit_versions(&HashSet::from([42]), &HashSet::new(), &HashMap::new()) + .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()); + 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 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(); + assert!( + guard + .commit_versions(&HashSet::from([42]), &HashSet::new(), &HashMap::new()) + .await + .unwrap() + .is_empty() + ); + 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(); + 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 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); + + 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()); + 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 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(); + assert_eq!( + store + .apply_reference_mutation_inner( + &admission.mutation, + false, + Some(&admission.operation_id), + ) + .await + .unwrap(), + ReferenceMutationOutcome::Published + ); + admission + .store + .object_store + .delete(&admission.path) + .await + .unwrap(); + + assert_eq!( + guard + .commit_versions(&HashSet::from([42]), &HashSet::new(), &HashMap::new()) + .await + .unwrap(), + HashSet::from([42]) + ); + assert!(!store.object_store.exists(&intent_path).await.unwrap()); + guard.cancel_all().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 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()); + } + + #[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 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() + .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()); + } + + #[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_inner(&mutation, false, None) + .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 manifest_path = Path::from("manifests/42.manifest"); + let canonical_path = Path::from("tags/deleted.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; + assert_eq!( + store + .apply_reference_mutation_inner(&admission.mutation, false, None,) + .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()); + assert!(matches!( + store + .reference_lifecycle_snapshot(&canonical_path) + .await + .unwrap() + .map(|snapshot| snapshot.state), + Some(ReferenceLifecycleState::Revoking { .. }) + )); + } + + #[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(), + 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 + .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 + ); + 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!( + serde_json::from_slice::(&fenced_payload).unwrap(), + serde_json::from_slice::(original_payload).unwrap() + ); + } + + #[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, + 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 + .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, None) + .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 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(); + 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(); + drop(guard); + + assert!( + store + .clone() + .recover_retirements() + .await + .unwrap() + .is_empty() + ); + assert!(store.version_marker_metadata(42).await.unwrap().is_empty()); + store.acquire(42, Duration::from_secs(60)).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()); + } + + #[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()); + } + + #[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 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 + .expired_reference_mutation_outcome(&mutation, &operation_id, 42, false) + .await + .unwrap(), + ReferenceMutationOutcome::Conflict + ); + 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 conditional rollback delete 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 + ); + assert!(!store.object_store.exists(&canonical_path).await.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] + 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}; + + 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_conditional_delete", + 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.object_store.exists(&canonical_path).await.unwrap()); + let snapshot = store + .reference_lifecycle_snapshot(&canonical_path) + .await + .unwrap() + .unwrap(); + assert!(matches!( + &snapshot.state, + ReferenceLifecycleState::Deleting { .. } + )); + assert!( + store + .retained_lifecycle_target(&snapshot.state) + .await + .unwrap() + .is_none(), + "an absent canonical reference must not retain its deleted version" + ); + } + + #[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( + "put", + "deletion-failure.json", + "injected conditional delete 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 conditional delete failure"), + "{error}" + ); + 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(); + 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()); + assert!( + store + .reference_lifecycle_snapshot(&canonical_path) + .await + .unwrap() + .is_none() + ); + } +}