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()); } }