From d29e6378908dd7cd32524d328ef8ce55cfb2c97c Mon Sep 17 00:00:00 2001 From: houseme Date: Fri, 10 Jul 2026 16:04:57 +0800 Subject: [PATCH] fix(object-data-cache): close identity-index races in fill/invalidate (#4657) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fix(object-data-cache): close identity-index races in fill/invalidate (backlog#1117/#1118/#1125/#1128/#1136) The object body cache's fill-vs-invalidation contract (index.insert before cache.insert, then re-check index membership and undo the fill on mismatch) had several races where the identity index and the cache could drift, either poisoning the race metric, breaking invalidation silently, or letting a principal evict a hot body on demand. ODC-12 (backlog#1117) — generation-aware index removal. The eviction listener removed (identity, key) by key equality with no way to tell an evicted old entry from a freshly refilled one under the same key. Verified against moka 0.12.15: an insert over a TTL/TTI-expired but not yet purged entry runs do_insert_with_hash inline during insert().await, and do_post_update_steps awaits notify_upsert INLINE with cause Expired (was_evicted() == true), so the listener deleted the index key the refill just registered and the post-insert recheck then self-destructed the fresh entry (SkippedInvalidationRace). A deferred Size notification for an old generation had the same effect on a cleanly refilled key. Chose Option A (generation token) over Option B (re-insert after cache.contains_key): B removes the key first and only reconciles afterward, which reopens a narrow window where a concurrent invalidate_object misses the key while it is transiently absent from the index. A never removes a key whose generation does not match, so the decision is atomic under the shard lock. The token is the cache entry's Arc pointer captured at fill time (readable from the value moka hands the listener), so no change to the entry type is needed. ODC-13 (backlog#1118) — cancellation-safe insert/recheck/undo. The recheck+undo sat after two awaits inside the S3 GET future; a client disconnect could cancel the fill at the recheck await after cache.insert made the entry visible, stranding a stale body with no index entry. The register/insert/recheck/undo sequence now runs in a spawned task whose JoinHandle the leader awaits: cancelling the await detaches the task, which still finishes the recheck and undo. The dropped leader unblocks waiters as before (SkippedSingleflightClosed), so no waiter is stranded. ODC-20 (backlog#1125) — do not prune in-flight sibling keys. prune_missing could remove another in-flight fill's index key (a different key of the same identity that had registered but not yet published its cache entry), making that fill's recheck fail and delete its own valid entry. Added a read-only ObjectDataCacheSingleflight::has_inflight and consult it in the prune predicate alongside cache.contains_key. ODC-23 (backlog#1128) — bounded eviction instead of clear-and-reject. Exceeding identity_keys_max drained the whole key set and rejected the incoming key, so a principal able to GET identity_keys_max+1 distinct versions could evict an object's hot body on demand, and identity_keys_max=1 on a versioned bucket meant a permanent 0% hit rate. The key set now evicts only the oldest key(s) (the Vec preserves insertion order) and inserts the new key. The insert result carries evicted_keys so the backend removes exactly those from the cache and still caches the newest body. The Overflow variant and the now-unused SkippedIdentityOverflow construction are gone. NOTE: the audit also asks for identity_keys_max >= 2 validation; that lives in config.rs (out of scope here) and is deferred to the config batch. ODC-31 (backlog#1136) — deterministic race coverage. Added a #[cfg(test)] fill barrier between the index insert and the cache insert so the fill-vs-invalidation recheck can be driven deterministically, plus a companion test asserting the identity-budget eviction drops the evicted keys' cache entries. New tests: - index: key_set_bounded_eviction_evicts_oldest_and_inserts_new, key_set_removes_only_matching_generation_token, key_set_duplicate_refreshes_generation_token - starshard_index: identity_index_bounded_eviction_evicts_oldest, identity_index_stale_eviction_preserves_refreshed_key, identity_index_matching_eviction_removes_key - moka_backend: moka_backend_refill_after_expiry_without_maintenance_inserts, moka_backend_concurrent_fills_same_identity_both_insert, moka_backend_identity_budget_evicts_oldest_and_caches_newest, moka_backend_fill_loses_race_to_invalidation, moka_backend_cancelled_fill_still_undoes_lost_race cargo fmt --all, cargo check/clippy/test -p rustfs-object-data-cache all pass (44 tests). Verified the refill-after-expiry and concurrent-fill tests fail without their respective fixes. Co-authored-by: heihutu --- crates/object-data-cache/src/index.rs | 127 +++++-- crates/object-data-cache/src/moka_backend.rs | 319 ++++++++++++++++-- crates/object-data-cache/src/singleflight.rs | 9 + .../object-data-cache/src/starshard_index.rs | 99 ++++-- 4 files changed, 474 insertions(+), 80 deletions(-) diff --git a/crates/object-data-cache/src/index.rs b/crates/object-data-cache/src/index.rs index f09f85e65..2dc1bb3ec 100644 --- a/crates/object-data-cache/src/index.rs +++ b/crates/object-data-cache/src/index.rs @@ -15,43 +15,72 @@ use crate::key::ObjectDataCacheKey; use crate::starshard_index::StarshardIdentityIndex; +/// Per-entry generation token used to tell a freshly refilled body apart from a +/// superseded one under the same key. The token is the cache entry's `Arc` +/// pointer captured at fill time: the eviction listener receives the evicted +/// value and can compare its pointer against the token currently tracked, so a +/// deferred or inline eviction of an old generation cannot remove the index key +/// registered for the current generation. +pub(crate) type ObjectDataCacheKeyToken = usize; + /// Result of inserting a cache key into the identity index. #[derive(Debug, Clone, PartialEq, Eq)] pub enum ObjectDataCacheIndexInsertResult { /// The key was inserted into the identity set. - Inserted, - /// The key was already tracked for the identity. - Duplicate, - /// The identity exceeded its configured key budget and was conservatively cleared. - Overflow { - /// Keys removed while clearing the identity. - cleared_keys: Vec, + Inserted { + /// Oldest keys evicted to keep the identity within its key budget. + /// + /// Empty unless the identity was already at `identity_keys_max`; the + /// caller must drop these keys from the cache. + evicted_keys: Vec, }, + /// The key was already tracked for the identity; its generation token was refreshed. + Duplicate, } -#[derive(Debug, Clone, Default, PartialEq, Eq)] +#[derive(Debug, Clone)] +struct TrackedKey { + key: ObjectDataCacheKey, + token: ObjectDataCacheKeyToken, +} + +#[derive(Debug, Clone, Default)] pub(crate) struct ObjectDataCacheKeySet { - keys: Vec, + keys: Vec, } impl ObjectDataCacheKeySet { - pub(crate) fn insert(&mut self, key: ObjectDataCacheKey, max_keys: usize) -> ObjectDataCacheIndexInsertResult { - if self.keys.iter().any(|existing| existing == &key) { + pub(crate) fn insert( + &mut self, + key: ObjectDataCacheKey, + token: ObjectDataCacheKeyToken, + max_keys: usize, + ) -> ObjectDataCacheIndexInsertResult { + if let Some(existing) = self.keys.iter_mut().find(|existing| existing.key == key) { + // Refresh the generation token so a later eviction of the superseded + // body cannot remove the key registered for this new body. + existing.token = token; return ObjectDataCacheIndexInsertResult::Duplicate; } - if self.keys.len() >= max_keys { - let cleared_keys = self.drain(); - return ObjectDataCacheIndexInsertResult::Overflow { cleared_keys }; + // Bounded eviction: the Vec preserves insertion order, so evicting from + // the front drops the oldest keys and keeps hot (recently filled) ones, + // rather than clearing the whole identity and rejecting the new key. + let mut evicted_keys = Vec::new(); + while self.keys.len() >= max_keys { + evicted_keys.push(self.keys.remove(0).key); } - self.keys.push(key); - ObjectDataCacheIndexInsertResult::Inserted + self.keys.push(TrackedKey { key, token }); + ObjectDataCacheIndexInsertResult::Inserted { evicted_keys } } - pub(crate) fn remove_key(&mut self, key: &ObjectDataCacheKey) -> bool { + /// Removes the key only when its tracked generation token matches, so an + /// eviction notification for an old generation leaves a refreshed key intact. + pub(crate) fn remove_evicted_key(&mut self, key: &ObjectDataCacheKey, token: ObjectDataCacheKeyToken) -> bool { let original_len = self.keys.len(); - self.keys.retain(|existing| existing != key); + self.keys + .retain(|existing| !(existing.key == *key && existing.token == token)); original_len != self.keys.len() } @@ -59,7 +88,7 @@ impl ObjectDataCacheKeySet { where F: FnMut(&ObjectDataCacheKey) -> bool, { - self.keys.retain(|key| keep(key)); + self.keys.retain(|tracked| keep(&tracked.key)); } pub(crate) fn is_empty(&self) -> bool { @@ -67,7 +96,7 @@ impl ObjectDataCacheKeySet { } pub(crate) fn contains(&self, key: &ObjectDataCacheKey) -> bool { - self.keys.iter().any(|existing| existing == key) + self.keys.iter().any(|existing| &existing.key == key) } #[cfg(test)] @@ -76,11 +105,7 @@ impl ObjectDataCacheKeySet { } pub(crate) fn cloned(&self) -> Vec { - self.keys.clone() - } - - pub(crate) fn drain(&mut self) -> Vec { - std::mem::take(&mut self.keys) + self.keys.iter().map(|tracked| tracked.key.clone()).collect() } } @@ -101,27 +126,63 @@ mod tests { let mut set = ObjectDataCacheKeySet::default(); let key = make_key("v1"); - let first = set.insert(key.clone(), 4); - let second = set.insert(key, 4); + let first = set.insert(key.clone(), 1, 4); + let second = set.insert(key, 2, 4); - assert_eq!(first, ObjectDataCacheIndexInsertResult::Inserted); + assert_eq!( + first, + ObjectDataCacheIndexInsertResult::Inserted { + evicted_keys: Vec::new() + } + ); assert_eq!(second, ObjectDataCacheIndexInsertResult::Duplicate); assert_eq!(set.len(), 1); } #[test] - fn key_set_overflow_clears_existing_keys() { + fn key_set_bounded_eviction_evicts_oldest_and_inserts_new() { let mut set = ObjectDataCacheKeySet::default(); let key_a = make_key("v1"); let key_b = make_key("v2"); - let _ = set.insert(key_a.clone(), 1); - let result = set.insert(key_b, 1); + let _ = set.insert(key_a.clone(), 1, 1); + let result = set.insert(key_b.clone(), 2, 1); assert!(matches!( result, - ObjectDataCacheIndexInsertResult::Overflow { cleared_keys } if cleared_keys == vec![key_a] + ObjectDataCacheIndexInsertResult::Inserted { evicted_keys } if evicted_keys == vec![key_a] )); - assert!(set.is_empty()); + // The new key replaces the evicted one instead of the identity being cleared. + assert!(set.contains(&key_b)); + assert_eq!(set.len(), 1); + } + + #[test] + fn key_set_removes_only_matching_generation_token() { + let mut set = ObjectDataCacheKeySet::default(); + let key = make_key("v1"); + + let _ = set.insert(key.clone(), 1, 4); + + // A stale eviction notification carrying the old token must not remove the key. + assert!(!set.remove_evicted_key(&key, 2)); + assert!(set.contains(&key)); + + // The matching token removes the key. + assert!(set.remove_evicted_key(&key, 1)); + assert!(!set.contains(&key)); + } + + #[test] + fn key_set_duplicate_refreshes_generation_token() { + let mut set = ObjectDataCacheKeySet::default(); + let key = make_key("v1"); + + let _ = set.insert(key.clone(), 1, 4); + let _ = set.insert(key.clone(), 2, 4); + + // After the refresh the old token no longer matches. + assert!(!set.remove_evicted_key(&key, 1)); + assert!(set.remove_evicted_key(&key, 2)); } } diff --git a/crates/object-data-cache/src/moka_backend.rs b/crates/object-data-cache/src/moka_backend.rs index 0f4d81fd2..dba17740f 100644 --- a/crates/object-data-cache/src/moka_backend.rs +++ b/crates/object-data-cache/src/moka_backend.rs @@ -15,7 +15,7 @@ use crate::cache::{ObjectDataCacheFillResult, ObjectDataCacheGetPlan, ObjectDataCacheInvalidationResult, ObjectDataCacheLookup}; use crate::config::ObjectDataCacheConfig; use crate::entry::ObjectDataCacheEntry; -use crate::index::{ObjectDataCacheIdentityIndex, ObjectDataCacheIndexInsertResult}; +use crate::index::{ObjectDataCacheIdentityIndex, ObjectDataCacheIndexInsertResult, ObjectDataCacheKeyToken}; use crate::key::{ObjectDataCacheIdentity, ObjectDataCacheKey}; use crate::memory::ObjectDataCacheMemoryGate; use crate::singleflight::{ObjectDataCacheSingleflight, ObjectDataCacheSingleflightAcquire}; @@ -31,6 +31,50 @@ pub struct MokaBackend { index: Arc, singleflight: ObjectDataCacheSingleflight, memory_gate: ObjectDataCacheMemoryGate, + /// Test-only barrier that pauses a fill between the index insert and the + /// cache insert so the fill-vs-invalidation race can be driven deterministically. + #[cfg(test)] + fill_barrier: std::sync::Mutex>>, +} + +/// Test-only rendezvous that pauses a fill between the index insert and the +/// cache insert. The fill signals `reached` and blocks on `release`; the test +/// waits for `reached`, drives an interleaving, then grants `release`. +#[cfg(test)] +#[derive(Debug)] +struct FillBarrier { + reached: tokio::sync::Semaphore, + release: tokio::sync::Semaphore, +} + +#[cfg(test)] +impl FillBarrier { + fn new() -> Self { + Self { + reached: tokio::sync::Semaphore::new(0), + release: tokio::sync::Semaphore::new(0), + } + } + + /// Called from inside the fill: announce arrival, then block until released. + async fn wait(&self) { + self.reached.add_permits(1); + if let Ok(permit) = self.release.acquire().await { + permit.forget(); + } + } + + /// Called from the test: block until a fill reaches the barrier. + async fn wait_until_reached(&self) { + if let Ok(permit) = self.reached.acquire().await { + permit.forget(); + } + } + + /// Called from the test: let the paused fill proceed. + fn release(&self) { + self.release.add_permits(1); + } } impl MokaBackend { @@ -58,7 +102,7 @@ impl MokaBackend { .weigher(|key, value: &Arc| value.estimated_weight(key)) .time_to_live(ttl) .time_to_idle(time_to_idle) - .async_eviction_listener(move |key, _value, cause| { + .async_eviction_listener(move |key, value, cause| { let index_for_eviction = index_for_eviction.clone(); async move { // Explicit removals and replacements are already reconciled @@ -70,7 +114,13 @@ impl MokaBackend { return; }; let identity = ObjectDataCacheIdentity::new(Arc::clone(&key.bucket), Arc::clone(&key.object)); - index.remove_key(&identity, &key).await; + // Prune by generation token: moka delivers the evicted value, + // so the index only drops the key when this evicted entry is + // still the one it tracks. An inline `Expired` upsert or a + // deferred `Size` notification for a superseded generation + // then cannot remove the key a fresh refill just registered. + let token = Arc::as_ptr(&value) as ObjectDataCacheKeyToken; + index.remove_evicted_key(&identity, &key, token).await; } .boxed() }) @@ -81,9 +131,19 @@ impl MokaBackend { index, singleflight: ObjectDataCacheSingleflight::new(Arc::clone(&stats)), memory_gate: ObjectDataCacheMemoryGate::new(config, stats), + #[cfg(test)] + fill_barrier: std::sync::Mutex::new(None), }) } + /// Installs a test-only fill barrier and returns it to the caller. + #[cfg(test)] + fn install_fill_barrier(&self) -> Arc { + let barrier = Arc::new(FillBarrier::new()); + *self.fill_barrier.lock().unwrap_or_else(|p| p.into_inner()) = Some(Arc::clone(&barrier)); + barrier + } + /// Returns the current cache entry count. pub fn entry_count(&self) -> u64 { self.cache.entry_count() @@ -125,34 +185,63 @@ impl MokaBackend { } let identity = ObjectDataCacheIdentity::new(Arc::clone(&key.bucket), Arc::clone(&key.object)); + // Keep keys that are still cached or still mid-fill: another in-flight + // fill for a different key of this identity may have registered in the + // index but not yet published its cache entry, and must not be pruned. self.index - .prune_missing(&identity, |candidate| self.cache.contains_key(candidate)) + .prune_missing(&identity, |candidate| { + self.cache.contains_key(candidate) || self.singleflight.has_inflight(candidate) + }) .await; - // Register the key in the identity index BEFORE the entry becomes - // visible in the cache, so a concurrent invalidation always finds it. - let result = match self.index.insert(identity.clone(), key.clone()).await { - ObjectDataCacheIndexInsertResult::Inserted | ObjectDataCacheIndexInsertResult::Duplicate => { - let entry = Arc::new(ObjectDataCacheEntry::new(bytes, key.size, Arc::clone(&key.etag))); - self.cache.insert(key.clone(), entry).await; + // Build the entry up front so its Arc pointer can serve as the + // generation token registered in the index alongside the key. + let entry = Arc::new(ObjectDataCacheEntry::new(bytes, key.size, Arc::clone(&key.etag))); + let token = Arc::as_ptr(&entry) as ObjectDataCacheKeyToken; - // An invalidation may have raced between the index and cache - // inserts; re-check the index and undo the fill so the stale - // body cannot outlive the invalidation. - if self.index.contains_key(&identity, key).await { - ObjectDataCacheFillResult::Inserted - } else { - self.cache.remove(key).await; - ObjectDataCacheFillResult::SkippedInvalidationRace + // Run the register/insert/recheck/undo sequence in a spawned task and + // await its handle. If the enclosing GET future is dropped (client + // disconnect) while we await, the JoinHandle is detached but the task + // still finishes the recheck and undo, so a stale body cannot survive + // with no index entry. Dropping the leader still unblocks waiters. + let cache = self.cache.clone(); + let index = Arc::clone(&self.index); + let fill_key = key.clone(); + let fill_identity = identity.clone(); + #[cfg(test)] + let fill_barrier = self.fill_barrier.lock().unwrap_or_else(|p| p.into_inner()).clone(); + + let handle = tokio::spawn(async move { + // Register the key in the identity index BEFORE the entry becomes + // visible in the cache, so a concurrent invalidation always finds it. + match index.insert(fill_identity.clone(), fill_key.clone(), token).await { + ObjectDataCacheIndexInsertResult::Inserted { evicted_keys } => { + for evicted in evicted_keys { + cache.remove(&evicted).await; + } } + ObjectDataCacheIndexInsertResult::Duplicate => {} } - ObjectDataCacheIndexInsertResult::Overflow { cleared_keys } => { - for stale_key in cleared_keys { - self.cache.remove(&stale_key).await; - } - ObjectDataCacheFillResult::SkippedIdentityOverflow + + #[cfg(test)] + if let Some(barrier) = fill_barrier { + barrier.wait().await; } - }; + + cache.insert(fill_key.clone(), entry).await; + + // An invalidation may have raced between the index and cache + // inserts; re-check the index and undo the fill so the stale + // body cannot outlive the invalidation. + if index.contains_key(&fill_identity, &fill_key).await { + ObjectDataCacheFillResult::Inserted + } else { + cache.remove(&fill_key).await; + ObjectDataCacheFillResult::SkippedInvalidationRace + } + }); + + let result = handle.await.unwrap_or(ObjectDataCacheFillResult::SkippedInvalidationRace); leader.finish(result).await } @@ -207,6 +296,12 @@ mod tests { } } + fn versioned_plan(object: &str, version: &str, etag: &str) -> ObjectDataCacheGetPlan { + ObjectDataCacheGetPlan::Cacheable { + key: ObjectDataCacheKey::new("bucket", object, Some(version), etag, 5, ObjectDataCacheBodyVariant::FullObjectPlainV1), + } + } + #[tokio::test] async fn moka_backend_round_trips_cached_body() { let backend = @@ -337,4 +432,180 @@ mod tests { assert_eq!(waiter_result, ObjectDataCacheFillResult::Inserted); assert!(matches!(lookup, ObjectDataCacheLookup::Hit(ref bytes) if bytes.as_ref() == b"hello")); } + + // ODC-12(1): the first refill after a TTL expiry must not self-destruct. + // The expired-but-resident entry is upserted, which fires moka's eviction + // listener inline with cause `Expired`; the generation token keeps that + // listener from removing the index key the refill just registered. + #[tokio::test] + async fn moka_backend_refill_after_expiry_without_maintenance_inserts() { + let backend = + MokaBackend::new(&enabled_config(), Arc::new(ObjectDataCacheStats::default())).expect("moka backend should build"); + let plan = cacheable_plan("object", "etag-a"); + + let first = backend.fill_body(&plan, Bytes::from_static(b"hello")).await; + assert_eq!(first, ObjectDataCacheFillResult::Inserted); + + // Let the entry expire but do NOT run pending tasks or look it up, so the + // expired entry is still resident when the refill upserts over it. + tokio::time::sleep(Duration::from_millis(150)).await; + + let refill = backend.fill_body(&plan, Bytes::from_static(b"world")).await; + let lookup = backend.lookup_body(&plan).await; + + assert_eq!(refill, ObjectDataCacheFillResult::Inserted, "refill after expiry must not self-destruct"); + assert!(matches!(lookup, ObjectDataCacheLookup::Hit(ref bytes) if bytes.as_ref() == b"world")); + } + + // ODC-20: two concurrent fills for different keys of one identity must both + // succeed. Fill A registers its key and pauses before publishing to the + // cache; fill B's prune must keep A's in-flight key so A's recheck holds. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn moka_backend_concurrent_fills_same_identity_both_insert() { + let backend = Arc::new( + MokaBackend::new(&enabled_config(), Arc::new(ObjectDataCacheStats::default())).expect("moka backend should build"), + ); + let plan_a = versioned_plan("object", "v1", "etag-a"); + let plan_b = versioned_plan("object", "v2", "etag-b"); + + let barrier = backend.install_fill_barrier(); + + let fill_a = { + let backend = Arc::clone(&backend); + let plan_a = plan_a.clone(); + tokio::spawn(async move { backend.fill_body(&plan_a, Bytes::from_static(b"aaaaa")).await }) + }; + + // Wait until A has registered its key and paused before the cache insert. + barrier.wait_until_reached().await; + // Clear the barrier so B runs to completion without pausing. + *backend.fill_barrier.lock().unwrap() = None; + + let fill_b = backend.fill_body(&plan_b, Bytes::from_static(b"bbbbb")).await; + assert_eq!(fill_b, ObjectDataCacheFillResult::Inserted); + + // Release A and let it finish its recheck. + barrier.release(); + let fill_a = fill_a.await.expect("fill A task should complete"); + + assert_eq!( + fill_a, + ObjectDataCacheFillResult::Inserted, + "an in-flight sibling fill must not phantom-race" + ); + assert!(matches!(backend.lookup_body(&plan_a).await, ObjectDataCacheLookup::Hit(_))); + assert!(matches!(backend.lookup_body(&plan_b).await, ObjectDataCacheLookup::Hit(_))); + } + + // ODC-23: exceeding the per-identity key budget evicts only the oldest key + // and still caches the newest body, instead of clearing the whole identity. + #[tokio::test] + async fn moka_backend_identity_budget_evicts_oldest_and_caches_newest() { + let mut config = enabled_config(); + config.ttl = Duration::from_secs(30); + config.time_to_idle = Duration::from_secs(30); + config.identity_keys_max = 2; + let backend = MokaBackend::new(&config, Arc::new(ObjectDataCacheStats::default())).expect("moka backend should build"); + + let plan_v1 = versioned_plan("object", "v1", "etag-1"); + let plan_v2 = versioned_plan("object", "v2", "etag-2"); + let plan_v3 = versioned_plan("object", "v3", "etag-3"); + + assert_eq!( + backend.fill_body(&plan_v1, Bytes::from_static(b"aaaaa")).await, + ObjectDataCacheFillResult::Inserted + ); + assert_eq!( + backend.fill_body(&plan_v2, Bytes::from_static(b"bbbbb")).await, + ObjectDataCacheFillResult::Inserted + ); + // The third fill exceeds the budget of 2 and evicts the oldest (v1). + assert_eq!( + backend.fill_body(&plan_v3, Bytes::from_static(b"ccccc")).await, + ObjectDataCacheFillResult::Inserted + ); + + assert!( + matches!(backend.lookup_body(&plan_v1).await, ObjectDataCacheLookup::Miss), + "oldest key must be evicted" + ); + assert!(matches!(backend.lookup_body(&plan_v2).await, ObjectDataCacheLookup::Hit(_))); + assert!( + matches!(backend.lookup_body(&plan_v3).await, ObjectDataCacheLookup::Hit(_)), + "newest body must be cached" + ); + } + + // ODC-31: the fill-vs-invalidation recheck. Pause the fill between the index + // insert and the cache insert, invalidate the identity, then resume; the fill + // must undo itself and report the race, leaving nothing cached. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn moka_backend_fill_loses_race_to_invalidation() { + let backend = Arc::new( + MokaBackend::new(&enabled_config(), Arc::new(ObjectDataCacheStats::default())).expect("moka backend should build"), + ); + let plan = cacheable_plan("object", "etag-a"); + let barrier = backend.install_fill_barrier(); + + let fill = { + let backend = Arc::clone(&backend); + let plan = plan.clone(); + tokio::spawn(async move { backend.fill_body(&plan, Bytes::from_static(b"hello")).await }) + }; + + barrier.wait_until_reached().await; + // Invalidate while the fill is parked after registering the index key. + let _ = backend + .invalidate_object(&ObjectDataCacheIdentity::new("bucket", "object")) + .await; + barrier.release(); + + let result = fill.await.expect("fill task should complete"); + let lookup = backend.lookup_body(&plan).await; + + assert_eq!(result, ObjectDataCacheFillResult::SkippedInvalidationRace); + assert!( + matches!(lookup, ObjectDataCacheLookup::Miss), + "a body must not survive a racing invalidation" + ); + } + + // ODC-13: cancelling the enclosing GET future at the recheck point must still + // undo a fill that lost the race, so no orphaned body survives with no index + // entry. The spawned fill task completes the recheck/undo after cancellation. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn moka_backend_cancelled_fill_still_undoes_lost_race() { + let backend = Arc::new( + MokaBackend::new(&enabled_config(), Arc::new(ObjectDataCacheStats::default())).expect("moka backend should build"), + ); + let plan = cacheable_plan("object", "etag-a"); + let barrier = backend.install_fill_barrier(); + + let fill = { + let backend = Arc::clone(&backend); + let plan = plan.clone(); + tokio::spawn(async move { backend.fill_body(&plan, Bytes::from_static(b"hello")).await }) + }; + + barrier.wait_until_reached().await; + // A concurrent invalidation removes the index key the fill registered. + let _ = backend + .invalidate_object(&ObjectDataCacheIdentity::new("bucket", "object")) + .await; + // Cancel the enclosing fill future before it reaches the recheck. + fill.abort(); + // Resume the detached fill task so it runs the recheck and undo. + barrier.release(); + + // The detached task finishes the undo; poll until nothing is cached. + let mut cached = true; + for _ in 0..200 { + if matches!(backend.lookup_body(&plan).await, ObjectDataCacheLookup::Miss) { + cached = false; + break; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + assert!(!cached, "a cancelled fill must not leave an orphaned body after a racing invalidation"); + } } diff --git a/crates/object-data-cache/src/singleflight.rs b/crates/object-data-cache/src/singleflight.rs index d02d954ce..32b89bb90 100644 --- a/crates/object-data-cache/src/singleflight.rs +++ b/crates/object-data-cache/src/singleflight.rs @@ -65,6 +65,15 @@ impl ObjectDataCacheSingleflight { finished: false, }) } + + /// Returns whether a leader fill is currently in flight for the key. + /// + /// Used by the fill hot path to keep another in-flight fill's index key + /// (a different key under the same identity that has registered in the index + /// but not yet published its cache entry) from being pruned as stale. + pub fn has_inflight(&self, key: &ObjectDataCacheKey) -> bool { + lock_fills(&self.fills).contains_key(key) + } } /// Leader or waiter result from the singleflight map. diff --git a/crates/object-data-cache/src/starshard_index.rs b/crates/object-data-cache/src/starshard_index.rs index c20dc6103..21c7f0a71 100644 --- a/crates/object-data-cache/src/starshard_index.rs +++ b/crates/object-data-cache/src/starshard_index.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::index::{ObjectDataCacheIndexInsertResult, ObjectDataCacheKeySet}; +use crate::index::{ObjectDataCacheIndexInsertResult, ObjectDataCacheKeySet, ObjectDataCacheKeyToken}; use crate::key::{ObjectDataCacheIdentity, ObjectDataCacheKey}; use starshard::{AsyncShardedHashMap, DEFAULT_SHARDS, SnapshotMode}; use std::collections::hash_map::RandomState; @@ -50,7 +50,12 @@ impl StarshardIdentityIndex { /// /// Runs the read-modify-write under the shard write lock so concurrent /// fills for the same identity cannot drop each other's keys. - pub async fn insert(&self, identity: ObjectDataCacheIdentity, key: ObjectDataCacheKey) -> ObjectDataCacheIndexInsertResult { + pub async fn insert( + &self, + identity: ObjectDataCacheIdentity, + key: ObjectDataCacheKey, + token: ObjectDataCacheKeyToken, + ) -> ObjectDataCacheIndexInsertResult { let max_keys = self.max_keys_per_identity; loop { let mut outcome = None; @@ -60,10 +65,12 @@ impl StarshardIdentityIndex { let _ = self .by_object .compute_if_present(&identity, move |mut key_set| { - let result = key_set.insert(key, max_keys); - let keep = !matches!(result, ObjectDataCacheIndexInsertResult::Overflow { .. }) && !key_set.is_empty(); + let result = key_set.insert(key, token, max_keys); + // Insert never empties the set (it either dedups or + // bounded-evicts and adds the new key), but keep the + // guard so an already-empty entry is not republished. *outcome = Some(result); - keep.then_some(key_set) + (!key_set.is_empty()).then_some(key_set) }) .await; } @@ -73,13 +80,10 @@ impl StarshardIdentityIndex { // Identity not tracked yet: publish a fresh single-key set. let mut fresh = ObjectDataCacheKeySet::default(); - let result = fresh.insert(key.clone(), max_keys); - if !matches!(result, ObjectDataCacheIndexInsertResult::Inserted) { - return result; - } + let result = fresh.insert(key.clone(), token, max_keys); let final_set = self.by_object.compute_if_absent(identity.clone(), move || fresh).await; if final_set.contains(&key) { - return ObjectDataCacheIndexInsertResult::Inserted; + return result; } // Lost the race to a concurrent insert; retry against the now // present entry. @@ -94,15 +98,26 @@ impl StarshardIdentityIndex { .map_or_else(Vec::new, |set| set.cloned()) } - /// Removes a single key tracked under an identity. - pub async fn remove_key(&self, identity: &ObjectDataCacheIdentity, key: &ObjectDataCacheKey) -> bool { + /// Removes a single evicted key tracked under an identity, but only when the + /// evicted entry's generation token still matches the tracked one. + /// + /// This lets the cache eviction listener prune keys of genuinely evicted + /// entries while leaving a key that was refilled under a new generation + /// intact, so a stale (inline `Expired` upsert or deferred `Size`) eviction + /// notification cannot silently break the invalidation contract. + pub async fn remove_evicted_key( + &self, + identity: &ObjectDataCacheIdentity, + key: &ObjectDataCacheKey, + token: ObjectDataCacheKeyToken, + ) -> bool { let mut removed = false; { let removed = &mut removed; let _ = self .by_object .compute_if_present(identity, move |mut key_set| { - *removed = key_set.remove_key(key); + *removed = key_set.remove_evicted_key(key, token); (!key_set.is_empty()).then_some(key_set) }) .await; @@ -160,8 +175,8 @@ mod tests { let key_a = key("v1"); let key_b = key("v2"); - let _ = index.insert(identity.clone(), key_a.clone()).await; - let _ = index.insert(identity.clone(), key_b.clone()).await; + let _ = index.insert(identity.clone(), key_a.clone(), 1).await; + let _ = index.insert(identity.clone(), key_b.clone(), 2).await; let removed = index.remove_identity(&identity).await; assert_eq!(removed, vec![key_a, key_b]); @@ -174,8 +189,8 @@ mod tests { let key_a = key("v1"); let key_b = key("v2"); - let _ = index.insert(identity.clone(), key_a.clone()).await; - let _ = index.insert(identity.clone(), key_b.clone()).await; + let _ = index.insert(identity.clone(), key_a.clone(), 1).await; + let _ = index.insert(identity.clone(), key_b.clone(), 2).await; index.prune_missing(&identity, |candidate| candidate == &key_b).await; let removed = index.remove_identity(&identity).await; @@ -191,7 +206,9 @@ mod tests { for i in 0..32 { let index = index.clone(); let identity = identity.clone(); - handles.push(tokio::spawn(async move { index.insert(identity, key(&format!("v{i}"))).await })); + handles.push(tokio::spawn( + async move { index.insert(identity, key(&format!("v{i}")), i as usize).await }, + )); } for handle in handles { handle.await.expect("insert task should complete"); @@ -202,20 +219,56 @@ mod tests { } #[tokio::test] - async fn identity_index_overflow_clears_identity() { + async fn identity_index_bounded_eviction_evicts_oldest() { let index = StarshardIdentityIndex::new(1); let identity = identity(); let key_a = key("v1"); let key_b = key("v2"); - let _ = index.insert(identity.clone(), key_a.clone()).await; - let result = index.insert(identity.clone(), key_b).await; + let _ = index.insert(identity.clone(), key_a.clone(), 1).await; + let result = index.insert(identity.clone(), key_b.clone(), 2).await; let removed = index.remove_identity(&identity).await; assert!(matches!( result, - ObjectDataCacheIndexInsertResult::Overflow { cleared_keys } if cleared_keys == vec![key_a] + ObjectDataCacheIndexInsertResult::Inserted { evicted_keys } if evicted_keys == vec![key_a] )); - assert!(removed.is_empty()); + // The newest key is retained rather than the identity being cleared. + assert_eq!(removed, vec![key_b]); + } + + #[tokio::test] + async fn identity_index_stale_eviction_preserves_refreshed_key() { + // Manifestation of ODC-12(2): a deferred eviction notification for an old + // generation must not break a later invalidate_object for the same key. + let index = StarshardIdentityIndex::new(4); + let identity = identity(); + let key_a = key("v1"); + + // First fill registers token 1; a refill (same key, new body) refreshes + // the token to 2. + let _ = index.insert(identity.clone(), key_a.clone(), 1).await; + let _ = index.insert(identity.clone(), key_a.clone(), 2).await; + + // A stale eviction notification for the old generation (token 1) arrives. + let removed = index.remove_evicted_key(&identity, &key_a, 1).await; + assert!(!removed, "stale-generation eviction must not remove the refreshed key"); + + // A later invalidate_object still finds and removes the key. + let invalidated = index.remove_identity(&identity).await; + assert_eq!(invalidated, vec![key_a]); + } + + #[tokio::test] + async fn identity_index_matching_eviction_removes_key() { + let index = StarshardIdentityIndex::new(4); + let identity = identity(); + let key_a = key("v1"); + + let _ = index.insert(identity.clone(), key_a.clone(), 7).await; + let removed = index.remove_evicted_key(&identity, &key_a, 7).await; + + assert!(removed, "an eviction carrying the tracked token must remove the key"); + assert!(index.remove_identity(&identity).await.is_empty()); } }