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