From 8379ee9251520be9c03ed78df86b5b05461d5b93 Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Mon, 26 Oct 2026 20:11:58 -0700 Subject: [PATCH] vector: snapshot lease + maintenance safety tests (P2) (#12107) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * vector: version-lease helpers + maintenance safety tests Pinned queries must survive Lance maintenance. maintenance::pin tags the current dataset version so cleanup refuses to reclaim it; release drops the tag. Tests prove the P2 gate: a pinned query returns identical hits across compact_files, a tagged version survives cleanup_old_versions (error_if_tagged protects it, the non-strict run skips it), and a missing data file fails the query instead of returning partial hits. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: lease the query's dataset version in the VectorService query_vectors pins the resolved version with a query tag for the duration of the RPC and releases it after, so concurrent maintenance cleanup cannot reclaim the files out from under an in-flight read. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: release the version lease on every query exit path The lease tag was only released on query_vectors' success path — get_vectors never released and query_vectors leaked its tag on early returns. pin now returns a guard whose drop deletes the tag, so every path frees the version. * vector: unique lease tag names + reap abandoned leases Nanos-since-epoch tag names can collide across concurrent readers and processes, so each lease now carries a unique suffix. pin() also drops query tags older than an hour, since a killed or cancelled request leaves its lease behind and would otherwise block cleanup forever. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: keep live lease tags + throttle abandoned-lease scans An expired tag is only reaped once no live lease in this process owns it, so a query running past LEASE_TTL keeps its pin. Tag scans run at most once per 60s per dataset instead of on every request. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: answer read-only queries without a version lease A namespace may hand the server credentials that can read a table but not update it, so the tag write in pin() fails and the whole query errors out. try_pin() degrades a denied tag write to an unleased query instead of failing the request. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * vector: fail a lease whose pinned version is already gone Cleanup lists tags once at the start, so a tag created mid-cleanup does not protect the version it names. After the tag lands, pin now verifies the version still exists and drops the lease when it does not, failing the read at pin time instead of mid-query. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- seaweed-vector/Cargo.lock | 2 + seaweed-vector/crates/core/Cargo.toml | 2 + seaweed-vector/crates/core/src/lib.rs | 1 + seaweed-vector/crates/core/src/maintenance.rs | 420 ++++++++++++++++++ seaweed-vector/crates/server/src/service.rs | 10 + 5 files changed, 435 insertions(+) create mode 100644 seaweed-vector/crates/core/src/maintenance.rs diff --git a/seaweed-vector/Cargo.lock b/seaweed-vector/Cargo.lock index e0b6355053..8d2f7b6e72 100644 --- a/seaweed-vector/Cargo.lock +++ b/seaweed-vector/Cargo.lock @@ -6465,6 +6465,7 @@ dependencies = [ "lance-index", "lance-linalg", "lance-table", + "object_store", "roaring", "serde", "serde_json", @@ -6472,6 +6473,7 @@ dependencies = [ "tokio", "tracing", "tracing-subscriber", + "uuid", "weed-lance-worker", ] diff --git a/seaweed-vector/crates/core/Cargo.toml b/seaweed-vector/crates/core/Cargo.toml index 7888bbe51b..4c645e60fe 100644 --- a/seaweed-vector/crates/core/Cargo.toml +++ b/seaweed-vector/crates/core/Cargo.toml @@ -18,6 +18,7 @@ path = "src/main.rs" # no common root manifest. weed-lance-worker = { path = "../../../seaweed-worker/crates/lance" } lance = { version = "10", default-features = false, features = ["aws"] } +object_store = "0.13" lance-index = "10" lance-linalg = "10" lance-table = "10" @@ -32,6 +33,7 @@ serde = { version = "1", features = ["derive"] } serde_json = "1" tokio.workspace = true tracing.workspace = true +uuid = { version = "1", features = ["v4"] } tracing-subscriber.workspace = true [dev-dependencies] diff --git a/seaweed-vector/crates/core/src/lib.rs b/seaweed-vector/crates/core/src/lib.rs index e4cb1453d2..f5ce66ad75 100644 --- a/seaweed-vector/crates/core/src/lib.rs +++ b/seaweed-vector/crates/core/src/lib.rs @@ -6,6 +6,7 @@ //! it does not, always covering the fragments no index has reached yet. pub mod coverage; +pub mod maintenance; pub mod query; pub mod table; #[cfg(test)] diff --git a/seaweed-vector/crates/core/src/maintenance.rs b/seaweed-vector/crates/core/src/maintenance.rs new file mode 100644 index 0000000000..04fc8029d6 --- /dev/null +++ b/seaweed-vector/crates/core/src/maintenance.rs @@ -0,0 +1,420 @@ +//! Snapshot discipline under Lance maintenance: queries pin a dataset version, +//! and version tags act as the lease maintenance must respect. + +use std::collections::{HashMap, HashSet}; +use std::sync::{LazyLock, Mutex}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use anyhow::Result; +use lance::Dataset; +use lance::dataset::refs::Ref; +use uuid::Uuid; + +const QUERY_TAG_PREFIX: &str = "weed-vector-query-"; +/// A lease outliving this is treated as abandoned by a dead or cancelled +/// request and reaped by the next [`pin`]. +const LEASE_TTL: Duration = Duration::from_secs(3600); +/// Expired-lease scans run at most once per interval per dataset, not per +/// request. +const REAP_INTERVAL: Duration = Duration::from_secs(60); +/// Nanos since the Unix epoch when each dataset last had its tags scanned, +/// keyed by dataset uri. +static LAST_REAP: LazyLock>> = LazyLock::new(Mutex::default); +/// Tags this process still holds. An expired tag is only reaped once no live +/// lease here owns it, so a long-running query keeps its pin. +static LIVE_TAGS: LazyLock>> = LazyLock::new(Mutex::default); + +/// A pinned dataset version. The lease tag is removed on drop; call +/// [`Lease::release`] to await removal and observe failures. +pub struct Lease { + dataset: Dataset, + tag: String, +} + +impl Lease { + /// Remove the lease tag now instead of on drop. + pub async fn release(mut self) -> Result<()> { + let tag = std::mem::take(&mut self.tag); + untrack(&tag); + release(&self.dataset, &tag).await + } +} + +impl Drop for Lease { + fn drop(&mut self) { + if self.tag.is_empty() { + return; + } + let tag = std::mem::take(&mut self.tag); + untrack(&tag); + if let Ok(handle) = tokio::runtime::Handle::try_current() { + let dataset = self.dataset.clone(); + handle.spawn(async move { + if let Err(err) = release(&dataset, &tag).await { + tracing::warn!("releasing query lease {tag}: {err:#}"); + } + }); + } + } +} + +/// Lease-pin `dataset`'s current version with a tag so cleanup skips it. +pub async fn pin(dataset: &Dataset) -> Result { + let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos(); + reap_expired(dataset, now).await; + let tag = format!("{QUERY_TAG_PREFIX}{now}-{}", Uuid::new_v4().simple()); + let version = dataset.version().version; + dataset + .tags() + .create(&tag, Ref::Version(None, Some(version))) + .await + .map_err(anyhow::Error::msg)?; + // A cleanup that listed tags before this create can still reclaim the + // version; fail the read now rather than mid-query. + if dataset.checkout_version(version).await.is_err() { + let _ = dataset.tags().delete(&tag).await; + anyhow::bail!("dataset version {version} was removed before it could be leased"); + } + if let Ok(mut live) = LIVE_TAGS.lock() { + live.insert(tag.clone()); + } + Ok(Lease { + dataset: dataset.clone(), + tag, + }) +} + +/// [`pin`] that answers the query anyway when the storage identity may read +/// but not write: a denied tag write degrades to an unleased query rather +/// than an error. +pub async fn try_pin(dataset: &Dataset) -> Result> { + match pin(dataset).await { + Ok(lease) => Ok(Some(lease)), + Err(err) if denied(&err) => { + tracing::warn!("query runs without a version lease: {err:#}"); + Ok(None) + } + Err(err) => Err(err), + } +} + +fn denied(err: &anyhow::Error) -> bool { + err.chain().any(|cause| { + matches!( + cause.downcast_ref::(), + Some(object_store::Error::PermissionDenied { .. }) + | Some(object_store::Error::Unauthenticated { .. }) + ) + }) +} + +async fn release(dataset: &Dataset, tag: &str) -> Result<()> { + dataset.tags().delete(tag).await.map_err(anyhow::Error::msg) +} + +fn untrack(tag: &str) { + if let Ok(mut live) = LIVE_TAGS.lock() { + live.remove(tag); + } +} + +/// Drop lease tags whose owner never released them, so an abandoned lease +/// cannot keep a version out of cleanup forever. +async fn reap_expired(dataset: &Dataset, now: u128) { + { + let Ok(mut scanned) = LAST_REAP.lock() else { + return; + }; + if now.saturating_sub(*scanned.get(dataset.uri()).unwrap_or(&0)) <= REAP_INTERVAL.as_nanos() + { + return; + } + if scanned.len() >= 4096 { + scanned.clear(); + } + scanned.insert(dataset.uri().to_string(), now); + } + let Ok(tags) = dataset.tags().list().await else { + return; + }; + for name in tags.keys() { + let held = LIVE_TAGS + .lock() + .map(|live| live.contains(name)) + .unwrap_or(false); + let expired = name + .strip_prefix(QUERY_TAG_PREFIX) + .and_then(|rest| rest.split('-').next()) + .and_then(|stamp| stamp.parse::().ok()) + .is_some_and(|stamp| now.saturating_sub(stamp) > LEASE_TTL.as_nanos()); + if expired && !held { + let _ = dataset.tags().delete(name).await; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{Query, VectorTable, query, testing}; + use lance::dataset::cleanup::{CleanupPolicy, cleanup_old_versions}; + use lance::dataset::optimize::{CompactionOptions, compact_files}; + use lance_linalg::distance::DistanceType; + use tempfile::TempDir; + + async fn pinned_hits(dataset: &Dataset, version: u64, query_vector: &[f32]) -> Vec { + let pinned = dataset.checkout_version(version).await.unwrap(); + let resolved = VectorTable::resolve(&pinned, None, None).unwrap(); + let q = Query { + vector: query_vector.to_vec(), + k: 3, + filter: None, + postfilter: false, + distance: DistanceType::L2, + nprobes: None, + }; + let mut ids: Vec = query::flat(&pinned, &resolved, &q) + .await + .unwrap() + .iter() + .map(|h| h.row_id) + .collect(); + ids.sort_unstable(); + ids + } + + /// Compaction commits a new version and rewrites fragments; a query pinned + /// at the old version still reads it. + #[tokio::test] + async fn pinned_query_survives_compaction() { + let dir = TempDir::new().unwrap(); + let mut dataset = + testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await; + testing::append_dataset(&mut dataset, &testing::vectors(4, 4), &["e", "f", "g", "h"]).await; + let v1 = dataset.version().version; + let before = pinned_hits(&dataset, v1, &[0.0, 0.0, 0.0, 1.0]).await; + + compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 64, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + assert!(dataset.version().version > v1); + assert_eq!( + pinned_hits(&dataset, v1, &[0.0, 0.0, 0.0, 1.0]).await, + before + ); + } + + /// A pinned version carrying a query tag is skipped by cleanup; removing + /// the lease first lets cleanup delete it. + #[tokio::test] + async fn tagged_version_survives_cleanup() { + let dir = TempDir::new().unwrap(); + let mut dataset = + testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await; + let v1 = dataset.version().version; + let before = pinned_hits(&dataset, v1, &[0.0, 0.0, 0.0, 1.0]).await; + let lease = pin(&dataset).await.unwrap(); + + testing::append_dataset(&mut dataset, &testing::vectors(4, 4), &["e", "f", "g", "h"]).await; + let v2 = dataset.version().version; + + // Tagged old versions refuse to clean by default. + let policy = CleanupPolicy { + before_version: Some(v2), + delete_unverified: true, + error_if_tagged_old_versions: true, + clean_referenced_branches: false, + delete_rate_limit: None, + ..Default::default() + }; + cleanup_old_versions(&dataset, policy.clone()) + .await + .unwrap_err(); + + let policy = CleanupPolicy { + error_if_tagged_old_versions: false, + ..policy + }; + cleanup_old_versions(&dataset, policy).await.unwrap(); + assert_eq!( + pinned_hits(&dataset, v1, &[0.0, 0.0, 0.0, 1.0]).await, + before + ); + + // Lease released: the old version is now fair game. + lease.release().await.unwrap(); + assert!( + dataset + .tags() + .list() + .await + .unwrap() + .keys() + .all(|t| !t.starts_with(QUERY_TAG_PREFIX)) + ); + cleanup_old_versions( + &dataset, + CleanupPolicy { + before_version: Some(v2), + delete_unverified: true, + error_if_tagged_old_versions: true, + clean_referenced_branches: false, + delete_rate_limit: None, + ..Default::default() + }, + ) + .await + .unwrap(); + assert!(dataset.checkout_version(v1).await.is_err()); + } + + /// Concurrent pins never collide on a tag name. + #[tokio::test] + async fn pins_get_distinct_tags() { + let dir = TempDir::new().unwrap(); + let dataset = + testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await; + let a = pin(&dataset).await.unwrap(); + let b = pin(&dataset).await.unwrap(); + assert_ne!(a.tag, b.tag); + assert_eq!(dataset.tags().list().await.unwrap().len(), 2); + a.release().await.unwrap(); + b.release().await.unwrap(); + } + + /// A stale lease tag left behind is reaped by the next pin; a stale-named + /// tag this process still holds is not. + #[tokio::test] + async fn pin_reaps_abandoned_leases() { + let dir = TempDir::new().unwrap(); + let dataset = + testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await; + let stale = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + - LEASE_TTL.as_nanos() + - 1; + dataset + .tags() + .create( + &format!("{QUERY_TAG_PREFIX}{stale}-deadbeef"), + Ref::Version(None, Some(dataset.version().version)), + ) + .await + .unwrap(); + let held = format!("{QUERY_TAG_PREFIX}{stale}-livebeef"); + dataset + .tags() + .create(&held, Ref::Version(None, Some(dataset.version().version))) + .await + .unwrap(); + LIVE_TAGS.lock().unwrap().insert(held.clone()); + + let lease = pin(&dataset).await.unwrap(); + let tags = dataset.tags().list().await.unwrap(); + assert_eq!(tags.len(), 2); + assert!(tags.contains_key(&held)); + assert!(tags.contains_key(&lease.tag)); + LIVE_TAGS.lock().unwrap().remove(&held); + lease.release().await.unwrap(); + } + + /// Pins within one reap interval do not rescan the dataset's tags. + #[tokio::test] + async fn pin_throttles_reap_scans() { + let dir = TempDir::new().unwrap(); + let dataset = + testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await; + let first = pin(&dataset).await.unwrap(); + let stale = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + - LEASE_TTL.as_nanos() + - 1; + let tag = format!("{QUERY_TAG_PREFIX}{stale}-slowbeef"); + dataset + .tags() + .create(&tag, Ref::Version(None, Some(dataset.version().version))) + .await + .unwrap(); + + let second = pin(&dataset).await.unwrap(); + assert!(dataset.tags().list().await.unwrap().contains_key(&tag)); + first.release().await.unwrap(); + second.release().await.unwrap(); + } + + /// A version reclaimed before its lease lands fails the pin instead of + /// dying mid-query, and leaves no lease tag behind. + #[tokio::test] + async fn pin_fails_when_version_already_reclaimed() { + let dir = TempDir::new().unwrap(); + let mut dataset = + testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await; + let v1 = dataset.version().version; + let old = dataset.checkout_version(v1).await.unwrap(); + testing::append_dataset(&mut dataset, &testing::vectors(4, 4), &["e", "f", "g", "h"]).await; + let v2 = dataset.version().version; + cleanup_old_versions( + &dataset, + CleanupPolicy { + before_version: Some(v2), + delete_unverified: true, + error_if_tagged_old_versions: false, + clean_referenced_branches: false, + delete_rate_limit: None, + ..Default::default() + }, + ) + .await + .unwrap(); + assert!(dataset.checkout_version(v1).await.is_err()); + + assert!(pin(&old).await.is_err()); + assert!( + dataset + .tags() + .list() + .await + .unwrap() + .keys() + .all(|t| !t.starts_with(QUERY_TAG_PREFIX)) + ); + } + + /// A data file deleted under a pinned query fails the request rather than + /// returning a partial result set. + #[tokio::test] + async fn missing_data_file_errors_instead_of_partial() { + let dir = TempDir::new().unwrap(); + let dataset = + testing::write_dataset(&dir, &testing::vectors(4, 4), &["a", "b", "c", "d"]).await; + let data_dir = dir.path().join("data"); + let file = std::fs::read_dir(&data_dir) + .unwrap() + .flatten() + .find(|e| e.file_name().to_string_lossy().ends_with(".lance")) + .unwrap(); + std::fs::remove_file(file.path()).unwrap(); + + let resolved = VectorTable::resolve(&dataset, None, None).unwrap(); + let q = Query { + vector: vec![0.0; 4], + k: 2, + filter: None, + postfilter: false, + distance: DistanceType::L2, + nprobes: None, + }; + assert!(query::flat(&dataset, &resolved, &q).await.is_err()); + } +} diff --git a/seaweed-vector/crates/server/src/service.rs b/seaweed-vector/crates/server/src/service.rs index 094eb932f8..376dea4edd 100644 --- a/seaweed-vector/crates/server/src/service.rs +++ b/seaweed-vector/crates/server/src/service.rs @@ -387,6 +387,11 @@ impl VectorService for VectorServer { .table .ok_or_else(|| Status::invalid_argument("missing table"))?; let dataset = self.open(&table).await.map_err(open_status)?; + // Lease the resolved version so maintenance cannot reclaim its files + // while this query runs against it. + let _lease = weed_vector_core::maintenance::try_pin(&dataset) + .await + .map_err(internal)?; let resolved = VectorTable::resolve(&dataset, None, Some(KEY_COLUMN)).map_err(invalid)?; let batches = query::get(&dataset, &resolved, &req.keys) .await @@ -481,6 +486,11 @@ impl VectorService for VectorServer { .table .ok_or_else(|| Status::invalid_argument("missing table"))?; let dataset = self.open(&table).await.map_err(open_status)?; + // Lease the resolved version so maintenance cannot reclaim its files + // while this query runs against it. + let _lease = weed_vector_core::maintenance::try_pin(&dataset) + .await + .map_err(internal)?; let resolved = VectorTable::resolve(&dataset, None, Some(KEY_COLUMN)).map_err(invalid)?; if req.query_vector.len() != resolved.dimension { return Err(Status::invalid_argument(format!(