pf_cache/
capture.rs

1// SPDX-License-Identifier: MIT
2//! High-level capture / restore one-shot helpers.
3//!
4//! Drives a [`CachePager`] through pause → walk-pages → serialize → resume,
5//! and the inverse on restore. Adapter authors normally call these instead
6//! of touching [`crate::serialize`] directly.
7
8use crate::pager::CachePager;
9use crate::serialize::{deserialize_pages, serialize_pages};
10use pf_core::cas::BlobStore;
11use pf_core::digest::Digest256;
12
13/// Capture the entire live state of `pager` into the supplied `blobs` store.
14///
15/// Returns the digest of the [`crate::format::PageManifest`] blob; this is
16/// the value that goes into the `.pfimg` manifest's `cache.manifest` field.
17///
18/// Pauses the pager for the duration of the capture and resumes on the way
19/// out — even if a page read errors midway through.
20pub fn capture_cache(
21    pager: &mut dyn CachePager,
22    blobs: &dyn BlobStore,
23) -> pf_core::Result<Digest256> {
24    pager.pause()?;
25    // Use a guard pattern to guarantee resume() runs.
26    let result = (|| {
27        let meta = pager.meta();
28        let occupied = pager.occupied_pages();
29        let logical_seqs = pager.logical_seqs();
30        let mut captured = Vec::with_capacity(occupied.len());
31        for ix in occupied {
32            let bytes = pager.read_page(ix)?;
33            captured.push((ix, bytes));
34        }
35        serialize_pages(blobs, meta, captured, &logical_seqs)
36    })();
37    let resume = pager.resume();
38    let cid = result?;
39    resume?;
40    Ok(cid)
41}
42
43/// Restore the cache described by `manifest_digest` into `pager`. The pager
44/// must be empty (no occupied pages) and the same engine config as the one
45/// that produced the snapshot. Mismatches surface as
46/// [`pf_core::Error::Integrity`].
47pub fn restore_cache(
48    pager: &mut dyn CachePager,
49    blobs: &dyn BlobStore,
50    manifest_digest: &Digest256,
51) -> pf_core::Result<()> {
52    let restored = deserialize_pages(blobs, manifest_digest)?;
53    if restored.meta != pager.meta() {
54        return Err(pf_core::Error::Integrity(format!(
55            "cache meta mismatch on restore: got {:?}, pager is {:?}",
56            restored.meta,
57            pager.meta()
58        )));
59    }
60    pager.pause()?;
61    let result = (|| {
62        // Allocate fresh slots; then, for each captured page, write into the
63        // pager at the captured *original* physical-page index — adapters
64        // are free to remap, but the synthetic pager assumes ix-stability.
65        let _ = pager.allocate_pages(restored.pages.len())?;
66        for (ix, bytes) in &restored.pages {
67            pager.write_page(*ix, bytes)?;
68        }
69        pager.install_logical_seqs(&restored.logical_seqs)
70    })();
71    let resume = pager.resume();
72    result?;
73    resume
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79    use crate::format::{CacheMeta, Dtype};
80    use crate::pager::{CachePager, SyntheticCachePager};
81    use pf_core::cas::MemBlobStore;
82
83    fn small_meta() -> CacheMeta {
84        CacheMeta {
85            page_size_tokens: 4,
86            n_layers: 2,
87            n_heads: 2,
88            head_dim: 4,
89            dtype: Dtype::Bf16,
90        }
91    }
92
93    #[test]
94    fn capture_then_restore_into_fresh_pager() {
95        let mut src = SyntheticCachePager::new(small_meta());
96        src.populate_synthetic(8, 99).unwrap();
97        let blobs = MemBlobStore::new();
98        let cid = capture_cache(&mut src, &blobs).unwrap();
99
100        let mut dst = SyntheticCachePager::new(small_meta());
101        restore_cache(&mut dst, &blobs, &cid).unwrap();
102
103        // Every page byte-identical?
104        for ix in src.occupied_pages() {
105            assert_eq!(src.read_page(ix).unwrap(), dst.read_page(ix).unwrap());
106        }
107        assert_eq!(src.logical_seqs(), dst.logical_seqs());
108    }
109
110    #[test]
111    fn restore_into_mismatched_meta_errors() {
112        let mut src = SyntheticCachePager::new(small_meta());
113        src.populate_synthetic(2, 0).unwrap();
114        let blobs = MemBlobStore::new();
115        let cid = capture_cache(&mut src, &blobs).unwrap();
116
117        let mut wider = small_meta();
118        wider.head_dim = 8; // diverges
119        let mut dst = SyntheticCachePager::new(wider);
120        let err = restore_cache(&mut dst, &blobs, &cid).unwrap_err();
121        assert!(matches!(err, pf_core::Error::Integrity(_)));
122    }
123}
pf_cache/capture.rs

pf_cache/
capture.rs
