litellm_rs/core/cache/
redis_cache.rs

1//! Redis cache implementation
2//!
3//! This module provides a Redis-based cache layer for distributed caching
4//! with serialization support using serde.
5
6use super::types::{
7    AtomicCacheStats, CacheEntry, CacheKey, DualCacheConfig, SerializableCacheEntry,
8};
9use crate::storage::redis::RedisPool;
10use crate::utils::error::gateway_error::{GatewayError, Result};
11use serde::{Serialize, de::DeserializeOwned};
12use std::marker::PhantomData;
13use std::sync::Arc;
14use std::time::Duration;
15use tracing::{trace, warn};
16
17/// Redis cache layer for distributed caching
18pub struct RedisCache<T> {
19    /// Redis connection pool
20    pool: Arc<RedisPool>,
21    /// Configuration
22    config: DualCacheConfig,
23    /// Statistics
24    stats: Arc<AtomicCacheStats>,
25    /// Phantom data for type parameter
26    _marker: PhantomData<T>,
27}
28
29impl<T> RedisCache<T>
30where
31    T: Serialize + DeserializeOwned + Clone + Send + Sync + 'static,
32{
33    /// Create a new Redis cache with the given pool and configuration
34    pub fn new(pool: Arc<RedisPool>, config: DualCacheConfig) -> Self {
35        Self {
36            pool,
37            config,
38            stats: Arc::new(AtomicCacheStats::new()),
39            _marker: PhantomData,
40        }
41    }
42
43    /// Create with shared statistics
44    pub fn with_stats(
45        pool: Arc<RedisPool>,
46        config: DualCacheConfig,
47        stats: Arc<AtomicCacheStats>,
48    ) -> Self {
49        Self {
50            pool,
51            config,
52            stats,
53            _marker: PhantomData,
54        }
55    }
56
57    /// Get a value from the cache
58    pub async fn get(&self, key: &CacheKey) -> Result<Option<T>> {
59        if self.pool.is_noop() {
60            self.stats.record_redis_miss();
61            return Ok(None);
62        }
63
64        let redis_key = self.make_redis_key(key);
65        match self.pool.get(&redis_key).await? {
66            Some(data) => {
67                match self.deserialize::<SerializableCacheEntry<T>>(&data) {
68                    Ok(entry) => {
69                        if entry.is_expired() {
70                            // Clean up expired entry
71                            let _ = self.pool.delete(&redis_key).await;
72                            self.stats.record_redis_miss();
73                            trace!(key = %key, "Redis cache entry expired");
74                            Ok(None)
75                        } else {
76                            self.stats.record_redis_hit();
77                            trace!(key = %key, "Redis cache hit");
78                            Ok(Some(entry.value))
79                        }
80                    }
81                    Err(e) => {
82                        warn!(key = %key, error = %e, "Failed to deserialize cache entry");
83                        // Clean up corrupted entry
84                        let _ = self.pool.delete(&redis_key).await;
85                        self.stats.record_redis_miss();
86                        Ok(None)
87                    }
88                }
89            }
90            None => {
91                self.stats.record_redis_miss();
92                trace!(key = %key, "Redis cache miss");
93                Ok(None)
94            }
95        }
96    }
97
98    /// Get an entry with metadata from the cache
99    pub async fn get_entry(&self, key: &CacheKey) -> Result<Option<CacheEntry<T>>> {
100        if self.pool.is_noop() {
101            self.stats.record_redis_miss();
102            return Ok(None);
103        }
104
105        let redis_key = self.make_redis_key(key);
106        match self.pool.get(&redis_key).await? {
107            Some(data) => match self.deserialize::<SerializableCacheEntry<T>>(&data) {
108                Ok(entry) => {
109                    if entry.is_expired() {
110                        let _ = self.pool.delete(&redis_key).await;
111                        self.stats.record_redis_miss();
112                        Ok(None)
113                    } else {
114                        self.stats.record_redis_hit();
115                        Ok(Some(entry.into_cache_entry()))
116                    }
117                }
118                Err(e) => {
119                    warn!(key = %key, error = %e, "Failed to deserialize cache entry");
120                    let _ = self.pool.delete(&redis_key).await;
121                    self.stats.record_redis_miss();
122                    Ok(None)
123                }
124            },
125            None => {
126                self.stats.record_redis_miss();
127                Ok(None)
128            }
129        }
130    }
131
132    /// Set a value in the cache with the default TTL
133    pub async fn set(&self, key: CacheKey, value: T) -> Result<()> {
134        self.set_with_ttl(key, value, self.config.default_ttl).await
135    }
136
137    /// Set a value in the cache with a specific TTL
138    pub async fn set_with_ttl(&self, key: CacheKey, value: T, ttl: Duration) -> Result<()> {
139        if self.pool.is_noop() {
140            return Ok(());
141        }
142
143        let entry = CacheEntry::new(value, ttl);
144        let serializable: SerializableCacheEntry<T> = (&entry).into();
145        let data = self.serialize(&serializable)?;
146        let redis_key = self.make_redis_key(&key);
147
148        self.pool
149            .set(&redis_key, &data, Some(ttl.as_secs()))
150            .await?;
151        self.stats.record_write();
152        trace!(key = %key, ttl_secs = ttl.as_secs(), "Redis cache set");
153
154        Ok(())
155    }
156
157    /// Set a value with size tracking
158    pub async fn set_with_size(
159        &self,
160        key: CacheKey,
161        value: T,
162        ttl: Duration,
163        size_bytes: usize,
164    ) -> Result<()> {
165        if self.pool.is_noop() {
166            return Ok(());
167        }
168
169        let entry = CacheEntry::with_size(value, ttl, size_bytes);
170        let serializable: SerializableCacheEntry<T> = (&entry).into();
171        let data = self.serialize(&serializable)?;
172        let redis_key = self.make_redis_key(&key);
173
174        self.pool
175            .set(&redis_key, &data, Some(ttl.as_secs()))
176            .await?;
177        self.stats.record_write();
178
179        Ok(())
180    }
181
182    /// Set a cache entry directly
183    pub async fn set_entry(&self, key: CacheKey, entry: CacheEntry<T>) -> Result<()> {
184        if self.pool.is_noop() {
185            return Ok(());
186        }
187
188        let ttl = entry.remaining_ttl().unwrap_or(Duration::from_secs(1));
189        let serializable: SerializableCacheEntry<T> = (&entry).into();
190        let data = self.serialize(&serializable)?;
191        let redis_key = self.make_redis_key(&key);
192
193        self.pool
194            .set(&redis_key, &data, Some(ttl.as_secs()))
195            .await?;
196        self.stats.record_write();
197
198        Ok(())
199    }
200
201    /// Delete a value from the cache
202    pub async fn delete(&self, key: &CacheKey) -> Result<bool> {
203        if self.pool.is_noop() {
204            return Ok(false);
205        }
206
207        let redis_key = self.make_redis_key(key);
208        let existed = self.pool.exists(&redis_key).await?;
209
210        if existed {
211            self.pool.delete(&redis_key).await?;
212            self.stats.record_deletion();
213            trace!(key = %key, "Redis cache delete");
214        }
215
216        Ok(existed)
217    }
218
219    /// Delete only while the live Redis value still satisfies `predicate`.
220    ///
221    /// Reads the raw blob, checks the deserialized value, then CAS-deletes with
222    /// Lua so a concurrent replacement is not removed. A CAS miss is re-read and
223    /// retried: an in-flight rewrite of still-matching poison changes the
224    /// serialized blob and would otherwise return false while L2 stays poisoned
225    /// (Dual can then report success after an L1-only delete).
226    /// Returns the deleted value when a matching live entry was removed.
227    pub async fn delete_if<F>(&self, key: &CacheKey, predicate: F) -> Result<Option<T>>
228    where
229        F: Fn(&T) -> bool,
230    {
231        if self.pool.is_noop() {
232            return Ok(None);
233        }
234
235        let redis_key = self.make_redis_key(key);
236        // Bound retries so a pathological rewrite livelock cannot hang forever.
237        const MAX_CAS_ATTEMPTS: u8 = 8;
238        for attempt in 1..=MAX_CAS_ATTEMPTS {
239            let Some(data) = self.pool.get(&redis_key).await? else {
240                return Ok(None);
241            };
242
243            let entry: SerializableCacheEntry<T> = match self.deserialize(&data) {
244                Ok(entry) => entry,
245                Err(e) => {
246                    // CAS-delete the exact malformed blob so a concurrent safe
247                    // replacement written after this GET is not removed by an
248                    // unconditional DEL (same guarantee as expired cleanup).
249                    warn!(key = %key, error = %e, "Failed to deserialize cache entry during conditional delete");
250                    let deleted = self.pool.delete_if_equals(&redis_key, &data).await?;
251                    if deleted {
252                        return Ok(None);
253                    }
254                    warn!(
255                        key = %key,
256                        attempt,
257                        "Redis malformed cleanup CAS miss; re-reading live value"
258                    );
259                    continue;
260                }
261            };
262
263            if entry.is_expired() {
264                // CAS-delete the exact expired blob so a concurrent safe
265                // replacement written in the logical-expiry/TTL race window is
266                // not removed by an unconditional DEL.
267                let deleted = self.pool.delete_if_equals(&redis_key, &data).await?;
268                if deleted {
269                    return Ok(None);
270                }
271                warn!(
272                    key = %key,
273                    attempt,
274                    "Redis expired cleanup CAS miss; re-reading live value"
275                );
276                continue;
277            }
278
279            if !predicate(&entry.value) {
280                return Ok(None);
281            }
282
283            let deleted = self.pool.delete_if_equals(&redis_key, &data).await?;
284            if deleted {
285                self.stats.record_deletion();
286                trace!(key = %key, "Redis cache conditional delete");
287                return Ok(Some(entry.value));
288            }
289
290            warn!(
291                key = %key,
292                attempt,
293                "Redis conditional delete CAS miss; re-reading live value"
294            );
295        }
296
297        Err(GatewayError::Internal(format!(
298            "Redis conditional delete exhausted CAS retries for key {redis_key}"
299        )))
300    }
301
302    /// Clear all Redis entries owned by this cache prefix.
303    pub async fn clear(&self) -> Result<usize> {
304        if self.pool.is_noop() {
305            return Ok(0);
306        }
307
308        let prefix = format!("{}:", self.config.key_prefix);
309        self.pool.delete_by_prefix(&prefix).await
310    }
311
312    /// Check if a key exists in the cache
313    pub async fn exists(&self, key: &CacheKey) -> Result<bool> {
314        if self.pool.is_noop() {
315            return Ok(false);
316        }
317
318        let redis_key = self.make_redis_key(key);
319        self.pool.exists(&redis_key).await
320    }
321
322    /// Get the remaining TTL for a key
323    pub async fn ttl(&self, key: &CacheKey) -> Result<Option<Duration>> {
324        if self.pool.is_noop() {
325            return Ok(None);
326        }
327
328        let redis_key = self.make_redis_key(key);
329        let ttl_secs = self.pool.ttl(&redis_key).await?;
330
331        if ttl_secs < 0 {
332            Ok(None)
333        } else {
334            Ok(Some(Duration::from_secs(ttl_secs as u64)))
335        }
336    }
337
338    /// Get cache statistics
339    pub fn stats(&self) -> Arc<AtomicCacheStats> {
340        Arc::clone(&self.stats)
341    }
342
343    /// Check if Redis is available
344    pub async fn is_available(&self) -> bool {
345        if self.pool.is_noop() {
346            return false;
347        }
348        self.pool.health_check().await.is_ok()
349    }
350
351    /// Perform a health check
352    pub async fn health_check(&self) -> Result<()> {
353        self.pool.health_check().await
354    }
355
356    // ==================== Private Methods ====================
357
358    /// Generate Redis key with prefix
359    fn make_redis_key(&self, key: &CacheKey) -> String {
360        format!("{}:{}", self.config.key_prefix, key.as_str())
361    }
362
363    /// Serialize a value to JSON string
364    fn serialize<S: Serialize>(&self, value: &S) -> Result<String> {
365        serde_json::to_string(value)
366            .map_err(|e| GatewayError::Config(format!("Failed to serialize cache value: {}", e)))
367    }
368
369    /// Deserialize a value from JSON string
370    fn deserialize<D: DeserializeOwned>(&self, data: &str) -> Result<D> {
371        serde_json::from_str(data)
372            .map_err(|e| GatewayError::Config(format!("Failed to deserialize cache value: {}", e)))
373    }
374}
375
376/// Batch operations for Redis cache
377impl<T> RedisCache<T>
378where
379    T: Serialize + DeserializeOwned + Clone + Send + Sync + 'static,
380{
381    /// Get multiple values from the cache using Redis MGET
382    pub async fn get_many(&self, keys: &[CacheKey]) -> Result<Vec<Option<T>>> {
383        if self.pool.is_noop() || keys.is_empty() {
384            return Ok(vec![None; keys.len()]);
385        }
386
387        let redis_keys: Vec<String> = keys.iter().map(|k| self.make_redis_key(k)).collect();
388        let raw_values = self.pool.mget(&redis_keys).await?;
389
390        let mut results = Vec::with_capacity(keys.len());
391        for (i, raw) in raw_values.into_iter().enumerate() {
392            match raw {
393                Some(data) => match self.deserialize::<SerializableCacheEntry<T>>(&data) {
394                    Ok(entry) if !entry.is_expired() => {
395                        self.stats.record_redis_hit();
396                        results.push(Some(entry.value));
397                    }
398                    Ok(_) => {
399                        let _ = self.pool.delete(&redis_keys[i]).await;
400                        self.stats.record_redis_miss();
401                        results.push(None);
402                    }
403                    Err(e) => {
404                        warn!(key = %keys[i], error = %e, "Failed to deserialize cache entry");
405                        let _ = self.pool.delete(&redis_keys[i]).await;
406                        self.stats.record_redis_miss();
407                        results.push(None);
408                    }
409                },
410                None => {
411                    self.stats.record_redis_miss();
412                    results.push(None);
413                }
414            }
415        }
416
417        Ok(results)
418    }
419
420    /// Set multiple values in the cache
421    pub async fn set_many(&self, entries: &[(CacheKey, T, Duration)]) -> Result<()> {
422        if self.pool.is_noop() || entries.is_empty() {
423            return Ok(());
424        }
425
426        for (key, value, ttl) in entries {
427            self.set_with_ttl(key.clone(), value.clone(), *ttl).await?;
428        }
429
430        Ok(())
431    }
432
433    /// Delete multiple keys from the cache
434    pub async fn delete_many(&self, keys: &[CacheKey]) -> Result<usize> {
435        if self.pool.is_noop() || keys.is_empty() {
436            return Ok(0);
437        }
438
439        let mut deleted = 0;
440        for key in keys {
441            if self.delete(key).await? {
442                deleted += 1;
443            }
444        }
445
446        Ok(deleted)
447    }
448}
449
450#[cfg(test)]
451mod tests {
452    use super::*;
453
454    // Helper to create a noop cache for testing
455    fn create_noop_cache() -> RedisCache<String> {
456        let pool = Arc::new(RedisPool::create_noop());
457        let config = DualCacheConfig::redis_only();
458        RedisCache::new(pool, config)
459    }
460
461    // ==================== Basic Noop Tests ====================
462    // These tests verify behavior when Redis is not available
463
464    #[tokio::test]
465    async fn test_redis_cache_noop_get() {
466        let cache = create_noop_cache();
467        let key = CacheKey::new("test-key");
468
469        let result = cache.get(&key).await.unwrap();
470        assert!(result.is_none());
471    }
472
473    #[tokio::test]
474    async fn test_redis_cache_noop_set() {
475        let cache = create_noop_cache();
476        let key = CacheKey::new("test-key");
477
478        // Should not error even in noop mode
479        let result = cache.set(key, "value".to_string()).await;
480        assert!(result.is_ok());
481    }
482
483    #[tokio::test]
484    async fn test_redis_cache_noop_delete() {
485        let cache = create_noop_cache();
486        let key = CacheKey::new("test-key");
487
488        let result = cache.delete(&key).await.unwrap();
489        assert!(!result);
490    }
491
492    #[tokio::test]
493    async fn test_redis_cache_noop_clear() {
494        let cache = create_noop_cache();
495
496        assert_eq!(cache.clear().await.ok(), Some(0));
497    }
498
499    #[tokio::test]
500    async fn test_redis_cache_noop_exists() {
501        let cache = create_noop_cache();
502        let key = CacheKey::new("test-key");
503
504        let result = cache.exists(&key).await.unwrap();
505        assert!(!result);
506    }
507
508    #[tokio::test]
509    async fn test_redis_cache_noop_ttl() {
510        let cache = create_noop_cache();
511        let key = CacheKey::new("test-key");
512
513        let result = cache.ttl(&key).await.unwrap();
514        assert!(result.is_none());
515    }
516
517    #[tokio::test]
518    async fn test_redis_cache_noop_availability() {
519        let cache = create_noop_cache();
520        assert!(!cache.is_available().await);
521    }
522
523    // ==================== Key Generation Tests ====================
524
525    #[test]
526    fn test_redis_key_generation() {
527        let pool = Arc::new(RedisPool::create_noop());
528        let config = DualCacheConfig::default();
529        let cache: RedisCache<String> = RedisCache::new(pool, config);
530
531        let key = CacheKey::new("my-key");
532        let redis_key = cache.make_redis_key(&key);
533
534        assert!(redis_key.starts_with("litellm:cache:"));
535        assert!(redis_key.ends_with("my-key"));
536    }
537
538    #[test]
539    fn test_redis_key_with_custom_prefix() {
540        let pool = Arc::new(RedisPool::create_noop());
541        let config = DualCacheConfig {
542            key_prefix: "custom:prefix".to_string(),
543            ..Default::default()
544        };
545        let cache: RedisCache<String> = RedisCache::new(pool, config);
546
547        let key = CacheKey::new("my-key");
548        let redis_key = cache.make_redis_key(&key);
549
550        assert_eq!(redis_key, "custom:prefix:my-key");
551    }
552
553    // ==================== Serialization Tests ====================
554
555    #[test]
556    fn test_serialization_roundtrip() {
557        let pool = Arc::new(RedisPool::create_noop());
558        let config = DualCacheConfig::default();
559        let cache: RedisCache<String> = RedisCache::new(pool, config);
560
561        let original = "test-value".to_string();
562        let serialized = cache.serialize(&original).unwrap();
563        let deserialized: String = cache.deserialize(&serialized).unwrap();
564
565        assert_eq!(original, deserialized);
566    }
567
568    #[test]
569    fn test_serialization_complex_type() {
570        use serde::{Deserialize, Serialize};
571
572        #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
573        struct ComplexType {
574            id: u64,
575            name: String,
576            tags: Vec<String>,
577        }
578
579        let pool = Arc::new(RedisPool::create_noop());
580        let config = DualCacheConfig::default();
581        let cache: RedisCache<ComplexType> = RedisCache::new(pool, config);
582
583        let original = ComplexType {
584            id: 123,
585            name: "test".to_string(),
586            tags: vec!["a".to_string(), "b".to_string()],
587        };
588
589        let serialized = cache.serialize(&original).unwrap();
590        let deserialized: ComplexType = cache.deserialize(&serialized).unwrap();
591
592        assert_eq!(original, deserialized);
593    }
594
595    // ==================== Statistics Tests ====================
596
597    #[tokio::test]
598    async fn test_redis_cache_stats() {
599        let cache = create_noop_cache();
600
601        // Generate some cache activity
602        let _ = cache.get(&CacheKey::new("miss1")).await;
603        let _ = cache.get(&CacheKey::new("miss2")).await;
604
605        let stats = cache.stats().snapshot();
606        assert_eq!(stats.redis_misses, 2);
607    }
608
609    // ==================== Batch Operations Tests ====================
610
611    #[tokio::test]
612    async fn test_redis_cache_get_many_noop() {
613        let cache = create_noop_cache();
614
615        let keys = vec![
616            CacheKey::new("key1"),
617            CacheKey::new("key2"),
618            CacheKey::new("key3"),
619        ];
620
621        let results = cache.get_many(&keys).await.unwrap();
622        assert_eq!(results.len(), 3);
623        assert!(results.iter().all(|r| r.is_none()));
624    }
625
626    #[tokio::test]
627    async fn test_redis_cache_set_many_noop() {
628        let cache = create_noop_cache();
629
630        let entries = vec![
631            (
632                CacheKey::new("key1"),
633                "value1".to_string(),
634                Duration::from_secs(60),
635            ),
636            (
637                CacheKey::new("key2"),
638                "value2".to_string(),
639                Duration::from_secs(60),
640            ),
641        ];
642
643        let result = cache.set_many(&entries).await;
644        assert!(result.is_ok());
645    }
646
647    #[tokio::test]
648    async fn test_redis_cache_delete_many_noop() {
649        let cache = create_noop_cache();
650
651        let keys = vec![CacheKey::new("key1"), CacheKey::new("key2")];
652
653        let deleted = cache.delete_many(&keys).await.unwrap();
654        assert_eq!(deleted, 0);
655    }
656
657    #[tokio::test]
658    async fn test_redis_cache_empty_batch() {
659        let cache = create_noop_cache();
660
661        assert!(cache.get_many(&[]).await.unwrap().is_empty());
662        assert!(cache.set_many(&[]).await.is_ok());
663        assert_eq!(cache.delete_many(&[]).await.unwrap(), 0);
664    }
665}
litellm_rs/core/cache/redis_cache.rs

litellm_rs/core/cache/
redis_cache.rs
