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.
This commit is contained in:
Wez Furlong
2025-03-04 16:23:16 -07:00
parent 68ff551aa9
commit 054508321b
15 changed files with 768 additions and 213 deletions
Generated
+28 -107
View File
@@ -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"
+1 -1
View File
@@ -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"
+1
View File
@@ -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<String>),
+1
View File
@@ -143,6 +143,7 @@ impl Default for SignerBuilder {
}
}
#[derive(Debug)]
pub struct Signer {
signed_headers: HeaderList,
private_key: Arc<DkimPrivateKey>,
+10 -8
View File
@@ -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<ByPreferenc
pub async fn ip_lookup(key: &str) -> anyhow::Result<(Arc<Vec<IpAddr>>, 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<Vec<IpAddr>>, 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<Vec<IpAddr>>, 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<Vec<IpAddr>>, 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<Vec<IpAddr>>, 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<Vec<IpAddr>>, 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))
}
@@ -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<LruCacheWithTtl<AuthKind, Result<bool, String>>> =
LazyLock::new(|| LruCacheWithTtl::new_named("http_server_auth", 128));
+1 -1
View File
@@ -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::Error>| anyhow::anyhow!("{err:#}"))?
+1 -1
View File
@@ -760,7 +760,7 @@ impl SmtpDispatcher {
.insert(
site_name.to_string(),
(),
std::time::Instant::now() + duration,
tokio::time::Instant::now() + duration,
)
.await;
}
+8 -1
View File
@@ -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
+703 -87
View File
@@ -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<IntCounterVec> = 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<IntCounterVec> = 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<IntCounterVec> = 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<IntCounterVec> = 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<IntCounterVec> = 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<IntCounterVec> = 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<IntCounterVec> = 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<IntCounterVec> = 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<IntGaugeVec> = 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<IntGaugeVec> = LazyLock::new(|| {
prometheus::register_int_gauge_vec!(
"lruttl_cache_size",
"number of items contained in an lruttl cache",
&["cache_name"]
)
.unwrap()
});
static CACHES: LazyLock<Mutex<Vec<Weak<dyn CachePurger + Send + Sync>>>> =
LazyLock::new(Mutex::default);
struct Inner<K: Clone + Hash + Eq, V: Clone + Send + Sync> {
struct Inner<K: Clone + Hash + Eq + Debug, V: Clone + Send + Sync + Debug> {
name: String,
cache: Cache<K, Item<V>>,
tick: AtomicUsize,
capacity: usize,
cache: DashMap<K, Item<V>>,
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<K: Clone + Send + Sync + Hash + Eq + 'static, V: Clone + Send + Sync + 'static> CachePurger
for Inner<K, V>
impl<
K: Clone + Debug + Send + Sync + Hash + Eq + 'static,
V: Clone + Debug + Send + Sync + 'static,
> CachePurger for Inner<K, V>
{
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<V>
where
V: Send,
V: Sync,
{
Present(V),
Pending(Arc<Semaphore>),
Failed(Arc<anyhow::Error>),
}
#[derive(Debug)]
struct Item<V>
where
V: Send,
V: Sync,
{
item: V,
item: ItemState<V>,
expiration: Instant,
last_tick: AtomicUsize,
}
struct PerItemExpiry<K, V> {
marker: PhantomData<(K, V)>,
}
impl<K: Send, V: Send + Sync> Expiry<K, Item<V>> for PerItemExpiry<K, V> {
fn expire_after_create(
&self,
_key: &K,
item: &Item<V>,
created_at: Instant,
) -> Option<Duration> {
Some(item.expiration - created_at)
impl<V: Clone + Send + Sync> Clone for Item<V> {
fn clone(&self) -> Self {
Self {
item: self.item.clone(),
expiration: self.expiration,
last_tick: self.last_tick.load(Ordering::Relaxed).into(),
}
}
}
pub struct ItemLookup<V> {
#[derive(Debug)]
pub struct ItemLookup<V: Debug> {
/// 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<V> {
pub expiration: Instant,
}
pub struct LruCacheWithTtl<K: Clone + Hash + Eq, V: Clone + Send + Sync> {
pub struct LruCacheWithTtl<K: Clone + Debug + Hash + Eq, V: Clone + Debug + Send + Sync> {
inner: Arc<Inner<K, V>>,
}
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<K, V>
{
#[deprecated = "use new_named instead"]
@@ -145,23 +293,56 @@ impl<
pub fn new_named<S: Into<String>>(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<V>) {
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 <https://redis.io/docs/latest/develop/reference/eviction/#apx-lru>
/// 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<Q: ?Sized>(&self, name: &Q) -> Option<ItemLookup<V>>
@@ -189,12 +567,48 @@ impl<
K: Borrow<Q>,
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<Q: ?Sized>(&self, name: &Q) -> Option<V>
@@ -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<V>, 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<E: Send + Sync + 'static, TTL: FnOnce(&V) -> Duration>(
pub async fn get_or_try_insert<E: Into<anyhow::Error>, TTL: FnOnce(&V) -> Duration>(
&self,
name: &K,
ttl: TTL,
ttl_func: TTL,
fut: impl Future<Output = Result<V, E>>,
) -> Result<ItemLookup<V>, Arc<E>> {
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<ItemLookup<V>, Arc<anyhow::Error>> {
// 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(&current_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");
}
}
+2 -1
View File
@@ -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<LruCacheWithTtl<SignerConfig, Arc<CFSigner>>> =
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,
}
+1 -1
View File
@@ -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.;
+7 -1
View File
@@ -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<Self> {
match value {
@@ -255,7 +261,7 @@ impl UserData for MemoizedTable {
}
}
#[derive(Clone)]
#[derive(Clone, Debug)]
enum CacheEntry {
Null,
Single(CacheValue),
+2 -2
View File
@@ -12,7 +12,7 @@ static CACHE: LazyLock<LruCacheWithTtl<Name, CachedPolicy>> =
pub mod dns;
pub mod policy;
#[derive(Clone)]
#[derive(Clone, Debug)]
struct CachedPolicy {
pub id: String,
pub policy: Arc<MtaStsPolicy>,
@@ -113,7 +113,7 @@ async fn get_policy_for_domain_impl(
id: record.id,
policy: Arc::clone(&policy),
},
expires,
expires.into(),
)
.await;
+1 -1
View File
@@ -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};