1use std::path::{Path, PathBuf};
10
11use redb::{Database, ReadableTable, ReadableTableMetadata, TableDefinition};
12use sha2::{Digest, Sha256};
13use tracing::{debug, info};
14
15const DRV_TABLE: TableDefinition<&str, &str> = TableDefinition::new("drv_paths");
17
18#[derive(Debug, thiserror::Error)]
20pub enum DrvCacheError {
21 #[error("redb error: {0}")]
22 Db(String),
23 #[error("I/O error: {0}")]
24 Io(#[from] std::io::Error),
25}
26
27#[derive(Debug, Clone)]
29pub struct DrvCacheEntry {
30 pub drv_path: String,
31 pub out_path: String,
32}
33
34pub struct DrvCache {
36 db: Database,
37}
38
39impl DrvCache {
40 pub fn open(path: &Path) -> Result<Self, DrvCacheError> {
42 if let Some(parent) = path.parent() {
43 std::fs::create_dir_all(parent)?;
44 }
45 let db = Database::create(path)
46 .map_err(|e| DrvCacheError::Db(format!("open: {e}")))?;
47
48 let txn = db
50 .begin_write()
51 .map_err(|e| DrvCacheError::Db(format!("txn: {e}")))?;
52 { let _ = txn.open_table(DRV_TABLE); }
53 txn.commit()
54 .map_err(|e| DrvCacheError::Db(format!("commit: {e}")))?;
55
56 info!(path = %path.display(), "Opened derivation cache");
57 Ok(Self { db })
58 }
59
60 pub fn default_path() -> PathBuf {
62 let base = std::env::var("XDG_CACHE_HOME")
63 .ok()
64 .filter(|s| !s.is_empty())
65 .map(PathBuf::from)
66 .or_else(|| {
67 std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".cache"))
68 })
69 .unwrap_or_else(|| PathBuf::from("/tmp"));
70 base.join("sui").join("drv-cache.redb")
71 }
72
73 pub fn get(
75 &self,
76 lock_hash: &str,
77 source_hash: &str,
78 attr_path: &str,
79 ) -> Option<DrvCacheEntry> {
80 let key = format!("{lock_hash}:{source_hash}:{attr_path}");
81 let txn = self.db.begin_read().ok()?;
82 let table = txn.open_table(DRV_TABLE).ok()?;
83 let value = table.get(key.as_str()).ok()??;
84 let s = value.value();
85 let (drv, out) = s.split_once('\n')?;
86 debug!(attr_path, "drv cache hit");
87 Some(DrvCacheEntry {
88 drv_path: drv.to_string(),
89 out_path: out.to_string(),
90 })
91 }
92
93 pub fn put(
95 &self,
96 lock_hash: &str,
97 source_hash: &str,
98 attr_path: &str,
99 entry: &DrvCacheEntry,
100 ) -> Result<(), DrvCacheError> {
101 let key = format!("{lock_hash}:{source_hash}:{attr_path}");
102 let value = format!("{}\n{}", entry.drv_path, entry.out_path);
103 let txn = self.db.begin_write()
104 .map_err(|e| DrvCacheError::Db(format!("txn: {e}")))?;
105 {
106 let mut table = txn.open_table(DRV_TABLE)
107 .map_err(|e| DrvCacheError::Db(format!("table: {e}")))?;
108 table.insert(key.as_str(), value.as_str())
109 .map_err(|e| DrvCacheError::Db(format!("insert: {e}")))?;
110 }
111 txn.commit()
112 .map_err(|e| DrvCacheError::Db(format!("commit: {e}")))?;
113 debug!(attr_path, "drv cache put");
114 Ok(())
115 }
116
117 pub fn len(&self) -> usize {
119 let Ok(txn) = self.db.begin_read() else { return 0 };
120 let Ok(table) = txn.open_table(DRV_TABLE) else { return 0 };
121 table.len().unwrap_or(0) as usize
122 }
123
124 pub fn is_empty(&self) -> bool {
126 self.len() == 0
127 }
128
129 pub fn hash_bytes(content: &[u8]) -> String {
131 let mut hasher = Sha256::new();
132 hasher.update(content);
133 format!("{:x}", hasher.finalize())
134 }
135}
136
137thread_local! {
140 static GLOBAL_CACHE: std::cell::RefCell<Option<DrvCache>> = const { std::cell::RefCell::new(None) };
141}
142
143pub fn init_global_cache() {
145 GLOBAL_CACHE.with(|cell| {
146 let mut cache = cell.borrow_mut();
147 if cache.is_none() {
148 let path = DrvCache::default_path();
149 match DrvCache::open(&path) {
150 Ok(c) => {
151 info!(entries = c.len(), "Derivation cache initialized");
152 *cache = Some(c);
153 }
154 Err(e) => {
155 tracing::warn!(error = %e, "Failed to open derivation cache (continuing without)");
156 }
157 }
158 }
159 });
160}
161
162pub fn with_cache<F, R>(f: F) -> Option<R>
164where
165 F: FnOnce(&DrvCache) -> Option<R>,
166{
167 GLOBAL_CACHE.with(|cell| {
168 let borrow = cell.borrow();
169 borrow.as_ref().and_then(f)
170 })
171}
172
173pub fn with_cache_mut<F>(f: F)
175where
176 F: FnOnce(&DrvCache),
177{
178 GLOBAL_CACHE.with(|cell| {
179 let borrow = cell.borrow();
180 if let Some(cache) = borrow.as_ref() {
181 f(cache);
182 }
183 });
184}
185
186#[cfg(test)]
189mod tests {
190 use super::*;
191
192 #[test]
193 fn roundtrip() {
194 let tmp = tempfile::tempdir().unwrap();
195 let cache = DrvCache::open(&tmp.path().join("test.redb")).unwrap();
196
197 assert!(cache.get("lock1", "src1", "packages.x86_64-linux.default").is_none());
198
199 cache
200 .put("lock1", "src1", "packages.x86_64-linux.default", &DrvCacheEntry {
201 drv_path: "/nix/store/abc-hello.drv".to_string(),
202 out_path: "/nix/store/xyz-hello-2.10".to_string(),
203 })
204 .unwrap();
205
206 let entry = cache.get("lock1", "src1", "packages.x86_64-linux.default").unwrap();
207 assert_eq!(entry.drv_path, "/nix/store/abc-hello.drv");
208 assert_eq!(entry.out_path, "/nix/store/xyz-hello-2.10");
209 assert_eq!(cache.len(), 1);
210 }
211
212 #[test]
213 fn different_keys_no_collision() {
214 let tmp = tempfile::tempdir().unwrap();
215 let cache = DrvCache::open(&tmp.path().join("test.redb")).unwrap();
216
217 cache.put("lock1", "src1", "attr.a", &DrvCacheEntry {
218 drv_path: "/nix/store/a.drv".into(),
219 out_path: "/nix/store/a".into(),
220 }).unwrap();
221
222 cache.put("lock1", "src1", "attr.b", &DrvCacheEntry {
223 drv_path: "/nix/store/b.drv".into(),
224 out_path: "/nix/store/b".into(),
225 }).unwrap();
226
227 cache.put("lock2", "src1", "attr.a", &DrvCacheEntry {
228 drv_path: "/nix/store/c.drv".into(),
229 out_path: "/nix/store/c".into(),
230 }).unwrap();
231
232 assert_eq!(cache.get("lock1", "src1", "attr.a").unwrap().out_path, "/nix/store/a");
233 assert_eq!(cache.get("lock1", "src1", "attr.b").unwrap().out_path, "/nix/store/b");
234 assert_eq!(cache.get("lock2", "src1", "attr.a").unwrap().out_path, "/nix/store/c");
235 assert_eq!(cache.len(), 3);
236 }
237
238 #[test]
239 fn hash_bytes_deterministic() {
240 let h1 = DrvCache::hash_bytes(b"hello world");
241 let h2 = DrvCache::hash_bytes(b"hello world");
242 assert_eq!(h1, h2);
243 assert_eq!(h1.len(), 64); }
245}