From 054508321bf8bfeece8b4bd6f5dfe6f158eccd41 Mon Sep 17 00:00:00 2001 From: Wez Furlong Date: Tue, 4 Mar 2025 15:30:00 -0700 Subject: [PATCH] lruttl: drop moka, roll our own cache on top of dashmap I've been having trouble with moka's eviction policy leading to some unwanted duplicate work, especially when the system is under load and the latency is higher. This commit swaps it out in favor of a relatively simple implementation on top of dashmap; that gives us a decent concurrent base and better respects the concurrency limit and LRU around the contentious insertion case. For LRU, since dashmap is a concurrent data structure, it is ~impossible to use the classic doubly linked LRU approach safely. The strategy used here is a proabalistic sampling LRU using a technique similar to that used in redis. We use atomics to tag read and write entries with a monotonic counter. Entries with smaller counter numbers are least-recent than entries with a larger number. During eviction we take a random sample of 10 entries from the cache map, remove any that have expired due to TTL, and if we still need space after that, we'll pick up to half of those to evict. That guarantees that we don't pick the most recent entry of the sample, and on aggregate this should be a reasonable approximation of the "true" LRU. In the background, every 30s, a maintenance task will remove expired entries from the caches. This is primarily to reduce memory utilization when the cache is otherwise idle. When a cache is hot, expiration is performed as part of the eviction logic described above. --- Cargo.lock | 135 +-- Cargo.toml | 2 +- crates/dkim/src/hash.rs | 1 + crates/dkim/src/sign.rs | 1 + crates/dns-resolver/src/lib.rs | 18 +- .../src/http_server/auth.rs | 2 +- crates/kumod/src/egress_source.rs | 2 +- crates/kumod/src/smtp_dispatcher.rs | 2 +- crates/lruttl/Cargo.toml | 9 +- crates/lruttl/src/lib.rs | 790 ++++++++++++++++-- crates/message/src/dkim.rs | 3 +- crates/mod-filesystem/src/lib.rs | 2 +- crates/mod-memoize/src/lib.rs | 8 +- crates/mta-sts/src/lib.rs | 4 +- crates/rfc5321/src/tls.rs | 2 +- 15 files changed, 768 insertions(+), 213 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5f3b1143..b219d3a5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2106,19 +2106,6 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d758ba1b47b00caf47f24925c0074ecb20d6dfcffe7f6d53395c0465674841a" -[[package]] -name = "generator" -version = "0.8.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc6bd114ceda131d3b1d665eba35788690ad37f5916457286b32ab6fd3c438dd" -dependencies = [ - "cfg-if", - "libc", - "log", - "rustversion", - "windows", -] - [[package]] name = "generic-array" version = "0.14.7" @@ -2538,7 +2525,7 @@ dependencies = [ "iana-time-zone-haiku", "js-sys", "wasm-bindgen", - "windows-core 0.52.0", + "windows-core", ] [[package]] @@ -3557,19 +3544,6 @@ version = "0.4.25" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04cbf5b083de1c7e0222a7a51dbfdba1cbe1c6ab0b15e29fff3f6c077fd9cd9f" -[[package]] -name = "loom" -version = "0.7.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "419e0dc8046cb947daa77eb95ae174acfbddb7673b4151f56d1eed8e93fbfaca" -dependencies = [ - "cfg-if", - "generator", - "scoped-tls", - "tracing", - "tracing-subscriber", -] - [[package]] name = "lru" version = "0.12.5" @@ -3592,9 +3566,13 @@ dependencies = [ name = "lruttl" version = "0.1.0" dependencies = [ + "anyhow", + "dashmap", "kumo-server-memory", - "moka", "parking_lot", + "prometheus", + "rand 0.8.5", + "test-log", "tokio", "tracing", ] @@ -4170,28 +4148,6 @@ dependencies = [ "uuid-helper", ] -[[package]] -name = "moka" -version = "0.12.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9321642ca94a4282428e6ea4af8cc2ca4eac48ac7a6a4ea8f33f76d0ce70926" -dependencies = [ - "async-lock 3.4.0", - "crossbeam-channel", - "crossbeam-epoch", - "crossbeam-utils", - "event-listener 5.4.0", - "futures-util", - "loom", - "parking_lot", - "portable-atomic", - "rustc_version", - "smallvec 1.13.2", - "tagptr", - "thiserror 1.0.69", - "uuid", -] - [[package]] name = "mta-sts" version = "0.1.0" @@ -5799,12 +5755,6 @@ dependencies = [ "windows-sys 0.59.0", ] -[[package]] -name = "scoped-tls" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e1cf6437eb19a8f4a6cc0f7dca544973b0b78843adbfeb3683d1a94a0024a294" - [[package]] name = "scopeguard" version = "1.2.0" @@ -6474,12 +6424,6 @@ dependencies = [ "termwiz", ] -[[package]] -name = "tagptr" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" - [[package]] name = "tailer" version = "0.1.0" @@ -6595,6 +6539,28 @@ dependencies = [ "xi-unicode", ] +[[package]] +name = "test-log" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7f46083d221181166e5b6f6b1e5f1d499f3a76888826e6cb1d057554157cd0f" +dependencies = [ + "env_logger", + "test-log-macros", + "tracing-subscriber", +] + +[[package]] +name = "test-log-macros" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "888d0c3c6db53c0fdab160d2ed5e12ba745383d3e85813f2ea0f2b1475ab553f" +dependencies = [ + "proc-macro2 1.0.93", + "quote 1.0.38", + "syn 2.0.98", +] + [[package]] name = "testcontainers" version = "0.23.2" @@ -7776,16 +7742,6 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" -[[package]] -name = "windows" -version = "0.58.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd04d41d93c4992d421894c18c8b43496aa748dd4c081bac0dc93eb0489272b6" -dependencies = [ - "windows-core 0.58.0", - "windows-targets 0.52.6", -] - [[package]] name = "windows-core" version = "0.52.0" @@ -7795,41 +7751,6 @@ dependencies = [ "windows-targets 0.52.6", ] -[[package]] -name = "windows-core" -version = "0.58.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ba6d44ec8c2591c134257ce647b7ea6b20335bf6379a27dac5f1641fcf59f99" -dependencies = [ - "windows-implement", - "windows-interface", - "windows-result", - "windows-strings", - "windows-targets 0.52.6", -] - -[[package]] -name = "windows-implement" -version = "0.58.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2bbd5b46c938e506ecbce286b6628a02171d56153ba733b6c741fc627ec9579b" -dependencies = [ - "proc-macro2 1.0.93", - "quote 1.0.38", - "syn 2.0.98", -] - -[[package]] -name = "windows-interface" -version = "0.58.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "053c4c462dc91d3b1504c6fe5a726dd15e216ba718e84a0e46a88fbe5ded3515" -dependencies = [ - "proc-macro2 1.0.93", - "quote 1.0.38", - "syn 2.0.98", -] - [[package]] name = "windows-registry" version = "0.2.0" diff --git a/Cargo.toml b/Cargo.toml index 9d4d3d40..0fe0c7dc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -125,7 +125,6 @@ metrics-tracing-context = "0.17" minijinja = {version="2.7.0",features=["loader", "builtins", "json", "loop_controls"]} minijinja-contrib = {version="2.7.0",features=["datetime", "pycompat", "rand", "textwrap", "unicode_wordwrap", "wordcount", "timezone"]} mlua = "0.10" -moka = {version="0.12", features=["future"]} nix = "0.28" nom = "7.1.0" nom_locate = "4.1" @@ -170,6 +169,7 @@ sqlite = "0.36" strum = { version = "0.26", features = ["derive"] } tabout = "0.3" tempfile = "3.10" +test-log = {version="0.2", features = ["trace"]} testcontainers-modules = { version = "0.11", default-features = false } textwrap = "0.16" thiserror = "1" diff --git a/crates/dkim/src/hash.rs b/crates/dkim/src/hash.rs index 307640d7..75e1e58c 100644 --- a/crates/dkim/src/hash.rs +++ b/crates/dkim/src/hash.rs @@ -117,6 +117,7 @@ pub(crate) fn compute_body_hash<'a>( } /// Holds a list of header names, normalized to lower case +#[derive(Debug)] pub(crate) enum HeaderList { /// A list of possibly duplicated header names MaybeMultiple(Vec), diff --git a/crates/dkim/src/sign.rs b/crates/dkim/src/sign.rs index ab2c5c3e..0bd20185 100644 --- a/crates/dkim/src/sign.rs +++ b/crates/dkim/src/sign.rs @@ -143,6 +143,7 @@ impl Default for SignerBuilder { } } +#[derive(Debug)] pub struct Signer { signed_headers: HeaderList, private_key: Arc, diff --git a/crates/dns-resolver/src/lib.rs b/crates/dns-resolver/src/lib.rs index 11e45b13..e38b5125 100644 --- a/crates/dns-resolver/src/lib.rs +++ b/crates/dns-resolver/src/lib.rs @@ -347,7 +347,7 @@ impl MailExchanger { .insert( name_fq, Err(error.clone()), - Instant::now() + get_mx_negative_ttl(), + tokio::time::Instant::now() + get_mx_negative_ttl(), ) .await; anyhow::bail!("{error}"); @@ -382,7 +382,9 @@ impl MailExchanger { }; let mx = Arc::new(mx); - let _ = MX_CACHE.insert(name_fq, Ok(mx.clone()), expires).await; + let _ = MX_CACHE + .insert(name_fq, Ok(mx.clone()), expires.into()) + .await; Ok(mx) } @@ -517,7 +519,7 @@ async fn lookup_mx_record(domain_name: &Name) -> anyhow::Result<(Vec anyhow::Result<(Arc>, Instant)> { let key_fq = fully_qualify(key)?; if let Some(lookup) = IP_CACHE.lookup(&key_fq).await { - return Ok((lookup.item, lookup.expiration)); + return Ok((lookup.item, lookup.expiration.into())); } let (v4, v6) = tokio::join!(ipv4_lookup(key), ipv6_lookup(key)); @@ -558,14 +560,14 @@ pub async fn ip_lookup(key: &str) -> anyhow::Result<(Arc>, Instant)> let addr = Arc::new(results); let exp = expires.take().unwrap_or_else(|| Instant::now()); - IP_CACHE.insert(key_fq, addr.clone(), exp).await; + IP_CACHE.insert(key_fq, addr.clone(), exp.into()).await; Ok((addr, exp)) } pub async fn ipv4_lookup(key: &str) -> anyhow::Result<(Arc>, Instant)> { let key_fq = fully_qualify(key)?; if let Some(lookup) = IPV4_CACHE.lookup(&key_fq).await { - return Ok((lookup.item, lookup.expiration)); + return Ok((lookup.item, lookup.expiration.into())); } let answer = RESOLVER @@ -576,14 +578,14 @@ pub async fn ipv4_lookup(key: &str) -> anyhow::Result<(Arc>, Instant let ips = Arc::new(ips); let expires = answer.expires; - IPV4_CACHE.insert(key_fq, ips.clone(), expires).await; + IPV4_CACHE.insert(key_fq, ips.clone(), expires.into()).await; Ok((ips, expires)) } pub async fn ipv6_lookup(key: &str) -> anyhow::Result<(Arc>, Instant)> { let key_fq = fully_qualify(key)?; if let Some(lookup) = IPV6_CACHE.lookup(&key_fq).await { - return Ok((lookup.item, lookup.expiration)); + return Ok((lookup.item, lookup.expiration.into())); } let answer = RESOLVER @@ -594,7 +596,7 @@ pub async fn ipv6_lookup(key: &str) -> anyhow::Result<(Arc>, Instant let ips = Arc::new(ips); let expires = answer.expires; - IPV6_CACHE.insert(key_fq, ips.clone(), expires).await; + IPV6_CACHE.insert(key_fq, ips.clone(), expires.into()).await; Ok((ips, expires)) } diff --git a/crates/kumo-server-common/src/http_server/auth.rs b/crates/kumo-server-common/src/http_server/auth.rs index 527da708..c16b121a 100644 --- a/crates/kumo-server-common/src/http_server/auth.rs +++ b/crates/kumo-server-common/src/http_server/auth.rs @@ -8,7 +8,7 @@ use config::{load_config, CallbackSignature}; use lruttl::LruCacheWithTtl; use std::net::{IpAddr, SocketAddr}; use std::sync::LazyLock; -use std::time::{Duration, Instant}; +use tokio::time::{Duration, Instant}; static AUTH_CACHE: LazyLock>> = LazyLock::new(|| LruCacheWithTtl::new_named("http_server_auth", 128)); diff --git a/crates/kumod/src/egress_source.rs b/crates/kumod/src/egress_source.rs index 8a503e7f..91bbd2ae 100644 --- a/crates/kumod/src/egress_source.rs +++ b/crates/kumod/src/egress_source.rs @@ -322,7 +322,7 @@ impl EgressPool { .with_context(|| format!("resolving egress pool '{name}'"))? }; - Ok(pool) + Ok::<_, anyhow::Error>(pool) }) .await .map_err(|err: Arc| anyhow::anyhow!("{err:#}"))? diff --git a/crates/kumod/src/smtp_dispatcher.rs b/crates/kumod/src/smtp_dispatcher.rs index 34f2cc96..8edc95e5 100644 --- a/crates/kumod/src/smtp_dispatcher.rs +++ b/crates/kumod/src/smtp_dispatcher.rs @@ -760,7 +760,7 @@ impl SmtpDispatcher { .insert( site_name.to_string(), (), - std::time::Instant::now() + duration, + tokio::time::Instant::now() + duration, ) .await; } diff --git a/crates/lruttl/Cargo.toml b/crates/lruttl/Cargo.toml index 99cb26d9..a1f7ea41 100644 --- a/crates/lruttl/Cargo.toml +++ b/crates/lruttl/Cargo.toml @@ -4,8 +4,15 @@ version = "0.1.0" edition = "2021" [dependencies] +anyhow.workspace = true +dashmap.workspace = true kumo-server-memory = {path="../kumo-server-memory"} -moka.workspace = true parking_lot = {workspace=true} +prometheus.workspace = true +rand.workspace = true tokio = {workspace=true, features=["sync"]} tracing = {workspace=true} + +[dev-dependencies] +tokio = {workspace=true, features=["test-util"]} +test-log.workspace = true diff --git a/crates/lruttl/src/lib.rs b/crates/lruttl/src/lib.rs index 4d82c2ea..4b4f831d 100644 --- a/crates/lruttl/src/lib.rs +++ b/crates/lruttl/src/lib.rs @@ -1,38 +1,171 @@ +use dashmap::DashMap; use kumo_server_memory::subscribe_to_memory_status_changes_async; -use moka::future::Cache; -use moka::policy::EvictionPolicy; -use moka::Expiry; use parking_lot::Mutex; +use prometheus::{IntCounter, IntCounterVec, IntGauge, IntGaugeVec}; use std::borrow::Borrow; +use std::fmt::Debug; use std::future::Future; use std::hash::Hash; -use std::marker::PhantomData; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, LazyLock, Weak}; -use std::time::{Duration, Instant}; +use tokio::sync::Semaphore; +use tokio::time::{timeout, Duration, Instant}; + +static CACHE_LOOKUP: LazyLock = LazyLock::new(|| { + prometheus::register_int_counter_vec!( + "lruttl_lookup_count", + "how many times a lruttl cache lookup was initiated for a given cache", + &["cache_name"] + ) + .unwrap() +}); +static CACHE_EVICT: LazyLock = LazyLock::new(|| { + prometheus::register_int_counter_vec!( + "lruttl_evict_count", + "how many times a lruttl cache evicted an item due to capacity constraints", + &["cache_name"] + ) + .unwrap() +}); +static CACHE_EXPIRE: LazyLock = LazyLock::new(|| { + prometheus::register_int_counter_vec!( + "lruttl_expire_count", + "how many times a lruttl cache removed an item due to ttl expiration", + &["cache_name"] + ) + .unwrap() +}); +static CACHE_HIT: LazyLock = LazyLock::new(|| { + prometheus::register_int_counter_vec!( + "lruttl_hit_count", + "how many times a lruttl cache lookup was a hit for a given cache", + &["cache_name"] + ) + .unwrap() +}); +static CACHE_MISS: LazyLock = LazyLock::new(|| { + prometheus::register_int_counter_vec!( + "lruttl_miss_count", + "how many times a lruttl cache lookup was a miss for a given cache", + &["cache_name"] + ) + .unwrap() +}); +static CACHE_INSERT: LazyLock = LazyLock::new(|| { + prometheus::register_int_counter_vec!( + "lruttl_insert_count", + "how many times a lruttl cache was populated via unconditional insert", + &["cache_name"] + ) + .unwrap() +}); +static CACHE_POPULATED: LazyLock = LazyLock::new(|| { + prometheus::register_int_counter_vec!( + "lruttl_populated_count", + "how many times a lruttl cache lookup resulted in performing the work to populate the entry", + &["cache_name"] + ) + .unwrap() +}); +static CACHE_ERROR: LazyLock = LazyLock::new(|| { + prometheus::register_int_counter_vec!( + "lruttl_error_count", + "how many times a lruttl cache population resulted in an error", + &["cache_name"] + ) + .unwrap() +}); +static CACHE_WAIT: LazyLock = LazyLock::new(|| { + prometheus::register_int_gauge_vec!( + "lruttl_waiting_populate", + "how many tasks are currently waiting for a cache entry to populate", + &["cache_name"] + ) + .unwrap() +}); +static CACHE_SIZE: LazyLock = LazyLock::new(|| { + prometheus::register_int_gauge_vec!( + "lruttl_cache_size", + "number of items contained in an lruttl cache", + &["cache_name"] + ) + .unwrap() +}); static CACHES: LazyLock>>> = LazyLock::new(Mutex::default); -struct Inner { +struct Inner { name: String, - cache: Cache>, + tick: AtomicUsize, + capacity: usize, + cache: DashMap>, + lru_samples: AtomicUsize, + lookup_counter: IntCounter, + evict_counter: IntCounter, + expire_counter: IntCounter, + hit_counter: IntCounter, + miss_counter: IntCounter, + populate_counter: IntCounter, + insert_counter: IntCounter, + error_counter: IntCounter, + wait_gauge: IntGauge, + size_gauge: IntGauge, } trait CachePurger { fn name(&self) -> &str; fn purge(&self) -> usize; + fn process_expirations(&self) -> usize; } -impl CachePurger - for Inner +impl< + K: Clone + Debug + Send + Sync + Hash + Eq + 'static, + V: Clone + Debug + Send + Sync + 'static, + > CachePurger for Inner { fn name(&self) -> &str { &self.name } fn purge(&self) -> usize { - let num_entries = self.cache.entry_count(); - self.cache.invalidate_all(); - num_entries as usize + let num_entries = self.cache.len(); + self.cache.clear(); + self.size_gauge.set(self.cache.len() as i64); + num_entries + } + fn process_expirations(&self) -> usize { + let now = Instant::now(); + let mut expired_keys = vec![]; + for map_entry in self.cache.iter() { + let item = map_entry.value(); + match &item.item { + ItemState::Pending(_) => { + // Cannot evict a pending lookup + } + ItemState::Present(_) | ItemState::Failed(_) => { + if now >= item.expiration { + expired_keys.push(map_entry.key().clone()); + } + } + } + } + + let mut num_removed = 0; + for key in expired_keys { + // Sanity check that it is still expired before removing it, + // because it would be a shame to remove it if another actor + // has just updated it + let removed = self + .cache + .remove_if(&key, |_k, entry| now >= entry.expiration) + .is_some(); + if removed { + num_removed += 1; + self.expire_counter.inc(); + self.size_gauge.set(self.cache.len() as i64); + } + } + num_removed } } @@ -69,7 +202,14 @@ async fn prune_expired_caches() { true } None => false, - }) + }); + + for p in purgers { + let n = p.process_expirations(); + if n > 0 { + tracing::debug!("expired {n} entries from cache {}", p.name()); + } + } } } } @@ -95,31 +235,39 @@ async fn purge_caches_on_memory_shortage() { } #[derive(Debug, Clone)] +enum ItemState +where + V: Send, + V: Sync, +{ + Present(V), + Pending(Arc), + Failed(Arc), +} + +#[derive(Debug)] struct Item where V: Send, V: Sync, { - item: V, + item: ItemState, expiration: Instant, + last_tick: AtomicUsize, } -struct PerItemExpiry { - marker: PhantomData<(K, V)>, -} - -impl Expiry> for PerItemExpiry { - fn expire_after_create( - &self, - _key: &K, - item: &Item, - created_at: Instant, - ) -> Option { - Some(item.expiration - created_at) +impl Clone for Item { + fn clone(&self) -> Self { + Self { + item: self.item.clone(), + expiration: self.expiration, + last_tick: self.last_tick.load(Ordering::Relaxed).into(), + } } } -pub struct ItemLookup { +#[derive(Debug)] +pub struct ItemLookup { /// A copy of the item pub item: V, /// If true, the get_or_try_insert operation populated the entry; @@ -129,13 +277,13 @@ pub struct ItemLookup { pub expiration: Instant, } -pub struct LruCacheWithTtl { +pub struct LruCacheWithTtl { inner: Arc>, } impl< - K: Clone + Hash + Eq + Send + Sync + std::fmt::Debug + 'static, - V: Clone + Send + Sync + 'static, + K: Clone + Debug + Hash + Eq + Send + Sync + std::fmt::Debug + 'static, + V: Clone + Debug + Send + Sync + 'static, > LruCacheWithTtl { #[deprecated = "use new_named instead"] @@ -145,23 +293,56 @@ impl< pub fn new_named>(name: S, capacity: usize) -> Self { let name = name.into(); + let cache = DashMap::new(); - let cache = Cache::builder() - .name(&name) - .eviction_policy(EvictionPolicy::lru()) - .eviction_listener({ - let name = name.clone(); - move |k, _v, reason| { - tracing::trace!("evicting {name} {k:?} {reason:?}"); - } - }) - .max_capacity(capacity as u64) - .expire_after(PerItemExpiry { - marker: PhantomData, - }) - .build(); + let lookup_counter = CACHE_LOOKUP + .get_metric_with_label_values(&[&name]) + .expect("failed to get counter"); + let hit_counter = CACHE_HIT + .get_metric_with_label_values(&[&name]) + .expect("failed to get counter"); + let evict_counter = CACHE_EVICT + .get_metric_with_label_values(&[&name]) + .expect("failed to get counter"); + let expire_counter = CACHE_EXPIRE + .get_metric_with_label_values(&[&name]) + .expect("failed to get counter"); + let miss_counter = CACHE_MISS + .get_metric_with_label_values(&[&name]) + .expect("failed to get counter"); + let populate_counter = CACHE_POPULATED + .get_metric_with_label_values(&[&name]) + .expect("failed to get counter"); + let insert_counter = CACHE_INSERT + .get_metric_with_label_values(&[&name]) + .expect("failed to get counter"); + let error_counter = CACHE_ERROR + .get_metric_with_label_values(&[&name]) + .expect("failed to get counter"); + let wait_gauge = CACHE_WAIT + .get_metric_with_label_values(&[&name]) + .expect("failed to get counter"); + let size_gauge = CACHE_SIZE + .get_metric_with_label_values(&[&name]) + .expect("failed to get counter"); - let inner = Arc::new(Inner { name, cache }); + let inner = Arc::new(Inner { + name, + cache, + tick: AtomicUsize::new(0), + capacity, + lru_samples: AtomicUsize::new(10), + lookup_counter, + evict_counter, + expire_counter, + hit_counter, + miss_counter, + populate_counter, + error_counter, + wait_gauge, + insert_counter, + size_gauge, + }); // Register with the global list of caches using a weak reference. // We need to "erase" the K/V types in order to do that, so we @@ -179,9 +360,206 @@ impl< } pub fn clear(&self) -> usize { - let num_entries = self.inner.cache.entry_count(); - self.inner.cache.invalidate_all(); - num_entries as usize + let num_entries = self.inner.cache.len(); + self.inner.cache.clear(); + self.inner.size_gauge.set(self.inner.cache.len() as i64); + num_entries + } + + fn inc_tick(&self) -> usize { + self.inner.tick.fetch_add(1, Ordering::Relaxed) + 1 + } + + fn update_tick(&self, item: &Item) { + let v = self.inc_tick(); + item.last_tick.store(v, Ordering::Relaxed); + } + + /// Evict up to target entries. + /// + /// We use a probablistic approach to the LRU, because + /// it is challenging to safely thread the classic doubly-linked-list + /// through dashmap. + /// + /// target is bounded to half of number of selected samples, in + /// order to ensure that we don't randomly pick the newest element + /// from the set when under pressure. + /// + /// Redis uses a similar technique for its LRU as described + /// in + /// which suggests that sampling 10 keys at random to them compare + /// their recency yields a reasonably close approximation to the + /// 100% precise LRU. + /// + /// Since we also support TTLs, we'll just go ahead and remove + /// any expired keys that show up in the sampled set. + fn evict_some(&self, target: usize) -> bool { + let now = Instant::now(); + + // Approximate (since it could change immediately after reading) + // cache size + let cache_size = self.inner.cache.len(); + // How many keys to sample + let num_samples = self + .inner + .lru_samples + .load(Ordering::Relaxed) + .min(cache_size); + + // a list of keys which have expired + let mut expired_keys = vec![]; + // a random selection of up to num_samples (key, tick) tuples + let mut samples = vec![]; + + // Pick some random keys. + // The rand crate has some helpers for working with iterators, + // but they appear to copy many elements into an internal buffer + // in order to make a selection, and we want to avoid directly + // considering every possible element because some users have + // very large capacity caches. + // + // The approach taken here is to produce a random list of iterator + // offsets so that we can skim across the iterator in a single + // pass and pull out a random selection of elements. + // The sample function provides a randomized list of indices that + // we can use for this; we need to sort it first, but the cost + // should be reasonably low as num_samples should be ~10 or so + // in the most common configuration. + { + let mut rng = rand::thread_rng(); + let mut indices = + rand::seq::index::sample(&mut rng, cache_size, num_samples).into_vec(); + indices.sort(); + let mut iter = self.inner.cache.iter(); + let mut current_idx = 0; + + /// Advance an iterator by skip_amount. + /// Ideally we'd use Iterator::advance_by for this, but at the + /// time of writing that method is nightly only. + /// Note that it also uses next() internally anyway + fn advance_by(iter: &mut impl Iterator, skip_amount: usize) { + for _ in 0..skip_amount { + if iter.next().is_none() { + return; + } + } + } + + for idx in indices { + // idx is the index we want to be on; we'll need to skip ahead + // by some number of slots based on the current one. skip_amount + // is that number. + let skip_amount = idx - current_idx; + advance_by(&mut iter, skip_amount); + + match iter.next() { + Some(map_entry) => { + current_idx = idx + 1; + let item = map_entry.value(); + match &item.item { + ItemState::Pending(_) => { + // Cannot evict a pending lookup + } + ItemState::Present(_) | ItemState::Failed(_) => { + if now >= item.expiration { + expired_keys.push(map_entry.key().clone()); + } else { + let last_tick = item.last_tick.load(Ordering::Relaxed); + samples.push((map_entry.key().clone(), last_tick)); + } + } + } + } + None => { + break; + } + } + } + } + + let mut num_removed = 0; + for key in expired_keys { + // Sanity check that it is still expired before removing it, + // because it would be a shame to remove it if another actor + // has just updated it + let removed = self + .inner + .cache + .remove_if(&key, |_k, entry| now >= entry.expiration) + .is_some(); + if removed { + tracing::trace!("{} expired {key:?}", self.inner.name); + num_removed += 1; + self.inner.expire_counter.inc(); + } + } + + // Since we're picking random elements, we want to ensure that + // we never pick the newest element from the set to evict because + // that is likely the wrong choice. We need enough samples to + // know that the lowest number we picked is representative + // of the eldest element in the map overall. + // We limit ourselves to half of the number of selected samples. + let target = target.min(samples.len() / 2).max(1); + + // If we met our target, skip the extra work below + if num_removed >= target { + self.inner.size_gauge.set(self.inner.cache.len() as i64); + tracing::trace!( + "{} expired {num_removed} of target {target}", + self.inner.name + ); + return true; + } + + // Sort by ascending tick, which is equivalent to having the + // LRU within that set towards the front of the vec + samples.sort_by(|(_ka, tick_a), (_kb, tick_b)| tick_a.cmp(&tick_b)); + + for (key, tick) in samples { + // Sanity check that the tick value is the same as we expect. + // If it has changed since we sampled it, then that element + // is no longer a good candidate for LRU eviction. + if self + .inner + .cache + .remove_if(&key, |_k, item| { + item.last_tick.load(Ordering::Relaxed) == tick + }) + .is_some() + { + tracing::debug!("{} evicted {key:?}", self.inner.name); + num_removed += 1; + self.inner.evict_counter.inc(); + self.inner.size_gauge.set(self.inner.cache.len() as i64); + if num_removed >= target { + return true; + } + } + } + + if num_removed == 0 { + tracing::warn!( + "{} did not find anything to evict, target was {target}", + self.inner.name + ); + } + + tracing::trace!( + "{} removed {num_removed} of target {target}", + self.inner.name + ); + + num_removed > 0 + } + + /// Potentially make some progress to get back under + /// budget on the cache capacity + fn maybe_evict(&self) { + let cache_size = self.inner.cache.len(); + if cache_size > self.inner.capacity { + self.evict_some(cache_size - self.inner.capacity); + } } pub async fn lookup(&self, name: &Q) -> Option> @@ -189,12 +567,48 @@ impl< K: Borrow, Q: Hash + Eq, { - let entry = self.inner.cache.get(name).await?; - Some(ItemLookup { - item: entry.item.clone(), - expiration: entry.expiration, - is_fresh: false, - }) + self.inner.lookup_counter.inc(); + match self.inner.cache.get_mut(name) { + None => { + self.inner.miss_counter.inc(); + return None; + } + Some(entry) => { + match &entry.item { + ItemState::Present(item) => { + let now = Instant::now(); + if now >= entry.expiration { + // Expired; remove it from the map. + // Take care to drop our ref first so that we don't + // self-deadlock + drop(entry); + if self + .inner + .cache + .remove_if(name, |_k, entry| now >= entry.expiration) + .is_some() + { + self.inner.expire_counter.inc(); + self.inner.size_gauge.set(self.inner.cache.len() as i64); + } + self.inner.miss_counter.inc(); + return None; + } + self.inner.hit_counter.inc(); + self.update_tick(&entry); + Some(ItemLookup { + item: item.clone(), + expiration: entry.expiration, + is_fresh: false, + }) + } + ItemState::Pending(_) | ItemState::Failed(_) => { + self.inner.miss_counter.inc(); + None + } + } + } + } } pub async fn get(&self, name: &Q) -> Option @@ -206,50 +620,252 @@ impl< } pub async fn insert(&self, name: K, item: V, expiration: Instant) -> V { - self.inner - .cache - .insert( - name, - Item { - item: item.clone(), - expiration, - }, - ) - .await; + self.inner.cache.insert( + name, + Item { + item: ItemState::Present(item.clone()), + expiration, + last_tick: self.inc_tick().into(), + }, + ); + + self.inner.insert_counter.inc(); + self.inner.size_gauge.set(self.inner.cache.len() as i64); + self.maybe_evict(); + item } + fn clone_item_state(&self, name: &K) -> (ItemState, Instant) { + let mut is_new = false; + let mut entry = self.inner.cache.entry(name.clone()).or_insert_with(|| { + is_new = true; + Item { + item: ItemState::Pending(Arc::new(Semaphore::new(1))), + expiration: Instant::now() + Duration::from_secs(60), + last_tick: self.inc_tick().into(), + } + }); + + match &entry.value().item { + ItemState::Pending(_) => {} + ItemState::Present(_) | ItemState::Failed(_) => { + let now = Instant::now(); + if now >= entry.expiration { + // Expired; we will need to fetch it + entry.value_mut().item = ItemState::Pending(Arc::new(Semaphore::new(1))); + } + } + } + + self.update_tick(&entry); + let item = entry.value(); + let result = (item.item.clone(), entry.expiration); + drop(entry); + + if is_new { + self.inner.size_gauge.set(self.inner.cache.len() as i64); + self.maybe_evict(); + } + + result + } + /// Get an existing item, but if that item doesn't already exist, /// execute the future `fut` to provide a value that will be inserted and then /// returned. This is done atomically wrt. other callers. /// The TTL parameter is a function that can extract the TTL from the value type, /// or just return a constant TTL. - pub async fn get_or_try_insert Duration>( + pub async fn get_or_try_insert, TTL: FnOnce(&V) -> Duration>( &self, name: &K, - ttl: TTL, + ttl_func: TTL, fut: impl Future>, - ) -> Result, Arc> { - let entry = self - .inner - .cache - .entry_by_ref(name) - .or_try_insert_with(async move { - let item = fut.await?; - let ttl = (ttl)(&item); - Ok(Item { - item, - expiration: Instant::now() + ttl, - }) - }) - .await?; - let is_fresh = entry.is_fresh(); - let item = entry.value(); + ) -> Result, Arc> { + // Fast path avoids cloning the key + if let Some(entry) = self.lookup(name).await { + return Ok(entry); + } - Ok(ItemLookup { - is_fresh, - item: item.item.clone(), - expiration: item.expiration, - }) + // Note: the lookup call increments lookup_counter and miss_counter + + match self.clone_item_state(name) { + (ItemState::Present(item), expiration) => { + return Ok(ItemLookup { + item, + expiration, + is_fresh: false, + }); + } + (ItemState::Failed(error), _) => { + return Err(error); + } + (ItemState::Pending(sema), _) => { + /// A little helper to ensure that we decrement the count + /// when we unwind, in the case that this future is cancelled + /// or abandoned prior to completion + struct DecOnDrop(IntGauge); + impl DecOnDrop { + /// Increment on acquire, decrement on drop + fn new(g: IntGauge) -> Self { + g.inc(); + Self(g) + } + } + impl Drop for DecOnDrop { + fn drop(&mut self) { + self.0.dec(); + } + } + + let wait_count = DecOnDrop::new(self.inner.wait_gauge.clone()); + let wait_result = + match timeout(Duration::from_secs(120), sema.acquire_owned()).await { + Err(_) => { + self.inner.error_counter.inc(); + tracing::error!( + "{} semaphore acquire for {name:?} timed out", + self.inner.name + ); + return Err(Arc::new(anyhow::anyhow!( + "{} lookup for {name:?} \ + timed out on semaphore acquire", + self.inner.name + ))); + } + Ok(r) => r, + }; + + drop(wait_count); + + // While we slept, someone else may have satisfied + // the lookup; check it + match self.clone_item_state(name) { + (ItemState::Present(item), expiration) => { + return Ok(ItemLookup { + item, + expiration, + is_fresh: false, + }); + } + (ItemState::Failed(error), _) => { + self.inner.hit_counter.inc(); + return Err(error); + } + (ItemState::Pending(current_sema), _) => { + // It's still outstanding + match wait_result { + Ok(permit) => { + // We're responsible for resolving it + + if !Arc::ptr_eq(¤t_sema, permit.semaphore()) { + self.inner.error_counter.inc(); + tracing::error!( + "{} mismatched semaphores for {name:?}", + self.inner.name + ); + + // sema is the one we started with, and + // we own the permit for it. Both us and + // anyone else waiting for this is going + // to be let down by this situation. + permit.semaphore().close(); + return Err(Arc::new(anyhow::anyhow!( + "{} lookup for {name:?} \ + but have mismatched semaphores", + self.inner.name + ))); + } + + self.inner.populate_counter.inc(); + let mut ttl = Duration::from_secs(60); + let future_result = fut.await; + let now = Instant::now(); + + let (item_result, return_value) = match future_result { + Ok(item) => { + ttl = ttl_func(&item); + ( + ItemState::Present(item.clone()), + Ok(ItemLookup { + item, + expiration: now + ttl, + is_fresh: true, + }), + ) + } + Err(err) => { + self.inner.error_counter.inc(); + let err = Arc::new(err.into()); + (ItemState::Failed(err.clone()), Err(err)) + } + }; + + self.inner.cache.insert( + name.clone(), + Item { + item: item_result, + expiration: Instant::now() + ttl, + last_tick: self.inc_tick().into(), + }, + ); + // Wake everybody up + permit.semaphore().close(); + self.maybe_evict(); + + return return_value; + } + Err(_) => { + self.inner.error_counter.inc(); + + // semaphore was closed, but the status is + // still somehow pending + tracing::error!( + "{} lookup for {name:?} woke up semas \ + but is still marked pending", + self.inner.name + ); + return Err(Arc::new(anyhow::anyhow!( + "{} lookup for {name:?} \ + sema was closed but state is still pending", + self.inner.name + ))); + } + } + } + } + } + } + } +} + +#[cfg(test)] +mod test { + use super::*; + use test_log::test; // run with RUST_LOG=lruttl=trace to trace + + #[test(tokio::test)] + async fn test_capacity() { + let cache = LruCacheWithTtl::new_named("test_capacity", 40); + + let expiration = Instant::now() + Duration::from_secs(60); + for i in 0..100 { + cache.insert(i, i, expiration).await; + } + + assert_eq!(cache.inner.cache.len(), 40, "capacity is respected"); + } + + #[test(tokio::test)] + async fn test_expiration() { + let cache = LruCacheWithTtl::new_named("test_expiration", 1); + + tokio::time::pause(); + let expiration = Instant::now() + Duration::from_secs(1); + cache.insert(0, 0, expiration).await; + + cache.get(&0).await.expect("still in cache"); + tokio::time::advance(Duration::from_secs(2)).await; + assert!(cache.get(&0).await.is_none(), "evicted due to ttl"); } } diff --git a/crates/message/src/dkim.rs b/crates/message/src/dkim.rs index a0676fdb..d866fca0 100644 --- a/crates/message/src/dkim.rs +++ b/crates/message/src/dkim.rs @@ -8,8 +8,8 @@ use mlua::{Lua, Value}; use prometheus::{Counter, Histogram}; use serde::Deserialize; use std::sync::{Arc, LazyLock, OnceLock}; -use std::time::{Duration, Instant}; use tokio::runtime::Runtime; +use tokio::time::{Duration, Instant}; static SIGNER_CACHE: LazyLock>> = LazyLock::new(|| LruCacheWithTtl::new_named("dkim_signer_cache", 1024)); @@ -302,6 +302,7 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> { Ok(()) } +#[derive(Debug)] pub struct CFSigner { signer: kumo_dkim::Signer, } diff --git a/crates/mod-filesystem/src/lib.rs b/crates/mod-filesystem/src/lib.rs index 7492ae3f..e6f74131 100644 --- a/crates/mod-filesystem/src/lib.rs +++ b/crates/mod-filesystem/src/lib.rs @@ -3,7 +3,7 @@ use config::get_or_create_module; use lruttl::LruCacheWithTtl; use mlua::Lua; use std::sync::{Arc, LazyLock}; -use std::time::{Duration, Instant}; +use tokio::time::{Duration, Instant}; const GLOB_CACHE_CAPACITY: usize = 32; const DEFAULT_CACHE_TTL: f32 = 60.; diff --git a/crates/mod-memoize/src/lib.rs b/crates/mod-memoize/src/lib.rs index 0e5192c6..ee7e8100 100644 --- a/crates/mod-memoize/src/lib.rs +++ b/crates/mod-memoize/src/lib.rs @@ -98,6 +98,12 @@ pub enum CacheValue { Memoized(Memoized), } +impl std::fmt::Debug for CacheValue { + fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result { + fmt.debug_struct("CacheValue").finish() + } +} + impl FromLua for CacheValue { fn from_lua(value: mlua::Value, lua: &Lua) -> mlua::Result { match value { @@ -255,7 +261,7 @@ impl UserData for MemoizedTable { } } -#[derive(Clone)] +#[derive(Clone, Debug)] enum CacheEntry { Null, Single(CacheValue), diff --git a/crates/mta-sts/src/lib.rs b/crates/mta-sts/src/lib.rs index e82fb668..41b1414b 100644 --- a/crates/mta-sts/src/lib.rs +++ b/crates/mta-sts/src/lib.rs @@ -12,7 +12,7 @@ static CACHE: LazyLock> = pub mod dns; pub mod policy; -#[derive(Clone)] +#[derive(Clone, Debug)] struct CachedPolicy { pub id: String, pub policy: Arc, @@ -113,7 +113,7 @@ async fn get_policy_for_domain_impl( id: record.id, policy: Arc::clone(&policy), }, - expires, + expires.into(), ) .await; diff --git a/crates/rfc5321/src/tls.rs b/crates/rfc5321/src/tls.rs index 59fcd3ba..6ae7d67a 100644 --- a/crates/rfc5321/src/tls.rs +++ b/crates/rfc5321/src/tls.rs @@ -3,7 +3,7 @@ use hickory_proto::rr::rdata::TLSA; use lruttl::LruCacheWithTtl; use openssl::ssl::SslOptions; use std::sync::{Arc, LazyLock}; -use std::time::{Duration, Instant}; +use tokio::time::{Duration, Instant}; use tokio_rustls::rustls::client::danger::ServerCertVerifier; use tokio_rustls::rustls::crypto::{aws_lc_rs as provider, CryptoProvider}; use tokio_rustls::rustls::{ClientConfig, SupportedCipherSuite};