diff --git a/litellm/rust/Cargo.lock b/litellm/rust/Cargo.lock index ea790ecfb61..8317c52b0e9 100644 --- a/litellm/rust/Cargo.lock +++ b/litellm/rust/Cargo.lock @@ -3610,6 +3610,7 @@ version = "0.1.0" dependencies = [ "litellm-cache", "litellm-cache-testing", + "parking_lot", "rstest", "tokio", ] @@ -4552,6 +4553,7 @@ dependencies = [ "litellm-redis-testing", "litellm-secret-value", "litellm-tracing", + "parking_lot", "rstest", "sha2 0.10.9", "strum", diff --git a/litellm/rust/crates/cache-memory/Cargo.toml b/litellm/rust/crates/cache-memory/Cargo.toml index 88124f5401e..e8c58839292 100644 --- a/litellm/rust/crates/cache-memory/Cargo.toml +++ b/litellm/rust/crates/cache-memory/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true +parking_lot.workspace = true [dev-dependencies] litellm-cache-testing.workspace = true diff --git a/litellm/rust/crates/cache-memory/src/cache.rs b/litellm/rust/crates/cache-memory/src/cache.rs index 809ffe3044e..c386ef26322 100644 --- a/litellm/rust/crates/cache-memory/src/cache.rs +++ b/litellm/rust/crates/cache-memory/src/cache.rs @@ -1,18 +1,15 @@ use std::{ - cmp::Reverse, - collections::{BinaryHeap, HashMap, HashSet}, - hash::Hash, sync::{ - Arc, Mutex, MutexGuard, + Arc, atomic::{AtomicUsize, Ordering}, }, time::{Duration, SystemTime, UNIX_EPOCH}, }; -use litellm_cache::{ - BaseCache, BatchCache, ClaimCache, CounterCache, DeleteCache, DisconnectCache, Error, - ExactCacheContext, FlushCache, SetCache, TtlCache, -}; +use litellm_cache::Error; +use parking_lot::Mutex; + +use crate::state::{Deadline, State}; const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; const DEFAULT_TTL: Duration = Duration::from_secs(600); @@ -26,14 +23,16 @@ pub enum CacheWrite { TooLarge, } -struct CacheState { - values: HashMap, - expirations: HashMap, - expiration_heap: BinaryHeap>, +/// The TTL a write gives its key, `None` being the cache's default. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Expiry { + /// Only a key the cache does not hold takes the TTL, so a counter keeps its window. + KeepExisting(Option), + Refresh(Option), } pub struct InMemoryCache { - state: Mutex>, + state: Mutex>, max_size_in_memory: AtomicUsize, default_ttl: Duration, max_entry_bytes: Option, @@ -72,11 +71,7 @@ impl InMemoryCache { now: impl Fn() -> Duration + Send + Sync + 'static, ) -> Self { Self { - state: Mutex::new(CacheState { - values: HashMap::new(), - expirations: HashMap::new(), - expiration_heap: BinaryHeap::new(), - }), + state: Mutex::new(State::new()), max_size_in_memory: AtomicUsize::new( max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), ), @@ -89,39 +84,39 @@ impl InMemoryCache { } } - fn lock(&self) -> Result>, Error> { - Ok(self.state.lock()?) - } - pub fn set_cache( &self, key: impl Into, value: V, - ttl: Option, + expiry: Expiry, ) -> Result { - if self.max_size_in_memory() == 0 { - return Ok(CacheWrite::Disabled); - } let now = (self.now)(); - let mut state = self.lock()?; - self.store(&mut state, key.into(), value, ttl, now) + self.write(&mut self.state.lock(), key.into(), value, expiry, now) } pub fn get_cache(&self, key: &str) -> Result, Error> { let now = (self.now)(); - let mut state = self.lock()?; - if state - .expirations - .get(key) - .is_some_and(|expiration| *expiration < now) - { - Self::remove(&mut state, key); - } - Ok(state.values.get(key).cloned()) + Ok(self.state.lock().get(key, now).cloned()) } - /// `check_value_size`: whether `value` fits `max_entry_bytes`. Always `true` without a - /// limit and a measure, since typed values have no generic size. + /// Reads `key` and writes back under one lock. `update` answers the value to write, or `None` + /// to leave the entry as it is, next to what the caller gets back. + pub fn update( + &self, + key: &str, + expiry: Expiry, + update: impl FnOnce(Option<&V>) -> Result<(Option, T), Error>, + ) -> Result { + let now = (self.now)(); + let mut state = self.state.lock(); + let (value, output) = update(state.get(key, now))?; + if let Some(value) = value { + self.write(&mut state, key.to_owned(), value, expiry, now)?; + } + Ok(output) + } + + /// Always `true` without a limit and a measure, since typed values have no generic size. pub fn check_value_size(&self, value: &V) -> Result { match (self.max_entry_bytes, &self.measure_value) { (Some(limit), Some(measure)) => Ok(measure(value)? <= limit), @@ -129,57 +124,21 @@ impl InMemoryCache { } } - /// `evict_cache`: drops expired entries, then the earliest-expiring ones until a new key - /// fits. - pub fn evict_cache(&self) -> Result<(), Error> { - let now = (self.now)(); - let mut state = self.lock()?; - Self::evict(&mut state, self.max_size_in_memory(), now, None); - Ok(()) - } - - /// `evict_element_if_expired`: `true` when `key` had expired and was removed. - pub fn evict_element_if_expired(&self, key: &str) -> Result { - let now = (self.now)(); - let mut state = self.lock()?; - let expired = state - .expirations - .get(key) - .is_some_and(|expiration| *expiration < now); - if expired { - Self::remove(&mut state, key); - } - Ok(expired) - } - - /// `allow_ttl_override`: a write may set the TTL when the key has none or it has passed. - pub fn allow_ttl_override(&self, key: &str) -> Result { - let now = (self.now)(); - Ok(self - .expires_at(key)? - .is_none_or(|expiration| expiration < now)) - } - - /// The number of stored entries, expired ones included until they are evicted. pub fn len(&self) -> Result { - Ok(self.lock()?.values.len()) + let now = (self.now)(); + Ok(self.state.lock().len(now)) } pub fn is_empty(&self) -> Result { Ok(self.len()? == 0) } - /// Entries in the expiration heap, stale ones included; bounded by eviction. - pub fn expiration_heap_len(&self) -> Result { - Ok(self.lock()?.expiration_heap.len()) - } - pub fn max_size_in_memory(&self) -> usize { self.max_size_in_memory.load(Ordering::Relaxed) } - /// `DualCache.update_in_memory_max_size`: `None` restores the default. Entries already held stay - /// until the next write evicts down to the new size. + /// `None` restores the default. Entries already held stay until the next write trims to the + /// new size. pub fn update_max_size_in_memory(&self, max_size_in_memory: Option) { self.max_size_in_memory.store( max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), @@ -187,12 +146,13 @@ impl InMemoryCache { ); } - pub fn max_entry_bytes(&self) -> Option { - self.max_entry_bytes + pub fn default_ttl(&self) -> Duration { + self.default_ttl } pub fn expires_at(&self, key: &str) -> Result, Error> { - Ok(self.lock()?.expirations.get(key).copied()) + let now = (self.now)(); + Ok(self.state.lock().expires_at(key, now)) } pub async fn async_get_ttl(&self, key: &str) -> Result, Error> { @@ -200,253 +160,406 @@ impl InMemoryCache { } pub async fn async_get_oldest_n_keys(&self, count: usize) -> Result, Error> { - let state = self.lock()?; - let mut expirations = state - .expirations - .iter() - .map(|(key, expiration)| (key.clone(), *expiration)) - .collect::>(); - expirations.sort_unstable_by_key(|(_, expiration)| *expiration); - Ok(expirations - .into_iter() + let now = (self.now)(); + Ok(self + .state + .lock() + .keys_by_expiry(now) .take(count) - .map(|(key, _)| key) + .map(str::to_owned) .collect()) } - /// Every live entry, for callers that inspect the whole tier. pub fn entries(&self) -> Result, Error> { let now = (self.now)(); - let state = self.lock()?; - Ok(state - .values - .iter() - .filter(|(key, _)| { - state - .expirations - .get(*key) - .is_none_or(|expiration| *expiration >= now) - }) - .map(|(key, value)| (key.clone(), value.clone())) + Ok(self + .state + .lock() + .live(now) + .map(|(key, value)| (key.to_owned(), value.clone())) .collect()) } pub fn delete_cache(&self, key: &str) -> Result<(), Error> { - let mut state = self.lock()?; - Self::remove(&mut state, key); + self.state.lock().remove(key); Ok(()) } pub fn flush_cache(&self) -> Result<(), Error> { - let mut state = self.lock()?; - state.values.clear(); - state.expirations.clear(); - state.expiration_heap.clear(); + self.state.lock().clear(); Ok(()) } - /// Writing an existing `key` never evicts another entry, unlike Python, which pops the - /// earliest-expiring entry whenever the cache is full. - fn evict(state: &mut CacheState, capacity: usize, now: Duration, key: Option<&str>) { - while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() { - if state.expirations.get(&key).copied() != Some(expiration) { - state.expiration_heap.pop(); - } else if expiration <= now { - state.expiration_heap.pop(); - Self::remove(state, &key); - } else { - break; - } - } - if key.is_some_and(|key| state.values.contains_key(key)) { - return; - } - while state.values.len() >= capacity { - let Some(Reverse((expiration, key))) = state.expiration_heap.pop() else { - break; - }; - if state.expirations.get(&key).copied() == Some(expiration) { - Self::remove(state, &key); - } - } - } - - fn set_expiration(state: &mut CacheState, key: &str, expiration: Duration) { - if state.expirations.get(key).copied() != Some(expiration) { - state.expirations.insert(key.into(), expiration); - state - .expiration_heap - .push(Reverse((expiration, key.into()))); - } - } - - fn remove(state: &mut CacheState, key: &str) { - state.values.remove(key); - state.expirations.remove(key); - } - - /// `get_cache` under the held lock: an expired entry is removed and reads as missing. - fn live(state: &mut CacheState, key: &str, now: Duration) -> Option { - if state - .expirations - .get(key) - .is_some_and(|expiration| *expiration < now) - { - Self::remove(state, key); - } - state.values.get(key).cloned() - } - - /// Python `set_cache` under the held lock: evict first (even when `key` already exists), - /// then skip oversized values, then write, keeping a live key's expiry. - fn store( + fn write( &self, - state: &mut CacheState, + state: &mut State, key: String, value: V, - ttl: Option, + expiry: Expiry, now: Duration, ) -> Result { - Self::evict(state, self.max_size_in_memory(), now, None); + let capacity = self.max_size_in_memory(); + if capacity == 0 { + return Ok(CacheWrite::Disabled); + } if !self.check_value_size(&value)? { return Ok(CacheWrite::TooLarge); } - let expiration = state.expirations.get(&key).copied(); - if expiration.is_none_or(|expiration| expiration < now) { - Self::set_expiration(state, &key, now + ttl.unwrap_or(self.default_ttl)); - } - state.values.insert(key, value); + let expires_at = + |ttl: Option| now.saturating_add(ttl.unwrap_or(self.default_ttl)); + let deadline = match expiry { + Expiry::KeepExisting(ttl) => Deadline::UnlessLive(expires_at(ttl)), + Expiry::Refresh(ttl) => Deadline::Always(expires_at(ttl)), + }; + state.write(key, value, deadline, capacity, now); Ok(CacheWrite::Stored) } } -impl ClaimCache for InMemoryCache -where - V: Clone + PartialEq + Send + Sync + 'static, -{ - fn claim_cache( - &self, - key: &str, - candidate: V, - eligible: &[V], - context: ExactCacheContext, - ) -> Result { - if self.max_size_in_memory() == 0 { - return Ok(candidate); - } - let now = (self.now)(); - let mut state = self.lock()?; - Self::evict(&mut state, self.max_size_in_memory(), now, Some(key)); - let existing = state - .values - .get(key) - .filter(|existing| eligible.is_empty() || eligible.contains(existing)) - .cloned(); - if let Some(existing) = &existing - && eligible.is_empty() - && *existing != candidate - { - return Ok(existing.clone()); - } - let winner = existing.unwrap_or(candidate); - Self::set_expiration( - &mut state, - key, - now + self.get_ttl(&context).unwrap_or(self.default_ttl), +#[cfg(test)] +mod tests { + use std::{ + panic::{AssertUnwindSafe, catch_unwind}, + sync::Arc, + time::Duration, + }; + + use litellm_cache::Error; + use rstest::rstest; + + use super::{CacheWrite, Expiry, InMemoryCache}; + use crate::test_support::{Clock, at, cache, clock, secs}; + + const DEFAULT: Expiry = Expiry::KeepExisting(None); + + fn keep(seconds: u64) -> Expiry { + Expiry::KeepExisting(secs(seconds)) + } + + fn measured(capacity: usize) -> InMemoryCache { + InMemoryCache::with_clock_and_size_measurement( + Some(capacity), + secs(60), + Some(4), + Some(Arc::new(|value: &String| { + if value.is_empty() { + return Err(Error::InvalidEntry(None)); + } + Ok(value.len()) + })), + || Duration::from_secs(100), + ) + } + + fn count(current: Option<&u64>) -> Result<(Option, u64), Error> { + let next = current.copied().unwrap_or_default() + 1; + Ok((Some(next), next)) + } + + #[rstest] + #[case::unset(None, secs(600))] + #[case::zero(Some(Duration::ZERO), secs(600))] + #[case::explicit(secs(5), secs(5))] + fn a_write_without_a_ttl_takes_the_default( + #[case] default_ttl: Option, + #[case] expected: Option, + ) { + let cache = InMemoryCache::::with_clock(None, default_ttl, || Duration::ZERO); + + cache.set_cache("key", "value".into(), DEFAULT).unwrap(); + + assert_eq!(cache.expires_at("key"), Ok(expected)); + assert_eq!(cache.max_size_in_memory(), 200); + } + + #[rstest] + #[case::before(109, Some("value"))] + #[case::at_the_deadline(110, None)] + #[case::after(111, None)] + fn an_entry_is_readable_until_its_deadline( + clock: Clock, + #[case] now: u64, + #[case] expected: Option<&str>, + ) { + let cache = cache::(&clock, 4); + cache.set_cache("key", "value".into(), keep(10)).unwrap(); + + at(&clock, now); + + assert_eq!(cache.get_cache("key"), Ok(expected.map(str::to_owned))); + } + + #[rstest] + #[case::keep_existing(keep(10), secs(160))] + #[case::refresh(Expiry::Refresh(secs(10)), secs(115))] + fn a_write_to_a_live_key_keeps_its_deadline_unless_it_refreshes( + clock: Clock, + #[case] expiry: Expiry, + #[case] expected: Option, + ) { + let cache = cache::(&clock, 4); + cache.set_cache("key", "first".into(), DEFAULT).unwrap(); + at(&clock, 105); + + cache.set_cache("key", "second".into(), expiry).unwrap(); + + assert_eq!(cache.get_cache("key"), Ok(Some("second".into()))); + assert_eq!(cache.expires_at("key"), Ok(expected)); + } + + #[rstest] + fn an_expired_key_takes_the_ttl_of_the_next_write(clock: Clock) { + let cache = cache::(&clock, 4); + cache.set_cache("key", "first".into(), keep(10)).unwrap(); + at(&clock, 110); + + cache.set_cache("key", "second".into(), keep(10)).unwrap(); + + assert_eq!(cache.expires_at("key"), Ok(secs(120))); + } + + #[rstest] + #[case::keep_existing(Expiry::KeepExisting(Some(Duration::MAX)))] + #[case::refresh(Expiry::Refresh(Some(Duration::MAX)))] + fn a_ttl_past_the_end_of_time_never_expires(clock: Clock, #[case] forever: Expiry) { + let cache = cache::(&clock, 3); + + cache.set_cache("forever", "a".into(), forever).unwrap(); + at(&clock, 1_000_000_000); + cache.set_cache("after", "b".into(), keep(1)).unwrap(); + + assert_eq!(cache.get_cache("forever"), Ok(Some("a".into()))); + assert_eq!(cache.get_cache("after"), Ok(Some("b".into()))); + } + + #[rstest] + fn a_panic_inside_an_update_leaves_the_cache_usable(clock: Clock) { + let cache = cache::(&clock, 4); + cache.set_cache("held", "a".into(), DEFAULT).unwrap(); + + let panicked = catch_unwind(AssertUnwindSafe(|| { + cache.update::<()>("held", DEFAULT, |_| panic!("the caller's update panics")) + })); + + assert!(panicked.is_err()); + assert_eq!( + cache.set_cache("next", "b".into(), DEFAULT), + Ok(CacheWrite::Stored) ); - state.values.insert(key.into(), winner.clone()); - Ok(winner) - } -} - -impl CounterCache for InMemoryCache { - fn increment_cache( - &self, - key: &str, - amount: f64, - context: ExactCacheContext, - ) -> Result { - if self.max_size_in_memory() == 0 { - return Ok(amount); - } - let now = (self.now)(); - let mut state = self.lock()?; - let value = Self::live(&mut state, key, now).unwrap_or_default() + amount; - self.store(&mut state, key.into(), value, self.get_ttl(&context), now)?; - Ok(value) - } -} - -impl BaseCache for InMemoryCache { - type Value = V; - type Context = ExactCacheContext; - - fn get_ttl(&self, context: &Self::Context) -> Option { - context.ttl.or(Some(self.default_ttl)) - } - - fn set_cache( - &self, - key: &str, - value: Self::Value, - context: &ExactCacheContext, - ) -> Result<(), Error> { - let ttl = self.get_ttl(context).unwrap_or(self.default_ttl); - self.set_cache(key, value, Some(ttl)).map(|_| ()) - } - - fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result, Error> { - self.get_cache(key) - } -} - -impl DisconnectCache for InMemoryCache { - async fn disconnect(&self) -> Result<(), Error> { - Ok(()) - } -} - -impl BatchCache for InMemoryCache {} - -impl DeleteCache for InMemoryCache { - fn delete_cache(&self, key: &str) -> Result<(), Error> { - InMemoryCache::delete_cache(self, key) - } -} - -impl FlushCache for InMemoryCache { - fn flush_cache(&self) -> Result<(), Error> { - InMemoryCache::flush_cache(self) - } -} - -impl TtlCache for InMemoryCache { - async fn async_get_ttl(&self, key: &str) -> Result, Error> { - InMemoryCache::async_get_ttl(self, key).await - } -} - -impl SetCache for InMemoryCache> -where - T: Clone + Eq + Hash + Send + Sync + 'static, -{ - type SetValue = T; - type SetResult = Vec; - - async fn async_set_cache_sadd( - &self, - key: &str, - values: Vec, - ttl: Option, - ) -> Result { - if self.max_size_in_memory() == 0 { - return Ok(values); - } - let now = (self.now)(); - let mut state = self.lock()?; - let mut stored = Self::live(&mut state, key, now).unwrap_or_default(); - stored.extend(values.iter().cloned()); - self.store(&mut state, key.into(), stored, ttl, now)?; - Ok(values) + assert_eq!(cache.get_cache("held"), Ok(Some("a".into()))); + } + + #[rstest] + fn a_full_cache_admits_a_new_key_by_evicting_the_soonest_to_expire(clock: Clock) { + let cache = cache::(&clock, 2); + cache.set_cache("early", "a".into(), keep(10)).unwrap(); + cache.set_cache("late", "b".into(), keep(20)).unwrap(); + + cache.set_cache("new", "c".into(), keep(5)).unwrap(); + + assert_eq!(cache.get_cache("early"), Ok(None)); + assert_eq!(cache.get_cache("late"), Ok(Some("b".into()))); + assert_eq!(cache.get_cache("new"), Ok(Some("c".into()))); + } + + #[rstest] + #[tokio::test] + async fn a_key_deleted_and_set_again_is_tracked_once(clock: Clock) { + let cache = cache::(&clock, 4); + cache.set_cache("older", "pin".into(), keep(3600)).unwrap(); + at(&clock, 101); + cache.set_cache("pinned", "pin".into(), keep(3600)).unwrap(); + at(&clock, 102); + cache.delete_cache("pinned").unwrap(); + cache.set_cache("pinned", "pin".into(), keep(3600)).unwrap(); + at(&clock, 103); + cache + .set_cache("pinned", "pin".into(), Expiry::Refresh(secs(3600))) + .unwrap(); + + assert_eq!( + cache.async_get_oldest_n_keys(10).await.unwrap(), + ["older", "pinned"] + ); + assert_eq!(cache.expires_at("pinned"), Ok(secs(3703))); + } + + #[rstest] + fn a_smaller_capacity_applies_on_the_next_write(clock: Clock) { + let cache = cache::(&clock, 3); + cache.set_cache("a", "1".into(), keep(10)).unwrap(); + cache.set_cache("b", "2".into(), keep(20)).unwrap(); + cache.set_cache("c", "3".into(), keep(30)).unwrap(); + + cache.update_max_size_in_memory(Some(1)); + let held = cache.len(); + cache.set_cache("d", "4".into(), keep(40)).unwrap(); + + assert_eq!(held, Ok(3)); + assert_eq!(cache.entries(), Ok(vec![("d".to_owned(), "4".to_owned())])); + } + + #[rstest] + fn an_unset_capacity_restores_the_default(clock: Clock) { + let cache = cache::(&clock, 3); + + cache.update_max_size_in_memory(None); + + assert_eq!(cache.max_size_in_memory(), 200); + } + + #[rstest] + fn a_cache_without_capacity_stores_nothing(clock: Clock) { + let cache = cache::(&clock, 0); + + let written = cache.set_cache("key", 1, DEFAULT); + let counted = cache.update("key", DEFAULT, count); + + assert_eq!(written, Ok(CacheWrite::Disabled)); + assert_eq!(counted, Ok(1)); + assert_eq!(cache.is_empty(), Ok(true)); + } + + #[rstest] + #[case::fits("ok", Ok(true))] + #[case::at_limit("four", Ok(true))] + #[case::too_large("oversized", Ok(false))] + #[case::measure_error("", Err(Error::InvalidEntry(None)))] + fn check_value_size_applies_the_entry_limit( + #[case] value: &str, + #[case] expected: Result, + ) { + assert_eq!(measured(2).check_value_size(&value.to_owned()), expected); + } + + #[rstest] + fn values_are_unbounded_without_a_measure() { + let cache = InMemoryCache::::default(); + + assert_eq!(cache.check_value_size(&"x".repeat(1 << 20)), Ok(true)); + } + + #[rstest] + #[case::too_large("oversized", Ok(CacheWrite::TooLarge))] + #[case::rejected_by_the_measure("", Err(Error::InvalidEntry(None)))] + fn a_value_that_does_not_fit_is_not_stored_and_keeps_the_held_one( + #[case] value: &str, + #[case] expected: Result, + ) { + let cache = measured(2); + cache.set_cache("key", "ok".into(), DEFAULT).unwrap(); + + let written = cache.set_cache("key", value.to_owned(), DEFAULT); + + assert_eq!(written, expected); + assert_eq!(cache.get_cache("key"), Ok(Some("ok".into()))); + } + + #[rstest] + fn an_update_writes_what_it_answers_and_keeps_the_window(clock: Clock) { + let cache = cache::(&clock, 4); + + let first = cache.update("count", keep(10), count); + at(&clock, 105); + let second = cache.update("count", keep(10), count); + + assert_eq!((first, second), (Ok(1), Ok(2))); + assert_eq!(cache.get_cache("count"), Ok(Some(2))); + assert_eq!(cache.expires_at("count"), Ok(secs(110))); + } + + #[rstest] + fn an_update_that_answers_no_value_leaves_the_entry_as_it_is(clock: Clock) { + let cache = cache::(&clock, 4); + cache.set_cache("count", 5, keep(10)).unwrap(); + at(&clock, 105); + + let read = cache.update("count", Expiry::Refresh(secs(50)), |current| { + Ok((None, current.copied())) + }); + + assert_eq!(read, Ok(Some(5))); + assert_eq!(cache.get_cache("count"), Ok(Some(5))); + assert_eq!(cache.expires_at("count"), Ok(secs(110))); + } + + #[rstest] + fn an_update_that_fails_writes_nothing(clock: Clock) { + let cache = cache::(&clock, 4); + + let failed = cache.update::<()>("count", DEFAULT, |_| Err(Error::InvalidEntry(None))); + + assert_eq!(failed, Err(Error::InvalidEntry(None))); + assert_eq!(cache.get_cache("count"), Ok(None)); + } + + #[rstest] + fn an_update_fails_when_its_value_cannot_be_measured() { + let updated = measured(2).update("key", DEFAULT, |_| Ok((Some(String::new()), ()))); + + assert_eq!(updated, Err(Error::InvalidEntry(None))); + } + + #[rstest] + fn only_live_entries_are_counted_and_listed(clock: Clock) { + let cache = cache::(&clock, 4); + cache.set_cache("short", "a".into(), keep(1)).unwrap(); + cache.set_cache("long", "b".into(), keep(50)).unwrap(); + + at(&clock, 101); + + assert_eq!(cache.len(), Ok(1)); + assert_eq!(cache.is_empty(), Ok(false)); + assert_eq!( + cache.entries(), + Ok(vec![("long".to_owned(), "b".to_owned())]) + ); + assert_eq!(cache.expires_at("short"), Ok(None)); + } + + #[rstest] + #[case::none(0, &[])] + #[case::one(1, &["first"])] + #[case::more_than_held(10, &["first", "later", "latest"])] + #[tokio::test] + async fn the_oldest_keys_are_the_soonest_to_expire( + clock: Clock, + #[case] count: usize, + #[case] expected: &[&str], + ) { + let cache = cache::(&clock, 3); + cache.set_cache("later", "2".into(), keep(20)).unwrap(); + cache.set_cache("first", "1".into(), keep(10)).unwrap(); + cache.set_cache("latest", "3".into(), keep(30)).unwrap(); + + assert_eq!( + cache.async_get_oldest_n_keys(count).await.unwrap(), + expected + ); + } + + #[rstest] + #[tokio::test] + async fn the_ttl_of_a_key_is_its_deadline(clock: Clock) { + let cache = cache::(&clock, 3); + cache.set_cache("key", "1".into(), keep(10)).unwrap(); + + assert_eq!(cache.async_get_ttl("key").await, Ok(secs(110))); + assert_eq!(cache.async_get_ttl("missing").await, Ok(None)); + } + + #[rstest] + fn delete_drops_one_key_and_flush_drops_the_rest(clock: Clock) { + let cache = cache::(&clock, 4); + cache.set_cache("a", "1".into(), DEFAULT).unwrap(); + cache.set_cache("b", "2".into(), DEFAULT).unwrap(); + cache.set_cache("c", "3".into(), DEFAULT).unwrap(); + + cache.delete_cache("a").unwrap(); + let after_delete = (cache.get_cache("a"), cache.len()); + cache.flush_cache().unwrap(); + + assert_eq!(after_delete, (Ok(None), Ok(2))); + assert_eq!(cache.is_empty(), Ok(true)); } } diff --git a/litellm/rust/crates/cache-memory/src/capabilities.rs b/litellm/rust/crates/cache-memory/src/capabilities.rs new file mode 100644 index 00000000000..6179732ce19 --- /dev/null +++ b/litellm/rust/crates/cache-memory/src/capabilities.rs @@ -0,0 +1,528 @@ +use std::{collections::HashSet, hash::Hash, time::Duration}; + +use litellm_cache::{ + BaseCache, BatchCache, ClaimCache, CounterCache, DeleteCache, DisconnectCache, Error, + ExactCacheContext, FlushCache, SetCache, TtlCache, +}; + +use crate::cache::{Expiry, InMemoryCache}; + +impl BaseCache for InMemoryCache { + type Value = V; + type Context = ExactCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl.or(Some(self.default_ttl())) + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &ExactCacheContext, + ) -> Result<(), Error> { + InMemoryCache::set_cache(self, key, value, Expiry::KeepExisting(context.ttl)).map(drop) + } + + fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result, Error> { + InMemoryCache::get_cache(self, key) + } +} + +impl ClaimCache for InMemoryCache +where + V: Clone + PartialEq + Send + Sync + 'static, +{ + /// A held value outside `eligible` gives way to `candidate`. With no `eligible` list the held + /// value always wins, and only a claim that agrees with it refreshes the TTL. + fn claim_cache( + &self, + key: &str, + candidate: V, + eligible: &[V], + context: ExactCacheContext, + ) -> Result { + self.update(key, Expiry::Refresh(context.ttl), |held| { + let held = held.filter(|held| eligible.is_empty() || eligible.contains(held)); + Ok(match held { + Some(held) if eligible.is_empty() && *held != candidate => (None, held.clone()), + Some(held) => (Some(held.clone()), held.clone()), + None => (Some(candidate.clone()), candidate), + }) + }) + } +} + +impl CounterCache for InMemoryCache { + fn increment_cache( + &self, + key: &str, + amount: f64, + context: ExactCacheContext, + ) -> Result { + self.update(key, Expiry::KeepExisting(context.ttl), |held| { + let total = held.copied().unwrap_or_default() + amount; + Ok((Some(total), total)) + }) + } +} + +impl DisconnectCache for InMemoryCache { + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } +} + +impl BatchCache for InMemoryCache {} + +impl DeleteCache for InMemoryCache { + fn delete_cache(&self, key: &str) -> Result<(), Error> { + InMemoryCache::delete_cache(self, key) + } +} + +impl FlushCache for InMemoryCache { + fn flush_cache(&self) -> Result<(), Error> { + InMemoryCache::flush_cache(self) + } +} + +impl TtlCache for InMemoryCache { + async fn async_get_ttl(&self, key: &str) -> Result, Error> { + InMemoryCache::async_get_ttl(self, key).await + } +} + +impl SetCache for InMemoryCache> +where + T: Clone + Eq + Hash + Send + Sync + 'static, +{ + type SetValue = T; + type SetResult = Vec; + + async fn async_set_cache_sadd( + &self, + key: &str, + values: Vec, + ttl: Option, + ) -> Result { + self.update(key, Expiry::KeepExisting(ttl), |held| { + let members = held.into_iter().flatten().chain(&values).cloned().collect(); + Ok((Some(members), ())) + })?; + Ok(values) + } +} + +#[cfg(test)] +mod tests { + use std::{collections::HashSet, sync::Arc, time::Duration}; + + use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CacheBackend, ClaimCache, CounterCache, DeleteCache, + DisconnectCache, ExactCacheContext, FlushCache, IncrementOperation, SetCache, TtlCache, + get_cache, set_cache, + }; + use rstest::rstest; + + use crate::{ + Expiry, InMemoryCache, + test_support::{Clock, at, cache, clock, secs}, + }; + + const DEFAULT: Expiry = Expiry::KeepExisting(None); + + fn keep(seconds: u64) -> Expiry { + Expiry::KeepExisting(secs(seconds)) + } + + fn ttl(seconds: u64) -> ExactCacheContext { + ExactCacheContext { ttl: secs(seconds) } + } + + #[rstest] + #[tokio::test] + async fn generic_consumers_share_typed_values_and_honor_expiration(clock: Clock) { + let cache: CacheBackend> = Arc::new(cache(&clock, 4)); + let reader = Arc::clone(&cache); + let context = ttl(5); + set_cache(cache.as_ref(), "sync", "first".into(), &context).unwrap(); + cache + .batch_cache_write("async", "second".into(), context.clone()) + .await + .unwrap(); + cache + .async_set_cache_pipeline(vec![("batch".into(), "third".into())], context.clone()) + .await + .unwrap(); + drop(cache); + + let written = ( + get_cache(reader.as_ref(), "sync", &context), + reader.async_get_cache("async", &context).await, + reader.async_get_cache("batch", &context).await, + ); + reader.async_delete_cache("async").await.unwrap(); + let deleted = reader.async_get_cache("async", &context).await; + at(&clock, 105); + let expired = ( + get_cache(reader.as_ref(), "sync", &context), + reader.async_get_cache("batch", &context).await, + ); + + assert_eq!( + written, + ( + Ok(Some("first".into())), + Ok(Some("second".into())), + Ok(Some("third".into())) + ) + ); + assert_eq!(deleted, Ok(None)); + assert_eq!(expired, (Ok(None), Ok(None))); + } + + #[rstest] + #[case::context_ttl(ttl(5), secs(105))] + #[case::default_ttl(ExactCacheContext::default(), secs(160))] + #[tokio::test] + async fn pipeline_writes_use_the_context_ttl_or_the_default( + clock: Clock, + #[case] context: ExactCacheContext, + #[case] expected: Option, + ) { + let cache = cache::(&clock, 4); + + cache + .async_set_cache_pipeline( + vec![("a".into(), "1".into()), ("b".into(), "2".into())], + context, + ) + .await + .unwrap(); + + assert_eq!(cache.expires_at("a"), Ok(expected)); + assert_eq!(cache.expires_at("b"), Ok(expected)); + } + + #[rstest] + #[tokio::test] + async fn batch_reads_return_one_entry_per_key_and_miss_expired_ones(clock: Clock) { + let cache = cache::(&clock, 4); + cache.set_cache("short", "a".into(), keep(1)).unwrap(); + cache.set_cache("long", "b".into(), keep(50)).unwrap(); + let keys = vec!["short".to_owned(), "missing".into(), "long".into()]; + + let live = cache.batch_get_cache(&keys, &ExactCacheContext::default()); + at(&clock, 101); + let later = cache + .async_batch_get_cache(keys, ExactCacheContext::default()) + .await; + + assert_eq!( + live, + Ok(vec![ + BatchEntry::Hit("a".to_owned()), + BatchEntry::Miss, + BatchEntry::Hit("b".into()), + ]) + ); + assert_eq!( + later, + Ok(vec![ + BatchEntry::Miss, + BatchEntry::Miss, + BatchEntry::Hit("b".into()), + ]) + ); + } + + #[rstest] + #[tokio::test] + async fn disconnect_keeps_entries_and_an_async_flush_drops_them() { + let cache = InMemoryCache::::default(); + cache.set_cache("key", "value".into(), DEFAULT).unwrap(); + + cache.disconnect().await.unwrap(); + let after_disconnect = cache.get_cache("key"); + FlushCache::async_flush_cache(&cache).await.unwrap(); + + assert_eq!(after_disconnect, Ok(Some("value".into()))); + assert_eq!(cache.is_empty(), Ok(true)); + } + + #[rstest] + #[tokio::test] + async fn the_ttl_capability_reads_the_deadline(clock: Clock) { + let cache = cache::(&clock, 3); + cache.set_cache("key", "1".into(), keep(20)).unwrap(); + + assert_eq!(TtlCache::async_get_ttl(&cache, "key").await, Ok(secs(120))); + } + + #[rstest] + fn a_claim_without_an_eligible_list_keeps_the_held_value_and_its_ttl(clock: Clock) { + let cache = cache::(&clock, 4); + let first = cache.claim_cache("affinity", "first".to_owned(), &[], ttl(10)); + at(&clock, 103); + + let second = cache.claim_cache("affinity", "second".to_owned(), &[], ttl(10)); + + assert_eq!((first, second), (Ok("first".into()), Ok("first".into()))); + assert_eq!(cache.expires_at("affinity"), Ok(secs(110))); + } + + #[rstest] + #[case::the_held_value_is_eligible(&["first", "second"], "first")] + #[case::the_held_value_is_not_eligible(&["second"], "second")] + fn a_claim_with_an_eligible_list_refreshes_the_ttl_of_the_winner( + clock: Clock, + #[case] eligible: &[&str], + #[case] winner: &str, + ) { + let cache = cache::(&clock, 4); + let eligible = eligible + .iter() + .copied() + .map(String::from) + .collect::>(); + cache + .set_cache("affinity", "first".into(), keep(10)) + .unwrap(); + at(&clock, 105); + + let claimed = cache.claim_cache("affinity", "second".to_owned(), &eligible, ttl(10)); + + assert_eq!(claimed, Ok(winner.to_owned())); + assert_eq!(cache.get_cache("affinity"), Ok(Some(winner.to_owned()))); + assert_eq!(cache.expires_at("affinity"), Ok(secs(115))); + } + + #[rstest] + fn claiming_a_held_key_at_capacity_keeps_other_entries(clock: Clock) { + let cache = cache::(&clock, 2); + cache.set_cache("hot", "1".into(), keep(10)).unwrap(); + cache.set_cache("cold", "2".into(), keep(20)).unwrap(); + + cache + .claim_cache("cold", "2".into(), &[], ExactCacheContext::default()) + .unwrap(); + + assert_eq!(cache.get_cache("hot"), Ok(Some("1".into()))); + assert_eq!(cache.expires_at("cold"), Ok(secs(160))); + } + + #[rstest] + fn counters_accumulate() { + let cache = InMemoryCache::::default(); + + let totals = ( + cache.increment_cache("counter", 1.5, ExactCacheContext::default()), + cache.increment_cache("counter", 2.0, ExactCacheContext::default()), + ); + + assert_eq!(totals, (Ok(1.5), Ok(3.5))); + } + + #[rstest] + fn concurrent_increments_lose_no_update() { + let cache = InMemoryCache::::default(); + + std::thread::scope(|scope| { + for _ in 0..8 { + scope.spawn(|| { + for _ in 0..10_000 { + cache + .increment_cache("counter", 1.0, ExactCacheContext::default()) + .unwrap(); + } + }); + } + }); + + assert_eq!(cache.get_cache("counter"), Ok(Some(80_000.0))); + } + + #[rstest] + fn a_trait_write_to_a_live_key_keeps_its_deadline(clock: Clock) { + let cache = cache::(&clock, 4); + set_cache(&cache, "key", "first".into(), &ttl(10)).unwrap(); + at(&clock, 105); + + set_cache(&cache, "key", "second".into(), &ttl(10)).unwrap(); + + assert_eq!(cache.get_cache("key"), Ok(Some("second".into()))); + assert_eq!(cache.expires_at("key"), Ok(secs(110))); + } + + #[rstest] + #[case::its_own(ttl(5), secs(5))] + #[case::the_default_of_the_cache(ExactCacheContext::default(), secs(60))] + fn a_context_resolves_to_its_ttl_or( + clock: Clock, + #[case] context: ExactCacheContext, + #[case] expected: Option, + ) { + assert_eq!(cache::(&clock, 4).get_ttl(&context), expected); + } + + #[rstest] + #[case::window_semantics(false)] + #[case::refresh_ttl_is_ignored(true)] + #[tokio::test] + async fn an_increment_keeps_the_window_of_a_live_counter( + clock: Clock, + #[case] refresh_ttl: bool, + ) { + let cache = cache::(&clock, 4); + + let first = cache + .async_increment("counter", 2.0, ttl(10), refresh_ttl) + .await; + at(&clock, 105); + let second = cache + .async_increment("counter", 3.0, ttl(10), refresh_ttl) + .await; + + assert_eq!((first, second), (Ok(2.0), Ok(5.0))); + assert_eq!(cache.expires_at("counter"), Ok(secs(110))); + } + + #[rstest] + fn an_expired_counter_restarts_from_zero_with_a_new_window(clock: Clock) { + let cache = cache::(&clock, 4); + cache.increment_cache("counter", 2.0, ttl(10)).unwrap(); + at(&clock, 110); + + let restarted = cache.increment_cache("counter", 1.0, ttl(10)); + + assert_eq!(restarted, Ok(1.0)); + assert_eq!(cache.expires_at("counter"), Ok(secs(120))); + } + + #[rstest] + fn incrementing_a_held_counter_at_capacity_evicts_nothing(clock: Clock) { + let cache = cache::(&clock, 2); + let increment = |key| cache.increment_cache(key, 1.0, ExactCacheContext::default()); + increment("a").unwrap(); + increment("b").unwrap(); + + let totals = (increment("a"), increment("b")); + + assert_eq!(totals, (Ok(2.0), Ok(2.0))); + } + + #[rstest] + #[tokio::test] + async fn a_cache_without_capacity_answers_claims_counters_and_sets_but_keeps_none() { + let claims = InMemoryCache::::new(Some(0), None); + let counters = InMemoryCache::::new(Some(0), None); + let sets = InMemoryCache::>::new(Some(0), None); + + let claimed = claims.claim_cache("key", "first".into(), &[], ExactCacheContext::default()); + let counted = counters.increment_cache("key", 2.0, ExactCacheContext::default()); + let added = sets + .async_set_cache_sadd("key", vec!["a".into()], None) + .await; + + assert_eq!(claimed, Ok("first".into())); + assert_eq!(counted, Ok(2.0)); + assert_eq!(added, Ok(vec!["a".to_owned()])); + assert_eq!( + (claims.is_empty(), counters.is_empty(), sets.is_empty()), + (Ok(true), Ok(true), Ok(true)) + ); + } + + #[rstest] + #[tokio::test] + async fn an_increment_pipeline_answers_in_operation_order(clock: Clock) { + let cache = cache::(&clock, 3); + let operation = |key: &str, amount, ttl| IncrementOperation { + key: key.into(), + amount, + ttl: secs(ttl), + }; + + let totals = cache + .async_increment_pipeline(vec![ + operation("a", 1.0, 10), + operation("b", 5.0, 30), + operation("a", 2.0, 20), + ]) + .await; + + assert_eq!(totals, Ok(vec![1.0, 5.0, 3.0])); + assert_eq!(cache.expires_at("a"), Ok(secs(110))); + assert_eq!(cache.expires_at("b"), Ok(secs(130))); + } + + #[rstest] + #[tokio::test] + async fn set_additions_answer_what_was_added_and_store_the_union(clock: Clock) { + let cache = cache::>(&clock, 4); + let added = vec!["a".to_owned(), "a".into(), "b".into()]; + + let first = cache + .async_set_cache_sadd("members", added.clone(), secs(10)) + .await; + let second = cache + .async_set_cache_sadd("members", vec!["c".into()], secs(99)) + .await; + + assert_eq!((first, second), (Ok(added), Ok(vec!["c".to_owned()]))); + assert_eq!( + cache.get_cache("members"), + Ok(Some(HashSet::from(["a".into(), "b".into(), "c".into()]))) + ); + assert_eq!(cache.expires_at("members"), Ok(secs(110))); + } + + #[rstest] + #[tokio::test] + async fn an_expired_set_restarts_empty_with_a_new_ttl(clock: Clock) { + let cache = cache::>(&clock, 4); + cache + .async_set_cache_sadd("members", vec!["a".into()], secs(10)) + .await + .unwrap(); + at(&clock, 110); + + cache + .async_set_cache_sadd("members", vec!["d".into()], None) + .await + .unwrap(); + + assert_eq!( + cache.get_cache("members"), + Ok(Some(HashSet::from(["d".into()]))) + ); + assert_eq!(cache.expires_at("members"), Ok(secs(170))); + } + + #[rstest] + #[tokio::test] + async fn a_set_addition_that_outgrows_the_entry_limit_is_not_stored() { + let cache = InMemoryCache::>::with_clock_and_size_measurement( + Some(4), + None, + Some(2), + Some(Arc::new(|value: &HashSet| Ok(value.len()))), + || Duration::ZERO, + ); + cache + .async_set_cache_sadd("members", vec!["a".into(), "b".into()], None) + .await + .unwrap(); + + let added = cache + .async_set_cache_sadd("members", vec!["c".into()], None) + .await; + + assert_eq!(added, Ok(vec!["c".to_owned()])); + assert_eq!( + cache.get_cache("members"), + Ok(Some(HashSet::from(["a".into(), "b".into()]))) + ); + } +} diff --git a/litellm/rust/crates/cache-memory/src/lib.rs b/litellm/rust/crates/cache-memory/src/lib.rs index c5b7fb6cb54..215e9d89edf 100644 --- a/litellm/rust/crates/cache-memory/src/lib.rs +++ b/litellm/rust/crates/cache-memory/src/lib.rs @@ -1,3 +1,7 @@ mod cache; +mod capabilities; +mod state; +#[cfg(test)] +mod test_support; -pub use cache::{CacheWrite, InMemoryCache}; +pub use cache::{CacheWrite, Expiry, InMemoryCache}; diff --git a/litellm/rust/crates/cache-memory/src/state.rs b/litellm/rust/crates/cache-memory/src/state.rs new file mode 100644 index 00000000000..06e42bc0308 --- /dev/null +++ b/litellm/rust/crates/cache-memory/src/state.rs @@ -0,0 +1,320 @@ +use std::{ + collections::{BTreeSet, HashMap}, + time::Duration, +}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Deadline { + UnlessLive(Duration), + Always(Duration), +} + +struct Entry { + value: V, + expires_at: Duration, +} + +/// Entries with an exact index of their deadlines: `by_expiry` holds one record per entry, so it +/// never outgrows `entries`. Every read and write first drops what has expired by `now`. +pub(crate) struct State { + entries: HashMap>, + by_expiry: BTreeSet<(Duration, String)>, +} + +impl State { + pub(crate) fn new() -> Self { + Self { + entries: HashMap::new(), + by_expiry: BTreeSet::new(), + } + } + + pub(crate) fn get(&mut self, key: &str, now: Duration) -> Option<&V> { + self.drop_expired(now); + self.entries.get(key).map(|entry| &entry.value) + } + + pub(crate) fn expires_at(&mut self, key: &str, now: Duration) -> Option { + self.drop_expired(now); + self.entries.get(key).map(|entry| entry.expires_at) + } + + pub(crate) fn len(&mut self, now: Duration) -> usize { + self.drop_expired(now); + self.entries.len() + } + + pub(crate) fn live(&mut self, now: Duration) -> impl Iterator { + self.drop_expired(now); + self.entries + .iter() + .map(|(key, entry)| (key.as_str(), &entry.value)) + } + + pub(crate) fn keys_by_expiry(&mut self, now: Duration) -> impl Iterator { + self.drop_expired(now); + self.by_expiry.iter().map(|(_, key)| key.as_str()) + } + + /// A key the cache does not hold evicts the entries soonest to expire until it fits `capacity`. + pub(crate) fn write( + &mut self, + key: String, + value: V, + deadline: Deadline, + capacity: usize, + now: Duration, + ) { + self.drop_expired(now); + let live = self.take(&key).map(|entry| entry.expires_at); + let expires_at = match (deadline, live) { + (Deadline::UnlessLive(_), Some(live)) => live, + (Deadline::UnlessLive(at) | Deadline::Always(at), _) => at, + }; + while self.entries.len() >= capacity { + let Some((_, evicted)) = self.by_expiry.pop_first() else { + break; + }; + self.entries.remove(&evicted); + } + self.by_expiry.insert((expires_at, key.clone())); + self.entries.insert(key, Entry { value, expires_at }); + } + + pub(crate) fn remove(&mut self, key: &str) { + self.take(key); + } + + pub(crate) fn clear(&mut self) { + self.entries.clear(); + self.by_expiry.clear(); + } + + fn take(&mut self, key: &str) -> Option> { + let entry = self.entries.remove(key)?; + self.by_expiry.remove(&(entry.expires_at, key.to_owned())); + Some(entry) + } + + fn drop_expired(&mut self, now: Duration) { + while self + .by_expiry + .first() + .is_some_and(|(expires_at, _)| *expires_at <= now) + { + if let Some((_, key)) = self.by_expiry.pop_first() { + self.entries.remove(&key); + } + } + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use rstest::{fixture, rstest}; + + use super::{Deadline, State}; + + const START: Duration = Duration::ZERO; + + fn secs(seconds: u64) -> Duration { + Duration::from_secs(seconds) + } + + fn put(state: &mut State<&'static str>, key: &str, value: &'static str, expires_at: u64) { + state.write( + key.to_owned(), + value, + Deadline::UnlessLive(secs(expires_at)), + 2, + START, + ); + } + + /// Full at capacity 2: `early` expires at 10 s and `late` at 20 s. + #[fixture] + fn full() -> State<&'static str> { + let mut state = State::new(); + put(&mut state, "early", "a", 10); + put(&mut state, "late", "b", 20); + state + } + + #[rstest] + #[case::before(9, Some("a"))] + #[case::at_the_deadline(10, None)] + #[case::after(11, None)] + fn an_entry_is_held_until_its_deadline( + mut full: State<&'static str>, + #[case] now: u64, + #[case] expected: Option<&'static str>, + ) { + assert_eq!(full.get("early", secs(now)).copied(), expected); + } + + #[rstest] + #[case::keeps_its_deadline(Deadline::UnlessLive(secs(50)), 10)] + #[case::takes_the_new_deadline(Deadline::Always(secs(50)), 50)] + fn a_write_to_a_live_key_replaces_the_value_and( + mut full: State<&'static str>, + #[case] deadline: Deadline, + #[case] expected: u64, + ) { + full.write("early".to_owned(), "new", deadline, 2, secs(5)); + + assert_eq!(full.get("early", secs(5)).copied(), Some("new")); + assert_eq!(full.expires_at("early", secs(5)), Some(secs(expected))); + } + + #[rstest] + fn an_expired_key_takes_the_deadline_of_the_next_write(mut full: State<&'static str>) { + full.write( + "early".to_owned(), + "new", + Deadline::UnlessLive(secs(30)), + 2, + secs(10), + ); + + assert_eq!(full.expires_at("early", secs(10)), Some(secs(30))); + } + + #[rstest] + fn a_new_key_at_capacity_evicts_the_entry_soonest_to_expire(mut full: State<&'static str>) { + put(&mut full, "new", "c", 5); + + assert_eq!( + full.keys_by_expiry(START).collect::>(), + ["new", "late"] + ); + } + + #[rstest] + fn entries_with_one_deadline_are_evicted_in_key_order() { + let mut state = State::new(); + put(&mut state, "b", "2", 10); + put(&mut state, "a", "1", 10); + + put(&mut state, "c", "3", 10); + + assert_eq!(state.keys_by_expiry(START).collect::>(), ["b", "c"]); + } + + #[rstest] + fn an_expired_entry_makes_room_before_a_live_one_is_evicted(mut full: State<&'static str>) { + full.write( + "new".to_owned(), + "c", + Deadline::UnlessLive(secs(15)), + 2, + secs(10), + ); + + assert_eq!( + full.keys_by_expiry(secs(10)).collect::>(), + ["new", "late"] + ); + } + + #[rstest] + #[case::keeping_its_deadline(Deadline::UnlessLive(secs(50)))] + #[case::taking_a_new_deadline(Deadline::Always(secs(50)))] + fn a_write_to_a_held_key_at_capacity_evicts_nothing( + mut full: State<&'static str>, + #[case] deadline: Deadline, + ) { + full.write("late".to_owned(), "new", deadline, 2, START); + + assert_eq!(full.get("early", START).copied(), Some("a")); + assert_eq!(full.len(START), 2); + } + + #[rstest] + fn a_key_given_a_later_deadline_is_evicted_by_that_deadline(mut full: State<&'static str>) { + full.write( + "early".to_owned(), + "a", + Deadline::Always(secs(30)), + 2, + START, + ); + + put(&mut full, "new", "c", 40); + + assert_eq!( + full.keys_by_expiry(START).collect::>(), + ["early", "new"] + ); + } + + #[rstest] + fn a_removed_key_written_again_is_indexed_once(mut full: State<&'static str>) { + full.remove("early"); + put(&mut full, "early", "a", 30); + full.remove("early"); + put(&mut full, "early", "a", 40); + + assert_eq!( + full.keys_by_expiry(START).collect::>(), + ["late", "early"] + ); + assert_eq!(full.len(START), 2); + } + + #[rstest] + #[case::a_new_key("new", &["new"])] + #[case::a_held_key("early", &["early"])] + fn a_smaller_capacity_applies_on_the_next_write_of( + mut full: State<&'static str>, + #[case] key: &str, + #[case] kept: &[&str], + ) { + full.write(key.to_owned(), "c", Deadline::Always(secs(40)), 1, START); + + assert_eq!(full.keys_by_expiry(START).collect::>(), kept); + } + + #[rstest] + fn an_expired_entry_is_not_listed(mut full: State<&'static str>) { + assert_eq!(full.live(secs(10)).collect::>(), [("late", &"b")]); + } + + #[rstest] + fn an_expired_entry_is_not_counted(mut full: State<&'static str>) { + assert_eq!(full.len(secs(10)), 1); + } + + #[rstest] + fn an_expired_entry_has_no_deadline(mut full: State<&'static str>) { + assert_eq!(full.expires_at("early", secs(10)), None); + } + + #[rstest] + fn an_expired_entry_is_not_among_the_keys_by_expiry(mut full: State<&'static str>) { + assert_eq!(full.keys_by_expiry(secs(10)).collect::>(), ["late"]); + } + + #[rstest] + fn clear_drops_every_entry_and_its_deadline(mut full: State<&'static str>) { + full.clear(); + + assert_eq!(full.len(START), 0); + assert_eq!(full.keys_by_expiry(START).count(), 0); + } + + #[rstest] + fn a_deadline_at_the_end_of_time_is_never_reached() { + let mut state = State::new(); + state.write( + "forever".to_owned(), + "a", + Deadline::Always(Duration::MAX), + 2, + START, + ); + + assert_eq!(state.get("forever", secs(u64::MAX)).copied(), Some("a")); + } +} diff --git a/litellm/rust/crates/cache-memory/src/test_support.rs b/litellm/rust/crates/cache-memory/src/test_support.rs new file mode 100644 index 00000000000..53f68e5e035 --- /dev/null +++ b/litellm/rust/crates/cache-memory/src/test_support.rs @@ -0,0 +1,33 @@ +use std::{ + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, +}; + +use rstest::fixture; + +use crate::InMemoryCache; + +pub(crate) type Clock = Arc; + +#[fixture] +pub(crate) fn clock() -> Clock { + Arc::new(AtomicU64::new(100)) +} + +pub(crate) fn cache(clock: &Clock, capacity: usize) -> InMemoryCache { + let clock = clock.clone(); + InMemoryCache::with_clock(Some(capacity), Some(Duration::from_secs(60)), move || { + Duration::from_secs(clock.load(Ordering::SeqCst)) + }) +} + +pub(crate) fn at(clock: &Clock, seconds: u64) { + clock.store(seconds, Ordering::SeqCst); +} + +pub(crate) fn secs(seconds: u64) -> Option { + Some(Duration::from_secs(seconds)) +} diff --git a/litellm/rust/crates/cache-memory/tests/cache.rs b/litellm/rust/crates/cache-memory/tests/cache.rs deleted file mode 100644 index e6c74e5ccda..00000000000 --- a/litellm/rust/crates/cache-memory/tests/cache.rs +++ /dev/null @@ -1,731 +0,0 @@ -use std::{ - collections::HashSet, - sync::{ - Arc, - atomic::{AtomicU64, Ordering}, - }, - time::Duration, -}; - -use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, CacheBackend, ClaimCache, CounterCache, DeleteCache, - DisconnectCache, Error, ExactCacheContext, FlushCache, IncrementOperation, SetCache, TtlCache, - get_cache, set_cache, -}; -use litellm_cache_memory::{CacheWrite, InMemoryCache}; -use rstest::{fixture, rstest}; - -type Clock = Arc; - -#[fixture] -fn clock() -> Clock { - Arc::new(AtomicU64::new(100)) -} - -fn cache_with(clock: &Clock, capacity: usize) -> InMemoryCache { - let clock = clock.clone(); - InMemoryCache::with_clock(Some(capacity), Some(Duration::from_secs(60)), move || { - Duration::from_millis(clock.load(Ordering::SeqCst) * 1000) - }) -} - -fn cache(clock: &Clock, capacity: usize) -> InMemoryCache { - cache_with(clock, capacity) -} - -fn at(clock: &Clock, seconds: u64) { - clock.store(seconds, Ordering::SeqCst); -} - -fn secs(seconds: u64) -> Option { - Some(Duration::from_secs(seconds)) -} - -fn ttl(seconds: u64) -> ExactCacheContext { - ExactCacheContext { ttl: secs(seconds) } -} - -fn measured(capacity: usize) -> InMemoryCache { - InMemoryCache::with_clock_and_size_measurement( - Some(capacity), - secs(60), - Some(4), - Some(Arc::new(|value: &String| { - if value.is_empty() { - return Err(Error::InvalidEntry(None)); - } - Ok(value.len()) - })), - || Duration::from_secs(100), - ) -} - -#[rstest] -fn default_explicit_and_override_ttls_follow_python_rules(clock: Clock) { - let cache = cache(&clock, 4); - cache.set_cache("key", "first".into(), None).unwrap(); - assert_eq!(cache.expires_at("key").unwrap(), secs(160)); - cache.set_cache("key", "second".into(), secs(10)).unwrap(); - assert_eq!(cache.expires_at("key").unwrap(), secs(160)); - at(&clock, 160); - assert_eq!(cache.get_cache("key").unwrap(), Some("second".into())); - at(&clock, 161); - assert_eq!(cache.get_cache("key").unwrap(), None); - assert_eq!(cache.expires_at("key").unwrap(), None); - cache.set_cache("key", "third".into(), secs(10)).unwrap(); - assert_eq!(cache.expires_at("key").unwrap(), secs(171)); -} - -#[rstest] -#[case::unset(None, secs(600))] -#[case::zero_falls_back_like_python_or(Some(Duration::ZERO), secs(600))] -#[case::explicit(secs(5), secs(5))] -fn default_ttl_falls_back_to_ten_minutes( - #[case] default_ttl: Option, - #[case] expected: Option, -) { - let cache = InMemoryCache::::with_clock(None, default_ttl, || Duration::ZERO); - assert_eq!(cache.get_ttl(&ExactCacheContext::default()), expected); - cache.set_cache("key", "value".into(), None).unwrap(); - assert_eq!(cache.expires_at("key").unwrap(), expected); - assert_eq!(cache.max_size_in_memory(), 200); -} - -#[rstest] -fn write_at_expiry_boundary_refreshes_ttl(clock: Clock) { - let cache = cache(&clock, 4); - cache.set_cache("key", "first".into(), secs(10)).unwrap(); - at(&clock, 110); - cache.set_cache("key", "second".into(), secs(10)).unwrap(); - assert_eq!(cache.expires_at("key").unwrap(), secs(120)); - at(&clock, 115); - assert_eq!(cache.get_cache("key").unwrap(), Some("second".into())); -} - -#[rstest] -fn expired_key_without_a_read_allows_a_ttl_override(clock: Clock) { - let cache = cache(&clock, 4); - cache.set_cache("key", "first".into(), secs(1)).unwrap(); - assert_eq!(cache.allow_ttl_override("key"), Ok(false)); - at(&clock, 102); - assert_eq!(cache.allow_ttl_override("key"), Ok(true)); - cache.set_cache("key", "second".into(), secs(1)).unwrap(); - assert_eq!(cache.expires_at("key").unwrap(), secs(103)); - assert_eq!(cache.allow_ttl_override("missing"), Ok(true)); -} - -#[rstest] -fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Clock) { - let cache = cache(&clock, 2); - cache.set_cache("early", "a".into(), secs(10)).unwrap(); - cache.set_cache("late", "b".into(), secs(20)).unwrap(); - cache.delete_cache("early").unwrap(); - cache.set_cache("new", "c".into(), secs(30)).unwrap(); - assert_eq!(cache.get_cache("late").unwrap(), Some("b".into())); - cache.set_cache("last", "d".into(), secs(40)).unwrap(); - assert_eq!(cache.get_cache("late").unwrap(), None); -} - -#[rstest] -fn max_size_is_respected_when_every_item_has_a_long_ttl(clock: Clock) { - let cache = cache(&clock, 3); - for index in 0..3 { - at(&clock, 100 + index); - cache - .set_cache( - format!("key_{index}"), - format!("value_{index}"), - secs(86_400), - ) - .unwrap(); - } - assert_eq!(cache.len(), Ok(3)); - cache - .set_cache("key_3", "value_3".into(), secs(86_400)) - .unwrap(); - assert_eq!(cache.len(), Ok(3)); - assert_eq!(cache.get_cache("key_0").unwrap(), None); - assert_eq!(cache.expires_at("key_0").unwrap(), None); - for key in ["key_1", "key_2", "key_3"] { - assert!(cache.get_cache(key).unwrap().is_some(), "{key}"); - } -} - -#[rstest] -fn expired_items_are_evicted_before_live_ones(clock: Clock) { - let cache = cache(&clock, 3); - cache.set_cache("expired_1", "1".into(), secs(1)).unwrap(); - cache.set_cache("expired_2", "2".into(), secs(1)).unwrap(); - cache - .set_cache("long_lived", "3".into(), secs(86_400)) - .unwrap(); - assert_eq!(cache.len(), Ok(3)); - at(&clock, 102); - cache - .set_cache("new_item", "4".into(), secs(86_400)) - .unwrap(); - assert_eq!(cache.len(), Ok(2)); - assert_eq!(cache.get_cache("long_lived").unwrap(), Some("3".into())); - assert_eq!(cache.get_cache("new_item").unwrap(), Some("4".into())); - for key in ["expired_1", "expired_2"] { - assert_eq!(cache.expires_at(key).unwrap(), None, "{key}"); - } -} - -#[rstest] -fn injected_clock_controls_expiry_and_eviction(clock: Clock) { - let cache = cache(&clock, 2); - at(&clock, 0); - cache - .set_cache("first", "original".into(), secs(10)) - .unwrap(); - at(&clock, 9); - cache.set_cache("second", "survivor".into(), None).unwrap(); - assert_eq!(cache.get_cache("first").unwrap(), Some("original".into())); - at(&clock, 11); - assert_eq!(cache.get_cache("first").unwrap(), None); - cache - .set_cache("third", "replacement".into(), None) - .unwrap(); - assert_eq!(cache.get_cache("second").unwrap(), Some("survivor".into())); - at(&clock, 70); - cache.set_cache("fourth", "new".into(), None).unwrap(); - assert_eq!(cache.get_cache("second").unwrap(), None); - assert_eq!( - cache.get_cache("third").unwrap(), - Some("replacement".into()) - ); - assert_eq!(cache.get_cache("fourth").unwrap(), Some("new".into())); -} - -#[rstest] -fn rewriting_one_key_keeps_one_heap_entry(clock: Clock) { - let cache = cache(&clock, 10); - for index in 0..1_000 { - cache - .set_cache("hot_key", format!("value_{index}"), secs(60)) - .unwrap(); - } - assert_eq!(cache.expiration_heap_len(), Ok(1)); -} - -#[rstest] -fn repeated_increments_keep_one_heap_entry_per_expiration() { - let cache = InMemoryCache::::new(Some(4), None); - for _ in 0..100 { - cache - .increment_cache("counter", 1.0, ExactCacheContext::default()) - .unwrap(); - } - assert_eq!(cache.expiration_heap_len(), Ok(1)); -} - -#[rstest] -fn reinserting_expired_keys_below_capacity_prunes_the_heap(clock: Clock) { - let cache = cache(&clock, 200); - for cycle in 0..3 { - for index in 0..5 { - cache - .set_cache(format!("key_{index}"), format!("value_{cycle}"), secs(1)) - .unwrap(); - } - at(&clock, 100 + 2 * (cycle + 1)); - } - for index in 0..5 { - cache - .set_cache(format!("key_{index}"), "final".into(), secs(1)) - .unwrap(); - } - assert_eq!(cache.len(), Ok(5)); - assert_eq!(cache.expiration_heap_len(), Ok(5)); -} - -#[rstest] -fn evict_cache_drops_expired_entries_then_makes_room(clock: Clock) { - let cache = cache(&clock, 2); - assert_eq!(cache.is_empty(), Ok(true)); - cache.set_cache("short", "a".into(), secs(1)).unwrap(); - cache.set_cache("long", "b".into(), secs(50)).unwrap(); - at(&clock, 102); - cache.evict_cache().unwrap(); - assert_eq!(cache.len(), Ok(1)); - assert_eq!(cache.expires_at("short").unwrap(), None); - cache.set_cache("longer", "c".into(), secs(90)).unwrap(); - cache.evict_cache().unwrap(); - assert_eq!(cache.len(), Ok(1)); - assert_eq!(cache.get_cache("long").unwrap(), None); - assert_eq!(cache.get_cache("longer").unwrap(), Some("c".into())); -} - -#[rstest] -fn evict_element_if_expired_reports_removal(clock: Clock) { - let cache = cache(&clock, 4); - cache.set_cache("key", "value".into(), secs(10)).unwrap(); - assert_eq!(cache.evict_element_if_expired("key"), Ok(false)); - assert_eq!(cache.evict_element_if_expired("missing"), Ok(false)); - at(&clock, 110); - assert_eq!(cache.evict_element_if_expired("key"), Ok(false)); - at(&clock, 111); - assert_eq!(cache.evict_element_if_expired("key"), Ok(true)); - assert_eq!(cache.len(), Ok(0)); - assert_eq!(cache.expires_at("key").unwrap(), None); -} - -#[rstest] -#[case::fits("ok", Ok(true))] -#[case::at_limit("four", Ok(true))] -#[case::too_large("oversized", Ok(false))] -#[case::measure_error("", Err(Error::InvalidEntry(None)))] -fn check_value_size_applies_the_entry_limit( - #[case] value: &str, - #[case] expected: Result, -) { - assert_eq!(measured(2).check_value_size(&value.to_string()), expected); -} - -#[rstest] -fn values_are_unbounded_without_a_measure() { - let cache = InMemoryCache::::default(); - assert_eq!(cache.max_entry_bytes(), None); - assert_eq!(cache.check_value_size(&"x".repeat(1 << 20)), Ok(true)); -} - -#[rstest] -fn disabled_size_limited_and_validated_writes_are_observable() { - assert_eq!( - measured(0).set_cache("a", "x".into(), None).unwrap(), - CacheWrite::Disabled - ); - let cache = measured(2); - assert_eq!(cache.max_entry_bytes(), Some(4)); - assert_eq!( - cache.set_cache("large", "oversized".into(), None).unwrap(), - CacheWrite::TooLarge - ); - assert_eq!(cache.get_cache("large").unwrap(), None); - assert_eq!( - cache.set_cache("small", "ok".into(), None).unwrap(), - CacheWrite::Stored - ); - assert_eq!(cache.get_cache("small").unwrap(), Some("ok".into())); - assert_eq!( - cache.set_cache("invalid", String::new(), None), - Err(Error::InvalidEntry(None)) - ); - assert_eq!(cache.get_cache("invalid").unwrap(), None); - cache.delete_cache("small").unwrap(); - assert_eq!(cache.get_cache("small").unwrap(), None); -} - -#[rstest] -#[tokio::test] -async fn disconnect_is_a_no_op_that_keeps_entries() { - let cache = InMemoryCache::::default(); - cache.set_cache("key", "value".into(), None).unwrap(); - cache.disconnect().await.unwrap(); - assert_eq!(cache.get_cache("key").unwrap(), Some("value".into())); -} - -#[rstest] -#[tokio::test] -async fn generic_consumers_share_typed_values_and_honor_expiration(clock: Clock) { - let cache: CacheBackend> = Arc::new(self::cache(&clock, 4)); - let reader = Arc::clone(&cache); - let context = ttl(5); - set_cache(cache.as_ref(), "sync", "first".into(), &context).unwrap(); - assert_eq!( - get_cache(reader.as_ref(), "sync", &context).unwrap(), - Some("first".into()) - ); - cache - .batch_cache_write("async", "second".into(), context.clone()) - .await - .unwrap(); - cache - .async_set_cache_pipeline(vec![("batch".into(), "third".into())], context.clone()) - .await - .unwrap(); - drop(cache); - for (key, value) in [("sync", "first"), ("async", "second"), ("batch", "third")] { - assert_eq!( - reader.async_get_cache(key, &context).await.unwrap(), - Some(value.into()) - ); - } - reader.async_delete_cache("async").await.unwrap(); - assert_eq!( - reader.async_get_cache("async", &context).await.unwrap(), - None - ); - at(&clock, 106); - assert_eq!(get_cache(reader.as_ref(), "sync", &context).unwrap(), None); - assert_eq!( - reader.async_get_cache("batch", &context).await.unwrap(), - None - ); -} - -#[rstest] -#[case::context_ttl(ttl(5), secs(105))] -#[case::default_ttl(ExactCacheContext::default(), secs(160))] -#[tokio::test] -async fn pipeline_writes_use_the_context_ttl_or_the_default( - clock: Clock, - #[case] context: ExactCacheContext, - #[case] expected: Option, -) { - let cache = cache(&clock, 4); - cache - .async_set_cache_pipeline( - vec![("a".into(), "1".into()), ("b".into(), "2".into())], - context, - ) - .await - .unwrap(); - assert_eq!(cache.expires_at("a").unwrap(), expected); - assert_eq!(cache.expires_at("b").unwrap(), expected); -} - -#[rstest] -#[tokio::test] -async fn batch_reads_return_one_entry_per_key_and_drop_expired_ones(clock: Clock) { - let cache = cache(&clock, 4); - cache.set_cache("short", "a".into(), secs(1)).unwrap(); - cache.set_cache("long", "b".into(), secs(50)).unwrap(); - let keys = vec!["short".to_string(), "missing".into(), "long".into()]; - assert_eq!( - cache - .batch_get_cache(&keys, &ExactCacheContext::default()) - .unwrap(), - [ - BatchEntry::Hit("a".to_string()), - BatchEntry::Miss, - BatchEntry::Hit("b".into()), - ] - ); - at(&clock, 102); - assert_eq!( - cache - .async_batch_get_cache(keys, ExactCacheContext::default()) - .await - .unwrap(), - [ - BatchEntry::Miss, - BatchEntry::Miss, - BatchEntry::Hit("b".into()) - ] - ); -} - -#[rstest] -#[tokio::test] -async fn flush_clears_values_and_expirations(clock: Clock) { - let cache = cache(&clock, 4); - cache.set_cache("a", "1".into(), None).unwrap(); - cache.set_cache("b", "2".into(), None).unwrap(); - cache.flush_cache().unwrap(); - assert_eq!(cache.len(), Ok(0)); - assert_eq!(cache.expiration_heap_len(), Ok(0)); - cache.set_cache("c", "3".into(), None).unwrap(); - FlushCache::async_flush_cache(&cache).await.unwrap(); - assert_eq!(cache.is_empty(), Ok(true)); - assert_eq!( - cache.async_get_oldest_n_keys(5).await.unwrap(), - Vec::::new() - ); -} - -#[rstest] -fn claims_are_atomic_and_refresh_eligible_winners(clock: Clock) { - let cache = cache(&clock, 4); - let context = ttl(10); - assert_eq!( - cache - .claim_cache("affinity", "first".to_string(), &[], context.clone()) - .unwrap(), - "first" - ); - at(&clock, 103); - assert_eq!( - cache - .claim_cache("affinity", "second".to_string(), &[], context.clone()) - .unwrap(), - "first" - ); - assert_eq!(cache.expires_at("affinity").unwrap(), secs(110)); - at(&clock, 105); - assert_eq!( - cache - .claim_cache( - "affinity", - "second".to_string(), - &["first".to_string(), "second".to_string()], - context, - ) - .unwrap(), - "first" - ); - assert_eq!(cache.expires_at("affinity").unwrap(), secs(115)); -} - -#[rstest] -fn counters_increment_under_one_lock() { - let cache = InMemoryCache::::default(); - assert_eq!( - CounterCache::increment_cache(&cache, "counter", 1.5, ExactCacheContext::default()) - .unwrap(), - 1.5 - ); - assert_eq!( - CounterCache::increment_cache(&cache, "counter", 2.0, ExactCacheContext::default()) - .unwrap(), - 3.5 - ); -} - -#[rstest] -fn concurrent_increments_are_atomic() { - let cache = Arc::new(InMemoryCache::::default()); - cache.set_cache("counter", 1000.0, None).unwrap(); - let threads = (0..8) - .map(|_| { - let cache = cache.clone(); - std::thread::spawn(move || { - cache - .increment_cache("counter", 1.0, ExactCacheContext::default()) - .unwrap() - }) - }) - .collect::>(); - for thread in threads { - thread.join().unwrap(); - } - assert_eq!(cache.get_cache("counter").unwrap(), Some(1008.0)); -} - -#[rstest] -#[case::window_semantics(false)] -#[case::refresh_ttl_is_ignored(true)] -#[tokio::test] -async fn async_increment_delegates_to_the_locked_sync_path( - clock: Clock, - #[case] refresh_ttl: bool, -) { - let cache = cache_with::(&clock, 4); - assert_eq!( - cache - .async_increment("counter", 2.0, ttl(10), refresh_ttl) - .await, - Ok(2.0) - ); - at(&clock, 105); - assert_eq!( - cache - .async_increment("counter", 3.0, ttl(10), refresh_ttl) - .await, - Ok(5.0) - ); - assert_eq!(cache.get_cache("counter").unwrap(), Some(5.0)); - assert_eq!(cache.expires_at("counter").unwrap(), secs(110)); -} - -#[rstest] -fn expired_counters_restart_from_zero_with_a_new_ttl(clock: Clock) { - let cache = cache_with::(&clock, 4); - cache.increment_cache("counter", 2.0, ttl(10)).unwrap(); - at(&clock, 111); - assert_eq!(cache.increment_cache("counter", 1.0, ttl(10)), Ok(1.0)); - assert_eq!(cache.expires_at("counter").unwrap(), secs(121)); -} - -/// Python `InMemoryCache.set_cache` runs `evict_cache()` before every insert, and step 2 evicts -/// the earliest expiry while `len(cache_dict) >= max_size_in_memory`, even when the key being -/// written already exists. -#[rstest] -fn overwriting_an_existing_key_at_capacity_evicts_the_earliest_expiry_like_python(clock: Clock) { - let cache = cache(&clock, 2); - cache.set_cache("hot", "1".into(), secs(10)).unwrap(); - cache.set_cache("cold", "2".into(), secs(20)).unwrap(); - - cache.set_cache("cold", "3".into(), None).unwrap(); - - assert_eq!(cache.get_cache("hot").unwrap(), None); - assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into())); -} - -/// `claim_cache` has no Python counterpart; it never evicts another entry for a key it holds. -#[rstest] -fn claiming_an_existing_key_at_capacity_keeps_other_entries(clock: Clock) { - let cache = cache(&clock, 2); - cache.set_cache("hot", "1".into(), secs(10)).unwrap(); - cache.set_cache("cold", "2".into(), secs(20)).unwrap(); - - cache - .claim_cache("cold", "4".into(), &[], ExactCacheContext::default()) - .unwrap(); - assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into())); - assert_eq!(cache.get_cache("cold").unwrap(), Some("2".into())); -} - -/// Python `increment_cache` is `get_cache` then `set_cache`, so at capacity the write evicts -/// the earliest expiry first: equal expiries tie-break on the key, and the value read before -/// eviction is the one written back. -#[rstest] -fn incrementing_at_capacity_evicts_the_earliest_expiry_like_python(clock: Clock) { - let cache = cache_with::(&clock, 2); - for key in ["a", "b", "a", "b"] { - cache - .increment_cache(key, 1.0, ExactCacheContext::default()) - .unwrap(); - } - assert_eq!(cache.get_cache("a").unwrap(), None); - assert_eq!(cache.get_cache("b").unwrap(), Some(2.0)); -} - -#[rstest] -#[tokio::test] -async fn disabled_cache_does_not_retain_claims_counters_or_sets() { - let claims = InMemoryCache::::new(Some(0), None); - assert_eq!( - claims - .claim_cache("key", "first".into(), &[], ExactCacheContext::default()) - .unwrap(), - "first" - ); - assert_eq!(claims.get_cache("key").unwrap(), None); - - let counters = InMemoryCache::::new(Some(0), None); - assert_eq!( - counters - .increment_cache("key", 2.0, ExactCacheContext::default()) - .unwrap(), - 2.0 - ); - assert_eq!(counters.get_cache("key").unwrap(), None); - - let sets = InMemoryCache::>::new(Some(0), None); - assert_eq!( - sets.async_set_cache_sadd("key", vec!["a".into()], None) - .await - .unwrap(), - ["a"] - ); - assert_eq!(sets.get_cache("key").unwrap(), None); -} - -#[rstest] -#[tokio::test] -async fn ttl_and_oldest_key_operations_use_the_stored_expirations(clock: Clock) { - let cache = cache(&clock, 3); - cache.set_cache("later", "2".into(), secs(20)).unwrap(); - cache.set_cache("first", "1".into(), secs(10)).unwrap(); - cache.set_cache("latest", "3".into(), secs(30)).unwrap(); - - assert_eq!(cache.async_get_ttl("first").await.unwrap(), secs(110)); - assert_eq!( - TtlCache::async_get_ttl(&cache, "later").await.unwrap(), - secs(120) - ); - assert_eq!(cache.async_get_oldest_n_keys(1).await.unwrap(), ["first"]); - assert_eq!( - cache.async_get_oldest_n_keys(10).await.unwrap(), - ["first", "later", "latest"] - ); - assert_eq!( - cache.async_get_oldest_n_keys(0).await.unwrap(), - Vec::::new() - ); - assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None); -} - -#[rstest] -#[tokio::test] -async fn increment_pipeline_preserves_operation_order(clock: Clock) { - let cache = cache_with::(&clock, 3); - let operation = |key: &str, amount, ttl| IncrementOperation { - key: key.into(), - amount, - ttl: secs(ttl), - }; - assert_eq!( - cache - .async_increment_pipeline(vec![ - operation("a", 1.0, 10), - operation("b", 5.0, 30), - operation("a", 2.0, 20), - ]) - .await - .unwrap(), - [1.0, 5.0, 3.0] - ); - assert_eq!(cache.get_cache("a").unwrap(), Some(3.0)); - assert_eq!(cache.expires_at("a").unwrap(), secs(110)); - assert_eq!(cache.expires_at("b").unwrap(), secs(130)); - assert_eq!( - cache.async_increment_pipeline(Vec::new()).await.unwrap(), - Vec::::new() - ); -} - -#[rstest] -#[tokio::test] -async fn set_capability_preserves_python_result_and_deduplicates_storage(clock: Clock) { - let cache = cache_with::>(&clock, 4); - let inserted = vec!["a".to_string(), "a".into(), "b".into()]; - assert_eq!( - cache - .async_set_cache_sadd("members", inserted.clone(), secs(10)) - .await - .unwrap(), - inserted - ); - assert_eq!( - cache - .async_set_cache_sadd("members", vec!["c".into()], secs(99)) - .await - .unwrap(), - ["c"] - ); - assert_eq!( - cache.get_cache("members").unwrap(), - Some(HashSet::from(["a".into(), "b".into(), "c".into()])) - ); - assert_eq!(cache.expires_at("members").unwrap(), secs(110)); - at(&clock, 111); - cache - .async_set_cache_sadd("members", vec!["d".into()], None) - .await - .unwrap(); - assert_eq!( - cache.get_cache("members").unwrap(), - Some(HashSet::from(["d".into()])) - ); - assert_eq!(cache.expires_at("members").unwrap(), secs(171)); -} - -#[rstest] -#[tokio::test] -async fn oversized_set_additions_are_not_stored() { - let cache = InMemoryCache::>::with_clock_and_size_measurement( - Some(4), - None, - Some(2), - Some(Arc::new(|value: &HashSet| Ok(value.len()))), - || Duration::ZERO, - ); - cache - .async_set_cache_sadd("members", vec!["a".into(), "b".into()], None) - .await - .unwrap(); - assert_eq!( - cache - .async_set_cache_sadd("members", vec!["c".into()], None) - .await - .unwrap(), - ["c"] - ); - assert_eq!( - cache.get_cache("members").unwrap(), - Some(HashSet::from(["a".into(), "b".into()])) - ); -} diff --git a/litellm/rust/crates/redis-store/Cargo.toml b/litellm/rust/crates/redis-store/Cargo.toml index cba7e992834..f363c8b67f8 100644 --- a/litellm/rust/crates/redis-store/Cargo.toml +++ b/litellm/rust/crates/redis-store/Cargo.toml @@ -11,6 +11,7 @@ litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-python-compat.workspace = true litellm-redis-client.workspace = true +parking_lot.workspace = true sha2.workspace = true strum.workspace = true thiserror.workspace = true diff --git a/litellm/rust/crates/redis-store/src/auth_cache/mod.rs b/litellm/rust/crates/redis-store/src/auth_cache/mod.rs index b5a83b5aec6..d058f337e03 100644 --- a/litellm/rust/crates/redis-store/src/auth_cache/mod.rs +++ b/litellm/rust/crates/redis-store/src/auth_cache/mod.rs @@ -1,20 +1,18 @@ mod keys; mod types; -use std::{ - sync::{Arc, Mutex}, - time::Duration, -}; +use std::{sync::Arc, time::Duration}; use litellm_python_compat::{Value, json}; use litellm_redis_client::{Declared, RedisBatch, commands}; +use parking_lot::Mutex; pub use self::keys::is_key_object; pub use self::types::Settings; use crate::{ Error, Number, Pending, RequestRedis, Send, dual_cache::DualCache, - memory::{Clock, lock, system_clock}, + memory::{Clock, system_clock}, remote::{Remote, RemoteSlot}, request::ready, }; @@ -54,13 +52,13 @@ impl AuthCache { } pub fn settings(&self) -> Settings { - *lock(&self.settings) + *self.settings.lock() } /// `UserApiKeyCache.update_cache_ttl` and `update_in_memory_max_size` on both partitions, which /// keep what they hold. pub fn reconfigure(&self, settings: Settings) { - *lock(&self.settings) = settings; + *self.settings.lock() = settings; self.key_objects.reconfigure(settings.memory()); self.objects.reconfigure(settings.memory()); } diff --git a/litellm/rust/crates/redis-store/src/dual_cache/batch_read.rs b/litellm/rust/crates/redis-store/src/dual_cache/batch_read.rs index abc49db428b..8196679b518 100644 --- a/litellm/rust/crates/redis-store/src/dual_cache/batch_read.rs +++ b/litellm/rust/crates/redis-store/src/dual_cache/batch_read.rs @@ -1,13 +1,13 @@ use std::{ collections::{HashMap, VecDeque}, - sync::Mutex, time::Duration, }; use litellm_redis_client::{Declared, Error as RedisError, RedisBatch, commands}; +use parking_lot::Mutex; use super::{DualCache, MemorySettings}; -use crate::{Error, FromValue, Send, codec, memory::lock, request::Misses}; +use crate::{Error, FromValue, Send, codec, request::Misses}; /// A batch-reading family's `DualCache` settings, as the host reads them from `litellm` globals. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -171,7 +171,7 @@ impl DualCache]) -> Reservation { let now = self.now(); let refresh_after = self.reads.refresh_after; - let mut recent = lock(&self.reads.recent); + let mut recent = self.reads.recent.lock(); let due = keys .iter() .zip(local) @@ -190,7 +190,7 @@ impl DualCache recent.insert(key, read), diff --git a/litellm/rust/crates/redis-store/src/dual_cache/mod.rs b/litellm/rust/crates/redis-store/src/dual_cache/mod.rs index 10e62076d41..36b090bc152 100644 --- a/litellm/rust/crates/redis-store/src/dual_cache/mod.rs +++ b/litellm/rust/crates/redis-store/src/dual_cache/mod.rs @@ -3,10 +3,9 @@ mod increment; mod read; mod write; -use std::{ - sync::{Arc, Mutex}, - time::Duration, -}; +use std::{sync::Arc, time::Duration}; + +use parking_lot::Mutex; pub use self::batch_read::TierPolicy; pub(crate) use self::{ @@ -14,7 +13,7 @@ pub(crate) use self::{ increment::PipelineIncrement, }; use crate::{ - memory::{Clock, Memory, lock}, + memory::{Clock, Memory}, remote::RemoteSlot, }; @@ -86,13 +85,13 @@ impl DualCache { } pub(crate) fn default_in_memory_ttl(&self) -> Option { - *lock(&self.default_in_memory_ttl) + *self.default_in_memory_ttl.lock() } /// `DualCache.update_cache_ttl` and `update_in_memory_max_size` together: both apply in place, /// and the entries memory already holds stay. pub(crate) fn reconfigure(&self, memory: MemorySettings) { - *lock(&self.default_in_memory_ttl) = memory.ttl; + *self.default_in_memory_ttl.lock() = memory.ttl; self.memory.update_max_size(memory.entries); } } @@ -100,10 +99,7 @@ impl DualCache { #[cfg(test)] mod tests { use std::{ - sync::{ - Arc, - atomic::{AtomicU64, Ordering}, - }, + sync::{Arc, atomic::Ordering}, time::Duration, }; @@ -111,7 +107,9 @@ mod tests { use tokio::runtime::Handle; use super::{DualCache, MemorySettings, PipelineIncrement}; - use crate::{Clock, Number, RequestRedis, memory::Counter, remote::RemoteSlot}; + use crate::{ + Clock, Number, RequestRedis, memory::Counter, remote::RemoteSlot, test_support::clock, + }; #[derive(Clone, Debug, PartialEq)] enum Entry { @@ -134,13 +132,6 @@ mod tests { } } - fn clock() -> (Clock, Arc) { - let seconds = Arc::new(AtomicU64::new(1_000)); - let read = seconds.clone(); - let clock: Clock = Arc::new(move || Duration::from_secs(read.load(Ordering::SeqCst))); - (clock, seconds) - } - fn without_redis( default_in_memory_ttl: Option, now: Clock, @@ -255,6 +246,23 @@ mod tests { assert_eq!(cache.memory().get("short").unwrap(), None); } + #[rstest] + #[tokio::test] + async fn reconfigure_applies_the_new_capacity_to_the_next_write() { + let (now, _) = clock(); + let cache = without_redis(None, now); + cache.set("first", Number::Int(1), None).await.unwrap(); + + cache.reconfigure(MemorySettings { + entries: Some(1), + ttl: None, + }); + cache.set("second", Number::Int(2), None).await.unwrap(); + + assert_eq!(cache.memory().len().unwrap(), 1); + assert_eq!(cache.memory().get("second").unwrap(), Some(Number::Int(2))); + } + // async_increment_cache_pipeline_post_call: InMemoryCache.async_increment_pipeline raises at the // first entry it cannot add to, so later increments never reach memory, and DualCache only logs #[tokio::test] diff --git a/litellm/rust/crates/redis-store/src/lib.rs b/litellm/rust/crates/redis-store/src/lib.rs index f1e3808fdce..46c72203599 100644 --- a/litellm/rust/crates/redis-store/src/lib.rs +++ b/litellm/rust/crates/redis-store/src/lib.rs @@ -14,6 +14,8 @@ pub mod router_session_pins; pub mod router_usage; pub mod sensitive_route_pins; pub mod spend_counters; +#[cfg(test)] +mod test_support; mod value; pub use codec::DecodeError; diff --git a/litellm/rust/crates/redis-store/src/memory.rs b/litellm/rust/crates/redis-store/src/memory.rs index f89d5926760..b67b6d7fe59 100644 --- a/litellm/rust/crates/redis-store/src/memory.rs +++ b/litellm/rust/crates/redis-store/src/memory.rs @@ -1,10 +1,10 @@ use std::{ - sync::{Arc, Mutex, MutexGuard, PoisonError}, + sync::Arc, time::{Duration, SystemTime, UNIX_EPOCH}, }; use litellm_cache::Error; -use litellm_cache_memory::InMemoryCache; +use litellm_cache_memory::{Expiry, InMemoryCache}; use crate::Number; @@ -18,11 +18,9 @@ pub fn system_clock() -> Clock { }) } -/// A store's memory tier: Python's `InMemoryCache`, with its 200 entries and 600 s TTL unless the -/// store names others, plus one lock for every read-modify-write on it. +/// A store's memory tier: 200 entries and a 600 s TTL unless the store names others. pub struct Memory { cache: InMemoryCache, - updates: Mutex<()>, } impl Memory { @@ -30,7 +28,6 @@ impl Memory { let now = now.clone(); Self { cache: InMemoryCache::with_clock(entries, None, move || now()), - updates: Mutex::new(()), } } @@ -38,8 +35,30 @@ impl Memory { self.cache.get_cache(key) } + /// A key memory already holds keeps its deadline, as a counter keeps its window. pub fn set(&self, key: &str, value: V, ttl: Option) -> Result<(), Error> { - self.cache.set_cache(key, value, ttl).map(drop) + self.cache + .set_cache(key, value, Expiry::KeepExisting(ttl)) + .map(drop) + } + + pub fn refresh(&self, key: &str, value: V, ttl: Duration) -> Result<(), Error> { + self.cache + .set_cache(key, value, Expiry::Refresh(Some(ttl))) + .map(drop) + } + + /// Writes `value` when memory holds nothing for `key` or `replaces` what it holds. + pub fn set_if( + &self, + key: &str, + value: V, + ttl: Option, + replaces: impl FnOnce(&V) -> bool, + ) -> Result<(), Error> { + self.cache.update(key, Expiry::KeepExisting(ttl), |held| { + Ok((held.is_none_or(replaces).then_some(value), ())) + }) } pub fn delete(&self, key: &str) -> Result<(), Error> { @@ -67,26 +86,21 @@ impl Memory { self.cache.update_max_size_in_memory(entries); } - /// Read-modify-write of one entry, atomic against every other update of this tier. + /// Read-modify-write of one entry, atomic against every other read and write of this tier. pub fn update( &self, key: &str, ttl: Option, update: impl FnOnce(Option) -> Result, ) -> Result { - let _guard = self.guard(); - let value = update(self.get(key)?)?; - self.set(key, value.clone(), ttl)?; - Ok(value) - } - - /// Held across a read and a conditional write that `update` cannot express. - pub fn guard(&self) -> MutexGuard<'_, ()> { - lock(&self.updates) + self.cache.update(key, Expiry::KeepExisting(ttl), |held| { + let value = update(held.cloned())?; + Ok((Some(value.clone()), value)) + }) } } -/// A value `InMemoryCache.increment_cache` can add to. +/// A value `Memory::increment` can add to. pub(crate) trait Counter: Clone + From { fn as_number(&self) -> Option; } @@ -98,25 +112,130 @@ impl Counter for Number { } impl Memory { - /// `InMemoryCache.increment_cache`: `(current or 0) + amount`, refusing an entry that is not a - /// number and leaving it in place. + /// `(current or 0) + amount`, refusing an entry that is not a number and leaving it in place. pub(crate) fn increment( &self, key: &str, amount: Number, ttl: Option, ) -> Result { - let _guard = self.guard(); - let current = match self.get(key)? { - None => Number::Int(0), - Some(value) => value.as_number().ok_or(Error::InvalidEntry(None))?, - }; - let total = current + amount; - self.set(key, V::from(total), ttl)?; - Ok(total) + self.cache.update(key, Expiry::KeepExisting(ttl), |held| { + let current = match held { + None => Number::Int(0), + Some(value) => value.as_number().ok_or(Error::InvalidEntry(None))?, + }; + let total = current + amount; + Ok((Some(V::from(total)), total)) + }) } } -pub(crate) fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { - mutex.lock().unwrap_or_else(PoisonError::into_inner) +#[cfg(test)] +mod tests { + use std::{sync::atomic::Ordering, time::Duration}; + + use rstest::rstest; + + use super::Memory; + use crate::{ + Number, + test_support::{START, clock}, + }; + + type Write = fn(&Memory); + + fn secs(seconds: u64) -> Duration { + Duration::from_secs(seconds) + } + + fn ttl(seconds: u64) -> Option { + Some(secs(seconds)) + } + + fn after(seconds: u64) -> Option { + ttl(START + seconds) + } + + fn set(memory: &Memory) { + memory.set("key", Number::Int(2), ttl(50)).unwrap(); + } + + fn set_if(memory: &Memory) { + memory + .set_if("key", Number::Int(2), ttl(50), |_| true) + .unwrap(); + } + + fn update(memory: &Memory) { + memory + .update("key", ttl(50), |_| Ok(Number::Int(2))) + .unwrap(); + } + + fn increment(memory: &Memory) { + memory.increment("key", Number::Int(2), ttl(50)).unwrap(); + } + + fn refresh(memory: &Memory) { + memory.refresh("key", Number::Int(2), secs(50)).unwrap(); + } + + #[rstest] + #[case::set(set, 10)] + #[case::set_if(set_if, 10)] + #[case::update(update, 10)] + #[case::increment(increment, 10)] + #[case::refresh(refresh, 55)] + fn a_write_to_a_held_key_keeps_its_deadline_unless_it_refreshes( + #[case] write: Write, + #[case] expires_after: u64, + ) { + let (now, seconds) = clock(); + let memory = Memory::new(None, &now); + memory.set("key", Number::Int(0), ttl(10)).unwrap(); + seconds.fetch_add(5, Ordering::SeqCst); + + write(&memory); + + assert_eq!(memory.get("key").unwrap(), Some(Number::Int(2))); + assert_eq!(memory.expires_at("key").unwrap(), after(expires_after)); + } + + #[rstest] + #[case::set(set)] + #[case::set_if(set_if)] + #[case::update(update)] + #[case::increment(increment)] + #[case::refresh(refresh)] + fn a_write_to_a_new_key_takes_its_ttl(#[case] write: Write) { + let (now, _) = clock(); + let memory = Memory::new(None, &now); + + write(&memory); + + assert_eq!(memory.get("key").unwrap(), Some(Number::Int(2))); + assert_eq!(memory.expires_at("key").unwrap(), after(50)); + } + + #[rstest] + #[case::nothing_held(None, false, 2)] + #[case::a_value_it_replaces(Some(Number::Int(1)), true, 2)] + #[case::a_value_it_does_not_replace(Some(Number::Int(1)), false, 1)] + fn set_if_writes_over( + #[case] held: Option, + #[case] replaces: bool, + #[case] expected: i64, + ) { + let (now, _) = clock(); + let memory = Memory::new(None, &now); + if let Some(held) = held { + memory.set("key", held, None).unwrap(); + } + + memory + .set_if("key", Number::Int(2), None, |_| replaces) + .unwrap(); + + assert_eq!(memory.get("key").unwrap(), Some(Number::Int(expected))); + } } diff --git a/litellm/rust/crates/redis-store/src/request.rs b/litellm/rust/crates/redis-store/src/request.rs index 93101ea01be..fbe0dc63b86 100644 --- a/litellm/rust/crates/redis-store/src/request.rs +++ b/litellm/rust/crates/redis-store/src/request.rs @@ -1,17 +1,11 @@ -use std::{ - collections::HashSet, - future::Future, - sync::{Arc, Mutex}, - time::Duration, -}; +use std::{collections::HashSet, future::Future, sync::Arc, time::Duration}; use futures_util::future::{BoxFuture, join_all}; use litellm_redis_client::{CallTag, Declared, RedisBatch, RedisClient, current_tag, tagged}; +use parking_lot::Mutex; use tokio::{runtime::Handle, sync::oneshot, task::JoinHandle}; use tracing::{Dispatch, Instrument, instrument::WithSubscriber}; -use crate::memory::lock; - /// One request's Redis batches: a pre-call batch for the reads it needs before answering and a /// post-call batch for the writes nobody waits on, one of each per Redis client. The host that owns /// the request builds one and passes it down; it knows nothing else about the request. @@ -43,15 +37,15 @@ pub(crate) struct Misses(Arc>>); impl Misses { pub(crate) fn contains(&self, key: &str) -> bool { - lock(&self.0).contains(key) + self.0.lock().contains(key) } pub(crate) fn note(&self, keys: impl IntoIterator) { - lock(&self.0).extend(keys); + self.0.lock().extend(keys); } pub(crate) fn forget(&self, key: &str) { - lock(&self.0).remove(key); + self.0.lock().remove(key); } } @@ -105,7 +99,7 @@ impl RequestRedis { .0 .runtime .spawn(self.logged(within(current_tag(), task))); - lock(&self.0.follow_ups).push(follow_up); + self.0.follow_ups.lock().push(follow_up); } /// Runs `task` on the request's runtime without `flush_post` waiting for it, for work that @@ -129,8 +123,8 @@ impl RequestRedis { /// Sends every post-call op declared so far and waits for the follow-ups spawned on them; the /// shutdown drain calls it too. pub async fn flush_post(&self) { - lock(&self.0.deadline).take(); - let follow_ups = std::mem::take(&mut *lock(&self.0.follow_ups)); + self.0.deadline.lock().take(); + let follow_ups = std::mem::take(&mut *self.0.follow_ups.lock()); join_all(self.batches().iter().map(|batches| batches.post.flush())).await; join_all(follow_ups) .await @@ -140,11 +134,11 @@ impl RequestRedis { } fn batches(&self) -> Vec { - lock(&self.0.batches).clone() + self.0.batches.lock().clone() } fn batches_for(&self, client: &RedisClient) -> ClientBatches { - let mut batches = lock(&self.0.batches); + let mut batches = self.0.batches.lock(); if let Some(existing) = batches .iter() .find(|batches| batches.client.same_client(client)) @@ -162,13 +156,13 @@ impl RequestRedis { } fn arm_deadline(&self) { - if lock(&self.0.deadline).is_some() { + if self.0.deadline.lock().is_some() { return; } // Asked before taking the lock: the host may answer under its interpreter lock, which // another declaring thread can hold while it waits for this one. let tag = (self.0.deadline_tag)(); - let mut deadline = lock(&self.0.deadline); + let mut deadline = self.0.deadline.lock(); if deadline.is_some() { return; } diff --git a/litellm/rust/crates/redis-store/src/router_cooldowns/mod.rs b/litellm/rust/crates/redis-store/src/router_cooldowns/mod.rs index 69493cbf7ce..ca958512c1a 100644 --- a/litellm/rust/crates/redis-store/src/router_cooldowns/mod.rs +++ b/litellm/rust/crates/redis-store/src/router_cooldowns/mod.rs @@ -191,14 +191,10 @@ impl RouterCooldowns { memory.delete(&key)?; return Ok(false); } - let remaining = Duration::from_secs_f64(remaining); - if memory.expires_at(&key)? > Some(now + remaining + CORRECTION_SLACK) { - memory.delete(&key)?; - memory.set( - &key, - cooldown.clone(), - Some(remaining.min(MAX_CORRECTED_TTL)), - )?; + let remaining = Duration::try_from_secs_f64(remaining).unwrap_or(Duration::MAX); + let ends = now.saturating_add(remaining); + if memory.expires_at(&key)? > Some(ends.saturating_add(CORRECTION_SLACK)) { + memory.refresh(&key, cooldown.clone(), remaining.min(MAX_CORRECTED_TTL))?; } Ok(true) } diff --git a/litellm/rust/crates/redis-store/src/router_session_pins/mod.rs b/litellm/rust/crates/redis-store/src/router_session_pins/mod.rs index 97e66d3a038..d3f3fa8263d 100644 --- a/litellm/rust/crates/redis-store/src/router_session_pins/mod.rs +++ b/litellm/rust/crates/redis-store/src/router_session_pins/mod.rs @@ -178,8 +178,7 @@ impl RouterSessionPins { } fn replace_local(&self, key: &str, value: Value, ttl: Duration) -> Result<(), Error> { - self.pins.memory().delete(key)?; - Ok(self.pins.memory().set(key, value, Some(ttl))?) + Ok(self.pins.memory().refresh(key, value, ttl)?) } } diff --git a/litellm/rust/crates/redis-store/src/spend_counters/mod.rs b/litellm/rust/crates/redis-store/src/spend_counters/mod.rs index 99ab9bea8f7..7ef4e0bb4a0 100644 --- a/litellm/rust/crates/redis-store/src/spend_counters/mod.rs +++ b/litellm/rust/crates/redis-store/src/spend_counters/mod.rs @@ -266,17 +266,9 @@ impl SpendCounters { spend: Number, ) -> Result, Error> { let key = counter.key(); - { - let _guard = self.cache.memory().guard(); - if self - .cache - .memory() - .get(&key)? - .is_none_or(|cached| cached.as_f64() < spend.as_f64()) - { - self.cache.memory().set(&key, spend, None)?; - } - } + self.cache + .memory() + .set_if(&key, spend, None, |cached| cached.as_f64() < spend.as_f64())?; let Some(remote) = self.cache.redis().get() else { return Ok(None); }; diff --git a/litellm/rust/crates/redis-store/src/test_support.rs b/litellm/rust/crates/redis-store/src/test_support.rs new file mode 100644 index 00000000000..092a9184481 --- /dev/null +++ b/litellm/rust/crates/redis-store/src/test_support.rs @@ -0,0 +1,19 @@ +use std::{ + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, +}; + +use crate::Clock; + +pub(crate) const START: u64 = 1_000; + +/// A clock at `START` seconds that a test moves by adding to the counter. +pub(crate) fn clock() -> (Clock, Arc) { + let seconds = Arc::new(AtomicU64::new(START)); + let read = seconds.clone(); + let clock: Clock = Arc::new(move || Duration::from_secs(read.load(Ordering::SeqCst))); + (clock, seconds) +} diff --git a/litellm/rust/crates/redis-store/tests/router_cooldowns.rs b/litellm/rust/crates/redis-store/tests/router_cooldowns.rs index 99b6b12bbf3..48ae6e2eb80 100644 --- a/litellm/rust/crates/redis-store/tests/router_cooldowns.rs +++ b/litellm/rust/crates/redis-store/tests/router_cooldowns.rs @@ -257,6 +257,87 @@ async fn a_router_without_redis_keeps_its_cooldowns_and_counts_in_memory() { ); } +#[rstest] +#[case::past_the_end_of_duration(1e300)] +#[case::infinite(f64::INFINITY)] +#[tokio::test] +async fn a_cooldown_too_long_to_measure_stays_active(#[case] cooldown_time: f64) { + let store = std::sync::Arc::new(RouterCooldowns::with_clock(None, POLICY, fixture_clock())); + let [case, _] = cooldowns(); + let cooldown = Cooldown { + cooldown_time: Number::Float(cooldown_time), + ..case.cooldown + }; + store + .put(case.model_id, &cooldown, Duration::from_secs(90)) + .await + .unwrap(); + + assert_eq!( + read_active(&store, &[case.model_id.to_owned()]).await, + vec![(case.model_id.to_owned(), cooldown)] + ); +} + +#[rstest] +#[case::outlasting_the_cooldown_is_shortened_to_its_end(60, 90.0, &[])] +#[case::within_the_slack_of_its_end_is_left_alone(60, 64.5, &["deployment:dep-a:cooldown"])] +#[case::of_a_long_cooldown_is_shortened_to_a_minute(600, 900.0, &[])] +#[tokio::test] +async fn reading_an_active_cooldown_corrects_a_memory_entry( + #[case] cooldown_time: i64, + #[case] memory_ttl: f64, + #[case] held_61_seconds_later: &[&str], +) { + let clock = StepClock::at(1_730_000_000.75); + let store = std::sync::Arc::new(RouterCooldowns::with_clock(None, POLICY, clock.clock())); + let [case, _] = cooldowns(); + let cooldown = Cooldown { + cooldown_time: Number::Int(cooldown_time), + ..case.cooldown + }; + store + .put( + case.model_id, + &cooldown, + Duration::from_secs_f64(memory_ttl), + ) + .await + .unwrap(); + + let active = read_active(&store, &[case.model_id.to_owned()]).await; + clock.advance(Duration::from_secs(61)); + let held = store + .snapshot() + .unwrap() + .into_iter() + .map(|(key, _)| key) + .collect::>(); + + assert_eq!(active, vec![(case.model_id.to_owned(), cooldown)]); + assert_eq!(held, held_61_seconds_later); +} + +#[rstest] +#[tokio::test] +async fn reading_an_ended_cooldown_drops_it_from_memory() { + let clock = StepClock::at(1_730_000_000.75); + let store = std::sync::Arc::new(RouterCooldowns::with_clock(None, POLICY, clock.clock())); + let [case, _] = cooldowns(); + store + .put(case.model_id, &case.cooldown, Duration::from_secs(90)) + .await + .unwrap(); + clock.advance(Duration::from_secs(61)); + let held = store.snapshot().unwrap().len(); + + let active = read_active(&store, &[case.model_id.to_owned()]).await; + + assert_eq!(held, 1); + assert_eq!(active, vec![]); + assert_eq!(store.snapshot().unwrap(), vec![]); +} + #[rstest] #[tokio::test] async fn a_redis_outage_serves_memory_and_counts_fails_locally() {