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 // --------------------------------------------------------------------------------------