1use std::collections::HashMap;
44use std::fs::File;
45use std::io::Write as IoWrite;
46use std::path::{Path, PathBuf};
47use std::sync::OnceLock;
48use std::time::{SystemTime, UNIX_EPOCH};
49
50use memmap2::Mmap;
51use parking_lot::Mutex;
52use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
53use std::os::unix::fs::MetadataExt;
54
55pub const SHARD_MAGIC: u32 = 0x5A525343;
58pub const SHARD_FORMAT_VERSION: u32 = 4;
64#[derive(Archive, RkyvDeserialize, RkyvSerialize, Debug, Clone)]
66#[archive(check_bytes)]
67pub struct ShardHeader {
68 pub magic: u32,
70 pub format_version: u32,
72 pub zshrs_version: String,
74 pub pointer_width: u32,
76 pub built_at_secs: u64,
78}
79#[derive(Archive, RkyvDeserialize, RkyvSerialize, Debug, Clone)]
81#[archive(check_bytes)]
82pub struct ScriptEntry {
83 pub mtime_secs: i64,
85 pub mtime_nsecs: i64,
87 pub binary_mtime_at_cache: i64,
89 pub binary_len_at_cache: u64,
95 pub cached_at_secs: i64,
97 pub chunk_blob: Vec<u8>,
99}
100#[derive(Archive, RkyvDeserialize, RkyvSerialize, Debug, Clone)]
102#[archive(check_bytes)]
103pub struct ScriptShard {
104 pub header: ShardHeader,
106 pub entries: HashMap<String, ScriptEntry>,
108}
109
110pub struct MmappedShard {
113 _mmap: Mmap,
115 archived: *const ArchivedScriptShard,
117}
118
119unsafe impl Send for MmappedShard {}
122unsafe impl Sync for MmappedShard {}
123
124impl MmappedShard {
125 pub fn open(path: &Path) -> Option<Self> {
127 let file = File::open(path).ok()?;
128 let mmap = unsafe { Mmap::map(&file).ok()? };
129 let archived = rkyv::check_archived_root::<ScriptShard>(&mmap[..]).ok()?;
130 let archived_ptr = archived as *const ArchivedScriptShard;
131 Some(Self {
132 _mmap: mmap,
133 archived: archived_ptr,
134 })
135 }
136
137 fn shard(&self) -> &ArchivedScriptShard {
138 unsafe { &*self.archived }
140 }
141
142 fn header_ok(&self) -> bool {
143 let h = &self.shard().header;
144 let magic: u32 = h.magic.into();
145 let fv: u32 = h.format_version.into();
146 let pw: u32 = h.pointer_width.into();
147 magic == SHARD_MAGIC
148 && fv == SHARD_FORMAT_VERSION
149 && pw as usize == std::mem::size_of::<usize>()
150 && h.zshrs_version.as_str() == env!("CARGO_PKG_VERSION")
151 }
152
153 fn lookup(&self, path: &str) -> Option<&ArchivedScriptEntry> {
154 self.shard().entries.get(path)
155 }
156
157 fn entry_count(&self) -> usize {
158 self.shard().entries.len()
159 }
160}
161
162pub struct ScriptCache {
164 path: PathBuf,
166 lock_path: PathBuf,
168 mmap: Mutex<Option<MmappedShard>>,
170}
171
172impl ScriptCache {
173 pub fn open(path: &Path) -> std::io::Result<Self> {
175 if let Some(parent) = path.parent() {
176 std::fs::create_dir_all(parent)?;
177 }
178 let parent = path.parent().unwrap_or_else(|| Path::new("/tmp"));
179 let lock_path = parent.join(format!(
180 "{}.lock",
181 path.file_name()
182 .and_then(|s| s.to_str())
183 .unwrap_or("scripts.rkyv")
184 ));
185 Ok(Self {
186 path: path.to_path_buf(),
187 lock_path,
188 mmap: Mutex::new(None),
189 })
190 }
191
192 fn ensure_mmap(&self) {
193 let mut guard = self.mmap.lock();
194 if guard.is_none() {
195 *guard = MmappedShard::open(&self.path);
196 }
197 }
198
199 fn invalidate_mmap(&self) {
200 let mut guard = self.mmap.lock();
201 *guard = None;
202 }
203
204 pub fn get(&self, path: &str, mtime_secs: i64, mtime_nsecs: i64) -> Option<Vec<u8>> {
207 self.ensure_mmap();
208 let guard = self.mmap.lock();
209 let shard = guard.as_ref()?;
210 if !shard.header_ok() {
211 return None;
212 }
213 let entry = shard.lookup(path)?;
214
215 let entry_mtime_s: i64 = entry.mtime_secs.into();
216 let entry_mtime_ns: i64 = entry.mtime_nsecs.into();
217 if entry_mtime_s != mtime_secs || entry_mtime_ns != mtime_nsecs {
218 return None;
219 }
220
221 match current_binary_identity() {
234 Some((bin_mtime, bin_len)) => {
235 let cached_bin_mtime: i64 = entry.binary_mtime_at_cache.into();
236 let cached_bin_len: u64 = entry.binary_len_at_cache.into();
237 if cached_bin_mtime != bin_mtime || cached_bin_len != bin_len {
238 return None;
239 }
240 }
241 None => return None,
243 }
244
245 Some(entry.chunk_blob.as_slice().to_vec())
246 }
247
248 pub fn put(
250 &self,
251 path: &str,
252 mtime_secs: i64,
253 mtime_nsecs: i64,
254 chunk_blob: Vec<u8>,
255 ) -> Result<(), String> {
256 let _lock = match acquire_lock(&self.lock_path) {
257 Some(l) => l,
258 None => return Ok(()),
259 };
260
261 let mut shard = match read_owned_shard(&self.path) {
262 Some(s)
263 if s.header.zshrs_version == env!("CARGO_PKG_VERSION")
264 && s.header.pointer_width as usize == std::mem::size_of::<usize>()
265 && s.header.format_version == SHARD_FORMAT_VERSION =>
266 {
267 s
268 }
269 _ => fresh_shard(),
270 };
271
272 let (bin_mtime, bin_len) = current_binary_identity().unwrap_or((0, 0));
273 let entry = ScriptEntry {
274 mtime_secs,
275 mtime_nsecs,
276 binary_mtime_at_cache: bin_mtime,
277 binary_len_at_cache: bin_len,
278 cached_at_secs: now_secs(),
279 chunk_blob,
280 };
281 shard.entries.insert(path.to_string(), entry);
282 shard.header.built_at_secs = now_secs() as u64;
283
284 write_shard_atomic(&self.path, &shard)?;
285 self.invalidate_mmap();
286 Ok(())
287 }
288
289 pub fn stats(&self) -> (i64, i64) {
291 self.ensure_mmap();
292 let guard = self.mmap.lock();
293 let Some(shard) = guard.as_ref() else {
294 return (0, 0);
295 };
296 let count = shard.entry_count() as i64;
297 let bytes: i64 = shard
298 .shard()
299 .entries
300 .values()
301 .map(|e| e.chunk_blob.len() as i64)
302 .sum();
303 (count, bytes)
304 }
305
306 pub fn list_scripts(&self) -> Vec<(String, f64, String, String)> {
309 self.ensure_mmap();
310 let guard = self.mmap.lock();
311 let Some(shard) = guard.as_ref() else {
312 return Vec::new();
313 };
314 let v = shard.shard().header.zshrs_version.as_str().to_string();
315 let mut out: Vec<(String, f64, String, String, i64)> = shard
316 .shard()
317 .entries
318 .iter()
319 .map(|(k, e)| {
320 let chunk_kb = e.chunk_blob.len() as f64 / 1024.0;
321 let cached_at: i64 = e.cached_at_secs.into();
322 let ts = format_local_ts(cached_at);
323 (k.as_str().to_string(), chunk_kb, v.clone(), ts, cached_at)
324 })
325 .collect();
326 out.sort_by_key(|x| std::cmp::Reverse(x.4));
327 out.into_iter()
328 .map(|(p, ck, ver, ts, _)| (p, ck, ver, ts))
329 .collect()
330 }
331
332 pub fn evict_stale(&self) -> usize {
334 let _lock = match acquire_lock(&self.lock_path) {
335 Some(l) => l,
336 None => return 0,
337 };
338 let mut shard = match read_owned_shard(&self.path) {
339 Some(s) => s,
340 None => return 0,
341 };
342 let before = shard.entries.len();
343 shard.entries.retain(|p, e| match file_mtime(Path::new(p)) {
344 Some((s, ns)) => s == e.mtime_secs && ns == e.mtime_nsecs,
345 None => false,
346 });
347 let evicted = before - shard.entries.len();
348 if evicted > 0 {
349 let _ = write_shard_atomic(&self.path, &shard);
350 self.invalidate_mmap();
351 }
352 evicted
353 }
354 pub fn clear(&self) -> std::io::Result<()> {
356 let _lock = acquire_lock(&self.lock_path);
357 let res = match std::fs::remove_file(&self.path) {
358 Ok(()) => Ok(()),
359 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
360 Err(e) => Err(e),
361 };
362 self.invalidate_mmap();
363 res
364 }
365}
366
367fn acquire_lock(path: &Path) -> Option<nix::fcntl::Flock<File>> {
368 let f = File::options()
369 .read(true)
370 .write(true)
371 .create(true)
372 .truncate(false)
373 .open(path)
374 .ok()?;
375 nix::fcntl::Flock::lock(f, nix::fcntl::FlockArg::LockExclusive).ok()
376}
377
378fn fresh_shard() -> ScriptShard {
379 ScriptShard {
380 header: ShardHeader {
381 magic: SHARD_MAGIC,
382 format_version: SHARD_FORMAT_VERSION,
383 zshrs_version: env!("CARGO_PKG_VERSION").to_string(),
384 pointer_width: std::mem::size_of::<usize>() as u32,
385 built_at_secs: now_secs() as u64,
386 },
387 entries: HashMap::new(),
388 }
389}
390
391fn read_owned_shard(path: &Path) -> Option<ScriptShard> {
392 let bytes = std::fs::read(path).ok()?;
393 let archived = rkyv::check_archived_root::<ScriptShard>(&bytes[..]).ok()?;
394 archived.deserialize(&mut rkyv::Infallible).ok()
395}
396
397fn write_shard_atomic(path: &Path, shard: &ScriptShard) -> Result<(), String> {
398 let bytes = rkyv::to_bytes::<_, 4096>(shard).map_err(|e| format!("rkyv serialize: {}", e))?;
399 crate::atomic_write::write_bytes_atomic(path, &bytes)
403}
404
405fn now_secs() -> i64 {
406 SystemTime::now()
407 .duration_since(UNIX_EPOCH)
408 .map(|d| d.as_secs() as i64)
409 .unwrap_or(0)
410}
411
412fn format_local_ts(secs: i64) -> String {
413 let dt = chrono::DateTime::<chrono::Local>::from(
414 UNIX_EPOCH + std::time::Duration::from_secs(secs.max(0) as u64),
415 );
416 dt.format("%Y-%m-%d %H:%M:%S").to_string()
417}
418pub fn file_mtime(path: &Path) -> Option<(i64, i64)> {
420 let meta = std::fs::metadata(path).ok()?;
421 Some((meta.mtime(), meta.mtime_nsec()))
422}
423
424fn current_binary_identity() -> Option<(i64, u64)> {
427 static BIN_ID: OnceLock<Option<(i64, u64)>> = OnceLock::new();
428 *BIN_ID.get_or_init(|| {
429 let exe = std::env::current_exe().ok()?;
430 let meta = std::fs::metadata(&exe).ok()?;
431 Some((meta.mtime(), meta.len()))
432 })
433}
434
435pub fn default_cache_path() -> PathBuf {
446 let root = if let Some(custom) = std::env::var_os("ZSHRS_HOME") {
447 PathBuf::from(custom)
448 } else {
449 dirs::home_dir()
450 .unwrap_or_else(|| PathBuf::from("/tmp"))
451 .join(".zshrs")
452 };
453 root.join("scripts.rkyv")
454}
455
456pub static CACHE_DISABLED: std::sync::atomic::AtomicBool =
461 std::sync::atomic::AtomicBool::new(false);
462
463pub fn cache_enabled() -> bool {
466 if CACHE_DISABLED.load(std::sync::atomic::Ordering::Relaxed) {
467 return false;
468 }
469 !matches!(
470 std::env::var("ZSHRS_CACHE").as_deref(),
471 Ok("0") | Ok("false") | Ok("no")
472 )
473}
474
475pub static CACHE: once_cell::sync::Lazy<Option<ScriptCache>> = once_cell::sync::Lazy::new(|| {
478 if !cache_enabled() {
479 return None;
480 }
481 ScriptCache::open(&default_cache_path()).ok()
482});
483
484pub fn try_load_bytes(path: &Path) -> Option<Vec<u8>> {
486 let cache = CACHE.as_ref()?;
487 let canonical = path.canonicalize().ok()?;
488 let path_str = canonical.to_string_lossy();
489 let (mtime_s, mtime_ns) = file_mtime(&canonical)?;
490 cache.get(&path_str, mtime_s, mtime_ns)
491}
492
493pub fn try_save_bytes(path: &Path, chunk_blob: &[u8]) -> Result<(), String> {
497 let Some(cache) = CACHE.as_ref() else {
498 return Ok(());
499 };
500 let canonical = match path.canonicalize() {
501 Ok(p) => p,
502 Err(_) => return Ok(()),
503 };
504 let path_str = canonical.to_string_lossy();
505 let (mtime_s, mtime_ns) = match file_mtime(&canonical) {
506 Some(m) => m,
507 None => return Ok(()),
508 };
509 cache.put(&path_str, mtime_s, mtime_ns, chunk_blob.to_vec())
510}
511pub fn stats() -> Option<(i64, i64)> {
513 CACHE.as_ref().map(|c| c.stats())
514}
515pub fn evict_stale() -> usize {
517 CACHE.as_ref().map(|c| c.evict_stale()).unwrap_or(0)
518}
519pub fn clear() -> bool {
521 CACHE.as_ref().map(|c| c.clear().is_ok()).unwrap_or(false)
522}
523
524#[cfg(test)]
525mod tests {
526 use super::*;
527 use tempfile::tempdir;
528
529 #[test]
530 fn round_trip() {
531 let _g = crate::test_util::global_state_lock();
532 let dir = tempdir().unwrap();
533 let cache_path = dir.path().join("scripts.rkyv");
534 let cache = ScriptCache::open(&cache_path).unwrap();
535
536 let script_path = dir.path().join("test.zsh");
537 std::fs::write(&script_path, "echo hi").unwrap();
538
539 let (mtime_s, mtime_ns) = file_mtime(&script_path).unwrap();
540 let path_str = script_path.to_string_lossy().to_string();
541
542 let blob = vec![1u8, 2, 3, 4, 5];
543 cache
544 .put(&path_str, mtime_s, mtime_ns, blob.clone())
545 .unwrap();
546
547 let loaded = cache.get(&path_str, mtime_s, mtime_ns).unwrap();
548 assert_eq!(loaded, blob);
549
550 let (count, _bytes) = cache.stats();
551 assert_eq!(count, 1);
552 }
553
554 #[test]
555 fn mtime_invalidation() {
556 let _g = crate::test_util::global_state_lock();
557 let dir = tempdir().unwrap();
558 let cache_path = dir.path().join("scripts.rkyv");
559 let cache = ScriptCache::open(&cache_path).unwrap();
560
561 let script_path = dir.path().join("test.zsh");
562 std::fs::write(&script_path, "echo hi").unwrap();
563
564 let (mtime_s, mtime_ns) = file_mtime(&script_path).unwrap();
565 let path_str = script_path.to_string_lossy().to_string();
566 cache.put(&path_str, mtime_s, mtime_ns, vec![9u8]).unwrap();
567
568 assert!(cache.get(&path_str, mtime_s + 1, mtime_ns).is_none());
569 }
570
571 #[test]
578 fn an_entry_from_another_binary_is_never_served() {
579 let _g = crate::test_util::global_state_lock();
580 let dir = tempdir().unwrap();
581 let cache_path = dir.path().join("scripts.rkyv");
582 let cache = ScriptCache::open(&cache_path).unwrap();
583
584 let script_path = dir.path().join("test.zsh");
585 std::fs::write(&script_path, "echo hi").unwrap();
586 let (mtime_s, mtime_ns) = file_mtime(&script_path).unwrap();
587 let path_str = script_path.to_string_lossy().to_string();
588 cache.put(&path_str, mtime_s, mtime_ns, vec![9u8]).unwrap();
589 assert_eq!(
590 cache.get(&path_str, mtime_s, mtime_ns),
591 Some(vec![9u8]),
592 "the emitting binary must hit its own entry",
593 );
594
595 let mut shard = read_owned_shard(&cache_path).expect("shard readable");
598 shard
599 .entries
600 .get_mut(&path_str)
601 .expect("entry present")
602 .binary_mtime_at_cache += 10_000;
603 write_shard_atomic(&cache_path, &shard).unwrap();
604 let reopened = ScriptCache::open(&cache_path).unwrap();
605 assert!(
606 reopened.get(&path_str, mtime_s, mtime_ns).is_none(),
607 "a chunk from a newer build was accepted",
608 );
609
610 let mut shard = read_owned_shard(&cache_path).expect("shard readable");
613 let entry = shard.entries.get_mut(&path_str).expect("entry present");
614 entry.binary_mtime_at_cache -= 10_000;
615 entry.binary_len_at_cache += 1;
616 write_shard_atomic(&cache_path, &shard).unwrap();
617 let reopened = ScriptCache::open(&cache_path).unwrap();
618 assert!(
619 reopened.get(&path_str, mtime_s, mtime_ns).is_none(),
620 "a chunk from a same-second build of a different size was accepted",
621 );
622 }
623
624 #[test]
625 fn second_put_replaces_first() {
626 let _g = crate::test_util::global_state_lock();
627 let dir = tempdir().unwrap();
628 let cache_path = dir.path().join("scripts.rkyv");
629 let cache = ScriptCache::open(&cache_path).unwrap();
630
631 let p1 = dir.path().join("a.zsh");
632 let p2 = dir.path().join("b.zsh");
633 std::fs::write(&p1, "1").unwrap();
634 std::fs::write(&p2, "2").unwrap();
635
636 let (m1s, m1n) = file_mtime(&p1).unwrap();
637 let (m2s, m2n) = file_mtime(&p2).unwrap();
638
639 cache
640 .put(&p1.to_string_lossy(), m1s, m1n, vec![1u8])
641 .unwrap();
642 cache
643 .put(&p2.to_string_lossy(), m2s, m2n, vec![2u8])
644 .unwrap();
645
646 let (count, _) = cache.stats();
647 assert_eq!(count, 2);
648 assert!(cache.get(&p1.to_string_lossy(), m1s, m1n).is_some());
649 assert!(cache.get(&p2.to_string_lossy(), m2s, m2n).is_some());
650 }
651
652 #[test]
653 fn corrupt_file_returns_no_mmap() {
654 let _g = crate::test_util::global_state_lock();
655 let dir = tempdir().unwrap();
656 let cache_path = dir.path().join("scripts.rkyv");
657 std::fs::write(&cache_path, b"this is not a valid rkyv archive").unwrap();
658 let cache = ScriptCache::open(&cache_path).unwrap();
659 assert!(cache.get("/nope", 0, 0).is_none());
660 }
661
662 #[test]
663 fn clear_removes_file() {
664 let _g = crate::test_util::global_state_lock();
665 let dir = tempdir().unwrap();
666 let cache_path = dir.path().join("scripts.rkyv");
667 let cache = ScriptCache::open(&cache_path).unwrap();
668
669 let script_path = dir.path().join("test.zsh");
670 std::fs::write(&script_path, "echo hi").unwrap();
671 let (mtime_s, mtime_ns) = file_mtime(&script_path).unwrap();
672 cache
673 .put(&script_path.to_string_lossy(), mtime_s, mtime_ns, vec![7u8])
674 .unwrap();
675 assert!(cache_path.exists());
676
677 cache.clear().unwrap();
678 assert!(!cache_path.exists());
679 }
680
681 #[test]
686 fn now_secs_is_positive_and_within_realistic_range() {
687 let _g = crate::test_util::global_state_lock();
688 let n = now_secs();
689 assert!(
691 (1_577_836_800..4_102_444_800).contains(&n),
692 "now_secs out of plausible range: {}",
693 n
694 );
695 }
696
697 #[test]
698 fn now_secs_does_not_go_backwards_in_quick_succession() {
699 let _g = crate::test_util::global_state_lock();
700 let a = now_secs();
701 let b = now_secs();
702 assert!(b >= a, "now_secs went backwards: {} -> {}", a, b);
703 }
704
705 #[test]
710 fn format_local_ts_includes_year_and_punctuation() {
711 let _g = crate::test_util::global_state_lock();
712 let s = format_local_ts(1_704_067_200);
715 assert!(s.starts_with("202"), "expected 21st century year: {}", s);
716 assert!(s.contains('-'), "expected dash separator: {}", s);
717 assert!(s.contains(':'), "expected colon separator: {}", s);
718 }
719
720 #[test]
721 fn format_local_ts_length_matches_pattern() {
722 let _g = crate::test_util::global_state_lock();
723 let s = format_local_ts(1_700_000_000);
724 assert_eq!(s.len(), 19, "unexpected width: {}", s);
726 }
727
728 #[test]
729 fn format_local_ts_handles_zero_secs_via_clamp() {
730 let _g = crate::test_util::global_state_lock();
731 let s = format_local_ts(0);
733 assert_eq!(s.len(), 19);
734 assert!(s.starts_with("19"), "expected 1970-ish year: {}", s);
735 }
736
737 #[test]
738 fn format_local_ts_negative_clamped_to_zero() {
739 let _g = crate::test_util::global_state_lock();
740 let s = format_local_ts(-1_000_000);
742 assert_eq!(s.len(), 19);
743 }
744
745 #[test]
750 fn file_mtime_returns_some_for_real_file() {
751 let _g = crate::test_util::global_state_lock();
752 let dir = tempdir().unwrap();
753 let p = dir.path().join("foo.zsh");
754 std::fs::write(&p, b"x").unwrap();
755 let (s, _ns) = file_mtime(&p).unwrap();
756 assert!(s > 0);
757 }
758
759 #[test]
760 fn file_mtime_returns_none_for_missing_path() {
761 let _g = crate::test_util::global_state_lock();
762 assert!(file_mtime(Path::new("/nonexistent/zshrs/script_cache_missing.bin")).is_none());
763 }
764
765 #[test]
770 fn default_cache_path_ends_in_scripts_rkyv() {
771 let _g = crate::test_util::global_state_lock();
772 let p = default_cache_path();
773 assert_eq!(p.file_name().and_then(|s| s.to_str()), Some("scripts.rkyv"));
774 }
775
776 #[test]
777 fn cache_enabled_true_when_env_unset() {
778 let _g = crate::test_util::global_state_lock();
779 let prev = std::env::var_os("ZSHRS_CACHE");
780 std::env::remove_var("ZSHRS_CACHE");
781 let on = cache_enabled();
782 if let Some(v) = prev {
783 std::env::set_var("ZSHRS_CACHE", v);
784 }
785 assert!(on, "cache should be enabled when ZSHRS_CACHE is unset");
786 }
787
788 #[test]
789 fn cache_enabled_false_when_env_is_zero_false_or_no() {
790 let _g = crate::test_util::global_state_lock();
791 let prev = std::env::var_os("ZSHRS_CACHE");
792 for v in ["0", "false", "no"] {
793 std::env::set_var("ZSHRS_CACHE", v);
794 assert!(!cache_enabled(), "ZSHRS_CACHE={} must disable cache", v);
795 }
796 if let Some(v) = prev {
797 std::env::set_var("ZSHRS_CACHE", v);
798 } else {
799 std::env::remove_var("ZSHRS_CACHE");
800 }
801 }
802
803 #[test]
804 fn cache_enabled_true_for_other_env_values() {
805 let _g = crate::test_util::global_state_lock();
808 let prev = std::env::var_os("ZSHRS_CACHE");
809 for v in ["1", "true", "yes", "on", ""] {
810 std::env::set_var("ZSHRS_CACHE", v);
811 assert!(
812 cache_enabled(),
813 "ZSHRS_CACHE={:?} must NOT disable cache",
814 v
815 );
816 }
817 if let Some(v) = prev {
818 std::env::set_var("ZSHRS_CACHE", v);
819 } else {
820 std::env::remove_var("ZSHRS_CACHE");
821 }
822 }
823}