zsh/extensions/
autoload_cache.rs

1//! rkyv-backed bytecode cache for autoload functions.
2//!
3//! Single-file shard at `~/.zshrs/autoloads.rkyv`. Keyed by function name
4//! (not file path) — autoload bytecode is identified by the resolved function
5//! name, regardless of which fpath dir or .zwc archive it came from.
6//!
7//! Storage layout (rkyv archived):
8//!   AutoloadShard {
9//!     header: { magic, format_version, zshrs_version, pointer_width, built_at_secs },
10//!     entries: HashMap<function_name, AutoloadEntry>,
11//!   }
12//!   AutoloadEntry { binary_mtime_at_cache, cached_at_secs, chunk_blob: `Vec<u8>` }
13//!
14//! Inner `chunk_blob` is bincode-encoded `fusevm::Chunk` (same constraint as
15//! [`script_cache`](crate::script_cache) module — `fusevm::Chunk` is upstream and only derives serde).
16//!
17//! Invalidation:
18//!   - zshrs binary mtime newer than `binary_mtime_at_cache` ⇒ entry stale
19//!     (any zshrs rebuild silently invalidates the whole shard).
20//!   - There is no per-source-file mtime check here. Autoload bodies live in
21//!     fpath dirs / .zwc archives and the existing `crate::compsys::cache::autoloads`
22//!     SQLite row tracks the source file/offset/size. Rebuild logic relies on
23//!     `compinit` clearing the whole rkyv shard at recompute time (see
24//!     `AutoloadShardWriter` — used by the compinit bulk-prewarm path).
25//!
26//! Bulk-write: compinit prewarms 16k+ autoload bytecodes in one go. Per-batch
27//! shard rewrites (the SQLite-era pattern) would re-serialize 16k entries
28//! 160 times. Instead the new API exposes `AutoloadShardWriter`: accumulate
29//! all `(name, blob)` pairs in memory, then `commit()` writes the shard once.
30//! The single-add `try_save_one` path remains for the cold-start case where
31//! one autoload at a time is compiled by the interactive shell.
32//!
33//! The on-disk shape mirrors [`ScriptShard`](crate::script_cache::ScriptShard) — same header,
34//! same magic-version-pointer_width discipline, same atomic-rename writes.
35
36use std::collections::HashMap;
37use std::fs::File;
38use std::io::Write as IoWrite;
39use std::path::{Path, PathBuf};
40use std::sync::OnceLock;
41use std::time::{SystemTime, UNIX_EPOCH};
42
43use memmap2::Mmap;
44use parking_lot::Mutex;
45use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
46use std::os::unix::fs::MetadataExt;
47
48/// "ZRAL" little-endian.
49pub const SHARD_MAGIC: u32 = 0x5A52414C;
50/// `SHARD_FORMAT_VERSION` constant.
51pub const SHARD_FORMAT_VERSION: u32 = 1;
52/// `ShardHeader` — see fields for layout.
53#[derive(Archive, RkyvDeserialize, RkyvSerialize, Debug, Clone)]
54#[archive(check_bytes)]
55pub struct ShardHeader {
56    /// `magic` field.
57    pub magic: u32,
58    /// `format_version` field.
59    pub format_version: u32,
60    /// `zshrs_version` field.
61    pub zshrs_version: String,
62    /// `pointer_width` field.
63    pub pointer_width: u32,
64    /// `built_at_secs` field.
65    pub built_at_secs: u64,
66}
67/// `AutoloadEntry` — see fields for layout.
68#[derive(Archive, RkyvDeserialize, RkyvSerialize, Debug, Clone)]
69#[archive(check_bytes)]
70pub struct AutoloadEntry {
71    /// `binary_mtime_at_cache` field.
72    pub binary_mtime_at_cache: i64,
73    /// `cached_at_secs` field.
74    pub cached_at_secs: i64,
75    /// `chunk_blob` field.
76    pub chunk_blob: Vec<u8>,
77}
78/// `AutoloadShard` — see fields for layout.
79#[derive(Archive, RkyvDeserialize, RkyvSerialize, Debug, Clone)]
80#[archive(check_bytes)]
81pub struct AutoloadShard {
82    /// `header` field.
83    pub header: ShardHeader,
84    /// `entries` field.
85    pub entries: HashMap<String, AutoloadEntry>,
86}
87/// `MmappedShard` — see fields for layout.
88pub struct MmappedShard {
89    /// `_mmap` field.
90    _mmap: Mmap,
91    /// `archived` field.
92    archived: *const ArchivedAutoloadShard,
93}
94
95unsafe impl Send for MmappedShard {}
96unsafe impl Sync for MmappedShard {}
97
98impl MmappedShard {
99    /// `open` — see implementation.
100    pub fn open(path: &Path) -> Option<Self> {
101        let file = File::open(path).ok()?;
102        let mmap = unsafe { Mmap::map(&file).ok()? };
103        let archived = rkyv::check_archived_root::<AutoloadShard>(&mmap[..]).ok()?;
104        let archived_ptr = archived as *const ArchivedAutoloadShard;
105        Some(Self {
106            _mmap: mmap,
107            archived: archived_ptr,
108        })
109    }
110
111    fn shard(&self) -> &ArchivedAutoloadShard {
112        unsafe { &*self.archived }
113    }
114
115    fn header_ok(&self) -> bool {
116        let h = &self.shard().header;
117        let magic: u32 = h.magic.into();
118        let fv: u32 = h.format_version.into();
119        let pw: u32 = h.pointer_width.into();
120        magic == SHARD_MAGIC
121            && fv == SHARD_FORMAT_VERSION
122            && pw as usize == std::mem::size_of::<usize>()
123            && h.zshrs_version.as_str() == env!("CARGO_PKG_VERSION")
124    }
125
126    fn lookup(&self, name: &str) -> Option<&ArchivedAutoloadEntry> {
127        self.shard().entries.get(name)
128    }
129}
130/// `AutoloadCache` — see fields for layout.
131pub struct AutoloadCache {
132    /// `path` field.
133    path: PathBuf,
134    /// `lock_path` field.
135    lock_path: PathBuf,
136    /// `mmap` field.
137    mmap: Mutex<Option<MmappedShard>>,
138}
139
140impl AutoloadCache {
141    /// `open` — see implementation.
142    pub fn open(path: &Path) -> std::io::Result<Self> {
143        if let Some(parent) = path.parent() {
144            std::fs::create_dir_all(parent)?;
145        }
146        let parent = path.parent().unwrap_or_else(|| Path::new("/tmp"));
147        let lock_path = parent.join(format!(
148            "{}.lock",
149            path.file_name()
150                .and_then(|s| s.to_str())
151                .unwrap_or("autoloads.rkyv")
152        ));
153        Ok(Self {
154            path: path.to_path_buf(),
155            lock_path,
156            mmap: Mutex::new(None),
157        })
158    }
159
160    fn ensure_mmap(&self) {
161        let mut guard = self.mmap.lock();
162        if guard.is_none() {
163            *guard = MmappedShard::open(&self.path);
164        }
165    }
166
167    fn invalidate_mmap(&self) {
168        let mut guard = self.mmap.lock();
169        *guard = None;
170    }
171    /// `get` — see implementation.
172    pub fn get(&self, name: &str) -> Option<Vec<u8>> {
173        self.ensure_mmap();
174        let guard = self.mmap.lock();
175        let shard = guard.as_ref()?;
176        if !shard.header_ok() {
177            return None;
178        }
179        let entry = shard.lookup(name)?;
180        if let Some(bin_mtime) = current_binary_mtime_secs() {
181            let cached_bin_mtime: i64 = entry.binary_mtime_at_cache.into();
182            if cached_bin_mtime < bin_mtime {
183                return None;
184            }
185        }
186        Some(entry.chunk_blob.as_slice().to_vec())
187    }
188
189    /// Single-write: read shard, insert one entry, write shard. Used by the
190    /// cold-start path when a function is autoloaded before compinit
191    /// pre-warm completes.
192    pub fn put_one(&self, name: &str, chunk_blob: Vec<u8>) -> Result<(), String> {
193        let _lock = match acquire_lock(&self.lock_path) {
194            Some(l) => l,
195            None => return Ok(()),
196        };
197        let mut shard = match read_owned_shard(&self.path) {
198            Some(s)
199                if s.header.zshrs_version == env!("CARGO_PKG_VERSION")
200                    && s.header.pointer_width as usize == std::mem::size_of::<usize>()
201                    && s.header.format_version == SHARD_FORMAT_VERSION =>
202            {
203                s
204            }
205            _ => fresh_shard(),
206        };
207        let bin_mtime = current_binary_mtime_secs().unwrap_or(0);
208        shard.entries.insert(
209            name.to_string(),
210            AutoloadEntry {
211                binary_mtime_at_cache: bin_mtime,
212                cached_at_secs: now_secs(),
213                chunk_blob,
214            },
215        );
216        shard.header.built_at_secs = now_secs() as u64;
217        write_shard_atomic(&self.path, &shard)?;
218        self.invalidate_mmap();
219        Ok(())
220    }
221
222    /// Merge `entries` into the existing shard, inserting/replacing each one.
223    /// Used by compinit's BACKFILL path — when an existing shard is missing
224    /// some entries (e.g. binary mtime bump invalidated a subset), the
225    /// caller computes the missing names + chunks and merges them in
226    /// without touching unrelated entries. Single read + single write,
227    /// even for 16k entries.
228    pub fn merge_in(&self, entries: HashMap<String, Vec<u8>>) -> Result<(), String> {
229        if entries.is_empty() {
230            return Ok(());
231        }
232        let _lock = match acquire_lock(&self.lock_path) {
233            Some(l) => l,
234            None => return Ok(()),
235        };
236        let mut shard = match read_owned_shard(&self.path) {
237            Some(s)
238                if s.header.zshrs_version == env!("CARGO_PKG_VERSION")
239                    && s.header.pointer_width as usize == std::mem::size_of::<usize>()
240                    && s.header.format_version == SHARD_FORMAT_VERSION =>
241            {
242                s
243            }
244            _ => fresh_shard(),
245        };
246        let bin_mtime = current_binary_mtime_secs().unwrap_or(0);
247        let now = now_secs();
248        for (name, chunk_blob) in entries {
249            shard.entries.insert(
250                name,
251                AutoloadEntry {
252                    binary_mtime_at_cache: bin_mtime,
253                    cached_at_secs: now,
254                    chunk_blob,
255                },
256            );
257        }
258        shard.header.built_at_secs = now as u64;
259        write_shard_atomic(&self.path, &shard)?;
260        self.invalidate_mmap();
261        Ok(())
262    }
263
264    /// Replace the entire shard with the given entries. Used by compinit's
265    /// bulk pre-warm — accumulate all (name, chunk_blob) pairs, then commit
266    /// once. Avoids re-serializing 16k entries on every batch flush.
267    pub fn replace_all(&self, entries: HashMap<String, Vec<u8>>) -> Result<(), String> {
268        let _lock = match acquire_lock(&self.lock_path) {
269            Some(l) => l,
270            None => return Ok(()),
271        };
272        let bin_mtime = current_binary_mtime_secs().unwrap_or(0);
273        let now = now_secs();
274        let mut shard = fresh_shard();
275        for (name, chunk_blob) in entries {
276            shard.entries.insert(
277                name,
278                AutoloadEntry {
279                    binary_mtime_at_cache: bin_mtime,
280                    cached_at_secs: now,
281                    chunk_blob,
282                },
283            );
284        }
285        write_shard_atomic(&self.path, &shard)?;
286        self.invalidate_mmap();
287        Ok(())
288    }
289    /// `entry_count` — see implementation.
290    pub fn entry_count(&self) -> usize {
291        self.ensure_mmap();
292        let guard = self.mmap.lock();
293        guard.as_ref().map(|s| s.shard().entries.len()).unwrap_or(0)
294    }
295
296    /// Set of cached function names — caller can subtract this from "all
297    /// known autoload names" to compute the missing-bytecode set without a
298    /// SQL JOIN.
299    pub fn cached_names(&self) -> std::collections::HashSet<String> {
300        self.ensure_mmap();
301        let guard = self.mmap.lock();
302        let Some(shard) = guard.as_ref() else {
303            return std::collections::HashSet::new();
304        };
305        shard
306            .shard()
307            .entries
308            .keys()
309            .map(|k| k.as_str().to_string())
310            .collect()
311    }
312    /// `stats` — see implementation.
313    pub fn stats(&self) -> (i64, i64) {
314        self.ensure_mmap();
315        let guard = self.mmap.lock();
316        let Some(shard) = guard.as_ref() else {
317            return (0, 0);
318        };
319        let count = shard.shard().entries.len() as i64;
320        let bytes: i64 = shard
321            .shard()
322            .entries
323            .values()
324            .map(|e| e.chunk_blob.len() as i64)
325            .sum();
326        (count, bytes)
327    }
328    /// `clear` — see implementation.
329    pub fn clear(&self) -> std::io::Result<()> {
330        let _lock = acquire_lock(&self.lock_path);
331        let res = match std::fs::remove_file(&self.path) {
332            Ok(()) => Ok(()),
333            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
334            Err(e) => Err(e),
335        };
336        self.invalidate_mmap();
337        res
338    }
339}
340
341fn acquire_lock(path: &Path) -> Option<nix::fcntl::Flock<File>> {
342    let f = File::options()
343        .read(true)
344        .write(true)
345        .create(true)
346        .truncate(false)
347        .open(path)
348        .ok()?;
349    nix::fcntl::Flock::lock(f, nix::fcntl::FlockArg::LockExclusive).ok()
350}
351
352fn fresh_shard() -> AutoloadShard {
353    AutoloadShard {
354        header: ShardHeader {
355            magic: SHARD_MAGIC,
356            format_version: SHARD_FORMAT_VERSION,
357            zshrs_version: env!("CARGO_PKG_VERSION").to_string(),
358            pointer_width: std::mem::size_of::<usize>() as u32,
359            built_at_secs: now_secs() as u64,
360        },
361        entries: HashMap::new(),
362    }
363}
364
365fn read_owned_shard(path: &Path) -> Option<AutoloadShard> {
366    let bytes = std::fs::read(path).ok()?;
367    let archived = rkyv::check_archived_root::<AutoloadShard>(&bytes[..]).ok()?;
368    archived.deserialize(&mut rkyv::Infallible).ok()
369}
370
371fn write_shard_atomic(path: &Path, shard: &AutoloadShard) -> Result<(), String> {
372    let bytes = rkyv::to_bytes::<_, 4096>(shard).map_err(|e| format!("rkyv serialize: {}", e))?;
373    let parent = path.parent().expect("cache path has parent");
374    let _ = std::fs::create_dir_all(parent);
375    let pid = std::process::id();
376    let nanos = SystemTime::now()
377        .duration_since(UNIX_EPOCH)
378        .map(|d| d.as_nanos())
379        .unwrap_or(0);
380    let tmp_path = parent.join(format!(
381        "{}.tmp.{}.{}",
382        path.file_name()
383            .and_then(|s| s.to_str())
384            .unwrap_or("autoloads.rkyv"),
385        pid,
386        nanos
387    ));
388    {
389        let mut f = File::create(&tmp_path).map_err(|e| e.to_string())?;
390        f.write_all(&bytes).map_err(|e| e.to_string())?;
391        f.sync_all().map_err(|e| e.to_string())?;
392    }
393    std::fs::rename(&tmp_path, path).map_err(|e| e.to_string())?;
394    Ok(())
395}
396
397fn now_secs() -> i64 {
398    SystemTime::now()
399        .duration_since(UNIX_EPOCH)
400        .map(|d| d.as_secs() as i64)
401        .unwrap_or(0)
402}
403
404fn file_mtime(path: &Path) -> Option<(i64, i64)> {
405    let meta = std::fs::metadata(path).ok()?;
406    Some((meta.mtime(), meta.mtime_nsec()))
407}
408
409fn current_binary_mtime_secs() -> Option<i64> {
410    static BIN_MTIME: OnceLock<Option<i64>> = OnceLock::new();
411    *BIN_MTIME.get_or_init(|| {
412        let exe = std::env::current_exe().ok()?;
413        let (secs, _) = file_mtime(&exe)?;
414        Some(secs)
415    })
416}
417/// `default_cache_path` — see implementation.
418///
419/// All zshrs state lives under `$ZSHRS_HOME` (default `~/.zshrs`),
420/// matching the daemon's `CachePaths` convention (daemon/paths.rs)
421/// and the `~/.zinit`/`~/.zpwr`/`~/.oh-my-zsh` precedent. Project
422/// policy forbids `~/.cache/zshrs/` and `~/Library/Caches/zshrs/`
423/// — both of which `dirs::cache_dir()` resolves to.
424pub fn default_cache_path() -> PathBuf {
425    let root = if let Some(custom) = std::env::var_os("ZSHRS_HOME") {
426        PathBuf::from(custom)
427    } else {
428        dirs::home_dir()
429            .unwrap_or_else(|| PathBuf::from("/tmp"))
430            .join(".zshrs")
431    };
432    root.join("autoloads.rkyv")
433}
434/// `cache_enabled` — see implementation. Honors the process-local
435/// `script_cache::CACHE_DISABLED` AtomicBool first so parity-mode
436/// init can disable caches without exporting `ZSHRS_CACHE=0` (which
437/// would otherwise leak into `${(k)parameters}`).
438pub fn cache_enabled() -> bool {
439    if crate::extensions::script_cache::CACHE_DISABLED.load(std::sync::atomic::Ordering::Relaxed) {
440        return false;
441    }
442    !matches!(
443        std::env::var("ZSHRS_CACHE").as_deref(),
444        Ok("0") | Ok("false") | Ok("no")
445    )
446}
447/// `CACHE` static.
448pub static CACHE: once_cell::sync::Lazy<Option<AutoloadCache>> = once_cell::sync::Lazy::new(|| {
449    if !cache_enabled() {
450        return None;
451    }
452    AutoloadCache::open(&default_cache_path()).ok()
453});
454/// `try_load` — see implementation.
455pub fn try_load(name: &str) -> Option<Vec<u8>> {
456    let cache = CACHE.as_ref()?;
457    cache.get(name)
458}
459/// `try_save_one` — see implementation.
460pub fn try_save_one(name: &str, chunk_blob: &[u8]) -> Result<(), String> {
461    let Some(cache) = CACHE.as_ref() else {
462        return Ok(());
463    };
464    cache.put_one(name, chunk_blob.to_vec())
465}
466
467/// Replace the entire autoload shard with the given entries. Use this from
468/// compinit's bulk pre-warm path — accumulates all `(name, chunk_blob)` in
469/// the `entries` HashMap and writes the shard exactly once.
470pub fn try_replace_all(entries: HashMap<String, Vec<u8>>) -> Result<(), String> {
471    let Some(cache) = CACHE.as_ref() else {
472        return Ok(());
473    };
474    cache.replace_all(entries)
475}
476
477/// Merge new entries into the existing shard. Use this from the compinit
478/// BACKFILL path (existing shard has most entries, just adding the missing
479/// ones).
480pub fn try_merge_in(entries: HashMap<String, Vec<u8>>) -> Result<(), String> {
481    let Some(cache) = CACHE.as_ref() else {
482        return Ok(());
483    };
484    cache.merge_in(entries)
485}
486/// `cached_names` — see implementation.
487pub fn cached_names() -> std::collections::HashSet<String> {
488    CACHE.as_ref().map(|c| c.cached_names()).unwrap_or_default()
489}
490/// `entry_count` — see implementation.
491pub fn entry_count() -> usize {
492    CACHE.as_ref().map(|c| c.entry_count()).unwrap_or(0)
493}
494/// `stats` — see implementation.
495pub fn stats() -> Option<(i64, i64)> {
496    CACHE.as_ref().map(|c| c.stats())
497}
498/// `clear` — see implementation.
499pub fn clear() -> bool {
500    CACHE.as_ref().map(|c| c.clear().is_ok()).unwrap_or(false)
501}
502
503#[cfg(test)]
504mod tests {
505    use super::*;
506    use tempfile::tempdir;
507
508    #[test]
509    fn round_trip_one() {
510        let _g = crate::test_util::global_state_lock();
511        let dir = tempdir().unwrap();
512        let cache_path = dir.path().join("autoloads.rkyv");
513        let cache = AutoloadCache::open(&cache_path).unwrap();
514        cache.put_one("foo", vec![1, 2, 3]).unwrap();
515        assert_eq!(cache.get("foo"), Some(vec![1, 2, 3]));
516        assert_eq!(cache.entry_count(), 1);
517    }
518
519    #[test]
520    fn replace_all_overwrites() {
521        let _g = crate::test_util::global_state_lock();
522        let dir = tempdir().unwrap();
523        let cache_path = dir.path().join("autoloads.rkyv");
524        let cache = AutoloadCache::open(&cache_path).unwrap();
525        cache.put_one("a", vec![10]).unwrap();
526        cache.put_one("b", vec![20]).unwrap();
527        assert_eq!(cache.entry_count(), 2);
528
529        let mut new_entries = HashMap::new();
530        new_entries.insert("c".to_string(), vec![30]);
531        new_entries.insert("d".to_string(), vec![40]);
532        cache.replace_all(new_entries).unwrap();
533
534        assert_eq!(cache.entry_count(), 2);
535        assert!(cache.get("a").is_none());
536        assert!(cache.get("b").is_none());
537        assert_eq!(cache.get("c"), Some(vec![30]));
538        assert_eq!(cache.get("d"), Some(vec![40]));
539    }
540
541    #[test]
542    fn cached_names_returns_keys() {
543        let _g = crate::test_util::global_state_lock();
544        let dir = tempdir().unwrap();
545        let cache_path = dir.path().join("autoloads.rkyv");
546        let cache = AutoloadCache::open(&cache_path).unwrap();
547        cache.put_one("alpha", vec![1]).unwrap();
548        cache.put_one("beta", vec![2]).unwrap();
549        let names = cache.cached_names();
550        assert!(names.contains("alpha"));
551        assert!(names.contains("beta"));
552        assert_eq!(names.len(), 2);
553    }
554
555    #[test]
556    fn corrupt_shard_returns_none() {
557        let _g = crate::test_util::global_state_lock();
558        let dir = tempdir().unwrap();
559        let cache_path = dir.path().join("autoloads.rkyv");
560        std::fs::write(&cache_path, b"garbage").unwrap();
561        let cache = AutoloadCache::open(&cache_path).unwrap();
562        assert!(cache.get("anything").is_none());
563        assert_eq!(cache.entry_count(), 0);
564    }
565
566    #[test]
567    fn clear_removes_file() {
568        let _g = crate::test_util::global_state_lock();
569        let dir = tempdir().unwrap();
570        let cache_path = dir.path().join("autoloads.rkyv");
571        let cache = AutoloadCache::open(&cache_path).unwrap();
572        cache.put_one("x", vec![1]).unwrap();
573        assert!(cache_path.exists());
574        cache.clear().unwrap();
575        assert!(!cache_path.exists());
576    }
577}
zsh/extensions/autoload_cache.rs

zsh/extensions/
autoload_cache.rs
