litellm_rs/core/cache/cloud/
s3.rs

1//! AWS S3 cache backend implementation
2//!
3//! Provides a cache backend using AWS S3 for persistent, distributed caching.
4
5#[cfg(feature = "s3")]
6mod implementation {
7    use async_trait::async_trait;
8    use aws_sdk_s3::Client;
9    use aws_sdk_s3::primitives::ByteStream;
10    use aws_sdk_s3::types::Object as S3Object;
11    use serde::{Serialize, de::DeserializeOwned};
12    use std::time::Duration;
13    use tracing::{debug, error, warn};
14
15    use crate::core::cache::cloud::{CacheMetadata, CloudCache, CloudCacheConfig};
16    use crate::core::cache::types::CacheKey;
17    use crate::utils::error::gateway_error::{GatewayError, Result};
18
19    /// S3 cache configuration
20    #[derive(Debug, Clone, Default)]
21    pub struct S3CacheConfig {
22        /// Base cloud cache config
23        pub base: CloudCacheConfig,
24        /// AWS region
25        pub region: Option<String>,
26        /// Custom endpoint URL (for LocalStack, MinIO, etc.)
27        pub endpoint_url: Option<String>,
28        /// Storage class for cache objects
29        pub storage_class: S3StorageClass,
30    }
31
32    /// S3 storage class options
33    #[derive(Debug, Clone, Copy, Default)]
34    pub enum S3StorageClass {
35        #[default]
36        Standard,
37        StandardIa,
38        OnezoneIa,
39        IntelligentTiering,
40        Glacier,
41        GlacierIr,
42    }
43
44    impl S3StorageClass {
45        /// Get the S3 storage class string
46        pub fn as_str(&self) -> &'static str {
47            match self {
48                S3StorageClass::Standard => "STANDARD",
49                S3StorageClass::StandardIa => "STANDARD_IA",
50                S3StorageClass::OnezoneIa => "ONEZONE_IA",
51                S3StorageClass::IntelligentTiering => "INTELLIGENT_TIERING",
52                S3StorageClass::Glacier => "GLACIER",
53                S3StorageClass::GlacierIr => "GLACIER_IR",
54            }
55        }
56    }
57
58    impl S3CacheConfig {
59        /// Create a new S3 cache configuration
60        pub fn new(bucket: impl Into<String>) -> Self {
61            Self {
62                base: CloudCacheConfig::new(bucket),
63                ..Default::default()
64            }
65        }
66
67        /// Set the AWS region
68        pub fn region(mut self, region: impl Into<String>) -> Self {
69            self.region = Some(region.into());
70            self
71        }
72
73        /// Set a custom endpoint URL
74        pub fn endpoint_url(mut self, url: impl Into<String>) -> Self {
75            self.endpoint_url = Some(url.into());
76            self
77        }
78
79        /// Set the storage class
80        pub fn storage_class(mut self, class: S3StorageClass) -> Self {
81            self.storage_class = class;
82            self
83        }
84
85        /// Set the key prefix
86        pub fn prefix(mut self, prefix: impl Into<String>) -> Self {
87            self.base.prefix = prefix.into();
88            self
89        }
90
91        /// Set the default TTL
92        pub fn default_ttl(mut self, ttl: Duration) -> Self {
93            self.base.default_ttl = ttl;
94            self
95        }
96
97        /// Create from environment variables
98        pub fn from_env() -> Option<Self> {
99            let bucket = std::env::var("S3_CACHE_BUCKET").ok()?;
100            Some(Self {
101                base: CloudCacheConfig::new(bucket),
102                region: std::env::var("AWS_REGION").ok(),
103                endpoint_url: std::env::var("S3_ENDPOINT_URL").ok(),
104                storage_class: S3StorageClass::default(),
105            })
106        }
107    }
108
109    /// AWS S3 cache backend
110    pub struct S3Cache {
111        client: Client,
112        config: S3CacheConfig,
113    }
114
115    impl S3Cache {
116        /// Create a new S3 cache
117        pub async fn new(config: S3CacheConfig) -> Result<Self> {
118            let mut aws_config_builder =
119                aws_config::defaults(aws_config::BehaviorVersion::latest());
120
121            if let Some(ref region) = config.region {
122                aws_config_builder =
123                    aws_config_builder.region(aws_config::Region::new(region.clone()));
124            }
125
126            let aws_config = aws_config_builder.load().await;
127
128            let mut s3_config = aws_sdk_s3::config::Builder::from(&aws_config);
129
130            if let Some(ref endpoint) = config.endpoint_url {
131                s3_config = s3_config.endpoint_url(endpoint).force_path_style(true);
132            }
133
134            let client = Client::from_conf(s3_config.build());
135
136            Ok(Self { client, config })
137        }
138
139        /// Create from environment variables
140        pub async fn from_env() -> Result<Self> {
141            let config = S3CacheConfig::from_env().ok_or_else(|| {
142                GatewayError::Config("S3_CACHE_BUCKET not set in environment".to_string())
143            })?;
144            Self::new(config).await
145        }
146
147        /// Build the full S3 key for a cache key
148        fn build_key(&self, key: &CacheKey) -> String {
149            format!("{}{}", self.config.base.prefix, key.as_str())
150        }
151
152        /// Build the metadata key for a cache key
153        fn build_metadata_key(&self, key: &CacheKey) -> String {
154            format!("{}{}._meta", self.config.base.prefix, key.as_str())
155        }
156    }
157
158    #[async_trait]
159    impl CloudCache for S3Cache {
160        async fn get<T: DeserializeOwned + Send>(&self, key: &CacheKey) -> Result<Option<T>> {
161            let s3_key = self.build_key(key);
162            let meta_key = self.build_metadata_key(key);
163            debug!(key = %s3_key, "Reading from S3 cache");
164
165            // First check metadata for expiration
166            match self
167                .client
168                .get_object()
169                .bucket(&self.config.base.bucket)
170                .key(&meta_key)
171                .send()
172                .await
173            {
174                Ok(meta_response) => {
175                    let meta_bytes = meta_response
176                        .body
177                        .collect()
178                        .await
179                        .map_err(|e| {
180                            GatewayError::Internal(format!("Failed to read metadata: {}", e))
181                        })?
182                        .into_bytes();
183
184                    let metadata: CacheMetadata =
185                        serde_json::from_slice(&meta_bytes).map_err(|e| {
186                            GatewayError::Internal(format!("Failed to parse metadata: {}", e))
187                        })?;
188
189                    if metadata.is_expired() {
190                        debug!(key = %s3_key, "Cache entry expired");
191                        // Clean up expired entry in background
192                        let _ = self.delete(key).await;
193                        return Ok(None);
194                    }
195                }
196                Err(err) => {
197                    let err_str = err.to_string();
198                    if err_str.contains("NoSuchKey") || err_str.contains("404") {
199                        debug!(key = %s3_key, "Cache miss - no metadata");
200                        return Ok(None);
201                    }
202                    warn!(key = %s3_key, error = %err, "Failed to read metadata");
203                    return Ok(None);
204                }
205            }
206
207            // Read the actual value
208            match self
209                .client
210                .get_object()
211                .bucket(&self.config.base.bucket)
212                .key(&s3_key)
213                .send()
214                .await
215            {
216                Ok(response) => {
217                    let bytes = response
218                        .body
219                        .collect()
220                        .await
221                        .map_err(|e| GatewayError::Internal(format!("Failed to read body: {}", e)))?
222                        .into_bytes();
223
224                    let value: T = serde_json::from_slice(&bytes).map_err(|e| {
225                        GatewayError::Internal(format!("Failed to deserialize: {}", e))
226                    })?;
227
228                    debug!(key = %s3_key, "S3 cache hit");
229                    Ok(Some(value))
230                }
231                Err(err) => {
232                    let err_str = err.to_string();
233                    if err_str.contains("NoSuchKey") || err_str.contains("404") {
234                        debug!(key = %s3_key, "Cache miss");
235                        Ok(None)
236                    } else {
237                        error!(key = %s3_key, error = %err, "Failed to read from S3");
238                        Err(GatewayError::Internal(format!("S3 read error: {}", err)))
239                    }
240                }
241            }
242        }
243
244        async fn set<T: Serialize + Send + Sync>(
245            &self,
246            key: &CacheKey,
247            value: &T,
248            ttl: Duration,
249        ) -> Result<()> {
250            let s3_key = self.build_key(key);
251            let meta_key = self.build_metadata_key(key);
252            debug!(key = %s3_key, ttl_secs = ttl.as_secs(), "Writing to S3 cache");
253
254            // Serialize the value
255            let bytes = serde_json::to_vec(value)
256                .map_err(|e| GatewayError::Internal(format!("Failed to serialize: {}", e)))?;
257
258            // Create metadata
259            let metadata = CacheMetadata::new(ttl, bytes.len(), false);
260            let meta_bytes = serde_json::to_vec(&metadata).map_err(|e| {
261                GatewayError::Internal(format!("Failed to serialize metadata: {}", e))
262            })?;
263
264            // Write the value
265            self.client
266                .put_object()
267                .bucket(&self.config.base.bucket)
268                .key(&s3_key)
269                .body(ByteStream::from(bytes))
270                .storage_class(self.config.storage_class.as_str().parse().map_err(|e| {
271                    GatewayError::Config(format!("Invalid S3 storage class: {}", e))
272                })?)
273                .content_type("application/json")
274                .send()
275                .await
276                .map_err(|e| GatewayError::Internal(format!("Failed to write to S3: {}", e)))?;
277
278            // Write metadata
279            self.client
280                .put_object()
281                .bucket(&self.config.base.bucket)
282                .key(&meta_key)
283                .body(ByteStream::from(meta_bytes))
284                .content_type("application/json")
285                .send()
286                .await
287                .map_err(|e| GatewayError::Internal(format!("Failed to write metadata: {}", e)))?;
288
289            debug!(key = %s3_key, "S3 cache write successful");
290            Ok(())
291        }
292
293        async fn delete(&self, key: &CacheKey) -> Result<bool> {
294            let s3_key = self.build_key(key);
295            let meta_key = self.build_metadata_key(key);
296            debug!(key = %s3_key, "Deleting from S3 cache");
297
298            // Delete both the value and metadata
299            let value_result = self
300                .client
301                .delete_object()
302                .bucket(&self.config.base.bucket)
303                .key(&s3_key)
304                .send()
305                .await;
306
307            let meta_result = self
308                .client
309                .delete_object()
310                .bucket(&self.config.base.bucket)
311                .key(&meta_key)
312                .send()
313                .await;
314
315            // S3 delete doesn't error if key doesn't exist
316            if value_result.is_err() && meta_result.is_err() {
317                return Ok(false);
318            }
319
320            debug!(key = %s3_key, "S3 cache delete successful");
321            Ok(true)
322        }
323
324        async fn exists(&self, key: &CacheKey) -> Result<bool> {
325            let s3_key = self.build_key(key);
326
327            match self
328                .client
329                .head_object()
330                .bucket(&self.config.base.bucket)
331                .key(&s3_key)
332                .send()
333                .await
334            {
335                Ok(_) => Ok(true),
336                Err(err) => {
337                    let err_str = err.to_string();
338                    if err_str.contains("NotFound") || err_str.contains("404") {
339                        Ok(false)
340                    } else {
341                        Err(GatewayError::Internal(format!("S3 head error: {}", err)))
342                    }
343                }
344            }
345        }
346
347        async fn list_keys(&self, prefix: &str) -> Result<Vec<String>> {
348            let full_prefix = format!("{}{}", self.config.base.prefix, prefix);
349
350            let response = self
351                .client
352                .list_objects_v2()
353                .bucket(&self.config.base.bucket)
354                .prefix(&full_prefix)
355                .send()
356                .await
357                .map_err(|e| GatewayError::Internal(format!("Failed to list objects: {}", e)))?;
358
359            let mut keys = Vec::new();
360            for obj in response.contents().iter() {
361                let obj: &S3Object = obj;
362                if let Some(k) = obj.key() {
363                    let key = k
364                        .strip_prefix(&self.config.base.prefix)
365                        .unwrap_or(k)
366                        .to_string();
367                    if !key.ends_with("._meta") {
368                        keys.push(key);
369                    }
370                }
371            }
372
373            Ok(keys)
374        }
375
376        async fn clear(&self) -> Result<()> {
377            warn!("Clearing all S3 cache entries");
378
379            let keys = self.list_keys("").await?;
380
381            for key in keys {
382                let cache_key = CacheKey::new(key);
383                let _ = self.delete(&cache_key).await;
384            }
385
386            Ok(())
387        }
388
389        fn name(&self) -> &'static str {
390            "s3"
391        }
392    }
393
394    impl std::fmt::Debug for S3Cache {
395        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
396            f.debug_struct("S3Cache")
397                .field("bucket", &self.config.base.bucket)
398                .field("prefix", &self.config.base.prefix)
399                .field("region", &self.config.region)
400                .finish()
401        }
402    }
403}
404
405#[cfg(feature = "s3")]
406pub use implementation::*;
407
408// Stub implementation when feature is disabled
409#[cfg(not(feature = "s3"))]
410mod stub {
411    use crate::utils::error::gateway_error::{GatewayError, Result};
412
413    /// S3 cache configuration (stub)
414    #[derive(Debug, Clone, Default)]
415    pub struct S3CacheConfig {
416        pub bucket: String,
417    }
418
419    impl S3CacheConfig {
420        pub fn new(bucket: impl Into<String>) -> Self {
421            Self {
422                bucket: bucket.into(),
423            }
424        }
425
426        pub fn from_env() -> Option<Self> {
427            None
428        }
429    }
430
431    /// S3 cache (stub)
432    #[derive(Debug)]
433    pub struct S3Cache;
434
435    impl S3Cache {
436        pub async fn new(_config: S3CacheConfig) -> Result<Self> {
437            Err(GatewayError::Config(
438                "S3 cache support not enabled. Enable the 's3' feature.".to_string(),
439            ))
440        }
441
442        pub async fn from_env() -> Result<Self> {
443            Err(GatewayError::Config(
444                "S3 cache support not enabled. Enable the 's3' feature.".to_string(),
445            ))
446        }
447    }
448}
449
450#[cfg(not(feature = "s3"))]
451pub use stub::*;
452
453#[cfg(all(test, feature = "s3"))]
454mod tests {
455    use super::*;
456
457    #[test]
458    fn test_s3_cache_config_builder() {
459        let config = S3CacheConfig::new("my-bucket")
460            .region("us-east-1")
461            .prefix("cache/")
462            .storage_class(S3StorageClass::StandardIa);
463
464        assert_eq!(config.base.bucket, "my-bucket");
465        assert_eq!(config.region, Some("us-east-1".to_string()));
466        assert_eq!(config.base.prefix, "cache/");
467    }
468
469    #[test]
470    fn test_s3_storage_class() {
471        assert_eq!(S3StorageClass::Standard.as_str(), "STANDARD");
472        assert_eq!(S3StorageClass::StandardIa.as_str(), "STANDARD_IA");
473        assert_eq!(S3StorageClass::Glacier.as_str(), "GLACIER");
474    }
475}
litellm_rs/core/cache/cloud/s3.rs

litellm_rs/core/cache/cloud/
s3.rs
