rs3gw 0.2.1

High-Performance AI/HPC Object Storage Gateway powered by scirs2-io
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

//! Preprocessing result cache.

use super::error::PreprocessingResult;
use bytes::Bytes;
use chrono::{DateTime, Utc};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// Cached preprocessing result
#[derive(Debug, Clone)]
pub(super) struct CachedResult {
    pub(super) data: Bytes,
    #[allow(dead_code)]
    pub(super) metadata: HashMap<String, String>,
    #[allow(dead_code)]
    pub(super) cached_at: DateTime<Utc>,
}

/// Cache for preprocessing results
pub struct PreprocessingCache {
    pub(super) cache: Arc<RwLock<HashMap<String, CachedResult>>>,
    pub(super) max_size_bytes: usize,
    pub(super) current_size: Arc<RwLock<usize>>,
}

impl PreprocessingCache {
    pub fn new(max_size_bytes: usize) -> Self {
        Self {
            cache: Arc::new(RwLock::new(HashMap::new())),
            max_size_bytes,
            current_size: Arc::new(RwLock::new(0)),
        }
    }

    pub async fn get(&self, key: &str) -> Option<Bytes> {
        let cache = self.cache.read().await;
        cache.get(key).map(|r| r.data.clone())
    }

    pub async fn put(
        &self,
        key: String,
        data: Bytes,
        metadata: HashMap<String, String>,
    ) -> PreprocessingResult<()> {
        let data_size = data.len();

        // Check if adding this would exceed max size
        let current_size = *self.current_size.read().await;
        if current_size + data_size > self.max_size_bytes {
            // Simple eviction: clear cache if too large
            let mut cache = self.cache.write().await;
            cache.clear();
            *self.current_size.write().await = 0;
        }

        let mut cache = self.cache.write().await;
        cache.insert(
            key,
            CachedResult {
                data,
                metadata,
                cached_at: Utc::now(),
            },
        );

        *self.current_size.write().await += data_size;
        Ok(())
    }

    pub async fn invalidate(&self, key: &str) {
        let mut cache = self.cache.write().await;
        if let Some(entry) = cache.remove(key) {
            let mut size = self.current_size.write().await;
            *size = size.saturating_sub(entry.data.len());
        }
    }

    pub async fn clear(&self) {
        let mut cache = self.cache.write().await;
        cache.clear();
        *self.current_size.write().await = 0;
    }

    pub async fn stats(&self) -> (usize, usize, usize) {
        let cache = self.cache.read().await;
        let current_size = *self.current_size.read().await;
        (cache.len(), current_size, self.max_size_bytes)
    }
}