From ae1f2c6d90d9f483901e864b4447710d0bd122cf Mon Sep 17 00:00:00 2001 From: Dan Lapid Date: Tue, 28 Jul 2026 04:09:58 +0000 Subject: [PATCH] Implement memory cache V2 in safe Rust Introduce a Rust memory-cache core behind the MEMORY_CACHE_V2 autogate, while retaining the existing C++ implementation as the rollback path. Keep the handwritten CXX bridge thin and isolate generated unsafe code from the core, which forbids unsafe code. Store cache entries in a LinkedHashMap for LRU ordering and maintain a separate expiration index for efficient cleanup. Use a per-key Tokio mutex to coalesce concurrent misses without a custom waiter queue. Represent registry membership with weak flight references so abandoned requests do not retain cache state, while flight destruction safely removes only its own registry entry. Make leader handoff cancellation-safe through mutex ownership. Promotion follows waiter poll order rather than request creation order, avoiding eager polling and custom scheduling machinery. Track waiter cardinality with lock-free, best-effort counters so metrics do not add contention to the cache-state lock. Expose the implementation through the existing memory-cache API and add coverage for eviction, expiration, coalescing, cancellation, abandoned reads, leader promotion, weak lifetime cleanup, and the autogated C++ integration. --- deps/rust/Cargo.lock | 65 +- deps/rust/Cargo.toml | 4 +- src/rust/memory-cache/BUILD.bazel | 12 + src/rust/memory-cache/ffi/BUILD.bazel | 13 + src/rust/memory-cache/ffi/lib.rs | 385 ++++ src/rust/memory-cache/lib.rs | 1763 +++++++++++++++++ src/workerd/api/BUILD.bazel | 24 +- src/workerd/api/memory-cache-test.c++ | 173 ++ src/workerd/api/memory-cache-v2-test.h | 17 + src/workerd/api/memory-cache-v2.c++ | 336 ++++ src/workerd/api/memory-cache.c++ | 22 +- src/workerd/api/memory-cache.h | 78 +- src/workerd/api/tests/memory-cache-test.js | 87 + .../api/tests/memory-cache-test.wd-test | 2 +- src/workerd/server/workerd-api.c++ | 13 +- src/workerd/util/autogate.h | 5 +- 16 files changed, 2936 insertions(+), 63 deletions(-) create mode 100644 src/rust/memory-cache/BUILD.bazel create mode 100644 src/rust/memory-cache/ffi/BUILD.bazel create mode 100644 src/rust/memory-cache/ffi/lib.rs create mode 100644 src/rust/memory-cache/lib.rs create mode 100644 src/workerd/api/memory-cache-v2-test.h create mode 100644 src/workerd/api/memory-cache-v2.c++ diff --git a/deps/rust/Cargo.lock b/deps/rust/Cargo.lock index a4fd4ee21d8..788a772807a 100644 --- a/deps/rust/Cargo.lock +++ b/deps/rust/Cargo.lock @@ -83,7 +83,7 @@ checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" dependencies = [ "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -231,9 +231,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.4.3" +version = "1.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "509591b7bcd67f4ef775afad7662703b4935daaa6ec0e5605cfb1090b32a2b6d" +checksum = "0ad534f4357a5264cce5019c989cf66a4f0dc4e0d1b1d15f8aacec0ff7360273" dependencies = [ "find-msvc-tools", "jobserver", @@ -292,7 +292,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -350,9 +350,9 @@ dependencies = [ [[package]] name = "crc32fast" -version = "1.5.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +checksum = "8498c871161e1742aaa9d52551b2d6ebdd4c3d45a3be423e3728f33b955be550" dependencies = [ "cfg-if", ] @@ -444,6 +444,7 @@ dependencies = [ "ada-url", "anyhow", "async-trait", + "bytes", "capnp", "capnp-rpc", "capnpc", @@ -455,6 +456,7 @@ dependencies = [ "flate2", "foldhash", "futures", + "hashlink", "libc", "libz-rs-sys", "lol_html_c_api", @@ -485,7 +487,7 @@ checksum = "c6232dd377dcc64799954cbd3a9bb882e9cdc1308ccd87b1c098f1fb2eaf82a8" dependencies = [ "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -511,9 +513,9 @@ dependencies = [ [[package]] name = "either" -version = "1.17.0" +version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" +checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" [[package]] name = "embedded-io" @@ -645,7 +647,7 @@ checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" dependencies = [ "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -739,6 +741,15 @@ dependencies = [ "foldhash", ] +[[package]] +name = "hashlink" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32069d97bb81e38fa67eab65e3393bf804bb85969f2bc06bf13f64aef5aba248" +dependencies = [ + "hashbrown 0.17.1", +] + [[package]] name = "heck" version = "0.5.0" @@ -835,9 +846,9 @@ checksum = "e590f038c1464a96894fd6d10127e90a8be4509f56ff7ecef851b15cee0b7caa" [[package]] name = "icu_provider" -version = "2.3.0" +version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92a7ed671a6aad807a8651a2e1782a6598fda9ce5185dd8158549e95a91c6428" +checksum = "d27bbb9d3abbefac45d55f647c9de1d44aafcd1186eb91879afef17c396c3e73" dependencies = [ "displaydoc", "icu_locale_core", @@ -887,9 +898,9 @@ dependencies = [ [[package]] name = "is-macro" -version = "0.3.7" +version = "0.3.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d57a3e447e24c22647738e4607f1df1e0ec6f72e16182c4cd199f647cdfb0e4" +checksum = "8267aa6001e25494f3015f9663bbd88a18240c74483afa5f0934a1b3e4c388e9" dependencies = [ "heck", "proc-macro2", @@ -962,9 +973,9 @@ checksum = "47d9d19d1d6efa0109d2f65ff4c85cddd50bd572e5a00127ab10987290bcefae" [[package]] name = "log" -version = "0.4.33" +version = "0.4.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" [[package]] name = "lol_html" @@ -1501,7 +1512,7 @@ checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -1950,9 +1961,9 @@ dependencies = [ [[package]] name = "syn" -version = "3.0.3" +version = "3.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f" dependencies = [ "proc-macro2", "quote", @@ -2002,7 +2013,7 @@ checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" dependencies = [ "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -2159,9 +2170,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.24.1" +version = "1.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2cefc03fd367c0c6d4305de1b312cf00248c4114f4a0418ce6a6af769e3b0bd9" +checksum = "f053576934f05a761a402421fbbe3d425d9366f75f978806a037b3ca481abecc" dependencies = [ "js-sys", "wasm-bindgen", @@ -2346,9 +2357,9 @@ dependencies = [ [[package]] name = "zerovec" -version = "0.11.7" +version = "0.11.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94b5c6b5976d66c1d703c4fd17d3f5e43c8cedaacf604961b171adc7130896d8" +checksum = "bb0464e17806c1d976d5cba29399c7f08e516e279e2ba493f63123b5fca67dd8" dependencies = [ "yoke", "zerofrom", @@ -2357,13 +2368,13 @@ dependencies = [ [[package]] name = "zerovec-derive" -version = "0.11.4" +version = "0.11.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47402523226a02bfe5230160dc3ccc089aa6f6f19e7fcbb4e6f824bbb1b4aa62" +checksum = "34df6fc39dbd26ddc9c10e6a2984476e13acce22e64e4487636ef494369225da" dependencies = [ "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] diff --git a/deps/rust/Cargo.toml b/deps/rust/Cargo.toml index 05b5b68d842..9aa2a97789c 100644 --- a/deps/rust/Cargo.toml +++ b/deps/rust/Cargo.toml @@ -32,6 +32,7 @@ syn = { version = "2", features = ["full"] } ada-url = { version = "4", default-features = false, features = ["std"] } anyhow = "1" async-trait = { version = "0", default-features = false } +bytes = { version = "1", default-features = false } capnp = "0" capnpc = "0" capnp-rpc = "0" @@ -45,6 +46,7 @@ libz-rs-sys = { version = "0.6", default-features = false, features = ["std", "r futures = "0" libc = "0" lol_html_c_api = { git = "https://github.com/cloudflare/lol-html" } +hashlink = "0" nix = "0" pico-args = "0" ruff_python_ast = { git = "https://github.com/astral-sh/ruff", tag = "0.12.1" } @@ -54,7 +56,7 @@ serde_json = { version = "1", features = ["unbounded_depth"] } serde = { version = "1", features = ["derive"] } thiserror = "2" # tokio is huge, let's enable only features when we actually need them. -tokio = { version = "1", default-features = false, features = ["net", "rt", "rt-multi-thread", "time"] } +tokio = { version = "1", default-features = false, features = ["net", "rt", "rt-multi-thread", "sync", "time"] } tracing = { version = "0", default-features = false, features = ["std"] } swc_common = "25" swc_ts_fast_strip = "57" diff --git a/src/rust/memory-cache/BUILD.bazel b/src/rust/memory-cache/BUILD.bazel new file mode 100644 index 00000000000..8d691ef9212 --- /dev/null +++ b/src/rust/memory-cache/BUILD.bazel @@ -0,0 +1,12 @@ +load("//:build/wd_rust_crate.bzl", "wd_rust_crate") + +wd_rust_crate( + name = "memory-cache", + visibility = ["//visibility:public"], + deps = [ + "@crates_vendor//:bytes", + "@crates_vendor//:futures", + "@crates_vendor//:hashlink", + "@crates_vendor//:tokio", + ], +) diff --git a/src/rust/memory-cache/ffi/BUILD.bazel b/src/rust/memory-cache/ffi/BUILD.bazel new file mode 100644 index 00000000000..ce38b9b5b27 --- /dev/null +++ b/src/rust/memory-cache/ffi/BUILD.bazel @@ -0,0 +1,13 @@ +load("//:build/wd_rust_crate.bzl", "wd_rust_crate") + +wd_rust_crate( + name = "memory-cache-ffi", + cxx_bridge_deps = [ + "@capnp-cpp//src/kj:kj-async", + ], + cxx_bridge_src = "lib.rs", + visibility = ["//visibility:public"], + deps = [ + "//src/rust/memory-cache", + ], +) diff --git a/src/rust/memory-cache/ffi/lib.rs b/src/rust/memory-cache/ffi/lib.rs new file mode 100644 index 00000000000..00cca5ef4fb --- /dev/null +++ b/src/rust/memory-cache/ffi/lib.rs @@ -0,0 +1,385 @@ +//! Generated workerd-cxx boundary for the safe memory-cache crate. + +#![expect( + clippy::needless_lifetimes, + clippy::needless_pass_by_value, + clippy::struct_field_names, + clippy::unnecessary_box_returns, + reason = "the CXX boundary requires these ownership and transport representations" +)] + +use kj_rs::KjMaybe; + +#[cxx::bridge(namespace = "workerd::rust::memory_cache")] +mod ffi { + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum ReadKind { + Miss, + Value, + Leader, + Waiter, + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum ReadMode { + CacheOnly, + WithFallback, + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum WaitKind { + Value, + Leader, + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum EvictionReason { + Expiration, + Lru, + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum WriteOutcome { + Success, + ValueTooLarge, + AlreadyExpired, + } + + struct Limits { + max_keys: u32, + max_value_size: u32, + max_total_value_size: u64, + } + + struct ReadTrace { + cache_hit: bool, + entry_size: usize, + total_value_size: usize, + entry_count: usize, + waiters_ahead: usize, + lock_wait_ns: u64, + } + + struct EvictionTrace { + reason: EvictionReason, + key: Vec, + value_size: usize, + total_before: usize, + entries_before: usize, + } + + struct WriteTrace { + value_size: usize, + has_expiration: bool, + outcome: WriteOutcome, + max_value_size: usize, + is_update: bool, + total_after: usize, + entries_after: usize, + evictions: Vec, + waiters_notified: usize, + } + + struct Stats { + bindings: usize, + in_flight_fallbacks: usize, + waiters: usize, + canceled_waiters: usize, + } + + extern "Rust" { + type Namespace; + type Binding; + type ReadDecision; + type Waiter; + type WaitOutcome; + type FallbackPermit; + type Value; + + fn namespace_new(max_total_value_size: KjMaybe) -> Box; + fn bind( + self: &Namespace, + id: &str, + is_private: bool, + limits: Limits, + ) -> Result>; + fn release(self: &mut Binding, now_ms: f64); + fn read( + self: &Binding, + key: &[u8], + now_ms: f64, + mode: ReadMode, + ) -> Result>; + fn remove(self: &Binding, key: &[u8]); + fn stats(self: &Binding) -> Stats; + + fn kind(self: &ReadDecision) -> ReadKind; + fn trace(self: &ReadDecision) -> ReadTrace; + fn take_value(self: &mut ReadDecision) -> Result>; + fn take_permit(self: &mut ReadDecision) -> Result>; + fn take_waiter(self: &mut ReadDecision) -> Result>; + + async fn waiter_wait(waiter: Box) -> Box; + fn kind(self: &WaitOutcome) -> WaitKind; + fn take_value(self: &mut WaitOutcome) -> Result>; + fn take_permit(self: &mut WaitOutcome) -> Result>; + + fn succeed( + self: &mut FallbackPermit, + bytes: &[u8], + expiration: KjMaybe, + now_ms: f64, + ) -> Result; + unsafe fn bytes<'a>(self: &'a Value) -> &'a [u8]; + } +} + +struct Namespace { + inner: memory_cache::Namespace, +} + +struct Binding { + inner: memory_cache::Binding, +} + +struct ReadDecision { + inner: memory_cache::ReadDecision, +} + +struct Waiter { + inner: memory_cache::Waiter, +} + +struct WaitOutcome { + inner: memory_cache::WaitOutcome, +} + +struct FallbackPermit { + inner: Option, +} + +struct Value { + inner: memory_cache::Value, +} + +fn limits_to_core(limits: ffi::Limits) -> memory_cache::Limits { + memory_cache::Limits { + max_keys: limits.max_keys, + max_value_size: limits.max_value_size, + max_total_value_size: limits.max_total_value_size, + } +} + +fn read_kind_from_core(kind: memory_cache::ReadKind) -> ffi::ReadKind { + match kind { + memory_cache::ReadKind::Miss => ffi::ReadKind::Miss, + memory_cache::ReadKind::Value => ffi::ReadKind::Value, + memory_cache::ReadKind::Leader => ffi::ReadKind::Leader, + memory_cache::ReadKind::Waiter => ffi::ReadKind::Waiter, + } +} + +fn read_mode_to_core( + mode: ffi::ReadMode, +) -> Result { + match mode { + ffi::ReadMode::CacheOnly => Ok(memory_cache::ReadMode::CacheOnly), + ffi::ReadMode::WithFallback => Ok(memory_cache::ReadMode::WithFallback), + _ => Err(memory_cache::CacheError::InvalidDecision( + "unknown memory cache read mode", + )), + } +} + +fn wait_kind_from_core(kind: memory_cache::WaitKind) -> ffi::WaitKind { + match kind { + memory_cache::WaitKind::Value => ffi::WaitKind::Value, + memory_cache::WaitKind::Leader => ffi::WaitKind::Leader, + } +} + +fn eviction_reason_from_core(reason: memory_cache::EvictionReason) -> ffi::EvictionReason { + match reason { + memory_cache::EvictionReason::Expiration => ffi::EvictionReason::Expiration, + memory_cache::EvictionReason::Lru => ffi::EvictionReason::Lru, + } +} + +fn namespace_new(max_total_value_size: KjMaybe) -> Box { + Box::new(Namespace { + inner: memory_cache::Namespace::new(max_total_value_size.into()), + }) +} + +impl Namespace { + fn bind( + &self, + id: &str, + is_private: bool, + limits: ffi::Limits, + ) -> Result, memory_cache::CacheError> { + Ok(Box::new(Binding { + inner: self + .inner + .bind((!is_private).then_some(id), limits_to_core(limits))?, + })) + } +} + +impl Binding { + fn release(&mut self, now_ms: f64) { + self.inner.release(now_ms); + } + + fn read( + &self, + key: &[u8], + now_ms: f64, + mode: ffi::ReadMode, + ) -> Result, memory_cache::CacheError> { + Ok(Box::new(ReadDecision { + inner: self.inner.read(key, now_ms, read_mode_to_core(mode)?)?, + })) + } + + fn remove(&self, key: &[u8]) { + self.inner.delete(key); + } + + fn stats(&self) -> ffi::Stats { + let stats = self.inner.stats(); + ffi::Stats { + bindings: stats.bindings, + in_flight_fallbacks: stats.in_flight_fallbacks, + waiters: stats.waiters, + canceled_waiters: stats.canceled_waiters, + } + } +} + +impl ReadDecision { + fn kind(&self) -> ffi::ReadKind { + read_kind_from_core(self.inner.kind()) + } + + fn trace(&self) -> ffi::ReadTrace { + let trace = self.inner.trace(); + ffi::ReadTrace { + cache_hit: trace.cache_hit, + entry_size: trace.entry_size, + total_value_size: trace.total_value_size, + entry_count: trace.entry_count, + waiters_ahead: trace.waiters_ahead, + lock_wait_ns: trace.lock_wait_ns, + } + } + + fn take_value(&mut self) -> Result, memory_cache::CacheError> { + Ok(Box::new(Value { + inner: self.inner.take_value()?, + })) + } + + fn take_permit(&mut self) -> Result, memory_cache::CacheError> { + Ok(Box::new(FallbackPermit { + inner: Some(self.inner.take_permit()?), + })) + } + + fn take_waiter(&mut self) -> Result, memory_cache::CacheError> { + Ok(Box::new(Waiter { + inner: self.inner.take_waiter()?, + })) + } +} + +async fn waiter_wait(waiter: Box) -> Box { + Box::new(WaitOutcome { + inner: waiter.inner.wait().await, + }) +} + +impl WaitOutcome { + fn kind(&self) -> ffi::WaitKind { + wait_kind_from_core(self.inner.kind()) + } + + fn take_value(&mut self) -> Result, memory_cache::CacheError> { + Ok(Box::new(Value { + inner: self.inner.take_value()?, + })) + } + + fn take_permit(&mut self) -> Result, memory_cache::CacheError> { + Ok(Box::new(FallbackPermit { + inner: Some(self.inner.take_permit()?), + })) + } +} + +impl FallbackPermit { + fn succeed( + &mut self, + bytes: &[u8], + expiration: KjMaybe, + now_ms: f64, + ) -> Result { + let trace = self + .inner + .take() + .ok_or(memory_cache::CacheError::InvalidDecision( + "fallback permit was already used", + ))? + .succeed(bytes.to_vec(), expiration.into(), now_ms)?; + let (outcome, max_value_size, is_update, total_after, entries_after) = match trace.outcome { + memory_cache::WriteOutcome::Success { + is_update, + total_after, + entries_after, + } => ( + ffi::WriteOutcome::Success, + 0, + is_update, + total_after, + entries_after, + ), + memory_cache::WriteOutcome::ValueTooLarge { max_value_size } => ( + ffi::WriteOutcome::ValueTooLarge, + max_value_size, + false, + 0, + 0, + ), + memory_cache::WriteOutcome::AlreadyExpired => { + (ffi::WriteOutcome::AlreadyExpired, 0, false, 0, 0) + } + }; + Ok(ffi::WriteTrace { + value_size: trace.value_size, + has_expiration: trace.has_expiration, + outcome, + max_value_size, + is_update, + total_after, + entries_after, + evictions: trace + .evictions + .into_iter() + .map(|eviction| ffi::EvictionTrace { + reason: eviction_reason_from_core(eviction.reason), + key: eviction.key.to_vec(), + value_size: eviction.value_size, + total_before: eviction.total_before, + entries_before: eviction.entries_before, + }) + .collect(), + waiters_notified: trace.waiters_notified, + }) + } +} + +impl Value { + fn bytes(&self) -> &[u8] { + self.inner.bytes() + } +} diff --git a/src/rust/memory-cache/lib.rs b/src/rust/memory-cache/lib.rs new file mode 100644 index 00000000000..b29b3ab8e30 --- /dev/null +++ b/src/rust/memory-cache/lib.rs @@ -0,0 +1,1763 @@ +#![forbid(unsafe_code)] + +//! Thread-safe storage and singleflight coordination for the memory cache. +//! +//! The C++ API owns a [`Namespace`] for each provider and creates a [`Binding`] for every configured +//! cache. Named bindings share a [`Cache`]; private bindings receive a dedicated one. Each binding +//! contributes limits, and the cache uses the component-wise maximum of all live bindings after +//! applying the namespace's policy cap. +//! +//! ```text +//! Namespace +//! (names + policy cap) +//! | +//! bind(name, limits) / private binding +//! | +//! v +//! +--------------------------------+ +//! | Cache: one logical cache | +//! | | +//! | Mutex | +//! | - LRU entries + expirations | +//! | - live binding limits | +//! | - weak in-flight fallbacks ---+----+ +//! +--------------------------------+ | +//! ^ v +//! Binding InFlightFallback(key) +//! / \ +//! FallbackPermit Waiters +//! (leader) / \ +//! | published value gate acquired +//! | | | +//! +-- succeed -->+ next leader +//! ``` +//! +//! A cache miss with fallback elects one caller as the leader. Its [`FallbackPermit`] holds the +//! fallback gate while C++ computes the value without holding the cache mutex. Other callers become +//! [`Waiter`] futures. A successful leader publishes one shared [`Bytes`] allocation to the cache +//! and all waiters. If the leader abandons its permit, the gate opens and exactly one waiter becomes +//! the next leader. Weak namespace and in-flight fallback indexes avoid keeping otherwise-unused +//! caches or abandoned work alive. + +use std::cmp::Ordering; +use std::collections::BTreeSet; +use std::collections::HashMap; +use std::collections::hash_map::RandomState; +use std::error::Error; +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::MutexGuard; +use std::sync::Weak; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering as AtomicOrdering; +use std::task::Context; +use std::task::Poll; +use std::time::Instant; + +use bytes::Bytes; +use futures::future::Either; +use futures::future::select; +use hashlink::LinkedHashMap; +use tokio::sync::Mutex as AsyncMutex; +use tokio::sync::OwnedMutexGuard; +use tokio::sync::watch; + +type OrderedMap = LinkedHashMap; +type SharedValue = Bytes; + +/// Identifies which payload a synchronous read decision contains. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ReadKind { + Miss, + Value, + Leader, + Waiter, +} + +/// Controls whether a miss is returned immediately or joins singleflight fallback coordination. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ReadMode { + CacheOnly, + WithFallback, +} + +/// Identifies whether a completed wait produced a value or promoted the waiter to leader. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum WaitKind { + Value, + Leader, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum EvictionReason { + Expiration, + Lru, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum WriteOutcome { + Success { + is_update: bool, + total_after: usize, + entries_after: usize, + }, + ValueTooLarge { + max_value_size: usize, + }, + AlreadyExpired, +} + +/// The configured capacity requested by one binding or computed for a cache cache. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct Limits { + pub max_keys: u32, + pub max_value_size: u32, + pub max_total_value_size: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CacheError { + IdExhausted(&'static str), + CounterOverflow(&'static str), + InvalidDecision(&'static str), +} + +impl fmt::Display for CacheError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::IdExhausted(kind) => write!(formatter, "memory cache {kind} IDs exhausted"), + Self::CounterOverflow(kind) => { + write!(formatter, "memory cache {kind} counter overflow") + } + Self::InvalidDecision(message) => formatter.write_str(message), + } + } +} + +impl Error for CacheError {} + +#[derive(Clone, Copy)] +pub struct ReadTrace { + pub cache_hit: bool, + pub entry_size: usize, + pub total_value_size: usize, + pub entry_count: usize, + pub waiters_ahead: usize, + pub lock_wait_ns: u64, +} + +pub struct EvictionTrace { + pub reason: EvictionReason, + pub key: Arc<[u8]>, + pub value_size: usize, + pub total_before: usize, + pub entries_before: usize, +} + +pub struct WriteTrace { + pub value_size: usize, + pub has_expiration: bool, + pub outcome: WriteOutcome, + pub evictions: Vec, + pub waiters_notified: usize, +} + +pub struct Stats { + pub limits: Limits, + pub bindings: usize, + pub entries: usize, + pub in_flight_fallbacks: usize, + pub waiters: usize, + pub canceled_waiters: usize, +} + +impl Limits { + fn normalize(mut self) -> Self { + if self.max_keys == 0 || self.max_value_size == 0 || self.max_total_value_size == 0 { + return Self::default(); + } + self.max_value_size = self + .max_value_size + .min(self.max_total_value_size.try_into().unwrap_or(u32::MAX)); + self + } + + fn component_max(self, other: Self) -> Self { + Self { + max_keys: self.max_keys.max(other.max_keys), + max_value_size: self.max_value_size.max(other.max_value_size), + max_total_value_size: self.max_total_value_size.max(other.max_total_value_size), + } + } +} + +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +fn elapsed_ns(start: Instant) -> u64 { + start.elapsed().as_nanos().try_into().unwrap_or(u64::MAX) +} + +fn saturating_increment(counter: &AtomicUsize) -> usize { + counter + .fetch_update(AtomicOrdering::Relaxed, AtomicOrdering::Relaxed, |value| { + Some(value.saturating_add(1)) + }) + .unwrap_or(usize::MAX) +} + +/// Defines the sharing boundary for named caches and applies provider-wide policy. +/// +/// Its name index contains weak references so the namespace does not extend a cache's lifetime. +pub struct Namespace { + inner: Arc, +} + +struct NamespaceInner { + state: Mutex, + max_total_value_size: Option, +} + +#[derive(Default)] +struct NamespaceState { + named: HashMap, Weak>, +} + +/// One private or named logical cache. +/// +/// All entry, binding, and fallback-index mutations are serialized by `state`. Waiter counters +/// remain atomic because cancellation and polling do not need to acquire that mutex merely for +/// telemetry. +struct Cache { + namespace: Weak, + name: Option>, + state: Mutex, + live_waiters: AtomicUsize, + canceled_waiters: AtomicUsize, +} + +/// Mutable state protected by a cache's synchronous mutex. +/// +/// `entries` supplies LRU order, while `expirations` supplies expiration order. The fallback map is +/// a weak index into asynchronous work whose lifetime is owned by leaders and waiters. +#[derive(Default)] +struct CacheState { + entries: OrderedMap, Entry>, + expirations: BTreeSet, + total_value_size: usize, + bindings: HashMap, + next_binding_id: u64, + effective_limits: Limits, + in_flight_fallbacks: HashMap, Weak>, +} + +/// A cached value and its optional absolute expiration time. +struct Entry { + value: SharedValue, + expiration: Option, +} + +/// Secondary index record used to find expired entries before falling back to LRU eviction. +#[derive(Clone)] +struct ExpirationRecord { + expiration: f64, + key: Arc<[u8]>, +} + +impl PartialEq for ExpirationRecord { + fn eq(&self, other: &Self) -> bool { + self.cmp(other) == Ordering::Equal + } +} + +impl Eq for ExpirationRecord {} + +impl Ord for ExpirationRecord { + fn cmp(&self, other: &Self) -> Ordering { + self.expiration + .total_cmp(&other.expiration) + .then_with(|| self.key.cmp(&other.key)) + } +} + +impl PartialOrd for ExpirationRecord { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +/// Coordinates one in-progress fallback for one key. +/// +/// The async mutex elects leaders, while the watch channel broadcasts a successful value without +/// requiring each waiter to reacquire the cache-state mutex. +struct InFlightFallback { + cache: Arc, + key: Arc<[u8]>, + gate: Arc>, + completion: watch::Sender>, + waiters: AtomicUsize, +} + +enum WaitResolution { + Value(SharedValue), + Leader(OwnedMutexGuard<()>), +} + +#[derive(Clone, Copy)] +enum WaiterExit { + Completed, + Canceled, +} + +/// One configured consumer's access to a cache and contribution to its effective limits. +/// +/// Releasing the binding removes its limits and immediately resizes the cache if the effective +/// limits shrink. +pub struct Binding { + namespace: Arc, + cache: Arc, + id: Option, +} + +/// An immutable shared cached value returned across the FFI boundary without another byte copy. +pub struct Value { + bytes: SharedValue, +} + +/// The immediate result of a read, pairing telemetry with exactly one outcome payload. +pub struct ReadDecision { + trace: ReadTrace, + outcome: ReadOutcome, +} + +/// The result of awaiting a coalesced fallback: either its value or leadership of the same operation. +pub struct WaitOutcome { + outcome: WaitResult, +} + +enum ReadOutcome { + Miss, + Value(Option), + Leader(Option), + Waiter(Option), +} + +enum WaitResult { + Value(Option), + Leader(Option), +} + +/// Exclusive permission to complete an in-flight fallback. +/// +/// Dropping the permit without succeeding releases the gate so a waiter can take over. +pub struct FallbackPermit { + // Fallback ownership drops before the guard can unlock an abandoned operation. + in_flight: Arc, + _guard: OwnedMutexGuard<()>, +} + +/// A cancellation-safe future waiting for either the published value or the fallback gate. +/// +/// Dropping it unregisters its waiter counters and removes its lock attempt from the async mutex. +pub struct Waiter { + // Fallback ownership drops before a canceled lock future leaves the mutex queue. + in_flight: Option>, + wait: Option + Send>>>, +} + +fn allocate_id( + next: &mut u64, + kind: &'static str, + mut occupied: impl FnMut(u64) -> bool, +) -> Result { + let start = *next; + loop { + let candidate = next.checked_add(1).unwrap_or(1); + *next = candidate; + if !occupied(candidate) { + return Ok(candidate); + } + if candidate == start { + return Err(CacheError::IdExhausted(kind)); + } + } +} + +impl NamespaceInner { + fn apply_policy(&self, mut limits: Limits) -> Limits { + if let Some(cap) = self.max_total_value_size { + limits.max_total_value_size = limits.max_total_value_size.min(cap); + } + limits.normalize() + } +} + +impl Drop for Cache { + fn drop(&mut self) { + let Some(name) = &self.name else { + return; + }; + let Some(namespace) = self.namespace.upgrade() else { + return; + }; + let mut state = lock(&namespace.state); + if state + .named + .get(name) + .is_some_and(|cache| std::ptr::eq(cache.as_ptr(), self)) + { + state.named.remove(name); + } + } +} + +impl Drop for InFlightFallback { + fn drop(&mut self) { + let mut state = lock(&self.cache.state); + if state + .in_flight_fallbacks + .get(&self.key) + .is_some_and(|current| std::ptr::eq(current.as_ptr(), self)) + { + state.remove_in_flight_fallback(&self.key); + } + } +} + +impl CacheState { + fn allocate_binding_id(&mut self) -> Result { + allocate_id(&mut self.next_binding_id, "binding", |id| { + self.bindings.contains_key(&id) + }) + } + + fn is_exact_in_flight_fallback(&self, in_flight: &Arc) -> bool { + self.in_flight_fallbacks + .get(&in_flight.key) + .is_some_and(|current| std::ptr::eq(current.as_ptr(), Arc::as_ptr(in_flight))) + } + + fn remove_in_flight_fallback(&mut self, key: &[u8]) -> Option> { + self.in_flight_fallbacks.remove(key) + } + + fn remove_indexes(&mut self, key: &Arc<[u8]>, entry: &Entry) { + if let Some(expiration) = entry.expiration { + let removed = self.expirations.remove(&ExpirationRecord { + expiration, + key: Arc::clone(key), + }); + debug_assert!(removed); + } + } + + fn insert_entry(&mut self, key: Arc<[u8]>, entry: Entry) -> Result<(), CacheError> { + let total_value_size = self + .total_value_size + .checked_add(entry.value.len()) + .ok_or(CacheError::CounterOverflow("value size"))?; + self.entries.reserve(1); + if let Some(expiration) = entry.expiration { + self.expirations.insert(ExpirationRecord { + expiration, + key: Arc::clone(&key), + }); + } + self.total_value_size = total_value_size; + let replaced = self.entries.insert(key, entry); + debug_assert!(replaced.is_none()); + Ok(()) + } + + fn remove_entry(&mut self, key: &[u8]) -> Option<(Arc<[u8]>, Entry)> { + let (key, entry) = self.entries.remove_entry(key)?; + self.remove_indexes(&key, &entry); + self.total_value_size = self + .total_value_size + .checked_sub(entry.value.len()) + .unwrap_or_else(|| unreachable!("memory cache size accounting underflow")); + Some((key, entry)) + } + + fn read_entry(&mut self, key: &[u8], now_ms: f64) -> Option { + if self.entries.get(key).is_some_and(|entry| { + entry + .expiration + .is_some_and(|expiration| expiration < now_ms) + }) { + self.remove_entry(key); + return None; + } + self.entries.to_back(key).map(|entry| entry.value.clone()) + } + + fn eviction_candidate(&self, now_ms: f64) -> Option<(Arc<[u8]>, EvictionReason)> { + if let Some(expiration) = self.expirations.first() + && expiration.expiration < now_ms + { + return Some((Arc::clone(&expiration.key), EvictionReason::Expiration)); + } + self.entries + .front() + .map(|(key, _)| (Arc::clone(key), EvictionReason::Lru)) + } + + fn evict_one(&mut self, now_ms: f64, traces: Option<&mut Vec>) -> bool { + let Some((key, reason)) = self.eviction_candidate(now_ms) else { + return false; + }; + let total_before = self.total_value_size; + let entries_before = self.entries.len(); + let Some((key, entry)) = self.remove_entry(&key) else { + return false; + }; + if let Some(traces) = traces { + traces.push(EvictionTrace { + reason, + key, + value_size: entry.value.len(), + total_before, + entries_before, + }); + } + true + } + + fn resize(&mut self, now_ms: f64) { + if self.effective_limits.max_keys == 0 { + self.entries.clear(); + self.expirations.clear(); + self.total_value_size = 0; + return; + } + let oversized: Vec<_> = self + .entries + .iter() + .filter(|(_, entry)| entry.value.len() > self.effective_limits.max_value_size as usize) + .map(|(key, _)| Arc::clone(key)) + .collect(); + for key in oversized { + self.remove_entry(&key); + } + while self.total_value_size > self.effective_limits.max_total_value_size as usize + || self.entries.len() > self.effective_limits.max_keys as usize + { + if !self.evict_one(now_ms, None) { + break; + } + } + } +} + +impl Future for Waiter { + type Output = WaitOutcome; + + fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll { + let result = match self + .wait + .as_mut() + .unwrap_or_else(|| unreachable!("completed memory cache waiter was polled")) + .as_mut() + .poll(context) + { + Poll::Ready(result) => result, + Poll::Pending => return Poll::Pending, + }; + self.unregister(WaiterExit::Completed); + self.wait = None; + let in_flight = self + .in_flight + .take() + .unwrap_or_else(|| unreachable!("completed memory cache waiter lost its fallback")); + Poll::Ready(WaitOutcome { + outcome: match result { + WaitResolution::Value(bytes) => WaitResult::Value(Some(Value { bytes })), + WaitResolution::Leader(guard) => WaitResult::Leader(Some(FallbackPermit { + in_flight, + _guard: guard, + })), + }, + }) + } +} + +impl Waiter { + fn unregister(&self, exit: WaiterExit) { + let in_flight = self + .in_flight + .as_ref() + .unwrap_or_else(|| unreachable!("completed memory cache waiter was still registered")); + let previous = in_flight.waiters.fetch_sub(1, AtomicOrdering::Relaxed); + debug_assert!(previous > 0); + let previous = in_flight + .cache + .live_waiters + .fetch_sub(1, AtomicOrdering::Relaxed); + debug_assert!(previous > 0); + if matches!(exit, WaiterExit::Canceled) { + saturating_increment(&in_flight.cache.canceled_waiters); + } + } + + pub async fn wait(self) -> WaitOutcome { + self.await + } +} + +impl Drop for Waiter { + fn drop(&mut self) { + if self.wait.is_some() { + self.unregister(WaiterExit::Canceled); + } + } +} + +async fn wait_for_in_flight_fallback( + gate: Arc>, + mut completion: watch::Receiver>, +) -> WaitResolution { + let published = completion.borrow().clone(); + if let Some(value) = published { + return WaitResolution::Value(value); + } + + let selected = { + let lock = gate.lock_owned(); + let completed = wait_for_completion(&mut completion); + futures::pin_mut!(lock, completed); + match select(lock, completed).await { + Either::Left((guard, _)) => WaitResolution::Leader(guard), + Either::Right((value, _)) => WaitResolution::Value(value), + } + }; + match selected { + WaitResolution::Leader(guard) => completion + .borrow() + .clone() + .map_or_else(|| WaitResolution::Leader(guard), WaitResolution::Value), + value @ WaitResolution::Value(_) => value, + } +} + +async fn wait_for_completion(completion: &mut watch::Receiver>) -> SharedValue { + completion + .wait_for(Option::is_some) + .await + .unwrap_or_else(|_| unreachable!("live memory cache fallback stopped publishing")); + completion + .borrow() + .clone() + .unwrap_or_else(|| unreachable!("memory cache fallback completion had no value")) +} + +impl Binding { + pub fn release(&mut self, now_ms: f64) { + let Some(id) = self.id.take() else { + return; + }; + let mut state = lock(&self.cache.state); + let removed = state.bindings.remove(&id); + debug_assert!(removed.is_some()); + if recompute_limits(&mut state, &self.namespace) { + state.resize(now_ms); + } + } +} + +impl Drop for Binding { + fn drop(&mut self) { + // C++ normally releases explicitly with its process clock. This fallback is only for + // exceptional bridge teardown; negative infinity preserves LRU behavior without guessing a clock. + self.release(f64::NEG_INFINITY); + } +} + +fn recompute_limits(state: &mut CacheState, namespace: &NamespaceInner) -> bool { + let limits = state + .bindings + .values() + .copied() + .fold(Limits::default(), Limits::component_max); + let effective_limits = namespace.apply_policy(limits); + if state.effective_limits == effective_limits { + false + } else { + state.effective_limits = effective_limits; + true + } +} + +fn new_cache(namespace: &Arc, name: Option>) -> Arc { + Arc::new(Cache { + namespace: Arc::downgrade(namespace), + name, + state: Mutex::new(CacheState::default()), + live_waiters: AtomicUsize::new(0), + canceled_waiters: AtomicUsize::new(0), + }) +} + +impl Namespace { + pub fn new(max_total_value_size: Option) -> Self { + Self { + inner: Arc::new(NamespaceInner { + state: Mutex::new(NamespaceState::default()), + max_total_value_size, + }), + } + } + + pub fn bind(&self, id: Option<&str>, limits: Limits) -> Result { + let cache = if let Some(id) = id { + let mut namespace_state = lock(&self.inner.state); + let cache = if let Some(cache) = namespace_state.named.get(id).and_then(Weak::upgrade) { + cache + } else { + let name: Arc = Arc::from(id); + let cache = new_cache(&self.inner, Some(Arc::clone(&name))); + namespace_state.named.reserve(1); + namespace_state.named.insert(name, Arc::downgrade(&cache)); + cache + }; + drop(namespace_state); + cache + } else { + new_cache(&self.inner, None) + }; + let mut state = lock(&cache.state); + let binding_id = state.allocate_binding_id()?; + state.bindings.reserve(1); + let limits = limits.normalize(); + state.bindings.insert(binding_id, limits); + state.effective_limits = self + .inner + .apply_policy(state.effective_limits.component_max(limits)); + drop(state); + Ok(Binding { + namespace: Arc::clone(&self.inner), + cache, + id: Some(binding_id), + }) + } +} + +impl Binding { + pub fn read( + &self, + key: &[u8], + now_ms: f64, + mode: ReadMode, + ) -> Result { + let lock_start = Instant::now(); + let mut state = lock(&self.cache.state); + let lock_wait_ns = elapsed_ns(lock_start); + let value = state.read_entry(key, now_ms); + let mut trace = ReadTrace { + cache_hit: value.is_some(), + entry_size: value.as_ref().map_or(0, Bytes::len), + total_value_size: state.total_value_size, + entry_count: state.entries.len(), + waiters_ahead: 0, + lock_wait_ns, + }; + if let Some(bytes) = value { + return Ok(ReadDecision { + trace, + outcome: ReadOutcome::Value(Some(Value { bytes })), + }); + } + if mode == ReadMode::CacheOnly { + return Ok(ReadDecision { + trace, + outcome: ReadOutcome::Miss, + }); + } + if let Some(in_flight) = state.in_flight_fallbacks.get(key).and_then(Weak::upgrade) { + trace.waiters_ahead = saturating_increment(&in_flight.waiters); + saturating_increment(&self.cache.live_waiters); + let wait = Box::pin(wait_for_in_flight_fallback( + Arc::clone(&in_flight.gate), + in_flight.completion.subscribe(), + )); + return Ok(ReadDecision { + trace, + outcome: ReadOutcome::Waiter(Some(Waiter { + in_flight: Some(in_flight), + wait: Some(wait), + })), + }); + } + + let gate = Arc::new(AsyncMutex::new(())); + let (completion, _) = watch::channel(None); + let guard = Arc::clone(&gate) + .try_lock_owned() + .unwrap_or_else(|_| unreachable!("new memory cache fallback gate was already locked")); + let in_flight = Arc::new(InFlightFallback { + cache: Arc::clone(&self.cache), + key: Arc::from(key), + gate, + completion, + waiters: AtomicUsize::new(0), + }); + let permit = FallbackPermit { + in_flight: Arc::clone(&in_flight), + _guard: guard, + }; + state.in_flight_fallbacks.reserve(1); + state + .in_flight_fallbacks + .insert(Arc::clone(&in_flight.key), Arc::downgrade(&in_flight)); + drop(state); + Ok(ReadDecision { + trace, + outcome: ReadOutcome::Leader(Some(permit)), + }) + } + + pub fn delete(&self, key: &[u8]) { + lock(&self.cache.state).remove_entry(key); + } + + pub fn stats(&self) -> Stats { + let state = lock(&self.cache.state); + Stats { + limits: state.effective_limits, + bindings: state.bindings.len(), + entries: state.entries.len(), + in_flight_fallbacks: state + .in_flight_fallbacks + .values() + .filter(|in_flight| in_flight.strong_count() > 0) + .count(), + waiters: self.cache.live_waiters.load(AtomicOrdering::Relaxed), + canceled_waiters: self.cache.canceled_waiters.load(AtomicOrdering::Relaxed), + } + } +} + +impl ReadDecision { + pub fn kind(&self) -> ReadKind { + match self.outcome { + ReadOutcome::Miss => ReadKind::Miss, + ReadOutcome::Value(_) => ReadKind::Value, + ReadOutcome::Leader(_) => ReadKind::Leader, + ReadOutcome::Waiter(_) => ReadKind::Waiter, + } + } + + pub fn trace(&self) -> ReadTrace { + self.trace + } + + pub fn take_value(&mut self) -> Result { + match &mut self.outcome { + ReadOutcome::Value(value) => value.take().ok_or(CacheError::InvalidDecision( + "read decision value was already taken", + )), + _ => Err(CacheError::InvalidDecision( + "read decision does not contain a value", + )), + } + } + + pub fn take_permit(&mut self) -> Result { + match &mut self.outcome { + ReadOutcome::Leader(permit) => permit.take().ok_or(CacheError::InvalidDecision( + "read decision fallback permit was already taken", + )), + _ => Err(CacheError::InvalidDecision( + "read decision does not contain a fallback permit", + )), + } + } + + pub fn take_waiter(&mut self) -> Result { + match &mut self.outcome { + ReadOutcome::Waiter(waiter) => waiter.take().ok_or(CacheError::InvalidDecision( + "read decision waiter was already taken", + )), + _ => Err(CacheError::InvalidDecision( + "read decision does not contain a waiter", + )), + } + } +} + +impl WaitOutcome { + pub fn kind(&self) -> WaitKind { + match self.outcome { + WaitResult::Value(_) => WaitKind::Value, + WaitResult::Leader(_) => WaitKind::Leader, + } + } + + pub fn take_value(&mut self) -> Result { + match &mut self.outcome { + WaitResult::Value(value) => value.take().ok_or(CacheError::InvalidDecision( + "wait outcome value was already taken", + )), + WaitResult::Leader(_) => Err(CacheError::InvalidDecision( + "wait outcome does not contain a value", + )), + } + } + + pub fn take_permit(&mut self) -> Result { + match &mut self.outcome { + WaitResult::Leader(permit) => permit.take().ok_or(CacheError::InvalidDecision( + "wait outcome fallback permit was already taken", + )), + WaitResult::Value(_) => Err(CacheError::InvalidDecision( + "wait outcome does not contain a fallback permit", + )), + } + } +} + +impl FallbackPermit { + pub fn succeed( + self, + bytes: Vec, + expiration: Option, + now_ms: f64, + ) -> Result { + let value = Bytes::from(bytes); + let is_update; + let mut trace = WriteTrace { + value_size: value.len(), + has_expiration: expiration.is_some(), + outcome: WriteOutcome::Success { + is_update: false, + total_after: 0, + entries_after: 0, + }, + evictions: Vec::new(), + waiters_notified: 0, + }; + { + let mut state = lock(&self.in_flight.cache.state); + if !state.is_exact_in_flight_fallback(&self.in_flight) { + return Ok(trace); + } + is_update = state.entries.contains_key(&self.in_flight.key); + if state.effective_limits.max_keys == 0 + || state.effective_limits.max_total_value_size == 0 + || value.len() > state.effective_limits.max_value_size as usize + { + trace.outcome = WriteOutcome::ValueTooLarge { + max_value_size: state.effective_limits.max_value_size as usize, + }; + state.remove_entry(&self.in_flight.key); + } else if expiration.is_some_and(|expiration| expiration < now_ms) { + trace.outcome = WriteOutcome::AlreadyExpired; + state.remove_entry(&self.in_flight.key); + } else { + state.remove_entry(&self.in_flight.key); + while state.entries.len() >= state.effective_limits.max_keys as usize + || state + .total_value_size + .checked_add(value.len()) + .is_none_or(|size| { + size > state.effective_limits.max_total_value_size as usize + }) + { + if !state.evict_one(now_ms, Some(&mut trace.evictions)) { + break; + } + } + state.insert_entry( + Arc::clone(&self.in_flight.key), + Entry { + value: value.clone(), + expiration, + }, + )?; + trace.outcome = WriteOutcome::Success { + is_update, + total_after: state.total_value_size, + entries_after: state.entries.len(), + }; + } + let previous = self.in_flight.completion.send_replace(Some(value)); + debug_assert!(previous.is_none()); + trace.waiters_notified = self.in_flight.waiters.load(AtomicOrdering::Relaxed); + state.remove_in_flight_fallback(&self.in_flight.key); + } + Ok(trace) + } +} + +impl Value { + pub fn bytes(&self) -> &[u8] { + self.bytes.as_ref() + } +} + +#[cfg(test)] +mod tests { + #![expect( + clippy::significant_drop_tightening, + reason = "tests keep decisions and permits alive to exercise ownership transitions" + )] + + use std::collections::HashMap; + use std::sync::Arc; + use std::sync::Barrier; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering as WakeOrdering; + use std::task::Context; + use std::task::Wake; + use std::task::Waker; + use std::thread; + + use super::*; + + struct CountingWake(AtomicUsize); + + impl Wake for CountingWake { + fn wake(self: Arc) { + self.0.fetch_add(1, WakeOrdering::Relaxed); + } + } + + fn test_namespace(max_total_value_size: Option) -> Namespace { + Namespace::new(max_total_value_size) + } + + fn read( + binding: &Binding, + key: &str, + now_ms: f64, + with_fallback: bool, + ) -> Result { + binding.read( + key.as_bytes(), + now_ms, + if with_fallback { + ReadMode::WithFallback + } else { + ReadMode::CacheOnly + }, + ) + } + + fn release(binding: &mut Binding, now_ms: f64) { + binding.release(now_ms); + } + + fn delete(binding: &Binding, key: &str) { + binding.delete(key.as_bytes()); + } + + fn stats(binding: &Binding) -> Stats { + binding.stats() + } + + fn fallback_succeed( + permit: FallbackPermit, + bytes: Vec, + has_expiration: bool, + expiration: f64, + now_ms: f64, + ) -> Result { + permit.succeed(bytes, has_expiration.then_some(expiration), now_ms) + } + + fn limits(keys: u32, value: u32, total: u64) -> Limits { + Limits { + max_keys: keys, + max_value_size: value, + max_total_value_size: total, + } + } + + fn binding(namespace: &Namespace, id: &str, limits: Limits) -> Binding { + namespace.bind(Some(id), limits).unwrap() + } + + fn poll(waiter: &mut Waiter) -> Poll { + let mut context = Context::from_waker(Waker::noop()); + Pin::new(waiter).poll(&mut context) + } + + fn poll_with_waker(waiter: &mut Waiter, wake: &Arc) -> Poll { + let waker = Arc::clone(wake).into(); + let mut context = Context::from_waker(&waker); + Pin::new(waiter).poll(&mut context) + } + + fn leader(binding: &Binding, key: &str) -> FallbackPermit { + let mut decision = read(binding, key, 1.0, true).unwrap(); + assert_eq!(decision.kind(), ReadKind::Leader); + decision.take_permit().unwrap() + } + + fn put(binding: &Binding, key: &str, value: Vec, expiration: Option, now: f64) { + fallback_succeed( + leader(binding, key), + value, + expiration.is_some(), + expiration.unwrap_or_default(), + now, + ) + .unwrap(); + } + + fn assert_consistent(binding: &Binding) { + let state = lock(&binding.cache.state); + assert_eq!( + state.expirations.len(), + state + .entries + .values() + .filter(|entry| entry.expiration.is_some()) + .count() + ); + assert_eq!( + state.total_value_size, + state + .entries + .values() + .map(|entry| entry.value.len()) + .sum::() + ); + + for (key, entry) in &state.entries { + if let Some(expiration) = entry.expiration { + assert!(state.expirations.contains(&ExpirationRecord { + expiration, + key: Arc::clone(key), + })); + } + } + for record in &state.expirations { + let entry = &state.entries[&record.key]; + assert_eq!(Some(record.expiration), entry.expiration); + } + + let mut waiter_count = 0; + for (key, in_flight) in &state.in_flight_fallbacks { + let in_flight = in_flight + .upgrade() + .unwrap_or_else(|| unreachable!("dead fallback remained discoverable")); + assert_eq!(&**key, &*in_flight.key); + waiter_count += in_flight.waiters.load(AtomicOrdering::Relaxed); + } + assert_eq!( + binding.cache.live_waiters.load(AtomicOrdering::Relaxed), + waiter_count + ); + + let mut expected_limits = Limits::default(); + for limits in state.bindings.values() { + expected_limits.max_keys = expected_limits.max_keys.max(limits.max_keys); + expected_limits.max_value_size = + expected_limits.max_value_size.max(limits.max_value_size); + expected_limits.max_total_value_size = expected_limits + .max_total_value_size + .max(limits.max_total_value_size); + } + if let Some(cap) = binding.namespace.max_total_value_size { + expected_limits.max_total_value_size = expected_limits.max_total_value_size.min(cap); + } + if expected_limits.max_keys == 0 + || expected_limits.max_value_size == 0 + || expected_limits.max_total_value_size == 0 + { + expected_limits = Limits::default(); + } else { + expected_limits.max_value_size = expected_limits.max_value_size.min( + expected_limits + .max_total_value_size + .try_into() + .unwrap_or(u32::MAX), + ); + } + assert_eq!(state.effective_limits, expected_limits); + assert!(state.entries.len() <= state.effective_limits.max_keys as usize); + assert!(state.total_value_size <= state.effective_limits.max_total_value_size as usize); + assert!( + state + .entries + .values() + .all(|entry| entry.value.len() <= state.effective_limits.max_value_size as usize) + ); + } + + fn insert_entry_for_test( + binding: &Binding, + key: &str, + value: Vec, + expiration: Option, + ) { + let mut state = lock(&binding.cache.state); + state + .insert_entry( + Arc::from(key.as_bytes()), + Entry { + value: Bytes::from(value), + expiration, + }, + ) + .unwrap(); + } + + #[test] + fn named_sharing_private_isolation_and_explicit_teardown() { + let namespace = test_namespace(None); + let named_a = binding(&namespace, "shared", limits(2, 8, 16)); + let named_b = binding(&namespace, "shared", limits(4, 4, 32)); + let private = namespace.bind(None, limits(2, 8, 16)).unwrap(); + assert_eq!(stats(&named_a).bindings, 2); + assert_eq!(stats(&private).bindings, 1); + assert_eq!(stats(&named_a).limits, limits(4, 8, 32)); + drop(named_b); + assert_eq!(stats(&named_a).bindings, 1); + drop(namespace); + assert_eq!(stats(&named_a).bindings, 1); + } + + #[test] + fn named_cache_drop_removes_weak_namespace_entry() { + let namespace = test_namespace(None); + let named = binding(&namespace, "shared", limits(2, 8, 16)); + let cache = Arc::downgrade(&named.cache); + assert_eq!(lock(&namespace.inner.state).named.len(), 1); + + drop(named); + + assert!(cache.upgrade().is_none()); + assert!(lock(&namespace.inner.state).named.is_empty()); + } + + #[test] + fn deterministic_operations_preserve_model_and_internal_indexes() { + let namespace = test_namespace(None); + let binding = binding(&namespace, "shared", limits(16, 8, 128)); + let mut model: HashMap, Option)> = HashMap::new(); + let mut seed = 0x6a09_e667_f3bc_c909_u64; + let mut now = 1.0; + let mut operation_counts = [0; 4]; + let mut model_hits = 0; + let mut expiration_reads = 0; + + for step in 0..10_000 { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + let key = format!("key-{}", (seed >> 32) % 8); + let operation = (seed % 4) as usize; + operation_counts[operation] += 1; + match operation { + 0 | 1 => { + delete(&binding, &key); + model.remove(&key); + let value = vec![(seed >> 8) as u8; (seed as usize % 8) + 1]; + let expiration = match (seed >> 16) % 4 { + 0 => None, + 1 => Some(now - 1.0), + 2 => Some(now), + _ => Some(now + 5.0), + }; + let trace = fallback_succeed( + leader(&binding, &key), + value.clone(), + expiration.is_some(), + expiration.unwrap_or_default(), + now, + ) + .unwrap(); + if expiration.is_some_and(|expiration| expiration < now) { + assert_eq!(trace.outcome, WriteOutcome::AlreadyExpired); + } else { + assert!(matches!(trace.outcome, WriteOutcome::Success { .. })); + model.insert(key, (value, expiration)); + } + } + 2 => { + if model + .get(&key) + .is_some_and(|(_, expiration)| expiration.is_some_and(|value| value < now)) + { + model.remove(&key); + expiration_reads += 1; + } + let mut decision = read(&binding, &key, now, false).unwrap(); + if let Some((expected, _)) = model.get(&key) { + model_hits += 1; + assert_eq!(decision.kind(), ReadKind::Value); + assert_eq!(decision.take_value().unwrap().bytes(), expected); + } else { + assert_eq!(decision.kind(), ReadKind::Miss); + } + } + _ => { + delete(&binding, &key); + model.remove(&key); + } + } + if step % 17 == 0 { + now += 1.0; + } + if step % 101 == 0 { + let temporary = self::binding(&namespace, "shared", limits(4, 4, 16)); + assert_eq!(stats(&temporary).bindings, 2); + drop(temporary); + } + assert_consistent(&binding); + } + assert!(operation_counts.into_iter().all(|count| count > 2_000)); + assert!(model_hits > 100); + assert!(expiration_reads > 10); + } + + #[test] + fn limits_expiration_rejections_and_decision_errors_cover_boundaries() { + let namespace = test_namespace(None); + let binding = binding(&namespace, "shared", limits(2, 4, 6)); + put(&binding, "a", vec![1; 4], None, 1.0); + put(&binding, "b", vec![2; 2], Some(5.0), 1.0); + assert_consistent(&binding); + + let mut equal_expiration = read(&binding, "b", 5.0, false).unwrap(); + assert_eq!(equal_expiration.kind(), ReadKind::Value); + assert_eq!(equal_expiration.take_value().unwrap().bytes(), [2; 2]); + assert!(matches!( + equal_expiration.take_value(), + Err(CacheError::InvalidDecision(_)) + )); + assert_eq!( + read(&binding, "b", 5.1, false).unwrap().kind(), + ReadKind::Miss + ); + + put(&binding, "b", vec![2; 2], None, 6.0); + put(&binding, "c", vec![3], None, 6.0); + assert_eq!( + read(&binding, "a", 6.0, false).unwrap().kind(), + ReadKind::Miss + ); + assert_eq!( + read(&binding, "b", 6.0, false).unwrap().kind(), + ReadKind::Value + ); + assert_eq!( + read(&binding, "c", 6.0, false).unwrap().kind(), + ReadKind::Value + ); + assert_consistent(&binding); + + let oversized = leader(&binding, "oversized"); + insert_entry_for_test(&binding, "oversized", vec![9], None); + let trace = fallback_succeed(oversized, vec![9; 5], false, 0.0, 6.0).unwrap(); + assert!(matches!(trace.outcome, WriteOutcome::ValueTooLarge { .. })); + assert_eq!( + read(&binding, "oversized", 6.0, false).unwrap().kind(), + ReadKind::Miss + ); + assert_consistent(&binding); + + let expired = leader(&binding, "expired"); + insert_entry_for_test(&binding, "expired", vec![8], None); + let trace = fallback_succeed(expired, vec![8], true, 5.0, 6.0).unwrap(); + assert_eq!(trace.outcome, WriteOutcome::AlreadyExpired); + assert_eq!( + read(&binding, "expired", 6.0, false).unwrap().kind(), + ReadKind::Miss + ); + assert_consistent(&binding); + + let capped_namespace = test_namespace(Some(6)); + let capped = self::binding(&capped_namespace, "capped", limits(4, 10, 10)); + assert_eq!(stats(&capped).limits, limits(4, 6, 6)); + put(&capped, "exact", vec![1; 6], None, 1.0); + let rejected = + fallback_succeed(leader(&capped, "too-large"), vec![1; 7], false, 0.0, 1.0).unwrap(); + assert!(matches!( + rejected.outcome, + WriteOutcome::ValueTooLarge { .. } + )); + assert_consistent(&capped); + + let disabled_namespace = test_namespace(None); + let disabled = self::binding(&disabled_namespace, "disabled", Limits::default()); + let rejected = + fallback_succeed(leader(&disabled, "key"), vec![1], false, 0.0, 1.0).unwrap(); + assert!(matches!( + rejected.outcome, + WriteOutcome::ValueTooLarge { .. } + )); + assert_consistent(&disabled); + } + + #[test] + fn fallback_transition_matrix_cleans_up() { + let namespace = test_namespace(None); + let binding = binding(&namespace, "shared", limits(8, 32, 128)); + + let permit = leader(&binding, "success"); + let mut first = read(&binding, "success", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(); + let mut second = read(&binding, "success", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(); + assert!(poll(&mut first).is_pending()); + assert!(poll(&mut second).is_pending()); + fallback_succeed(permit, vec![1, 2, 3], false, 0.0, 1.0).unwrap(); + for waiter in [&mut first, &mut second] { + let Poll::Ready(mut outcome) = poll(waiter) else { + panic!("successful fallback did not notify waiter"); + }; + assert_eq!(outcome.kind(), WaitKind::Value); + assert_eq!(outcome.take_value().unwrap().bytes(), [1, 2, 3]); + } + assert_consistent(&binding); + + let permit = leader(&binding, "failure"); + let mut first = read(&binding, "failure", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(); + let mut second = read(&binding, "failure", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(); + assert!(poll(&mut first).is_pending()); + assert!(poll(&mut second).is_pending()); + drop(permit); + let Poll::Ready(mut promoted) = poll(&mut first) else { + panic!("first waiter was not promoted"); + }; + assert_eq!(promoted.kind(), WaitKind::Leader); + drop(promoted.take_permit().unwrap()); + let Poll::Ready(mut promoted) = poll(&mut second) else { + panic!("second waiter was not promoted"); + }; + assert_eq!(promoted.kind(), WaitKind::Leader); + drop(promoted.take_permit().unwrap()); + assert_eq!(stats(&binding).in_flight_fallbacks, 0); + assert_consistent(&binding); + } + + #[test] + fn successful_completion_broadcasts_to_all_waiters() { + const COUNT: usize = 64; + let namespace = test_namespace(None); + let binding = binding(&namespace, "shared", limits(4, 32, 64)); + let permit = leader(&binding, "key"); + let mut waiters = Vec::with_capacity(COUNT); + let mut wakes = Vec::with_capacity(COUNT); + for _ in 0..COUNT { + let mut waiter = read(&binding, "key", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(); + let wake = Arc::new(CountingWake(AtomicUsize::new(0))); + assert!(poll_with_waker(&mut waiter, &wake).is_pending()); + waiters.push(waiter); + wakes.push(wake); + } + let mut late = read(&binding, "key", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(); + let value = vec![1, 2, 3]; + let allocation = value.as_ptr(); + + fallback_succeed(permit, value, false, 0.0, 1.0).unwrap(); + + assert!( + wakes + .iter() + .all(|wake| wake.0.load(WakeOrdering::Relaxed) > 0) + ); + for mut waiter in waiters { + let Poll::Ready(mut outcome) = poll(&mut waiter) else { + panic!("broadcast waiter remained pending"); + }; + assert_eq!(outcome.kind(), WaitKind::Value); + assert_eq!(outcome.take_value().unwrap().bytes().as_ptr(), allocation); + } + let Poll::Ready(mut outcome) = poll(&mut late) else { + panic!("waiter first polled after completion remained pending"); + }; + assert_eq!(outcome.take_value().unwrap().bytes().as_ptr(), allocation); + } + + #[test] + fn fallback_success_retains_the_input_allocation() { + let namespace = test_namespace(None); + let binding = binding(&namespace, "shared", limits(4, 32, 64)); + let value = vec![1, 2, 3, 4]; + let allocation = value.as_ptr(); + + fallback_succeed(leader(&binding, "key"), value, false, 0.0, 1.0).unwrap(); + + let mut hit = read(&binding, "key", 1.0, false).unwrap(); + assert_eq!(hit.take_value().unwrap().bytes().as_ptr(), allocation); + } + + #[test] + fn repeated_completion_cancellation_races_preserve_invariants() { + let namespace = test_namespace(None); + let binding = Arc::new(binding(&namespace, "shared", limits(128, 32, 4096))); + for round in 0..100 { + let key = format!("race-{round}"); + let permit = leader(&binding, &key); + let mut waiters = Vec::new(); + for _ in 0..64 { + waiters.push( + read(&binding, &key, 1.0, true) + .unwrap() + .take_waiter() + .unwrap(), + ); + } + let canceled_before = stats(&binding).canceled_waiters; + let barrier = Arc::new(Barrier::new(3)); + let complete_barrier = Arc::clone(&barrier); + let complete = thread::spawn(move || { + complete_barrier.wait(); + fallback_succeed(permit, vec![1, 2, 3], false, 0.0, 1.0).unwrap() + }); + let cancel_barrier = Arc::clone(&barrier); + let cancel = thread::spawn(move || { + cancel_barrier.wait(); + drop(waiters); + }); + barrier.wait(); + let trace = complete.join().unwrap(); + cancel.join().unwrap(); + let canceled_after = stats(&binding).canceled_waiters; + assert!(trace.waiters_notified <= 64); + assert!(canceled_after - canceled_before <= 64); + let stats = stats(&binding); + assert_eq!(stats.in_flight_fallbacks, 0); + assert_eq!(stats.waiters, 0); + assert_consistent(&binding); + } + } + + #[test] + fn fifty_thousand_reverse_and_random_cancellations_clean_up() { + const COUNT: usize = 50_000; + let namespace = test_namespace(None); + let binding = binding(&namespace, "shared", limits(4, 32, 64)); + let permit = leader(&binding, "key"); + let mut waiters = Vec::with_capacity(COUNT); + for _ in 0..COUNT { + waiters.push( + read(&binding, "key", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(), + ); + } + for index in (COUNT / 2..COUNT).rev() { + drop(waiters.swap_remove(index)); + } + let mut seed = 0x9e37_79b9_u64; + while !waiters.is_empty() { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + let index = seed as usize % waiters.len(); + drop(waiters.swap_remove(index)); + } + let stats = stats(&binding); + assert_eq!(stats.waiters, 0); + assert_eq!(stats.canceled_waiters, COUNT); + drop(permit); + } + + #[test] + fn abandonment_promotes_a_live_waiter_and_skips_canceled_waiters() { + let namespace = test_namespace(None); + let binding = binding(&namespace, "shared", limits(4, 32, 64)); + let first = leader(&binding, "key"); + let mut canceled = read(&binding, "key", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(); + let mut next = read(&binding, "key", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(); + assert!(poll(&mut canceled).is_pending()); + assert!(poll(&mut next).is_pending()); + drop(canceled); + drop(first); + let Poll::Ready(promoted) = poll(&mut next) else { + panic!("live waiter was not promoted"); + }; + assert_eq!(promoted.kind(), WaitKind::Leader); + } + + #[test] + fn poll_order_controls_promotion_and_post_failure_cancellation_cleans_up() { + let namespace = test_namespace(None); + let binding = binding(&namespace, "shared", limits(4, 32, 64)); + let permit = leader(&binding, "fifo"); + let mut first = read(&binding, "fifo", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(); + let mut second = read(&binding, "fifo", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(); + + assert!(poll(&mut second).is_pending()); + assert!(poll(&mut first).is_pending()); + drop(permit); + assert!(poll(&mut first).is_pending()); + let Poll::Ready(mut promoted) = poll(&mut second) else { + panic!("first polled waiter was not promoted"); + }; + drop(promoted.take_permit().unwrap()); + drop(first); + assert_eq!(stats(&binding).in_flight_fallbacks, 0); + + let permit = leader(&binding, "cleanup"); + let waiter = read(&binding, "cleanup", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(); + drop(permit); + drop(waiter); + assert_eq!(stats(&binding).in_flight_fallbacks, 0); + } + + #[test] + fn read_during_abandonment_becomes_leader() { + let namespace = test_namespace(None); + let binding = binding(&namespace, "shared", limits(4, 32, 64)); + let permit = leader(&binding, "key"); + let in_flight = Arc::clone(&permit.in_flight); + + drop(permit); + let mut waiter = read(&binding, "key", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(); + let Poll::Ready(mut outcome) = poll(&mut waiter) else { + panic!("waiter did not acquire the abandoned fallback"); + }; + let replacement = outcome.take_permit().unwrap(); + + assert!(Arc::ptr_eq(&in_flight, &replacement.in_flight)); + drop(in_flight); + drop(replacement); + assert_eq!(stats(&binding).in_flight_fallbacks, 0); + } + + #[test] + fn successful_fanout_shares_arc_and_races_with_cancellation() { + let namespace = test_namespace(None); + let binding = Arc::new(binding(&namespace, "shared", limits(4, 32, 64))); + let permit = leader(&binding, "key"); + let mut waiters = Vec::new(); + for _ in 0..1000 { + waiters.push( + read(&binding, "key", 1.0, true) + .unwrap() + .take_waiter() + .unwrap(), + ); + } + let cancel = thread::spawn(move || drop(waiters)); + let trace = fallback_succeed(permit, vec![1, 2, 3], false, 0.0, 1.0).unwrap(); + cancel.join().unwrap(); + assert!(trace.waiters_notified <= 1000); + assert_eq!(stats(&binding).waiters, 0); + let mut hit = read(&binding, "key", 2.0, false).unwrap(); + assert_eq!(hit.take_value().unwrap().bytes(), [1, 2, 3]); + } + + #[test] + fn indexed_eviction_and_large_limit_reduction_preserve_order() { + let namespace = test_namespace(None); + let mut large_binding = binding(&namespace, "shared", limits(10_000, 64, 640_000)); + for index in 0..5000 { + let expiration = (index % 10 == 0).then_some(10.0 + f64::from(index)); + put( + &large_binding, + &format!("key-{index:05}"), + vec![0; 32], + expiration, + 1.0, + ); + } + let small = binding(&namespace, "shared", limits(10, 8, 80)); + release(&mut large_binding, 100_000.0); + let stats = stats(&small); + assert!(stats.entries <= 10); + assert!(stats.limits.max_value_size >= 8); + drop(small); + } + + #[test] + fn expired_entries_win_before_lru_and_ties_use_key() { + let namespace = test_namespace(None); + let binding = binding(&namespace, "shared", limits(2, 8, 16)); + put(&binding, "permanent", vec![1], None, 1.0); + put(&binding, "expired-b", vec![2], Some(2.0), 1.0); + let trace = + fallback_succeed(leader(&binding, "incoming"), vec![3], false, 0.0, 3.0).unwrap(); + assert_eq!(trace.evictions[0].reason, EvictionReason::Expiration); + assert_eq!(&*trace.evictions[0].key, b"expired-b"); + } + + #[test] + fn reads_refresh_lru_order() { + let namespace = test_namespace(None); + let binding = binding(&namespace, "shared", limits(2, 8, 16)); + put(&binding, "a", vec![1], None, 1.0); + put(&binding, "b", vec![2], None, 1.0); + read(&binding, "a", 2.0, false).unwrap(); + + let trace = fallback_succeed(leader(&binding, "c"), vec![3], false, 0.0, 2.0).unwrap(); + + assert_eq!(trace.evictions[0].reason, EvictionReason::Lru); + assert_eq!(&*trace.evictions[0].key, b"b"); + } + + #[test] + fn stale_fallback_cannot_mutate_replacement() { + let namespace = test_namespace(None); + let binding = binding(&namespace, "shared", limits(4, 8, 32)); + let stale = leader(&binding, "key"); + { + let mut cache_state = lock(&binding.cache.state); + cache_state.remove_in_flight_fallback(b"key"); + } + let replacement = leader(&binding, "key"); + fallback_succeed(stale, vec![1], false, 0.0, 1.0).unwrap(); + assert_eq!(stats(&binding).in_flight_fallbacks, 1); + assert_eq!(stats(&binding).entries, 0); + drop(replacement); + } + + #[test] + fn binding_teardown_during_fallback_keeps_fallback_and_clears_entries() { + let namespace = test_namespace(None); + let mut original = binding(&namespace, "shared", limits(2, 8, 16)); + put(&original, "stored", vec![1], None, 1.0); + let permit = leader(&original, "in_flight"); + release(&mut original, 2.0); + let replacement = binding(&namespace, "shared", limits(2, 8, 16)); + assert_eq!(stats(&replacement).in_flight_fallbacks, 1); + assert_eq!(stats(&replacement).entries, 0); + drop(permit); + } +} diff --git a/src/workerd/api/BUILD.bazel b/src/workerd/api/BUILD.bazel index a63724821f4..ce05f89a460 100644 --- a/src/workerd/api/BUILD.bazel +++ b/src/workerd/api/BUILD.bazel @@ -338,20 +338,35 @@ wd_cc_library( wd_cc_library( name = "memory-cache", - srcs = ["memory-cache.c++"], - hdrs = ["memory-cache.h"], + srcs = [ + "memory-cache.c++", + "memory-cache-v2.c++", + "memory-cache-v2-test.h", + ], + hdrs = [ + "memory-cache.h", + ], implementation_deps = [ + "//src/rust/memory-cache/ffi:memory-cache-ffi", "//src/workerd/io", + "//src/workerd/util:autogate", + "//src/workerd/util:thread-scopes", ], visibility = ["//visibility:public"], deps = [ "//src/workerd/io:compatibility-date_capnp", "//src/workerd/jsg", "//src/workerd/util:checked-queue", - "//src/workerd/util:uuid", ], ) +wd_cc_library( + name = "memory-cache-test-support", + testonly = True, + hdrs = ["memory-cache-v2-test.h"], + deps = [":memory-cache"], +) + wd_cc_library( name = "rtti", srcs = [ @@ -627,8 +642,11 @@ kj_test( src = "memory-cache-test.c++", deps = [ ":memory-cache", + ":memory-cache-test-support", "//src/workerd/io", "//src/workerd/io:trace", + "//src/workerd/tests:test-fixture", + "//src/workerd/util:autogate", ], ) diff --git a/src/workerd/api/memory-cache-test.c++ b/src/workerd/api/memory-cache-test.c++ index 6d2947e80fc..d17d6c7cd31 100644 --- a/src/workerd/api/memory-cache-test.c++ +++ b/src/workerd/api/memory-cache-test.c++ @@ -12,11 +12,15 @@ // the callback is still live. This test simulates that sequence: obtain a // callback, destroy the Use, then invoke it. +#include "memory-cache-v2-test.h" #include "memory-cache.h" #include +#include +#include #include +#include namespace workerd::api { namespace { @@ -29,6 +33,175 @@ static SharedMemoryCache::Limits testLimits() { }; } +static bool memoryCacheV2Enabled() { + return util::Autogate::isEnabled(util::AutogateKey::MEMORY_CACHE_V2); +} + +KJ_TEST("MemoryCacheProvider captures its implementation at construction") { + const auto& clock = kj::systemCoarseMonotonicClock(); + MemoryCacheProvider provider(clock); + KJ_EXPECT(isMemoryCacheV2ForTest(provider) == memoryCacheV2Enabled()); +} + +KJ_TEST("V2 serializes concurrent final release and acquisition") { + auto cacheNamespace = MemoryCacheNamespace::create(MemoryCachePolicy{kj::none}); + auto run = [&cacheNamespace]() { + for (size_t i = 0; i < 1000; ++i) { + auto binding = cacheNamespace->getBinding("shared"_kj, testLimits()); + } + }; + { + kj::Thread first(run); + kj::Thread second(run); + } + auto first = cacheNamespace->getBinding("shared"_kj, testLimits()); + auto second = cacheNamespace->getBinding("shared"_kj, testLimits()); + KJ_EXPECT(getMemoryCacheV2StatsForTest(*first).bindings == 2); +} + +KJ_TEST("V2 provider teardown does not invalidate a live binding") { + if (!memoryCacheV2Enabled()) return; + + TestFixture fixture; + fixture.runInIoContext([&](const TestFixture::Environment&) -> kj::Promise { + kj::Own use; + kj::Own privateUse; + { + MemoryCacheProvider provider(kj::systemCoarseMonotonicClock()); + use = provider.getUse("shared"_kj, testLimits()); + privateUse = provider.getUse(kj::none, testLimits()); + } + + SpanBuilder span(nullptr); + auto result = use->getWithFallback(kj::str("key"), span); + KJ_ASSERT(result.is>()); + return kj::mv(result.get>()) + .then([use = kj::mv(use), privateUse = kj::mv(privateUse)]( + SharedMemoryCache::Use::GetWithFallbackOutcome outcome) mutable { + (void)use; + (void)privateUse; + KJ_ASSERT(outcome.is()); + SpanBuilder span(nullptr); + outcome.get()(kj::none, span); + }); + }); +} + +KJ_TEST("V2 canceled waiters unlink immediately") { + TestFixture fixture; + fixture.runInIoContext([&](const TestFixture::Environment&) -> kj::Promise { + auto cacheNamespace = MemoryCacheNamespace::create(MemoryCachePolicy{kj::none}); + auto cache = cacheNamespace->getBinding("shared"_kj, testLimits()); + auto key = kj::str("key"); + SpanBuilder span(nullptr); + auto leader = cache->getWithFallback(key, span); + KJ_ASSERT(leader.is>()); + { + auto follower = cache->getWithFallback(key, span); + KJ_ASSERT(follower.is>()); + KJ_EXPECT(getMemoryCacheV2StatsForTest(*cache).waiters == 1); + } + KJ_EXPECT(getMemoryCacheV2StatsForTest(*cache).waiters == 0); + KJ_EXPECT(getMemoryCacheV2StatsForTest(*cache).canceledWaiters == 1); + + return kj::mv(leader.get>()) + .then([cache = kj::mv(cache)]( + SharedMemoryCache::Use::GetWithFallbackOutcome outcome) mutable { + KJ_ASSERT(outcome.is()); + SpanBuilder span(nullptr); + outcome.get()(kj::none, span); + auto stats = getMemoryCacheV2StatsForTest(*cache); + KJ_EXPECT(stats.inFlightFallbacks == 0); + KJ_EXPECT(stats.waiters == 0); + KJ_EXPECT(stats.canceledWaiters == 1); + }); + }); +} + +KJ_TEST("V2 abandoned fallback token promotes the next waiter") { + TestFixture fixture; + fixture.runInIoContext([&](const TestFixture::Environment&) -> kj::Promise { + auto cacheNamespace = MemoryCacheNamespace::create(MemoryCachePolicy{kj::none}); + auto cache = cacheNamespace->getBinding("shared"_kj, testLimits()); + auto key = kj::str("key"); + SpanBuilder span(nullptr); + auto leader = cache->getWithFallback(key, span); + auto follower = cache->getWithFallback(key, span); + KJ_ASSERT(leader.is>()); + KJ_ASSERT(follower.is>()); + auto followerPromise = + kj::mv(follower.get>()); + return kj::mv(leader.get>()) + .then([follower = kj::mv(followerPromise), cache = kj::mv(cache)]( + SharedMemoryCache::Use::GetWithFallbackOutcome outcome) mutable { + KJ_ASSERT(outcome.is()); + { auto abandoned = kj::mv(outcome.get()); } + return kj::mv(follower).then( + [cache = kj::mv(cache)](SharedMemoryCache::Use::GetWithFallbackOutcome outcome) { + KJ_ASSERT(outcome.is()); + SpanBuilder span(nullptr); + outcome.get()(kj::none, span); + KJ_EXPECT(getMemoryCacheV2StatsForTest(*cache).inFlightFallbacks == 0); + }); + }); + }); +} + +KJ_TEST("V2 fallback callback is one-shot") { + TestFixture fixture; + fixture.runInIoContext([&](const TestFixture::Environment&) -> kj::Promise { + auto cacheNamespace = MemoryCacheNamespace::create(MemoryCachePolicy{kj::none}); + auto cache = cacheNamespace->getBinding("shared"_kj, testLimits()); + auto key = kj::str("key"); + SpanBuilder span(nullptr); + auto leader = cache->getWithFallback(key, span); + KJ_ASSERT(leader.is>()); + return kj::mv(leader.get>()) + .then([](SharedMemoryCache::Use::GetWithFallbackOutcome outcome) { + KJ_ASSERT(outcome.is()); + auto callback = kj::mv(outcome.get()); + SpanBuilder span(nullptr); + callback(kj::none, span); + KJ_EXPECT_THROW_MESSAGE( + "memory cache fallback callback invoked more than once", callback(kj::none, span)); + }); + }); +} + +KJ_TEST("V2 canceled fallback waiters do not overflow the stack") { + TestFixture fixture; + fixture.runInIoContext([&](const TestFixture::Environment&) -> kj::Promise { + auto cacheNamespace = MemoryCacheNamespace::create(MemoryCachePolicy{kj::none}); + auto cache = cacheNamespace->getBinding("shared"_kj, testLimits()); + auto key = kj::str("test-key"); + SpanBuilder span(nullptr); + + auto leader = cache->getWithFallback(key, span); + KJ_ASSERT(leader.is>()); + auto leaderCallback = kj::mv(leader.get>()) + .then([](MemoryCacheUse::GetWithFallbackOutcome outcome) { + KJ_ASSERT(outcome.is()); + return kj::mv(outcome.get()); + }); + + constexpr size_t waiterCount = 50'000; + for (size_t i = 0; i < waiterCount; ++i) { + auto waiter = cache->getWithFallback(key, span); + KJ_ASSERT(waiter.is>()); + } + + KJ_EXPECT(getMemoryCacheV2StatsForTest(*cache).waiters == 0); + KJ_EXPECT(getMemoryCacheV2StatsForTest(*cache).canceledWaiters == waiterCount); + + return leaderCallback.then( + [cache = kj::mv(cache)](MemoryCacheUse::FallbackDoneCallback callback) mutable { + SpanBuilder span(nullptr); + callback(kj::none, span); + KJ_EXPECT(getMemoryCacheV2StatsForTest(*cache).inFlightFallbacks == 0); + }); + }); +} + KJ_TEST("regression: FallbackDoneCallback survives Use destruction") { kj::EventLoop loop; kj::WaitScope waitScope(loop); diff --git a/src/workerd/api/memory-cache-v2-test.h b/src/workerd/api/memory-cache-v2-test.h new file mode 100644 index 00000000000..68b6f6d3b20 --- /dev/null +++ b/src/workerd/api/memory-cache-v2-test.h @@ -0,0 +1,17 @@ +#pragma once + +#include "memory-cache.h" + +namespace workerd::api { + +struct MemoryCacheV2TestStats { + size_t bindings; + size_t inFlightFallbacks; + size_t waiters; + size_t canceledWaiters; +}; + +MemoryCacheV2TestStats getMemoryCacheV2StatsForTest(const MemoryCacheUse& use); +bool isMemoryCacheV2ForTest(const MemoryCacheProvider& provider); + +} // namespace workerd::api diff --git a/src/workerd/api/memory-cache-v2.c++ b/src/workerd/api/memory-cache-v2.c++ new file mode 100644 index 00000000000..85823826716 --- /dev/null +++ b/src/workerd/api/memory-cache-v2.c++ @@ -0,0 +1,336 @@ +#include "memory-cache-v2-test.h" +#include "memory-cache.h" + +#include +#include +#include +#include +#include +#include + +#include + +namespace workerd::api { +namespace { + +namespace rustCache = workerd::rust::memory_cache; + +using Limits = SharedMemoryCache::Limits; +using Outcome = MemoryCacheUse::GetWithFallbackOutcome; +using FallbackResult = MemoryCacheUse::FallbackResult; +using FallbackDoneCallback = MemoryCacheUse::FallbackDoneCallback; + +class MemoryCacheUseV1 final: public MemoryCacheUse { + public: + MemoryCacheUseV1(kj::Own cache, Limits limits) + : use(kj::mv(cache), limits) {} + + kj::Maybe> getWithoutFallback( + const kj::String& key, SpanBuilder& readSpan) const override { + return use.getWithoutFallback(key, readSpan); + } + + kj::OneOf, kj::Promise> getWithFallback( + const kj::String& key, SpanBuilder& readSpan) const override { + return use.getWithFallback(key, readSpan); + } + + void delete_(const kj::String& key) const override { + use.delete_(key); + } + + private: + SharedMemoryCache::Use use; +}; + +static ::rust::Str asRustStr(kj::StringPtr value) { + return ::rust::Str(value.begin(), value.size()); +} + +static ::rust::Slice asRustBytes(kj::ArrayPtr value) { + return ::rust::Slice( + reinterpret_cast(value.begin()), value.size()); +} + +static ::rust::Slice asRustBytes(kj::StringPtr value) { + return ::rust::Slice( + reinterpret_cast(value.begin()), value.size()); +} + +static rustCache::Limits toRustLimits(Limits limits) { + return { + .max_keys = limits.maxKeys, + .max_value_size = limits.maxValueSize, + .max_total_value_size = limits.maxTotalValueSize, + }; +} + +static double cacheNow() { + if (IoContext::tryCurrent() != kj::none) { + return dateNow(); + } + return (kj::systemPreciseCalendarClock().now() - kj::UNIX_EPOCH) / kj::MILLISECONDS; +} + +class RustCacheValueBacking final: public CacheValueBacking { + public: + explicit RustCacheValueBacking(::rust::Box value): value(kj::mv(value)) {} + + kj::ArrayPtr asBytes() const override { + auto bytes = value->bytes(); + return kj::arrayPtr(reinterpret_cast(bytes.data()), bytes.size()); + } + + private: + ::rust::Box value; +}; + +static kj::Own makeCacheValue(::rust::Box value) { + return kj::atomicRefcounted(kj::heap(kj::mv(value))); +} + +class FallbackPermitOwner final { + public: + explicit FallbackPermitOwner(::rust::Box permit) + : permit(kj::mv(permit)) {} + + ::rust::Box take() { + KJ_IF_SOME(current, permit) { + auto result = kj::mv(current); + permit = kj::none; + return result; + } + KJ_FAIL_REQUIRE("memory cache fallback callback invoked more than once"); + } + + private: + kj::Maybe<::rust::Box> permit; +}; + +static int64_t lockWaitNsForTrace(uint64_t lockWaitNs) { + if (isPredictableModeForTest()) { + return 0; + } + return static_cast(kj::min(lockWaitNs, static_cast(INT64_MAX))); +} + +static void emitReadTrace(SpanBuilder& span, const rustCache::ReadTrace& trace) { + span.setTag("memory_cache_lock_wait_time_ns"_kjc, lockWaitNsForTrace(trace.lock_wait_ns)); + span.setTag("cache_hit"_kjc, trace.cache_hit); + if (trace.cache_hit) { + span.setTag("entry_size"_kjc, static_cast(trace.entry_size)); + } + span.setTag("cache_total_size"_kjc, static_cast(trace.total_value_size)); + span.setTag("cache_entry_count"_kjc, static_cast(trace.entry_count)); +} + +static void emitWriteTrace(kj::StringPtr key, const rustCache::WriteTrace& trace) { + auto writeSpan = IoContext::current().makeTraceSpan("memory_cache_write"_kjc); + writeSpan.setTag("key"_kjc, key); + writeSpan.setTag("value_size"_kjc, static_cast(trace.value_size)); + writeSpan.setTag("has_expiration"_kjc, trace.has_expiration); + switch (trace.outcome) { + case rustCache::WriteOutcome::Success: + writeSpan.setTag("write_success"_kjc, true); + writeSpan.setTag("is_update"_kjc, trace.is_update); + writeSpan.setTag("evictions_triggered"_kjc, static_cast(trace.evictions.size())); + writeSpan.setTag("cache_total_size_after"_kjc, static_cast(trace.total_after)); + writeSpan.setTag("cache_entry_count_after"_kjc, static_cast(trace.entries_after)); + break; + case rustCache::WriteOutcome::ValueTooLarge: + writeSpan.setTag("write_rejected"_kjc, true); + writeSpan.setTag("rejection_reason"_kjc, "value_too_large"_kjc); + writeSpan.setTag("max_value_size"_kjc, static_cast(trace.max_value_size)); + break; + case rustCache::WriteOutcome::AlreadyExpired: + writeSpan.setTag("write_rejected"_kjc, true); + writeSpan.setTag("rejection_reason"_kjc, "already_expired"_kjc); + break; + } + + for (const auto& eviction: trace.evictions) { + auto span = IoContext::current().makeTraceSpan("memory_cache_eviction"_kjc); + switch (eviction.reason) { + case rustCache::EvictionReason::Expiration: + span.setTag("eviction_reason"_kjc, "expiration"_kjc); + break; + case rustCache::EvictionReason::Lru: + span.setTag("eviction_reason"_kjc, "lru"_kjc); + break; + } + span.setTag("evicted_key"_kjc, + kj::str( + kj::arrayPtr(reinterpret_cast(eviction.key.data()), eviction.key.size()))); + span.setTag("evicted_size"_kjc, static_cast(eviction.value_size)); + span.setTag("cache_size_before"_kjc, static_cast(eviction.total_before)); + span.setTag("cache_entries_before"_kjc, static_cast(eviction.entries_before)); + } +} + +static FallbackDoneCallback makeFallback( + ::rust::Box permit, kj::String key) { + return [permit = kj::heap(kj::mv(permit)), key = kj::mv(key)]( + kj::Maybe result, SpanBuilder& fallbackSpan) mutable { + auto currentPermit = permit->take(); + KJ_IF_SOME(value, result) { + auto source = value.value->asBytes(); + KJ_IF_SOME(exception, kj::runCatchingExceptions([&]() { + auto trace = currentPermit->succeed(asRustBytes(source), value.expiration, cacheNow()); + emitWriteTrace(key, trace); + fallbackSpan.setTag("waiters_notified"_kjc, static_cast(trace.waiters_notified)); + })) { + KJ_LOG(ERROR, "memory cache fallback completion failed", exception); + } + } + }; +} + +static Outcome convertWaitOutcome(::rust::Box outcome, kj::String key) { + switch (outcome->kind()) { + case rustCache::WaitKind::Value: + return makeCacheValue(outcome->take_value()); + case rustCache::WaitKind::Leader: + return makeFallback(outcome->take_permit(), kj::mv(key)); + default: + KJ_UNREACHABLE; + } +} + +class MemoryCacheUseV2 final: public MemoryCacheUse { + public: + explicit MemoryCacheUseV2(::rust::Box binding): binding(kj::mv(binding)) {} + ~MemoryCacheUseV2() noexcept override { + binding->release(cacheNow()); + } + + kj::Maybe> getWithoutFallback( + const kj::String& key, SpanBuilder& readSpan) const override; + kj::OneOf, kj::Promise> getWithFallback( + const kj::String& key, SpanBuilder& readSpan) const override; + void delete_(const kj::String& key) const override; + MemoryCacheV2TestStats getStatsForTest() const; + + private: + ::rust::Box binding; +}; + +} // namespace + +MemoryCacheProvider::MemoryCacheProvider(const kj::MonotonicClock& timer) + : MemoryCacheProvider(timer, MemoryCachePolicy{}) {} + +MemoryCacheProvider::MemoryCacheProvider(const kj::MonotonicClock& timer, MemoryCachePolicy policy) + : additionalResizeMemoryLimitHandler([policy](SharedMemoryCache::ThreadUnsafeData& data) { + KJ_IF_SOME(cap, policy.maxTotalValueSize) { + data.effectiveLimits.maxTotalValueSize = + kj::min(data.effectiveLimits.maxTotalValueSize, cap); + } + data.effectiveLimits = data.effectiveLimits.normalize(); + }), + timer(timer) { + if (util::Autogate::isEnabled(util::AutogateKey::MEMORY_CACHE_V2)) { + namespaceV2 = MemoryCacheNamespace::create(policy); + } +} + +kj::Own MemoryCacheProvider::getUse( + kj::Maybe cacheId, SharedMemoryCache::Limits limits) const { + KJ_IF_SOME(cacheNamespace, namespaceV2) { + return cacheNamespace->getBinding(cacheId, limits); + } + return kj::heap(getInstance(cacheId), limits); +} + +kj::Own MemoryCacheNamespace::create(MemoryCachePolicy policy) { + class V2 final: public MemoryCacheNamespace { + public: + explicit V2(MemoryCachePolicy policy) + : cacheNamespace(rustCache::namespace_new(policy.maxTotalValueSize)) {} + + kj::Own getBinding(kj::Maybe id, Limits limits) const override { + ::rust::Str name; + bool isPrivate = id == kj::none; + KJ_IF_SOME(value, id) { + name = asRustStr(value); + } + auto binding = cacheNamespace->bind(name, isPrivate, toRustLimits(limits)); + return kj::heap(kj::mv(binding)); + } + + private: + ::rust::Box cacheNamespace; + }; + + return kj::heap(policy); +} + +kj::Maybe> MemoryCacheUseV2::getWithoutFallback( + const kj::String& key, SpanBuilder& readSpan) const { + auto decision = binding->read(asRustBytes(key), dateNow(), rustCache::ReadMode::CacheOnly); + auto trace = decision->trace(); + emitReadTrace(readSpan, trace); + switch (decision->kind()) { + case rustCache::ReadKind::Miss: + return kj::none; + case rustCache::ReadKind::Value: + return makeCacheValue(decision->take_value()); + default: + KJ_FAIL_ASSERT("unexpected Rust memory cache decision without fallback"); + } +} + +kj::OneOf, kj::Promise> MemoryCacheUseV2::getWithFallback( + const kj::String& key, SpanBuilder& readSpan) const { + auto decision = binding->read(asRustBytes(key), dateNow(), rustCache::ReadMode::WithFallback); + auto trace = decision->trace(); + switch (decision->kind()) { + case rustCache::ReadKind::Value: + emitReadTrace(readSpan, trace); + return makeCacheValue(decision->take_value()); + case rustCache::ReadKind::Leader: + readSpan.setTag("memory_cache_lock_wait_time_ns"_kjc, lockWaitNsForTrace(trace.lock_wait_ns)); + readSpan.setTag("cache_hit"_kjc, false); + readSpan.setTag("coalesced_request"_kjc, false); + readSpan.setTag("initiating_fallback"_kjc, true); + readSpan.setTag("cache_total_size"_kjc, static_cast(trace.total_value_size)); + readSpan.setTag("cache_entry_count"_kjc, static_cast(trace.entry_count)); + return kj::Promise(makeFallback(decision->take_permit(), kj::str(key))); + case rustCache::ReadKind::Waiter: { + readSpan.setTag("memory_cache_lock_wait_time_ns"_kjc, lockWaitNsForTrace(trace.lock_wait_ns)); + readSpan.setTag("cache_hit"_kjc, false); + readSpan.setTag("coalesced_request"_kjc, true); + readSpan.setTag("waiting_on_inflight"_kjc, true); + readSpan.setTag("inflight_waiters_count"_kjc, static_cast(trace.waiters_ahead + 1)); + auto waitSpan = kj::rc(readSpan.newChild("memory_cache_coalesce_wait"_kjc)); + waitSpan->setTag("key"_kjc, key.asPtr()); + waitSpan->setTag("waiters_ahead"_kjc, static_cast(trace.waiters_ahead)); + return rustCache::waiter_wait(decision->take_waiter()) + .then([key = kj::str(key)](::rust::Box outcome) mutable { + return convertWaitOutcome(kj::mv(outcome), kj::mv(key)); + }).attach(IoContext::current().registerPendingEvent(), waitSpan.addRef()); + } + case rustCache::ReadKind::Miss: + KJ_FAIL_ASSERT("unexpected Rust memory cache miss with fallback"); + default: + KJ_UNREACHABLE; + } +} + +void MemoryCacheUseV2::delete_(const kj::String& key) const { + binding->remove(asRustBytes(key)); +} + +MemoryCacheV2TestStats MemoryCacheUseV2::getStatsForTest() const { + auto stats = binding->stats(); + return {stats.bindings, stats.in_flight_fallbacks, stats.waiters, stats.canceled_waiters}; +} + +MemoryCacheV2TestStats getMemoryCacheV2StatsForTest(const MemoryCacheUse& use) { + return static_cast(use).getStatsForTest(); +} + +bool isMemoryCacheV2ForTest(const MemoryCacheProvider& provider) { + return provider.namespaceV2 != kj::none; +} + +} // namespace workerd::api diff --git a/src/workerd/api/memory-cache.c++ b/src/workerd/api/memory-cache.c++ index 03dc8bf6007..93ecb531e7e 100644 --- a/src/workerd/api/memory-cache.c++ +++ b/src/workerd/api/memory-cache.c++ @@ -132,7 +132,7 @@ void SharedMemoryCache::putWhileLocked(ThreadUnsafeData& data, const kj::String& key, kj::Own&& value, kj::Maybe expiration) const { - size_t valueSize = value->bytes.size(); + size_t valueSize = value->size(); auto writeSpan = IoContext::current().makeTraceSpan("memory_cache_write"_kjc); writeSpan.setTag("key"_kjc, key.asPtr()); @@ -292,7 +292,7 @@ kj::Maybe> SharedMemoryCache::Use::getWithoutFallback( // Track cache hit/miss readSpan.setTag("cache_hit"_kjc, result != kj::none); KJ_IF_SOME(value, result) { - readSpan.setTag("entry_size"_kjc, static_cast(value->bytes.size())); + readSpan.setTag("entry_size"_kjc, static_cast(value->size())); } readSpan.setTag("cache_total_size"_kjc, static_cast(data->totalValueSize)); readSpan.setTag("cache_entry_count"_kjc, static_cast(data->cache.size())); @@ -310,7 +310,7 @@ SharedMemoryCache::Use::getWithFallback(const kj::String& key, SpanBuilder& read KJ_IF_SOME(existingValue, cache->getWhileLocked(*data, key)) { // Cache hit readSpan.setTag("cache_hit"_kjc, true); - readSpan.setTag("entry_size"_kjc, static_cast(existingValue->bytes.size())); + readSpan.setTag("entry_size"_kjc, static_cast(existingValue->size())); readSpan.setTag("cache_total_size"_kjc, static_cast(data->totalValueSize)); readSpan.setTag("cache_entry_count"_kjc, static_cast(data->cache.size())); return kj::mv(existingValue); @@ -469,10 +469,10 @@ jsg::Promise> MemoryCache::read(jsg::Lock& js, auto userReadSpan = IoContext::current().makeUserTraceSpan("memory_cache_read"_kjc); KJ_IF_SOME(fallback, optionalFallback) { - KJ_SWITCH_ONEOF(cacheUse.getWithFallback(key.value, readSpan)) { + KJ_SWITCH_ONEOF(cacheUse->getWithFallback(key.value, readSpan)) { KJ_CASE_ONEOF(result, kj::Own) { // Optimization: Don't even release the isolate lock if the value is already in cache. - jsg::Deserializer deserializer(js, result->bytes.asPtr()); + jsg::Deserializer deserializer(js, result->asBytes()); auto value = jsg::JsRef(js, deserializer.readValue(js)); return js.resolvedPromise(kj::mv(value)); @@ -486,9 +486,9 @@ jsg::Promise> MemoryCache::read(jsg::Lock& js, KJ_SWITCH_ONEOF(cacheResult) { KJ_CASE_ONEOF(serialized, kj::Own) { readSpan.setTag("fallback_cache_hit"_kjc, true); - readSpan.setTag("entry_size"_kjc, static_cast(serialized->bytes.size())); + readSpan.setTag("entry_size"_kjc, static_cast(serialized->size())); - jsg::Deserializer deserializer(js, serialized->bytes.asPtr()); + jsg::Deserializer deserializer(js, serialized->asBytes()); return js.resolvedPromise(jsg::JsRef(js, deserializer.readValue(js))); } KJ_CASE_ONEOF(callback, SharedMemoryCache::Use::FallbackDoneCallback) { @@ -515,7 +515,7 @@ jsg::Promise> MemoryCache::read(jsg::Lock& js, auto serialized = hackySerialize(js, result.value); fallbackSpan->setTag( - "fallback_result_size"_kjc, static_cast(serialized->bytes.size())); + "fallback_result_size"_kjc, static_cast(serialized->size())); KJ_IF_SOME(expiration, result.expiration) { JSG_REQUIRE( @@ -549,8 +549,8 @@ jsg::Promise> MemoryCache::read(jsg::Lock& js, } KJ_UNREACHABLE; } else { - KJ_IF_SOME(cacheValue, cacheUse.getWithoutFallback(key.value, readSpan)) { - jsg::Deserializer deserializer(js, cacheValue->bytes.asPtr()); + KJ_IF_SOME(cacheValue, cacheUse->getWithoutFallback(key.value, readSpan)) { + jsg::Deserializer deserializer(js, cacheValue->asBytes()); return js.resolvedPromise(jsg::JsRef(js, deserializer.readValue(js))); } return js.resolvedPromise(jsg::JsRef(js, js.undefined())); @@ -567,7 +567,7 @@ void MemoryCache::delete_(jsg::Lock& js, jsg::NonCoercible key) { auto deleteSpan = IoContext::current().makeTraceSpan("memory_cache_delete"_kjc); deleteSpan.setTag("key"_kjc, key.value.asPtr()); - cacheUse.delete_(key.value); + cacheUse->delete_(key.value); deleteSpan.setTag("delete_completed"_kjc, true); } diff --git a/src/workerd/api/memory-cache.h b/src/workerd/api/memory-cache.h index fdbef22c1b7..0730d9de894 100644 --- a/src/workerd/api/memory-cache.h +++ b/src/workerd/api/memory-cache.h @@ -38,10 +38,34 @@ namespace workerd::api { // instances, etc). Objects that represent i/o (like streams or promises are // explicitly not supported. +class CacheValueBacking { + public: + virtual kj::ArrayPtr asBytes() const = 0; + virtual ~CacheValueBacking() noexcept(false) = default; +}; + struct CacheValue: kj::AtomicRefcounted { - CacheValue(kj::Array&& bytes): bytes(kj::mv(bytes)) {} + CacheValue(kj::Array&& bytes): data(kj::mv(bytes)) {} + CacheValue(kj::Own backing): data(kj::mv(backing)) {} - kj::Array bytes; + kj::ArrayPtr asBytes() const { + KJ_SWITCH_ONEOF(data) { + KJ_CASE_ONEOF(d, kj::Array) { + return d.asPtr(); + } + KJ_CASE_ONEOF(b, kj::Own) { + return b->asBytes(); + } + } + KJ_UNREACHABLE; + } + + size_t size() const { + return asBytes().size(); + } + + private: + kj::OneOf, kj::Own> data; }; struct MemoryCacheEntry { @@ -68,7 +92,7 @@ struct MemoryCacheEntry { kj::Own value; inline size_t size() const { - return value->bytes.size(); + return value->size(); } // The expiration timestamp of this cache entry, usually the time at which the @@ -158,7 +182,6 @@ class SharedMemoryCache: public kj::AtomicRefcounted { kj::StringPtr id, kj::Maybe additionalResizeMemoryLimitHandler, const kj::MonotonicClock& timer); - ~SharedMemoryCache() noexcept(false); kj::StringPtr getId() const { @@ -472,13 +495,37 @@ class SharedMemoryCache: public kj::AtomicRefcounted { const kj::MonotonicClock& timer; }; +struct MemoryCachePolicy { + kj::Maybe maxTotalValueSize; +}; + +class MemoryCacheUse { + public: + using FallbackResult = SharedMemoryCache::Use::FallbackResult; + using FallbackDoneCallback = SharedMemoryCache::Use::FallbackDoneCallback; + using GetWithFallbackOutcome = SharedMemoryCache::Use::GetWithFallbackOutcome; + + virtual ~MemoryCacheUse() noexcept(false) = default; + virtual kj::Maybe> getWithoutFallback( + const kj::String& key, SpanBuilder& readSpan) const = 0; + virtual kj::OneOf, kj::Promise> getWithFallback( + const kj::String& key, SpanBuilder& readSpan) const = 0; + virtual void delete_(const kj::String& key) const = 0; +}; + +class MemoryCacheNamespace { + public: + static kj::Own create(MemoryCachePolicy policy); + virtual ~MemoryCacheNamespace() noexcept(false) = default; + virtual kj::Own getBinding( + kj::Maybe id, SharedMemoryCache::Limits limits) const = 0; +}; + // JavaScript class that allows accessing an in-memory cache. -// Each instance of this class holds a SharedMemoryCache::Use object and -// all calls from JavaScript are essentially forwarded to that object, which -// manages interaction with the shared cache in a thread-safe manner. +// Each instance forwards JavaScript calls to the selected backend lease. class MemoryCache: public jsg::Object { public: - MemoryCache(SharedMemoryCache::Use&& use): cacheUse(kj::mv(use)) {} + MemoryCache(kj::Own use): cacheUse(kj::mv(use)) {} using FallbackFunction = jsg::Function(kj::String)>; @@ -499,24 +546,27 @@ class MemoryCache: public jsg::Object { } private: - SharedMemoryCache::Use cacheUse; + kj::Own cacheUse; }; // The MemoryCacheProvider provides the internal implementation of the MemoryCache mechanism. // It is responsible for owning the SharedMemoryCache instances and providing them to the -// bindings as needed. The default implementation (created and returned by createDefault()) -// uses a simple in-memory map to store the SharedMemoryCache instances. +// bindings as needed. The selected implementation is fixed when the provider is constructed. // TODO(later): It may be worth considering some kind of metrics observer for the provider // that can be passed along to the individual cache instances so we can monitor just how much // the in memory cache is being used. class MemoryCacheProvider { public: + explicit MemoryCacheProvider(const kj::MonotonicClock& timer); MemoryCacheProvider(const kj::MonotonicClock& timer, kj::Maybe - additionalResizeMemoryLimitHandler = kj::none); + additionalResizeMemoryLimitHandler); + MemoryCacheProvider(const kj::MonotonicClock& timer, MemoryCachePolicy policy); KJ_DISALLOW_COPY_AND_MOVE(MemoryCacheProvider); ~MemoryCacheProvider() noexcept(false); + kj::Own getUse( + kj::Maybe cacheId, SharedMemoryCache::Limits limits) const; kj::Own getInstance(kj::Maybe cacheId = kj::none) const; void removeInstance(const SharedMemoryCache& instance) const; @@ -532,7 +582,11 @@ class MemoryCacheProvider { // is destroyed, it will remove itself from this cache by calling removeInstance. kj::MutexGuarded> caches; + kj::Maybe> namespaceV2; + const kj::MonotonicClock& timer; + + friend bool isMemoryCacheV2ForTest(const MemoryCacheProvider& provider); }; // clang-format off diff --git a/src/workerd/api/tests/memory-cache-test.js b/src/workerd/api/tests/memory-cache-test.js index bde51efd23f..dc8aa8fea36 100644 --- a/src/workerd/api/tests/memory-cache-test.js +++ b/src/workerd/api/tests/memory-cache-test.js @@ -2,6 +2,7 @@ // Licensed under the Apache 2.0 license found in the LICENSE file or at: // https://opensource.org/licenses/Apache-2.0 import { strictEqual } from 'node:assert'; +import unsafe from 'workerd:unsafe'; export const basic = { async test(ctrl, env) { @@ -20,6 +21,35 @@ export const basic = { }, }; +export const loneSurrogateKey = { + async test(ctrl, env) { + const key = '\ud800'; + const replacement = '\ufffd'; + strictEqual( + await env.CACHE.read(key, async () => ({ value: 'surrogate value' })), + 'surrogate value' + ); + strictEqual( + await env.CACHE.read(replacement, async () => ({ + value: 'replacement value', + })), + 'replacement value' + ); + strictEqual(await env.CACHE.read(key), 'surrogate value'); + strictEqual(await env.CACHE.read(replacement), 'replacement value'); + env.CACHE.delete(key); + strictEqual(await env.CACHE.read(key), undefined); + strictEqual(await env.CACHE.read(replacement), 'replacement value'); + + await env.CACHE.read(key, async () => ({ value: 'surrogate value' })); + await env.CACHE.read(replacement); + await env.CACHE.read('eviction trigger', async () => ({ + value: 'trigger value', + })); + strictEqual(await env.CACHE.read(key), undefined); + }, +}; + export const keysTooLarge = { async test(ctrl, env) { // Keys that are fewer than 2048 bytes should work. @@ -226,6 +256,63 @@ export const fallbackChainingOnError = { }, }; +export const fallbackChainingOnErrorManyWaiters = { + async test(ctrl, env) { + if (!unsafe.isTestAutogateEnabled()) return; + + const waiters = 64; + const promises = []; + for (let i = 0; i < waiters - 1; i++) { + promises.push( + env.CACHE.read('manyWaiters', () => { + throw new Error(`fallback ${i} failed`); + }) + ); + } + promises.push( + env.CACHE.read('manyWaiters', () => { + return { value: 'last' }; + }) + ); + + const results = await Promise.allSettled(promises); + strictEqual(results.length, waiters); + for (let i = 0; i < waiters - 1; i++) { + strictEqual(results[i].status, 'rejected'); + strictEqual(results[i].reason.message, `fallback ${i} failed`); + } + strictEqual(results[waiters - 1].status, 'fulfilled'); + strictEqual(results[waiters - 1].value, 'last'); + strictEqual(await env.CACHE.read('manyWaiters'), 'last'); + }, +}; + +export const fallbackChainingAllWaitersFail = { + async test(ctrl, env) { + if (!unsafe.isTestAutogateEnabled()) return; + + const waiters = 32; + const promises = []; + for (let i = 0; i < waiters; i++) { + promises.push( + env.CACHE.read('allFail', () => { + throw new Error(`nope ${i}`); + }) + ); + } + const results = await Promise.allSettled(promises); + for (let i = 0; i < waiters; i++) { + strictEqual(results[i].status, 'rejected'); + strictEqual(results[i].reason.message, `nope ${i}`); + } + strictEqual(await env.CACHE.read('allFail'), undefined); + strictEqual( + await env.CACHE.read('allFail', () => ({ value: 'recovered' })), + 'recovered' + ); + }, +}; + export const fallbackNotLocked = { async test(ctrl, env) { // Test that one long running fallback does not block another one. diff --git a/src/workerd/api/tests/memory-cache-test.wd-test b/src/workerd/api/tests/memory-cache-test.wd-test index 1958f0914e6..70e058e7b66 100644 --- a/src/workerd/api/tests/memory-cache-test.wd-test +++ b/src/workerd/api/tests/memory-cache-test.wd-test @@ -7,7 +7,7 @@ const unitTests :Workerd.Config = ( modules = [ (name = "worker", esModule = embed "memory-cache-test.js") ], - compatibilityFlags = ["nodejs_compat"], + compatibilityFlags = ["nodejs_compat", "unsafe_module", "memory_cache_delete"], bindings = [ (name = "CACHE", memoryCache = ( limits = ( diff --git a/src/workerd/server/workerd-api.c++ b/src/workerd/server/workerd-api.c++ index 8e83b8bccdd..dd969f0fab7 100644 --- a/src/workerd/server/workerd-api.c++ +++ b/src/workerd/server/workerd-api.c++ @@ -644,13 +644,12 @@ static v8::Local createBindingValue(JsgWorkerdIsolate::Lock& lock, KJ_CASE_ONEOF(cache, Global::MemoryCache) { value = lock.wrap(context, - lock.alloc( - api::SharedMemoryCache::Use(memoryCacheProvider.getInstance(cache.cacheId), - { - .maxKeys = cache.maxKeys, - .maxValueSize = cache.maxValueSize, - .maxTotalValueSize = cache.maxTotalValueSize, - }))); + lock.alloc(memoryCacheProvider.getUse(cache.cacheId, + { + .maxKeys = cache.maxKeys, + .maxValueSize = cache.maxValueSize, + .maxTotalValueSize = cache.maxTotalValueSize, + }))); } KJ_CASE_ONEOF(ns, Global::EphemeralActorNamespace) { diff --git a/src/workerd/util/autogate.h b/src/workerd/util/autogate.h index 8cdf5e1363c..f3f9da6655f 100644 --- a/src/workerd/util/autogate.h +++ b/src/workerd/util/autogate.h @@ -123,7 +123,10 @@ namespace workerd::util { compression utils. Chromium zlib remains the default. */ \ V(COMPRESSION_RS) \ /* Enables per-call JSRPC tracing, trace-context propagation, and related Fetcher spans. */ \ - V(JSRPC_TRACING) + V(JSRPC_TRACING) \ + /* Selects the redesigned memory cache implementation. The legacy implementation remains \ + available for rollback while this gate is rolled out. */ \ + V(MEMORY_CACHE_V2) // clang-format on // --------------------------------------------------------------------------------------