use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::OnceLock;
use std::time::{SystemTime, UNIX_EPOCH};
use memmap2::Mmap;
use parking_lot::Mutex;
use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
use std::os::unix::fs::MetadataExt;
pub const SHARD_MAGIC: u32 = 0x5A52414C;
pub const SHARD_FORMAT_VERSION: u32 = 3;
#[derive(Archive, RkyvDeserialize, RkyvSerialize, Debug, Clone)]
#[archive(check_bytes)]
pub struct ShardHeader {
pub magic: u32,
pub format_version: u32,
pub zshrs_version: String,
pub pointer_width: u32,
pub built_at_secs: u64,
}
#[derive(Archive, RkyvDeserialize, RkyvSerialize, Debug, Clone)]
#[archive(check_bytes)]
pub struct AutoloadEntry {
pub binary_mtime_at_cache: i64,
pub binary_len_at_cache: u64,
pub cached_at_secs: i64,
pub source_dir: String,
pub source_sha: [u8; 32],
pub chunk_blob: Vec<u8>,
}
#[derive(Archive, RkyvDeserialize, RkyvSerialize, Debug, Clone)]
#[archive(check_bytes)]
pub struct AutoloadShard {
pub header: ShardHeader,
pub entries: HashMap<String, AutoloadEntry>,
}
pub struct MmappedShard {
_mmap: Mmap,
archived: *const ArchivedAutoloadShard,
}
unsafe impl Send for MmappedShard {}
unsafe impl Sync for MmappedShard {}
impl MmappedShard {
pub fn open(path: &Path) -> Option<Self> {
let file = File::open(path).ok()?;
let mmap = unsafe { Mmap::map(&file).ok()? };
let archived = rkyv::check_archived_root::<AutoloadShard>(&mmap[..]).ok()?;
let archived_ptr = archived as *const ArchivedAutoloadShard;
Some(Self {
_mmap: mmap,
archived: archived_ptr,
})
}
fn shard(&self) -> &ArchivedAutoloadShard {
unsafe { &*self.archived }
}
fn header_ok(&self) -> bool {
let h = &self.shard().header;
let magic: u32 = h.magic.into();
let fv: u32 = h.format_version.into();
let pw: u32 = h.pointer_width.into();
magic == SHARD_MAGIC
&& fv == SHARD_FORMAT_VERSION
&& pw as usize == std::mem::size_of::<usize>()
&& h.zshrs_version.as_str() == env!("CARGO_PKG_VERSION")
}
fn lookup(&self, name: &str) -> Option<&ArchivedAutoloadEntry> {
self.shard().entries.get(name)
}
}
fn entry_binary_matches(entry: &ArchivedAutoloadEntry) -> bool {
let Some((mtime, len)) = current_binary_identity() else {
return false;
};
let cached_mtime: i64 = entry.binary_mtime_at_cache.into();
let cached_len: u64 = entry.binary_len_at_cache.into();
cached_mtime == mtime && cached_len == len
}
pub struct AutoloadCache {
path: PathBuf,
lock_path: PathBuf,
mmap: Mutex<Option<MmappedShard>>,
}
impl AutoloadCache {
pub fn open(path: &Path) -> std::io::Result<Self> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let parent = path.parent().unwrap_or_else(|| Path::new("/tmp"));
let lock_path = parent.join(format!(
"{}.lock",
path.file_name()
.and_then(|s| s.to_str())
.unwrap_or("autoloads.rkyv")
));
Ok(Self {
path: path.to_path_buf(),
lock_path,
mmap: Mutex::new(None),
})
}
fn ensure_mmap(&self) {
let mut guard = self.mmap.lock();
if guard.is_none() {
*guard = MmappedShard::open(&self.path);
}
}
fn invalidate_mmap(&self) {
let mut guard = self.mmap.lock();
*guard = None;
}
pub fn get(&self, name: &str) -> Option<Vec<u8>> {
self.ensure_mmap();
let guard = self.mmap.lock();
let shard = guard.as_ref()?;
if !shard.header_ok() {
return None;
}
let entry = shard.lookup(name)?;
if !entry_binary_matches(entry) {
return None;
}
Some(entry.chunk_blob.as_slice().to_vec())
}
pub fn get_for_source(
&self,
name: &str,
source_dir: &str,
source_sha: &[u8; 32],
) -> Option<Vec<u8>> {
self.ensure_mmap();
let guard = self.mmap.lock();
let shard = guard.as_ref()?;
if !shard.header_ok() {
return None;
}
let entry = shard.lookup(name)?;
if !entry_binary_matches(entry) {
return None;
}
if entry.source_dir.as_str() != source_dir {
return None;
}
if entry.source_sha.as_slice() != source_sha.as_slice() {
return None;
}
Some(entry.chunk_blob.as_slice().to_vec())
}
fn owned_shard_for_write(&self) -> AutoloadShard {
match read_owned_shard(&self.path) {
Some(s)
if s.header.zshrs_version == env!("CARGO_PKG_VERSION")
&& s.header.pointer_width as usize == std::mem::size_of::<usize>()
&& s.header.format_version == SHARD_FORMAT_VERSION =>
{
s
}
_ => fresh_shard(),
}
}
pub fn put_one(
&self,
name: &str,
chunk_blob: Vec<u8>,
source_dir: &str,
source_sha: [u8; 32],
) -> Result<(), String> {
self.put_many(&[(
name.to_string(),
chunk_blob,
source_dir.to_string(),
source_sha,
)])
}
pub fn put_many(&self, entries: &[(String, Vec<u8>, String, [u8; 32])]) -> Result<(), String> {
if entries.is_empty() {
return Ok(());
}
let _lock = match acquire_lock(&self.lock_path) {
Some(l) => l,
None => return Ok(()),
};
let mut shard = self.owned_shard_for_write();
let (bin_mtime, bin_len) = current_binary_identity().unwrap_or((0, 0));
let now = now_secs();
for (name, chunk_blob, source_dir, source_sha) in entries {
shard.entries.insert(
name.clone(),
AutoloadEntry {
binary_mtime_at_cache: bin_mtime,
binary_len_at_cache: bin_len,
cached_at_secs: now,
source_dir: source_dir.clone(),
source_sha: *source_sha,
chunk_blob: chunk_blob.clone(),
},
);
}
shard.header.built_at_secs = now as u64;
write_shard_atomic(&self.path, &shard)?;
self.invalidate_mmap();
Ok(())
}
pub fn remove(&self, name: &str) -> Result<(), String> {
let _lock = match acquire_lock(&self.lock_path) {
Some(l) => l,
None => return Ok(()),
};
let mut shard = self.owned_shard_for_write();
if shard.entries.remove(name).is_none() {
return Ok(());
}
shard.header.built_at_secs = now_secs() as u64;
write_shard_atomic(&self.path, &shard)?;
self.invalidate_mmap();
Ok(())
}
pub fn entry_count(&self) -> usize {
self.ensure_mmap();
let guard = self.mmap.lock();
guard.as_ref().map(|s| s.shard().entries.len()).unwrap_or(0)
}
pub fn cached_names(&self) -> std::collections::HashSet<String> {
self.ensure_mmap();
let guard = self.mmap.lock();
let Some(shard) = guard.as_ref() else {
return std::collections::HashSet::new();
};
shard
.shard()
.entries
.keys()
.map(|k| k.as_str().to_string())
.collect()
}
pub fn stats(&self) -> (i64, i64) {
self.ensure_mmap();
let guard = self.mmap.lock();
let Some(shard) = guard.as_ref() else {
return (0, 0);
};
let count = shard.shard().entries.len() as i64;
let bytes: i64 = shard
.shard()
.entries
.values()
.map(|e| e.chunk_blob.len() as i64)
.sum();
(count, bytes)
}
pub fn clear(&self) -> std::io::Result<()> {
let _lock = acquire_lock(&self.lock_path);
let res = match std::fs::remove_file(&self.path) {
Ok(()) => Ok(()),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(e) => Err(e),
};
crate::atomic_write::reap_orphan_temps(&self.path);
self.invalidate_mmap();
res
}
}
fn acquire_lock(path: &Path) -> Option<nix::fcntl::Flock<File>> {
let f = File::options()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(path)
.ok()?;
nix::fcntl::Flock::lock(f, nix::fcntl::FlockArg::LockExclusive).ok()
}
fn fresh_shard() -> AutoloadShard {
AutoloadShard {
header: ShardHeader {
magic: SHARD_MAGIC,
format_version: SHARD_FORMAT_VERSION,
zshrs_version: env!("CARGO_PKG_VERSION").to_string(),
pointer_width: std::mem::size_of::<usize>() as u32,
built_at_secs: now_secs() as u64,
},
entries: HashMap::new(),
}
}
fn read_owned_shard(path: &Path) -> Option<AutoloadShard> {
let bytes = std::fs::read(path).ok()?;
let archived = rkyv::check_archived_root::<AutoloadShard>(&bytes[..]).ok()?;
archived.deserialize(&mut rkyv::Infallible).ok()
}
fn write_shard_atomic(path: &Path, shard: &AutoloadShard) -> Result<(), String> {
let bytes = rkyv::to_bytes::<_, 4096>(shard).map_err(|e| format!("rkyv serialize: {}", e))?;
crate::atomic_write::write_bytes_atomic(path, &bytes)
}
fn now_secs() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs() as i64)
.unwrap_or(0)
}
fn current_binary_identity() -> Option<(i64, u64)> {
static BIN_ID: OnceLock<Option<(i64, u64)>> = OnceLock::new();
*BIN_ID.get_or_init(|| {
let exe = std::env::current_exe().ok()?;
let meta = std::fs::metadata(&exe).ok()?;
Some((meta.mtime(), meta.len()))
})
}
pub fn source_digest(text: &str) -> [u8; 32] {
use sha2::Digest;
let mut hasher = sha2::Sha256::new();
hasher.update(text.as_bytes());
hasher.finalize().into()
}
pub fn default_cache_path() -> PathBuf {
let root = if let Some(custom) = std::env::var_os("ZSHRS_HOME") {
PathBuf::from(custom)
} else {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("/tmp"))
.join(".zshrs")
};
root.join("autoloads.rkyv")
}
pub fn cache_enabled() -> bool {
if crate::extensions::script_cache::CACHE_DISABLED.load(std::sync::atomic::Ordering::Relaxed) {
return false;
}
!matches!(
std::env::var("ZSHRS_CACHE").as_deref(),
Ok("0") | Ok("false") | Ok("no")
)
}
pub static CACHE: once_cell::sync::Lazy<Option<AutoloadCache>> = once_cell::sync::Lazy::new(|| {
if !cache_enabled() {
return None;
}
AutoloadCache::open(&default_cache_path()).ok()
});
pub fn try_load(name: &str) -> Option<Vec<u8>> {
let cache = CACHE.as_ref()?;
cache.get(name)
}
pub fn try_load_for_source(name: &str, source_dir: &str, source_sha: &[u8; 32]) -> Option<Vec<u8>> {
let cache = CACHE.as_ref()?;
cache.get_for_source(name, source_dir, source_sha)
}
pub fn try_save_one(
name: &str,
chunk_blob: &[u8],
source_dir: &str,
source_sha: [u8; 32],
) -> Result<(), String> {
let Some(cache) = CACHE.as_ref() else {
return Ok(());
};
cache.put_one(name, chunk_blob.to_vec(), source_dir, source_sha)
}
pub fn try_put_many(entries: &[(String, Vec<u8>, String, [u8; 32])]) -> Result<(), String> {
let Some(cache) = CACHE.as_ref() else {
return Ok(());
};
cache.put_many(entries)
}
pub fn try_remove(name: &str) {
if let Some(cache) = CACHE.as_ref() {
if let Err(e) = cache.remove(name) {
tracing::warn!(name, error = %e, "autoload: could not drop bad cache entry");
}
}
}
pub fn cached_names() -> std::collections::HashSet<String> {
CACHE.as_ref().map(|c| c.cached_names()).unwrap_or_default()
}
pub fn entry_count() -> usize {
CACHE.as_ref().map(|c| c.entry_count()).unwrap_or(0)
}
pub fn stats() -> Option<(i64, i64)> {
CACHE.as_ref().map(|c| c.stats())
}
pub fn clear() -> bool {
CACHE.as_ref().map(|c| c.clear().is_ok()).unwrap_or(false)
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
const DIR: &str = "/some/fpath/dir";
#[test]
fn round_trip_one() {
let _g = crate::test_util::global_state_lock();
let dir = tempdir().unwrap();
let cache_path = dir.path().join("autoloads.rkyv");
let cache = AutoloadCache::open(&cache_path).unwrap();
cache
.put_one("foo", vec![1, 2, 3], DIR, source_digest("body"))
.unwrap();
assert_eq!(cache.get("foo"), Some(vec![1, 2, 3]));
assert_eq!(cache.entry_count(), 1);
}
#[test]
fn source_text_mismatch_is_a_miss() {
let _g = crate::test_util::global_state_lock();
let dir = tempdir().unwrap();
let cache_path = dir.path().join("autoloads.rkyv");
let cache = AutoloadCache::open(&cache_path).unwrap();
let sha = source_digest("foo() {\nprint one\n}");
cache.put_one("foo", vec![1, 2, 3], DIR, sha).unwrap();
assert_eq!(cache.get_for_source("foo", DIR, &sha), Some(vec![1, 2, 3]));
let edited = source_digest("foo() {\nprint two\n}");
assert!(cache.get_for_source("foo", DIR, &edited).is_none());
assert!(cache.get_for_source("foo", "/other/dir", &sha).is_none());
}
#[test]
fn an_entry_from_another_binary_is_never_served() {
let _g = crate::test_util::global_state_lock();
let dir = tempdir().unwrap();
let cache_path = dir.path().join("autoloads.rkyv");
let cache = AutoloadCache::open(&cache_path).unwrap();
let sha = source_digest("foo() {\nprint one\n}");
cache.put_one("foo", vec![1, 2, 3], DIR, sha).unwrap();
let mut shard = read_owned_shard(&cache_path).expect("shard readable");
let entry = shard.entries.get_mut("foo").expect("entry present");
entry.binary_mtime_at_cache += 10_000;
write_shard_atomic(&cache_path, &shard).unwrap();
let reopened = AutoloadCache::open(&cache_path).unwrap();
assert!(
reopened.get_for_source("foo", DIR, &sha).is_none(),
"a chunk from a newer build was accepted",
);
assert!(reopened.get("foo").is_none());
}
#[test]
fn remove_drops_the_entry() {
let _g = crate::test_util::global_state_lock();
let dir = tempdir().unwrap();
let cache_path = dir.path().join("autoloads.rkyv");
let cache = AutoloadCache::open(&cache_path).unwrap();
let sha = source_digest("body");
cache.put_one("foo", vec![1], DIR, sha).unwrap();
cache.put_one("bar", vec![2], DIR, sha).unwrap();
cache.remove("foo").unwrap();
assert!(cache.get("foo").is_none());
assert_eq!(cache.get("bar"), Some(vec![2]));
}
#[test]
fn cached_names_returns_keys() {
let _g = crate::test_util::global_state_lock();
let dir = tempdir().unwrap();
let cache_path = dir.path().join("autoloads.rkyv");
let cache = AutoloadCache::open(&cache_path).unwrap();
let sha = source_digest("body");
cache.put_one("alpha", vec![1], DIR, sha).unwrap();
cache.put_one("beta", vec![2], DIR, sha).unwrap();
let names = cache.cached_names();
assert!(names.contains("alpha"));
assert!(names.contains("beta"));
assert_eq!(names.len(), 2);
}
#[test]
fn corrupt_shard_returns_none() {
let _g = crate::test_util::global_state_lock();
let dir = tempdir().unwrap();
let cache_path = dir.path().join("autoloads.rkyv");
std::fs::write(&cache_path, b"garbage").unwrap();
let cache = AutoloadCache::open(&cache_path).unwrap();
assert!(cache.get("anything").is_none());
assert_eq!(cache.entry_count(), 0);
}
#[test]
fn clear_removes_file() {
let _g = crate::test_util::global_state_lock();
let dir = tempdir().unwrap();
let cache_path = dir.path().join("autoloads.rkyv");
let cache = AutoloadCache::open(&cache_path).unwrap();
cache
.put_one("x", vec![1], DIR, source_digest("body"))
.unwrap();
assert!(cache_path.exists());
cache.clear().unwrap();
assert!(!cache_path.exists());
}
#[test]
fn a_write_leaves_no_temp_file_behind() {
let _g = crate::test_util::global_state_lock();
let dir = tempdir().unwrap();
let cache_path = dir.path().join("autoloads.rkyv");
let cache = AutoloadCache::open(&cache_path).unwrap();
cache
.put_one("x", vec![1], DIR, source_digest("body"))
.unwrap();
let temps: Vec<_> = std::fs::read_dir(dir.path())
.unwrap()
.flatten()
.map(|e| e.file_name().into_string().unwrap())
.filter(|n| n.contains(".tmp."))
.collect();
assert!(temps.is_empty(), "temp files left behind: {temps:?}");
}
}