scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(cache-memory): rebuild the in-memory cache on an exact expiry index (#51388)

Browse source 
The cache copied Python's lazy expiry heap. Deleting or re-arming a key
left a stale heap record, and a longer-lived key at the top kept those
records in memory without bound. A TTL past the end of Duration also
panicked while the lock was held, which poisoned it for every later call

The cache now keeps one index record per entry in a BTreeSet and removes
it with the entry, so there is nothing to compact. A write picks
Expiry::KeepExisting or Expiry::Refresh, update() reads and writes under
one lock, and deadlines saturate

The Redis stores' memory tier uses update() and drops its second mutex.
RouterCooldowns no longer panics on a cooldown too long to measure, and
both crates lock with parking_lot
This commit is contained in:
yujonglee 2026-11-14 06:16:33 +09:00 • committed by GitHub
parent a920b5be6e
commit 9986dab8a1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 1610 additions and 1131 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/rust/Cargo.lock generated
View file

@ -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",

1
litellm/rust/crates/cache-memory/Cargo.toml
View file

@ -7,6 +7,7 @@ repository.workspace = true
[dependencies]
litellm-cache.workspace = true
parking_lot.workspace = true
[dev-dependencies]
litellm-cache-testing.workspace = true

713
litellm/rust/crates/cache-memory/src/cache.rs
View file

@ -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<V> {
values: HashMap<String, V>,
expirations: HashMap<String, Duration>,
expiration_heap: BinaryHeap<Reverse<(Duration, String)>>,
/// 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<Duration>),
Refresh(Option<Duration>),
}
pub struct InMemoryCache<V: Clone> {
state: Mutex<CacheState<V>>,
state: Mutex<State<V>>,
max_size_in_memory: AtomicUsize,
default_ttl: Duration,
max_entry_bytes: Option<usize>,
@ -72,11 +71,7 @@ impl<V: Clone> InMemoryCache<V> {
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<V: Clone> InMemoryCache<V> {
}
}
fn lock(&self) -> Result<MutexGuard<'_, CacheState<V>>, Error> {
Ok(self.state.lock()?)
}
pub fn set_cache(
&self,
key: impl Into<String>,
value: V,
ttl: Option<Duration>,
expiry: Expiry,
) -> Result<CacheWrite, Error> {
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<Option<V>, 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<T>(
&self,
key: &str,
expiry: Expiry,
update: impl FnOnce(Option<&V>) -> Result<(Option<V>, T), Error>,
) -> Result<T, Error> {
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<bool, Error> {
match (self.max_entry_bytes, &self.measure_value) {
(Some(limit), Some(measure)) => Ok(measure(value)? <= limit),
@ -129,57 +124,21 @@ impl<V: Clone> InMemoryCache<V> {
}
}
/// `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<bool, Error> {
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<bool, Error> {
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<usize, Error> {
Ok(self.lock()?.values.len())
let now = (self.now)();
Ok(self.state.lock().len(now))
}
pub fn is_empty(&self) -> Result<bool, Error> {
Ok(self.len()? == 0)
}
/// Entries in the expiration heap, stale ones included; bounded by eviction.
pub fn expiration_heap_len(&self) -> Result<usize, Error> {
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<usize>) {
self.max_size_in_memory.store(
max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
@ -187,12 +146,13 @@ impl<V: Clone> InMemoryCache<V> {
);
}
pub fn max_entry_bytes(&self) -> Option<usize> {
self.max_entry_bytes
pub fn default_ttl(&self) -> Duration {
self.default_ttl
}
pub fn expires_at(&self, key: &str) -> Result<Option<Duration>, 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<Option<Duration>, Error> {
@ -200,253 +160,406 @@ impl<V: Clone> InMemoryCache<V> {
}
pub async fn async_get_oldest_n_keys(&self, count: usize) -> Result<Vec<String>, Error> {
let state = self.lock()?;
let mut expirations = state
.expirations
.iter()
.map(|(key, expiration)| (key.clone(), *expiration))
.collect::<Vec<_>>();
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<Vec<(String, V)>, 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<V>, 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<V>, 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<V>, 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<V>, key: &str, now: Duration) -> Option<V> {
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<V>,
state: &mut State<V>,
key: String,
value: V,
ttl: Option<Duration>,
expiry: Expiry,
now: Duration,
) -> Result<CacheWrite, Error> {
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<Duration>| 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<V> ClaimCache for InMemoryCache<V>
where
V: Clone + PartialEq + Send + Sync + 'static,
{
fn claim_cache(
&self,
key: &str,
candidate: V,
eligible: &[V],
context: ExactCacheContext,
) -> Result<V, Error> {
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<String> {
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>, 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<Duration>,
#[case] expected: Option<Duration>,
) {
let cache = InMemoryCache::<String>::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::<String>(&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<Duration>,
) {
let cache = cache::<String>(&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::<String>(&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::<String>(&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::<String>(&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<f64> {
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
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<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
type Value = V;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
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<Option<Self::Value>, Error> {
self.get_cache(key)
}
}
impl<V: Clone + Send + Sync + 'static> DisconnectCache for InMemoryCache<V> {
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
}
impl<V: Clone + Send + Sync + 'static> BatchCache for InMemoryCache<V> {}
impl<V: Clone + Send + Sync + 'static> DeleteCache for InMemoryCache<V> {
fn delete_cache(&self, key: &str) -> Result<(), Error> {
InMemoryCache::delete_cache(self, key)
}
}
impl<V: Clone + Send + Sync + 'static> FlushCache for InMemoryCache<V> {
fn flush_cache(&self) -> Result<(), Error> {
InMemoryCache::flush_cache(self)
}
}
impl<V: Clone + Send + Sync + 'static> TtlCache for InMemoryCache<V> {
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
InMemoryCache::async_get_ttl(self, key).await
}
}
impl<T> SetCache for InMemoryCache<HashSet<T>>
where
T: Clone + Eq + Hash + Send + Sync + 'static,
{
type SetValue = T;
type SetResult = Vec<T>;
async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
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::<String>(&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::<String>(&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::<String>(&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::<String>(&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::<u64>(&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<bool, Error>,
) {
assert_eq!(measured(2).check_value_size(&value.to_owned()), expected);
}
#[rstest]
fn values_are_unbounded_without_a_measure() {
let cache = InMemoryCache::<String>::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<CacheWrite, Error>,
) {
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::<u64>(&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::<u64>(&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::<u64>(&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::<String>(&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::<String>(&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::<String>(&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::<String>(&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));
}
}

528
litellm/rust/crates/cache-memory/src/capabilities.rs Normal file
View file

@ -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<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
type Value = V;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
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<Option<Self::Value>, Error> {
InMemoryCache::get_cache(self, key)
}
}
impl<V> ClaimCache for InMemoryCache<V>
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<V, Error> {
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<f64> {
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
self.update(key, Expiry::KeepExisting(context.ttl), |held| {
let total = held.copied().unwrap_or_default() + amount;
Ok((Some(total), total))
})
}
}
impl<V: Clone + Send + Sync + 'static> DisconnectCache for InMemoryCache<V> {
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
}
impl<V: Clone + Send + Sync + 'static> BatchCache for InMemoryCache<V> {}
impl<V: Clone + Send + Sync + 'static> DeleteCache for InMemoryCache<V> {
fn delete_cache(&self, key: &str) -> Result<(), Error> {
InMemoryCache::delete_cache(self, key)
}
}
impl<V: Clone + Send + Sync + 'static> FlushCache for InMemoryCache<V> {
fn flush_cache(&self) -> Result<(), Error> {
InMemoryCache::flush_cache(self)
}
}
impl<V: Clone + Send + Sync + 'static> TtlCache for InMemoryCache<V> {
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
InMemoryCache::async_get_ttl(self, key).await
}
}
impl<T> SetCache for InMemoryCache<HashSet<T>>
where
T: Clone + Eq + Hash + Send + Sync + 'static,
{
type SetValue = T;
type SetResult = Vec<T>;
async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
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<InMemoryCache<String>> = 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<Duration>,
) {
let cache = cache::<String>(&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::<String>(&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::<String>::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::<String>(&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::<String>(&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::<String>(&clock, 4);
let eligible = eligible
.iter()
.copied()
.map(String::from)
.collect::<Vec<_>>();
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::<String>(&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::<f64>::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::<f64>::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::<String>(&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<Duration>,
) {
assert_eq!(cache::<String>(&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::<f64>(&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::<f64>(&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::<f64>(&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::<String>::new(Some(0), None);
let counters = InMemoryCache::<f64>::new(Some(0), None);
let sets = InMemoryCache::<HashSet<String>>::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::<f64>(&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::<HashSet<String>>(&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::<HashSet<String>>(&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::<HashSet<String>>::with_clock_and_size_measurement(
Some(4),
None,
Some(2),
Some(Arc::new(|value: &HashSet<String>| 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()])))
);
}
}

6
litellm/rust/crates/cache-memory/src/lib.rs
View file

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

320
litellm/rust/crates/cache-memory/src/state.rs Normal file
View file

@ -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<V> {
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<V> {
entries: HashMap<String, Entry<V>>,
by_expiry: BTreeSet<(Duration, String)>,
}
impl<V> State<V> {
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<Duration> {
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<Item = (&str, &V)> {
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<Item = &str> {
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<Entry<V>> {
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::<Vec<_>>(),
["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::<Vec<_>>(), ["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::<Vec<_>>(),
["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::<Vec<_>>(),
["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::<Vec<_>>(),
["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::<Vec<_>>(), kept);
}
#[rstest]
fn an_expired_entry_is_not_listed(mut full: State<&'static str>) {
assert_eq!(full.live(secs(10)).collect::<Vec<_>>(), [("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::<Vec<_>>(), ["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"));
}
}

33
litellm/rust/crates/cache-memory/src/test_support.rs Normal file
View file

@ -0,0 +1,33 @@
use std::{
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::Duration,
};
use rstest::fixture;
use crate::InMemoryCache;
pub(crate) type Clock = Arc<AtomicU64>;
#[fixture]
pub(crate) fn clock() -> Clock {
Arc::new(AtomicU64::new(100))
}
pub(crate) fn cache<V: Clone>(clock: &Clock, capacity: usize) -> InMemoryCache<V> {
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<Duration> {
Some(Duration::from_secs(seconds))
}

731
litellm/rust/crates/cache-memory/tests/cache.rs
View file

@ -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<AtomicU64>;
#[fixture]
fn clock() -> Clock {
Arc::new(AtomicU64::new(100))
}
fn cache_with<V: Clone>(clock: &Clock, capacity: usize) -> InMemoryCache<V> {
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<String> {
cache_with(clock, capacity)
}
fn at(clock: &Clock, seconds: u64) {
clock.store(seconds, Ordering::SeqCst);
}
fn secs(seconds: u64) -> Option<Duration> {
Some(Duration::from_secs(seconds))
}
fn ttl(seconds: u64) -> ExactCacheContext {
ExactCacheContext { ttl: secs(seconds) }
}
fn measured(capacity: usize) -> InMemoryCache<String> {
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<Duration>,
#[case] expected: Option<Duration>,
) {
let cache = InMemoryCache::<String>::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::<f64>::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<bool, Error>,
) {
assert_eq!(measured(2).check_value_size(&value.to_string()), expected);
}
#[rstest]
fn values_are_unbounded_without_a_measure() {
let cache = InMemoryCache::<String>::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::<String>::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<InMemoryCache<String>> = 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<Duration>,
) {
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::<String>::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::<f64>::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::<f64>::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::<Vec<_>>();
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::<f64>(&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::<f64>(&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::<f64>(&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::<String>::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::<f64>::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::<HashSet<String>>::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::<String>::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::<f64>(&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::<f64>::new()
);
}
#[rstest]
#[tokio::test]
async fn set_capability_preserves_python_result_and_deduplicates_storage(clock: Clock) {
let cache = cache_with::<HashSet<String>>(&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::<HashSet<String>>::with_clock_and_size_measurement(
Some(4),
None,
Some(2),
Some(Arc::new(|value: &HashSet<String>| 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()]))
);
}

1
litellm/rust/crates/redis-store/Cargo.toml
View file

@ -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

12
litellm/rust/crates/redis-store/src/auth_cache/mod.rs
View file

@ -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());
}

8
litellm/rust/crates/redis-store/src/dual_cache/batch_read.rs
View file

@ -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<V: Clone + FromValue + std::marker::Send + Sync + 'static> DualCache<V, Bat
fn reserve(&self, keys: &[String], local: &[Option<V>]) -> 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<V: Clone + FromValue + std::marker::Send + Sync + 'static> DualCache<V, Bat
}
pub(crate) fn release(&self, reservation: Reservation) {
let mut recent = lock(&self.reads.recent);
let mut recent = self.reads.recent.lock();
for (key, previous) in reservation.keys.into_iter().zip(reservation.previous) {
match previous {
Some(read) => recent.insert(key, read),

46
litellm/rust/crates/redis-store/src/dual_cache/mod.rs
View file

@ -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<V: Clone, R> DualCache<V, R> {
}
pub(crate) fn default_in_memory_ttl(&self) -> Option<Duration> {
*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<V: Clone, R> DualCache<V, R> {
#[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<AtomicU64>) {
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<V: Clone>(
default_in_memory_ttl: Option<Duration>,
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]

2
litellm/rust/crates/redis-store/src/lib.rs
View file

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

179
litellm/rust/crates/redis-store/src/memory.rs
View file

@ -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<V: Clone> {
cache: InMemoryCache<V>,
updates: Mutex<()>,
}
impl<V: Clone> Memory<V> {
@ -30,7 +28,6 @@ impl<V: Clone> Memory<V> {
let now = now.clone();
Self {
cache: InMemoryCache::with_clock(entries, None, move || now()),
updates: Mutex::new(()),
}
}
@ -38,8 +35,30 @@ impl<V: Clone> Memory<V> {
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<Duration>) -> 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<Duration>,
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<V: Clone> Memory<V> {
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<Duration>,
update: impl FnOnce(Option<V>) -> Result<V, Error>,
) -> Result<V, Error> {
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<Number> {
fn as_number(&self) -> Option<Number>;
}
@ -98,25 +112,130 @@ impl Counter for Number {
}
impl<V: Counter> Memory<V> {
/// `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<Duration>,
) -> Result<Number, Error> {
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<T>(mutex: &Mutex<T>) -> 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<Number>);
fn secs(seconds: u64) -> Duration {
Duration::from_secs(seconds)
}
fn ttl(seconds: u64) -> Option<Duration> {
Some(secs(seconds))
}
fn after(seconds: u64) -> Option<Duration> {
ttl(START + seconds)
}
fn set(memory: &Memory<Number>) {
memory.set("key", Number::Int(2), ttl(50)).unwrap();
}
fn set_if(memory: &Memory<Number>) {
memory
.set_if("key", Number::Int(2), ttl(50), |_| true)
.unwrap();
}
fn update(memory: &Memory<Number>) {
memory
.update("key", ttl(50), |_| Ok(Number::Int(2)))
.unwrap();
}
fn increment(memory: &Memory<Number>) {
memory.increment("key", Number::Int(2), ttl(50)).unwrap();
}
fn refresh(memory: &Memory<Number>) {
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<Number>,
#[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)));
}
}

30
litellm/rust/crates/redis-store/src/request.rs
View file

@ -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<Mutex<HashSet<String>>>);
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<Item = String>) {
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<ClientBatches> {
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;
}

12
litellm/rust/crates/redis-store/src/router_cooldowns/mod.rs
View file

@ -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)
}

3
litellm/rust/crates/redis-store/src/router_session_pins/mod.rs
View file

@ -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)?)
}
}

14
litellm/rust/crates/redis-store/src/spend_counters/mod.rs
View file

@ -266,17 +266,9 @@ impl SpendCounters {
spend: Number,
) -> Result<Option<Number>, 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);
};

19
litellm/rust/crates/redis-store/src/test_support.rs Normal file
View file

@ -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<AtomicU64>) {
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)
}

81
litellm/rust/crates/redis-store/tests/router_cooldowns.rs
View file

@ -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::<Vec<_>>();
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() {