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