mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
refactor(rust): remove generic core cache
This commit is contained in:
Yujong Lee
parent
5b7ba86c30
commit
b1310a1bf8
4 changed files with 50 additions and 269 deletions
|
|
@ -1,258 +0,0 @@
|
|||
use std::cmp::Reverse;
|
||||
use std::collections::{BinaryHeap, HashMap};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
|
||||
const DEFAULT_TTL: Duration = Duration::from_secs(600);
|
||||
|
||||
pub struct InMemoryCache<V: Clone> {
|
||||
pub cache_dict: HashMap<String, V>,
|
||||
pub ttl_dict: HashMap<String, Duration>,
|
||||
pub expiration_heap: BinaryHeap<Reverse<(Duration, String)>>,
|
||||
pub max_size_in_memory: usize,
|
||||
pub default_ttl: Duration,
|
||||
now: Box<dyn Fn() -> Duration + Send + Sync>,
|
||||
}
|
||||
|
||||
impl<V: Clone> Default for InMemoryCache<V> {
|
||||
fn default() -> Self {
|
||||
Self::new(None, None)
|
||||
}
|
||||
}
|
||||
|
||||
impl<V: Clone> InMemoryCache<V> {
|
||||
pub fn new(max_size_in_memory: Option<usize>, default_ttl: Option<Duration>) -> Self {
|
||||
Self::with_clock(max_size_in_memory, default_ttl, || {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn with_clock(
|
||||
max_size_in_memory: Option<usize>,
|
||||
default_ttl: Option<Duration>,
|
||||
now: impl Fn() -> Duration + Send + Sync + 'static,
|
||||
) -> Self {
|
||||
Self {
|
||||
cache_dict: HashMap::new(),
|
||||
ttl_dict: HashMap::new(),
|
||||
expiration_heap: BinaryHeap::new(),
|
||||
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
|
||||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
now: Box::new(now),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn evict_cache(&mut self) {
|
||||
if self.max_size_in_memory == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
let current_time = (self.now)();
|
||||
while let Some(Reverse((expiration_time, key))) = self.expiration_heap.peek().cloned() {
|
||||
if self.ttl_dict.get(&key).copied() != Some(expiration_time) {
|
||||
self.expiration_heap.pop();
|
||||
} else if expiration_time <= current_time {
|
||||
self.expiration_heap.pop();
|
||||
self.remove_key(&key);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
while self.cache_dict.len() >= self.max_size_in_memory {
|
||||
let Some(Reverse((expiration_time, key))) = self.expiration_heap.pop() else {
|
||||
break;
|
||||
};
|
||||
if self.ttl_dict.get(&key).copied() == Some(expiration_time) {
|
||||
self.remove_key(&key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn allow_ttl_override(&self, key: &str) -> bool {
|
||||
match self.ttl_dict.get(key).copied() {
|
||||
None => true,
|
||||
Some(expiration_time) => expiration_time < (self.now)(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_cache(&mut self, key: impl Into<String>, value: V, ttl: Option<Duration>) {
|
||||
if self.max_size_in_memory == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
self.evict_cache();
|
||||
let key = key.into();
|
||||
self.cache_dict.insert(key.clone(), value);
|
||||
if self.allow_ttl_override(&key) {
|
||||
let expiration_time = (self.now)() + ttl.unwrap_or(self.default_ttl);
|
||||
self.ttl_dict.insert(key.clone(), expiration_time);
|
||||
self.expiration_heap.push(Reverse((expiration_time, key)));
|
||||
}
|
||||
}
|
||||
|
||||
// Generic values intentionally omit Python's per-item size check.
|
||||
pub fn get_cache(&mut self, key: &str) -> Option<V> {
|
||||
if self.cache_dict.contains_key(key) {
|
||||
if self.is_key_expired(key) {
|
||||
self.remove_key(key);
|
||||
return None;
|
||||
}
|
||||
return self.cache_dict.get(key).cloned();
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
pub fn get_ttl(&self, key: &str) -> Option<Duration> {
|
||||
self.ttl_dict.get(key).copied()
|
||||
}
|
||||
|
||||
pub fn delete_cache(&mut self, key: &str) {
|
||||
self.remove_key(key);
|
||||
}
|
||||
|
||||
pub fn flush_cache(&mut self) {
|
||||
self.cache_dict.clear();
|
||||
self.ttl_dict.clear();
|
||||
self.expiration_heap.clear();
|
||||
}
|
||||
|
||||
fn is_key_expired(&self, key: &str) -> bool {
|
||||
self.ttl_dict
|
||||
.get(key)
|
||||
.is_some_and(|expiration_time| *expiration_time < (self.now)())
|
||||
}
|
||||
|
||||
fn remove_key(&mut self, key: &str) {
|
||||
self.cache_dict.remove(key);
|
||||
self.ttl_dict.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
};
|
||||
|
||||
use super::InMemoryCache;
|
||||
use std::time::Duration;
|
||||
|
||||
fn cache(now: Arc<AtomicU64>, max_size: usize, default_ttl: Duration) -> InMemoryCache<String> {
|
||||
InMemoryCache::with_clock(Some(max_size), Some(default_ttl), move || {
|
||||
Duration::from_secs(now.load(Ordering::Relaxed))
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ttl_expiry_is_deterministic() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
|
||||
cache.set_cache("key", "value".to_string(), None);
|
||||
assert_eq!(cache.get_cache("key"), Some("value".to_string()));
|
||||
now.store(161, Ordering::Relaxed);
|
||||
assert_eq!(cache.get_cache("key"), None);
|
||||
assert_eq!(cache.get_ttl("key"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_and_per_set_ttl_are_applied() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
|
||||
cache.set_cache("default", "value".to_string(), None);
|
||||
cache.set_cache("custom", "value".to_string(), Some(Duration::from_secs(20)));
|
||||
assert_eq!(cache.get_ttl("default"), Some(Duration::from_secs(160)));
|
||||
assert_eq!(cache.get_ttl("custom"), Some(Duration::from_secs(120)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unexpired_entries_do_not_allow_ttl_override() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
|
||||
cache.set_cache("key", "first".to_string(), Some(Duration::from_secs(20)));
|
||||
cache.set_cache("key", "second".to_string(), Some(Duration::from_secs(80)));
|
||||
assert_eq!(cache.get_cache("key"), Some("second".to_string()));
|
||||
assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(120)));
|
||||
now.store(121, Ordering::Relaxed);
|
||||
cache.set_cache("key", "third".to_string(), Some(Duration::from_secs(80)));
|
||||
assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(201)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn max_size_evicts_earliest_expiration() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now, 2, Duration::from_secs(60));
|
||||
cache.set_cache("early", "value".to_string(), Some(Duration::from_secs(10)));
|
||||
cache.set_cache("late", "value".to_string(), Some(Duration::from_secs(20)));
|
||||
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
|
||||
assert_eq!(cache.get_cache("early"), None);
|
||||
assert!(cache.get_cache("late").is_some());
|
||||
assert!(cache.get_cache("new").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expired_entries_are_evicted_before_live_entries() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now.clone(), 3, Duration::from_secs(60));
|
||||
cache.set_cache(
|
||||
"expired-one",
|
||||
"value".to_string(),
|
||||
Some(Duration::from_secs(10)),
|
||||
);
|
||||
cache.set_cache(
|
||||
"expired-two",
|
||||
"value".to_string(),
|
||||
Some(Duration::from_secs(20)),
|
||||
);
|
||||
cache.set_cache("live", "value".to_string(), Some(Duration::from_secs(100)));
|
||||
now.store(121, Ordering::Relaxed);
|
||||
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(100)));
|
||||
assert_eq!(cache.get_cache("expired-one"), None);
|
||||
assert_eq!(cache.get_cache("expired-two"), None);
|
||||
assert!(cache.get_cache("live").is_some());
|
||||
assert!(cache.get_cache("new").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_heap_entries_are_skipped() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now, 1, Duration::from_secs(60));
|
||||
cache.set_cache(
|
||||
"removed",
|
||||
"value".to_string(),
|
||||
Some(Duration::from_secs(10)),
|
||||
);
|
||||
cache.delete_cache("removed");
|
||||
cache.set_cache("kept", "value".to_string(), Some(Duration::from_secs(20)));
|
||||
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
|
||||
assert_eq!(cache.get_cache("removed"), None);
|
||||
assert_eq!(cache.get_cache("kept"), None);
|
||||
assert!(cache.get_cache("new").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn delete_and_flush_remove_values_and_ttls() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now, 10, Duration::from_secs(60));
|
||||
cache.set_cache("one", "value".to_string(), None);
|
||||
cache.set_cache("two", "value".to_string(), None);
|
||||
cache.delete_cache("one");
|
||||
assert_eq!(cache.get_cache("one"), None);
|
||||
cache.flush_cache();
|
||||
assert!(cache.cache_dict.is_empty());
|
||||
assert!(cache.ttl_dict.is_empty());
|
||||
assert!(cache.expiration_heap.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_max_size_does_not_cache() {
|
||||
let now = Arc::new(AtomicU64::new(100));
|
||||
let mut cache = cache(now, 0, Duration::from_secs(60));
|
||||
cache.set_cache("key", "value".to_string(), None);
|
||||
assert_eq!(cache.get_cache("key"), None);
|
||||
assert!(cache.cache_dict.is_empty());
|
||||
}
|
||||
}
|
||||
|
|
@ -1 +0,0 @@
|
|||
pub mod in_memory_cache;
|
||||
|
|
@ -1,5 +1,4 @@
|
|||
pub mod audio_transcription;
|
||||
pub mod caching;
|
||||
pub mod chat_completions;
|
||||
pub mod constants;
|
||||
pub mod error;
|
||||
|
|
|
|||
|
|
@ -1,9 +1,8 @@
|
|||
use std::collections::BTreeMap;
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::{BTreeMap, BinaryHeap, HashMap};
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
use std::time::Duration;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::caching::in_memory_cache::InMemoryCache;
|
||||
use crate::error::Error;
|
||||
use aws_credential_types::Credentials;
|
||||
use aws_credential_types::provider::ProvideCredentials;
|
||||
|
|
@ -25,8 +24,50 @@ use super::constants::{
|
|||
|
||||
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
|
||||
const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600);
|
||||
const MAX_CACHED_CREDENTIALS: usize = 200;
|
||||
|
||||
static IAM_CREDENTIALS_CACHE: OnceLock<Mutex<InMemoryCache<Credentials>>> = OnceLock::new();
|
||||
#[derive(Default)]
|
||||
struct CredentialCache {
|
||||
entries: HashMap<String, (Credentials, Duration)>,
|
||||
expirations: BinaryHeap<Reverse<(Duration, String)>>,
|
||||
}
|
||||
|
||||
impl CredentialCache {
|
||||
fn get(&mut self, key: &str) -> Option<Credentials> {
|
||||
let now = unix_time();
|
||||
let (credentials, expiration) = self.entries.get(key)?;
|
||||
if *expiration > now {
|
||||
return Some(credentials.clone());
|
||||
}
|
||||
self.entries.remove(key);
|
||||
None
|
||||
}
|
||||
|
||||
fn insert(&mut self, key: String, credentials: Credentials, ttl: Duration) {
|
||||
let now = unix_time();
|
||||
while let Some(Reverse((expiration, key))) = self.expirations.peek().cloned() {
|
||||
if self.entries.get(&key).map(|(_, current)| *current) != Some(expiration) {
|
||||
self.expirations.pop();
|
||||
} else if expiration <= now || self.entries.len() >= MAX_CACHED_CREDENTIALS {
|
||||
self.expirations.pop();
|
||||
self.entries.remove(&key);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let expiration = now + ttl;
|
||||
self.entries.insert(key.clone(), (credentials, expiration));
|
||||
self.expirations.push(Reverse((expiration, key)));
|
||||
}
|
||||
}
|
||||
|
||||
static IAM_CREDENTIALS_CACHE: OnceLock<Mutex<CredentialCache>> = OnceLock::new();
|
||||
|
||||
fn unix_time() -> Duration {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option<Duration> {
|
||||
match flow {
|
||||
|
|
@ -108,15 +149,15 @@ fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String {
|
|||
}
|
||||
|
||||
fn get_cached_credentials(key: &str) -> Option<Credentials> {
|
||||
let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default()));
|
||||
let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(CredentialCache::default()));
|
||||
let mut entries = cache.lock().ok()?;
|
||||
entries.get_cache(key)
|
||||
entries.get(key)
|
||||
}
|
||||
|
||||
fn set_cached_credentials(key: String, credentials: Credentials, ttl: Duration) {
|
||||
let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default()));
|
||||
let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(CredentialCache::default()));
|
||||
if let Ok(mut entries) = cache.lock() {
|
||||
entries.set_cache(key, credentials, Some(ttl));
|
||||
entries.insert(key, credentials, ttl);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue