openjd_snapshots/
s3_check_cache.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// Copyright by contributors to this project.
3// SPDX-License-Identifier: (Apache-2.0 OR MIT)
4
5use std::path::{Path, PathBuf};
6use std::sync::Mutex;
7use std::time::{SystemTime, UNIX_EPOCH};
8
9const ENTRY_EXPIRY_DAYS: u64 = 30;
10
11pub struct S3CheckCache {
12    conn: Mutex<rusqlite::Connection>,
13}
14
15impl S3CheckCache {
16    pub fn new(cache_dir: impl AsRef<Path>) -> crate::Result<Self> {
17        let dir = cache_dir.as_ref();
18        std::fs::create_dir_all(dir)?;
19        let db_path = dir.join("s3_check_cache.db");
20        let conn = rusqlite::Connection::open(&db_path)
21            .map_err(|e| crate::SnapshotError::Cache(e.to_string()))?;
22        conn.pragma_update(None, "journal_mode", "WAL")
23            .map_err(|e| crate::SnapshotError::Cache(e.to_string()))?;
24        conn.execute_batch(
25            "CREATE TABLE IF NOT EXISTS s3checkV1(
26                s3_key text primary key,
27                last_seen_time timestamp
28            );",
29        )
30        .map_err(|e| crate::SnapshotError::Cache(e.to_string()))?;
31
32        // Prune expired entries (older than 30 days) on open
33        let cutoff = SystemTime::now()
34            .duration_since(UNIX_EPOCH)
35            .unwrap_or_default()
36            .as_secs_f64()
37            - (ENTRY_EXPIRY_DAYS as f64 * 86400.0);
38        conn.execute(
39            "DELETE FROM s3checkV1 WHERE CAST(last_seen_time AS REAL) < ?1",
40            rusqlite::params![cutoff],
41        )
42        .map_err(|e| crate::SnapshotError::Cache(e.to_string()))?;
43
44        Ok(Self {
45            conn: Mutex::new(conn),
46        })
47    }
48
49    pub fn open_default() -> crate::Result<Self> {
50        let home = std::env::var("HOME").map_err(|_| {
51            crate::SnapshotError::Cache("$HOME is not set, cannot locate default cache".into())
52        })?;
53        Self::new(PathBuf::from(home).join(".deadline/job_attachments"))
54    }
55
56    /// Check if an entry exists and hasn't expired (30 days).
57    pub fn get_entry(&self, s3_key: &str) -> Option<String> {
58        let conn = self.conn.lock().unwrap();
59        let result: String = conn
60            .query_row(
61                "SELECT last_seen_time FROM s3checkV1 WHERE s3_key = ?1",
62                rusqlite::params![s3_key],
63                |row| match row.get_ref(0)?.as_str() {
64                    Ok(s) => Ok(s.to_string()),
65                    Err(_) => {
66                        let f: f64 = row.get(0)?;
67                        Ok(f.to_string())
68                    }
69                },
70            )
71            .ok()?;
72
73        let last_seen: f64 = result.parse().ok()?;
74        let now = SystemTime::now()
75            .duration_since(UNIX_EPOCH)
76            .ok()?
77            .as_secs_f64();
78        if (now - last_seen) / 86400.0 < ENTRY_EXPIRY_DAYS as f64 {
79            Some(result)
80        } else {
81            None
82        }
83    }
84
85    /// Insert or replace an entry with the current timestamp.
86    pub fn put_entry(&self, s3_key: &str) -> crate::Result<()> {
87        let now = SystemTime::now()
88            .duration_since(UNIX_EPOCH)
89            .map_err(|e| crate::SnapshotError::Cache(e.to_string()))?
90            .as_secs_f64();
91        let time_str = rusqlite::types::Value::Text(now.to_string());
92        let conn = self.conn.lock().unwrap();
93        conn.execute(
94            "INSERT OR REPLACE INTO s3checkV1 VALUES(?1, ?2)",
95            rusqlite::params![s3_key, time_str],
96        )
97        .map_err(|e| crate::SnapshotError::Cache(e.to_string()))?;
98        Ok(())
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105    use tempfile::TempDir;
106
107    #[test]
108    fn put_and_get() {
109        let tmp = TempDir::new().unwrap();
110        let cache = S3CheckCache::new(tmp.path()).unwrap();
111        cache.put_entry("bucket/Data/abc123.xxh128").unwrap();
112        assert!(cache.get_entry("bucket/Data/abc123.xxh128").is_some());
113    }
114
115    #[test]
116    fn missing_entry_returns_none() {
117        let tmp = TempDir::new().unwrap();
118        let cache = S3CheckCache::new(tmp.path()).unwrap();
119        assert!(cache.get_entry("nonexistent").is_none());
120    }
121
122    #[test]
123    fn expired_entry_returns_none() {
124        let tmp = TempDir::new().unwrap();
125        let cache = S3CheckCache::new(tmp.path()).unwrap();
126        let old_time = SystemTime::now()
127            .duration_since(UNIX_EPOCH)
128            .unwrap()
129            .as_secs_f64()
130            - (31.0 * 86400.0);
131        let time_str = rusqlite::types::Value::Text(old_time.to_string());
132        let conn = cache.conn.lock().unwrap();
133        conn.execute(
134            "INSERT OR REPLACE INTO s3checkV1 VALUES(?1, ?2)",
135            rusqlite::params!["old_key", time_str],
136        )
137        .unwrap();
138        drop(conn);
139        assert!(cache.get_entry("old_key").is_none());
140    }
141
142    #[test]
143    fn recent_entry_not_expired() {
144        let tmp = TempDir::new().unwrap();
145        let cache = S3CheckCache::new(tmp.path()).unwrap();
146        let recent_time = SystemTime::now()
147            .duration_since(UNIX_EPOCH)
148            .unwrap()
149            .as_secs_f64()
150            - 86400.0;
151        let time_str = rusqlite::types::Value::Text(recent_time.to_string());
152        let conn = cache.conn.lock().unwrap();
153        conn.execute(
154            "INSERT OR REPLACE INTO s3checkV1 VALUES(?1, ?2)",
155            rusqlite::params!["recent_key", time_str],
156        )
157        .unwrap();
158        drop(conn);
159        assert!(cache.get_entry("recent_key").is_some());
160    }
161}
openjd_snapshots/s3_check_cache.rs

openjd_snapshots/
s3_check_cache.rs
