Skip to main content

aft/
symbol_cache_disk.rs

1use std::fs::{self, File};
2use std::io::{BufReader, BufWriter, Read, Write};
3use std::path::{Path, PathBuf};
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::sync::Mutex;
6use std::time::{Duration, SystemTime, UNIX_EPOCH};
7
8use crate::fs_lock;
9use crate::parser::SymbolCache;
10use crate::search_index::{cache_relative_path, validate_cached_relative_path};
11use crate::symbols::Symbol;
12use crate::{slog_info, slog_warn};
13
14const MAGIC: &[u8; 8] = b"AFTSYM1\0";
15const FORMAT_VERSION: u32 = 3;
16
17/// Version of the symbol extraction schema stored in the disk cache.
18///
19/// Bump this whenever symbol-extraction logic changes: tree-sitter grammar
20/// upgrades, query updates, extractor behavior, or symbol shape changes. A
21/// mismatch rejects persisted symbols so they are regenerated on next access.
22pub const SCHEMA_VERSION: u32 = 3;
23
24const MAX_ENTRIES: usize = 2_000_000;
25const MAX_PATH_BYTES: usize = 16 * 1024;
26const MAX_SYMBOL_BYTES: usize = 16 * 1024 * 1024;
27static TMP_COUNTER: AtomicU64 = AtomicU64::new(0);
28static SYMBOL_LOCK_ACQUIRE_MUTEX: Mutex<()> = Mutex::new(());
29
30pub struct SymbolCacheLock {
31    _guard: Option<fs_lock::LockGuard>,
32}
33
34impl SymbolCacheLock {
35    pub fn acquire(
36        storage_dir: &Path,
37        project_key: &str,
38        project_root: &Path,
39    ) -> std::io::Result<Self> {
40        let dir = storage_dir.join("symbols").join(project_key);
41        let path = dir.join("symbols.lock");
42        let access = crate::root_cache::ArtifactAccess::for_root(project_root);
43        if !access.allows_write(project_key, &path) {
44            return Ok(Self { _guard: None });
45        }
46        fs::create_dir_all(&dir)?;
47        let _acquire_guard = SYMBOL_LOCK_ACQUIRE_MUTEX
48            .lock()
49            .map_err(|_| std::io::Error::other("symbol cache lock acquisition mutex poisoned"))?;
50        fs_lock::try_acquire(&path, Duration::from_secs(2))
51            .map(|guard| Self {
52                _guard: Some(guard),
53            })
54            .map_err(|error| match error {
55                fs_lock::AcquireError::Timeout => {
56                    std::io::Error::other("timed out acquiring symbol cache lock")
57                }
58                fs_lock::AcquireError::Io(error) => error,
59            })
60    }
61}
62
63#[derive(Debug, Clone)]
64pub struct DiskSymbolCache {
65    pub(crate) entries: Vec<DiskSymbolEntry>,
66}
67
68#[derive(Debug, Clone)]
69pub(crate) struct DiskSymbolEntry {
70    pub(crate) relative_path: PathBuf,
71    pub(crate) mtime: SystemTime,
72    pub(crate) size: u64,
73    pub(crate) content_hash: blake3::Hash,
74    pub(crate) symbols: Vec<Symbol>,
75}
76
77impl DiskSymbolCache {
78    pub fn len(&self) -> usize {
79        self.entries.len()
80    }
81
82    pub fn is_empty(&self) -> bool {
83        self.entries.is_empty()
84    }
85}
86
87pub(crate) fn cache_path(storage_dir: &Path, project_key: &str) -> PathBuf {
88    storage_dir
89        .join("symbols")
90        .join(project_key)
91        .join("symbols.bin")
92}
93
94pub fn read_from_disk(storage_dir: &Path, project_key: &str) -> Option<DiskSymbolCache> {
95    let data_path = cache_path(storage_dir, project_key);
96    if !data_path.exists() {
97        return None;
98    }
99
100    match read_cache_file(&data_path) {
101        Ok(cache) => Some(cache),
102        Err(error) => {
103            slog_warn!(
104                "corrupt symbol cache at {}: {}, rebuilding",
105                data_path.display(),
106                error
107            );
108            None
109        }
110    }
111}
112
113pub fn write_to_disk(
114    cache: &SymbolCache,
115    storage_dir: &Path,
116    project_key: &str,
117) -> std::io::Result<()> {
118    if cache.len() == 0 {
119        slog_info!("skipping symbol cache persistence (0 entries)");
120        return Ok(());
121    }
122
123    let project_root = cache.project_root().ok_or_else(|| {
124        std::io::Error::other("symbol cache project root is not set; cannot persist relative paths")
125    })?;
126
127    let dir = storage_dir.join("symbols").join(project_key);
128    let data_path = dir.join("symbols.bin");
129    let access = crate::root_cache::ArtifactAccess::for_root(&project_root);
130    if !access.allows_write(project_key, &data_path) {
131        return Ok(());
132    }
133    let _cache_lock = SymbolCacheLock::acquire(storage_dir, project_key, &project_root)?;
134    fs::create_dir_all(&dir)?;
135    let tmp_path = dir.join(format!(
136        "symbols.bin.tmp.{}.{}.{}",
137        std::process::id(),
138        SystemTime::now()
139            .duration_since(UNIX_EPOCH)
140            .unwrap_or(Duration::ZERO)
141            .as_nanos(),
142        TMP_COUNTER.fetch_add(1, Ordering::Relaxed)
143    ));
144    let write_result = write_cache_file(cache, &project_root, &tmp_path).and_then(|()| {
145        fs::rename(&tmp_path, &data_path)?;
146        if let Ok(dir_file) = File::open(&dir) {
147            let _ = dir_file.sync_all();
148        }
149        Ok(())
150    });
151
152    if write_result.is_err() {
153        let _ = fs::remove_file(&tmp_path);
154    }
155
156    write_result
157}
158
159fn read_cache_file(path: &Path) -> Result<DiskSymbolCache, String> {
160    let mut reader = BufReader::new(File::open(path).map_err(|error| error.to_string())?);
161
162    let mut magic = [0u8; 8];
163    reader
164        .read_exact(&mut magic)
165        .map_err(|error| format!("failed to read symbol cache magic: {error}"))?;
166    if &magic != MAGIC {
167        return Err("invalid symbol cache magic".to_string());
168    }
169
170    let format_version = read_u32(&mut reader)?;
171    if format_version != FORMAT_VERSION {
172        return Err(format!(
173            "unsupported symbol cache format version: {format_version} (expected {FORMAT_VERSION})"
174        ));
175    }
176
177    let schema_version = read_u32(&mut reader)?;
178    if schema_version != SCHEMA_VERSION {
179        return Err(format!(
180            "unsupported symbol cache schema version: {schema_version} (expected {SCHEMA_VERSION})"
181        ));
182    }
183
184    let root_len = read_u32(&mut reader)? as usize;
185    let entry_count = read_u32(&mut reader)? as usize;
186    if root_len > MAX_PATH_BYTES {
187        return Err(format!("project root path too large: {root_len} bytes"));
188    }
189    if entry_count > MAX_ENTRIES {
190        return Err(format!("too many symbol cache entries: {entry_count}"));
191    }
192
193    let _project_root = PathBuf::from(read_string_with_len(&mut reader, root_len)?);
194    let mut entries = Vec::with_capacity(entry_count);
195
196    for _ in 0..entry_count {
197        let path_len = read_u32(&mut reader)? as usize;
198        if path_len > MAX_PATH_BYTES {
199            return Err(format!("cached path too large: {path_len} bytes"));
200        }
201        let relative_path = validate_cached_relative_path(&PathBuf::from(read_string_with_len(
202            &mut reader,
203            path_len,
204        )?))
205        .ok_or_else(|| "cached symbol path escapes project root".to_string())?;
206        let mtime_secs = read_i64(&mut reader)?;
207        let mtime_nanos = read_u32(&mut reader)?;
208        let size = read_u64(&mut reader)?;
209        let mut hash_bytes = [0u8; 32];
210        reader
211            .read_exact(&mut hash_bytes)
212            .map_err(|error| format!("failed to read symbol content hash: {error}"))?;
213        let content_hash = blake3::Hash::from_bytes(hash_bytes);
214        let symbol_bytes_len = read_u32(&mut reader)? as usize;
215        if symbol_bytes_len > MAX_SYMBOL_BYTES {
216            return Err(format!(
217                "cached symbol payload too large: {symbol_bytes_len} bytes"
218            ));
219        }
220
221        let mut symbol_bytes = vec![0u8; symbol_bytes_len];
222        reader
223            .read_exact(&mut symbol_bytes)
224            .map_err(|error| format!("failed to read symbol payload: {error}"))?;
225        let symbols: Vec<Symbol> = serde_json::from_slice(&symbol_bytes)
226            .map_err(|error| format!("failed to decode cached symbols: {error}"))?;
227
228        entries.push(DiskSymbolEntry {
229            relative_path,
230            mtime: system_time_from_parts(mtime_secs, mtime_nanos)?,
231            size,
232            content_hash,
233            symbols,
234        });
235    }
236
237    Ok(DiskSymbolCache { entries })
238}
239
240fn write_cache_file(
241    cache: &SymbolCache,
242    project_root: &Path,
243    tmp_path: &Path,
244) -> std::io::Result<()> {
245    let mut writer = BufWriter::new(File::create(tmp_path)?);
246    let entries = cache
247        .disk_entries()
248        .into_iter()
249        .map(|(path, mtime, size, content_hash, symbols)| {
250            cache_relative_path(project_root, path)
251                .map(|relative_path| (relative_path, mtime, size, content_hash, symbols))
252        })
253        .collect::<Option<Vec<_>>>()
254        .ok_or_else(|| std::io::Error::other("refusing to cache path outside project root"))?;
255    let root = project_root.to_string_lossy();
256    let root_len = u32::try_from(root.len())
257        .map_err(|_| std::io::Error::other("project root too large to cache"))?;
258    let entry_count = u32::try_from(entries.len())
259        .map_err(|_| std::io::Error::other("too many symbol cache entries"))?;
260
261    writer.write_all(MAGIC)?;
262    write_u32(&mut writer, FORMAT_VERSION)?;
263    write_u32(&mut writer, SCHEMA_VERSION)?;
264    write_u32(&mut writer, root_len)?;
265    write_u32(&mut writer, entry_count)?;
266    writer.write_all(root.as_bytes())?;
267
268    for (relative_path, mtime, size, content_hash, symbols) in entries {
269        let path_bytes = relative_path.to_string_lossy();
270        let path_len = u32::try_from(path_bytes.len())
271            .map_err(|_| std::io::Error::other("cached path too large"))?;
272        let (secs, nanos) = system_time_parts(mtime);
273        let symbol_bytes = serde_json::to_vec(symbols).map_err(|error| {
274            std::io::Error::other(format!("symbol serialization failed: {error}"))
275        })?;
276        let symbol_len = u32::try_from(symbol_bytes.len())
277            .map_err(|_| std::io::Error::other("cached symbol payload too large"))?;
278
279        write_u32(&mut writer, path_len)?;
280        writer.write_all(path_bytes.as_bytes())?;
281        write_i64(&mut writer, secs)?;
282        write_u32(&mut writer, nanos)?;
283        write_u64(&mut writer, size)?;
284        writer.write_all(content_hash.as_bytes())?;
285        write_u32(&mut writer, symbol_len)?;
286        writer.write_all(&symbol_bytes)?;
287    }
288
289    writer.flush()?;
290    writer.get_ref().sync_all()?;
291    Ok(())
292}
293
294fn system_time_parts(time: SystemTime) -> (i64, u32) {
295    match time.duration_since(UNIX_EPOCH) {
296        Ok(duration) => (
297            i64::try_from(duration.as_secs()).unwrap_or(i64::MAX),
298            duration.subsec_nanos(),
299        ),
300        Err(error) => {
301            let duration = error.duration();
302            let nanos = duration.subsec_nanos();
303            if nanos == 0 {
304                (-(duration.as_secs() as i64), 0)
305            } else {
306                (-(duration.as_secs() as i64) - 1, 1_000_000_000 - nanos)
307            }
308        }
309    }
310}
311
312fn system_time_from_parts(secs: i64, nanos: u32) -> Result<SystemTime, String> {
313    if nanos >= 1_000_000_000 {
314        return Err(format!(
315            "invalid symbol cache mtime nanos: {nanos} >= 1_000_000_000"
316        ));
317    }
318
319    if secs >= 0 {
320        let duration = Duration::new(secs as u64, nanos);
321        UNIX_EPOCH
322            .checked_add(duration)
323            .ok_or_else(|| format!("symbol cache mtime overflows SystemTime: {secs}.{nanos}"))
324    } else {
325        let whole = Duration::new(secs.unsigned_abs(), 0);
326        let base = UNIX_EPOCH.checked_sub(whole).ok_or_else(|| {
327            format!("symbol cache negative mtime overflows SystemTime: {secs}.{nanos}")
328        })?;
329        base.checked_add(Duration::new(0, nanos)).ok_or_else(|| {
330            format!("symbol cache negative mtime overflows SystemTime: {secs}.{nanos}")
331        })
332    }
333}
334
335fn read_string_with_len<R: Read>(reader: &mut R, len: usize) -> Result<String, String> {
336    let mut bytes = vec![0u8; len];
337    reader
338        .read_exact(&mut bytes)
339        .map_err(|error| format!("failed to read string: {error}"))?;
340    String::from_utf8(bytes).map_err(|error| format!("invalid utf-8 string: {error}"))
341}
342
343fn read_u32<R: Read>(reader: &mut R) -> Result<u32, String> {
344    let mut bytes = [0u8; 4];
345    reader
346        .read_exact(&mut bytes)
347        .map_err(|error| format!("failed to read u32: {error}"))?;
348    Ok(u32::from_le_bytes(bytes))
349}
350
351fn read_i64<R: Read>(reader: &mut R) -> Result<i64, String> {
352    let mut bytes = [0u8; 8];
353    reader
354        .read_exact(&mut bytes)
355        .map_err(|error| format!("failed to read i64: {error}"))?;
356    Ok(i64::from_le_bytes(bytes))
357}
358
359fn read_u64<R: Read>(reader: &mut R) -> Result<u64, String> {
360    let mut bytes = [0u8; 8];
361    reader
362        .read_exact(&mut bytes)
363        .map_err(|error| format!("failed to read u64: {error}"))?;
364    Ok(u64::from_le_bytes(bytes))
365}
366
367fn write_u32<W: Write>(writer: &mut W, value: u32) -> std::io::Result<()> {
368    writer.write_all(&value.to_le_bytes())
369}
370
371fn write_i64<W: Write>(writer: &mut W, value: i64) -> std::io::Result<()> {
372    writer.write_all(&value.to_le_bytes())
373}
374
375fn write_u64<W: Write>(writer: &mut W, value: u64) -> std::io::Result<()> {
376    writer.write_all(&value.to_le_bytes())
377}
378
379#[cfg(test)]
380mod tests {
381    use super::*;
382    use crate::symbols::{Range, SymbolKind};
383
384    fn test_symbol(name: &str) -> Symbol {
385        Symbol {
386            name: name.to_string(),
387            kind: SymbolKind::Function,
388            range: Range {
389                start_line: 0,
390                start_col: 0,
391                end_line: 0,
392                end_col: 1,
393            },
394            signature: None,
395            scope_chain: Vec::new(),
396            exported: false,
397            parent: None,
398        }
399    }
400
401    fn test_cache(project: &Path, file_name: &str) -> SymbolCache {
402        let file = project.join(file_name);
403        fs::write(&file, format!("fn {file_name}() {{}}\n")).expect("write file");
404        let metadata = fs::metadata(&file).expect("metadata");
405        let content_hash = blake3::hash(&fs::read(&file).expect("read file"));
406        let mut cache = SymbolCache::new();
407        cache.set_project_root(project.to_path_buf());
408        cache.insert(
409            file,
410            metadata.modified().expect("mtime"),
411            metadata.len(),
412            content_hash,
413            vec![test_symbol(file_name)],
414        );
415        cache
416    }
417
418    #[test]
419    fn borrow_only_root_skips_symbol_lock_and_persist() {
420        let project = tempfile::tempdir().expect("project");
421        let storage = tempfile::tempdir().expect("storage");
422        let project_key = "shared-artifact-key".to_string();
423        crate::root_cache::configure_artifact_access(project.path(), &project_key, true);
424        let cache = test_cache(project.path(), "lib.rs");
425
426        let _lock = SymbolCacheLock::acquire(storage.path(), &project_key, project.path())
427            .expect("borrow-only lock downgrade");
428        let cache_dir = storage.path().join("symbols").join(&project_key);
429        assert!(!cache_dir.join("symbols.lock").exists());
430
431        write_to_disk(&cache, storage.path(), &project_key).expect("borrow-only persist downgrade");
432
433        assert!(!cache_dir.join("symbols.bin").exists());
434        assert!(!cache_dir.exists());
435    }
436
437    #[test]
438    fn concurrent_symbol_cache_writes_do_not_share_temp_file() {
439        let dir = tempfile::tempdir().expect("create temp dir");
440        let project = dir.path().join("project");
441        fs::create_dir_all(&project).expect("create project");
442        let storage = dir.path().join("storage");
443
444        let cache_a = test_cache(&project, "a");
445        let cache_b = test_cache(&project, "b");
446        let storage_a = storage.clone();
447        let writer_a = std::thread::spawn(move || {
448            write_to_disk(&cache_a, &storage_a, "unit-project").expect("write a");
449        });
450        let storage_b = storage.clone();
451        let writer_b = std::thread::spawn(move || {
452            write_to_disk(&cache_b, &storage_b, "unit-project").expect("write b");
453        });
454
455        writer_a.join().expect("writer a");
456        writer_b.join().expect("writer b");
457
458        let loaded = read_from_disk(&storage, "unit-project").expect("load symbol cache");
459        assert_eq!(loaded.len(), 1);
460        assert!(fs::read_dir(storage.join("symbols").join("unit-project"))
461            .expect("read symbol cache dir")
462            .all(|entry| !entry
463                .expect("cache entry")
464                .file_name()
465                .to_string_lossy()
466                .contains(".tmp.")));
467    }
468
469    #[test]
470    fn symbol_cache_rejects_mismatched_schema_version() {
471        let storage = tempfile::tempdir().expect("create storage dir");
472        let path = cache_path(storage.path(), "schema-project");
473        fs::create_dir_all(path.parent().expect("cache parent")).expect("create cache dir");
474
475        let mut bytes = Vec::new();
476        bytes.extend_from_slice(MAGIC);
477        bytes.extend_from_slice(&FORMAT_VERSION.to_le_bytes());
478        bytes.extend_from_slice(&SCHEMA_VERSION.wrapping_add(1).to_le_bytes());
479        bytes.extend_from_slice(&0u32.to_le_bytes());
480        bytes.extend_from_slice(&0u32.to_le_bytes());
481        fs::write(&path, bytes).expect("write wrong-schema cache");
482
483        assert!(read_from_disk(storage.path(), "schema-project").is_none());
484    }
485
486    #[test]
487    fn symbol_cache_rejects_paths_outside_project_root_on_write() {
488        let dir = tempfile::tempdir().expect("create temp dir");
489        let project = dir.path().join("project");
490        fs::create_dir_all(&project).expect("create project");
491        let outside = dir.path().join("outside.rs");
492        fs::write(&outside, "fn outside() {}\n").expect("write outside");
493        let metadata = fs::metadata(&outside).expect("metadata");
494
495        let mut cache = SymbolCache::new();
496        cache.set_project_root(project);
497        cache.insert(
498            outside.clone(),
499            metadata.modified().expect("mtime"),
500            metadata.len(),
501            blake3::hash(&fs::read(&outside).expect("read outside")),
502            vec![test_symbol("outside")],
503        );
504
505        let error = write_to_disk(&cache, dir.path(), "escape-project").expect_err("reject escape");
506        assert!(error.to_string().contains("outside project root"));
507    }
508}