diff --git a/crates/dns-resolver/src/lib.rs b/crates/dns-resolver/src/lib.rs index eced7f44..11e45b13 100644 --- a/crates/dns-resolver/src/lib.rs +++ b/crates/dns-resolver/src/lib.rs @@ -516,8 +516,8 @@ async fn lookup_mx_record(domain_name: &Name) -> anyhow::Result<(Vec anyhow::Result<(Arc>, Instant)> { let key_fq = fully_qualify(key)?; - if let Some(value) = IP_CACHE.get_with_expiry(&key_fq).await { - return Ok(value); + if let Some(lookup) = IP_CACHE.lookup(&key_fq).await { + return Ok((lookup.item, lookup.expiration)); } let (v4, v6) = tokio::join!(ipv4_lookup(key), ipv6_lookup(key)); @@ -564,8 +564,8 @@ pub async fn ip_lookup(key: &str) -> anyhow::Result<(Arc>, Instant)> pub async fn ipv4_lookup(key: &str) -> anyhow::Result<(Arc>, Instant)> { let key_fq = fully_qualify(key)?; - if let Some(value) = IPV4_CACHE.get_with_expiry(&key_fq).await { - return Ok(value); + if let Some(lookup) = IPV4_CACHE.lookup(&key_fq).await { + return Ok((lookup.item, lookup.expiration)); } let answer = RESOLVER @@ -582,8 +582,8 @@ pub async fn ipv4_lookup(key: &str) -> anyhow::Result<(Arc>, Instant pub async fn ipv6_lookup(key: &str) -> anyhow::Result<(Arc>, Instant)> { let key_fq = fully_qualify(key)?; - if let Some(value) = IPV6_CACHE.get_with_expiry(&key_fq).await { - return Ok(value); + if let Some(lookup) = IPV6_CACHE.lookup(&key_fq).await { + return Ok((lookup.item, lookup.expiration)); } let answer = RESOLVER diff --git a/crates/kumod/src/egress_source.rs b/crates/kumod/src/egress_source.rs index 795920b7..8a503e7f 100644 --- a/crates/kumod/src/egress_source.rs +++ b/crates/kumod/src/egress_source.rs @@ -7,7 +7,7 @@ use data_loader::KeySource; use gcd::Gcd; use kumo_log_types::MaybeProxiedSourceAddress; use kumo_server_common::config_handle::ConfigHandle; -use lruttl::{ItemTtl, LruCacheWithTtl}; +use lruttl::LruCacheWithTtl; use mlua::prelude::LuaUserData; use parking_lot::FairMutex as Mutex; use prometheus::IntCounter; @@ -69,16 +69,10 @@ pub struct EgressSource { impl LuaUserData for EgressSource {} -impl ItemTtl for EgressSource { - fn get_ttl(&self) -> Duration { - self.ttl - } -} - impl EgressSource { pub async fn resolve(name: &str, config: &mut LuaConfig) -> anyhow::Result { SOURCES - .get_or_try_insert_embedded_ttl(name.to_string(), async { + .get_or_try_insert(&name.to_string(), |source| source.ttl, async { if name == "unspecified" { Ok(EgressSource { name: name.to_string(), @@ -103,6 +97,7 @@ impl EgressSource { }) .await .map_err(|err| anyhow::anyhow!("{err:#}")) + .map(|lookup| lookup.item) } fn resolve_proxy_protocol(&self, address: SocketAddr) -> anyhow::Result { @@ -303,18 +298,12 @@ pub struct EgressPool { impl LuaUserData for EgressPool {} -impl ItemTtl for EgressPool { - fn get_ttl(&self) -> Duration { - self.ttl - } -} - impl EgressPool { pub async fn resolve(name: Option<&str>, config: &mut LuaConfig) -> anyhow::Result { let name = name.unwrap_or("unspecified"); let pool = POOLS - .get_or_try_insert_embedded_ttl(name.to_string(), async { + .get_or_try_insert(&name.to_string(), |pool| pool.ttl, async { let pool = if name == "unspecified" { EgressPool { name: "unspecified".to_string(), @@ -336,7 +325,8 @@ impl EgressPool { Ok(pool) }) .await - .map_err(|err: Arc| anyhow::anyhow!("{err:#}"))?; + .map_err(|err: Arc| anyhow::anyhow!("{err:#}"))? + .item; // Validate each of the sources for entry in &pool.entries { diff --git a/crates/lruttl/src/lib.rs b/crates/lruttl/src/lib.rs index 403642f5..4d82c2ea 100644 --- a/crates/lruttl/src/lib.rs +++ b/crates/lruttl/src/lib.rs @@ -119,8 +119,14 @@ impl Expiry> for PerItemExpiry { } } -pub trait ItemTtl { - fn get_ttl(&self) -> Duration; +pub struct ItemLookup { + /// A copy of the item + pub item: V, + /// If true, the get_or_try_insert operation populated the entry; + /// the operation was a cache miss + pub is_fresh: bool, + /// The instant at which this entry will expire + pub expiration: Instant, } pub struct LruCacheWithTtl { @@ -178,13 +184,17 @@ impl< num_entries as usize } - pub async fn get_with_expiry(&self, name: &Q) -> Option<(V, Instant)> + pub async fn lookup(&self, name: &Q) -> Option> where K: Borrow, Q: Hash + Eq, { let entry = self.inner.cache.get(name).await?; - Some((entry.item.clone(), entry.expiration)) + Some(ItemLookup { + item: entry.item.clone(), + expiration: entry.expiration, + is_fresh: false, + }) } pub async fn get(&self, name: &Q) -> Option @@ -192,8 +202,7 @@ impl< K: Borrow, Q: Hash + Eq, { - let entry = self.inner.cache.get(name).await?; - entry.item.clone().into() + self.lookup(name).await.map(|lookup| lookup.item) } pub async fn insert(&self, name: K, item: V, expiration: Instant) -> V { @@ -213,63 +222,21 @@ impl< /// 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. - pub async fn get_or_insert(&self, name: K, ttl: Duration, fut: impl Future) -> V { - let item = self - .inner - .cache - .get_with(name, async move { - let item = fut.await; - Item { - item, - expiration: Instant::now() + ttl, - } - }) - .await; - item.item - } - - /// 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. - /// This variant allows the future to be fallible. - pub async fn get_or_try_insert( - &self, - name: K, - ttl: Duration, - fut: impl Future>, - ) -> Result> { - let item = self - .inner - .cache - .try_get_with(name, async move { - let item = fut.await?; - Ok(Item { - item, - expiration: Instant::now() + ttl, - }) - }) - .await?; - Ok(item.item) - } - - /// 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. - /// This variant allows the future to be fallible, as well as returns a freshness - /// flag that can be used to update cache metrics. - /// if the freshness flag is true, it was a cache miss. - pub async fn get_or_try_insert_with_freshness_status( + /// 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>( &self, name: &K, - ttl: Duration, + ttl: TTL, fut: impl Future>, - ) -> Result<(bool, V), Arc> { + ) -> 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, @@ -277,32 +244,12 @@ impl< }) .await?; let is_fresh = entry.is_fresh(); - Ok((is_fresh, entry.value().item.clone())) - } -} + let item = entry.value(); -impl< - K: Clone + Hash + Eq + Send + Sync + std::fmt::Debug + 'static, - V: Clone + Send + Sync + ItemTtl + 'static, - > LruCacheWithTtl -{ - pub async fn get_or_try_insert_embedded_ttl( - &self, - name: K, - fut: impl Future>, - ) -> Result> { - let item = self - .inner - .cache - .try_get_with(name, async move { - let item = fut.await?; - let ttl = item.get_ttl(); - Ok(Item { - item, - expiration: Instant::now() + ttl, - }) - }) - .await?; - Ok(item.item) + Ok(ItemLookup { + is_fresh, + item: item.item.clone(), + expiration: item.expiration, + }) } } diff --git a/crates/mod-memoize/src/lib.rs b/crates/mod-memoize/src/lib.rs index 73abffa2..0e5192c6 100644 --- a/crates/mod-memoize/src/lib.rs +++ b/crates/mod-memoize/src/lib.rs @@ -443,7 +443,7 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> { let key = (epoch_key, key); let value_result = cache - .get_or_try_insert_with_freshness_status(&key, ttl, async { + .get_or_try_insert(&key, |_| ttl, async { tracing::trace!("populate {key:?}"); populate_counter.inc(); let result: MultiValue = (func?).call_async(params).await?; @@ -452,13 +452,13 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> { .await; match value_result { - Ok((is_fresh, value)) => { - if is_fresh { + Ok(lookup) => { + if lookup.is_fresh { miss_counter.inc(); } else { hit_counter.inc(); } - value.to_value(&lua) + lookup.item.to_value(&lua) } Err(err) => { tracing::error!("{cache_name} {key:?} failed: {err:#}");