cranpose_render_wgpu/
pipeline_disk_cache.rs

1use std::{
2    hash::Hasher,
3    io::Write,
4    path::{Path, PathBuf},
5    sync::{
6        Arc, Condvar, Mutex, MutexGuard, OnceLock, PoisonError,
7        atomic::{AtomicBool, Ordering},
8    },
9};
10
11use cranpose_ui_graphics::FxHasher;
12use web_time::{Duration, Instant};
13
14use crate::debug_toggles::DebugToggle;
15
16static DISK_CACHE: DebugToggle = DebugToggle::new("CRANPOSE_PIPELINE_DISK_CACHE");
17
18fn disk_cache_enabled() -> bool {
19    !DISK_CACHE.equals("0")
20}
21
22static HOST_FILE: OnceLock<PathBuf> = OnceLock::new();
23
24/// Names the file the driver's compiled pipelines are kept in between runs.
25///
26/// This crate renders and does not know where an application may write, so the
27/// platform backend passes the path in once, before the first device is
28/// created. Without it the cache lives and dies with the process, and every
29/// launch pays the driver's compile again. A later call is ignored: the
30/// loaded blob and the watcher that writes it back must name one file.
31pub fn set_file_path(path: PathBuf) {
32    let _ = HOST_FILE.set(path);
33}
34
35pub(crate) fn file_path() -> Option<PathBuf> {
36    if !disk_cache_enabled() {
37        return None;
38    }
39    match crate::debug_toggles::debug_toggle_os("CRANPOSE_PIPELINE_CACHE_FILE") {
40        Some(path) if !path.is_empty() => Some(PathBuf::from(path)),
41        _ => HOST_FILE.get().cloned(),
42    }
43}
44
45/// What the last launch left for this one: the driver's compiled
46/// pipelines, where the device keeps them, and the shape pipelines that
47/// launch drew its first screen with.
48pub(crate) struct Loaded {
49    pub(crate) cache: Option<wgpu::PipelineCache>,
50    pub(crate) first_screen: Vec<u64>,
51}
52
53pub(crate) fn load(device: &wgpu::Device) -> Loaded {
54    let path = file_path();
55    let file = path.as_deref().and_then(|path| match std::fs::read(path) {
56        Ok(bytes) => Some(bytes),
57        Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
58        Err(error) => {
59            log::warn!("[pipeline-cache] unreadable {path:?}: {error}");
60            None
61        }
62    });
63    let contents = file.as_deref().and_then(current_contents);
64    let first_screen = contents
65        .as_ref()
66        .map_or_else(Vec::new, |contents| contents.first_screen.clone().collect());
67    if !device.features().contains(wgpu::Features::PIPELINE_CACHE) {
68        log::info!(
69            "[pipeline-cache] not offered by {:?}; compiled pipelines persist only as far \
70             as the driver's own cache does",
71            device.adapter_info().backend
72        );
73        return Loaded {
74            cache: None,
75            first_screen,
76        };
77    }
78    let data = contents
79        .as_ref()
80        .and_then(|contents| contents.blob)
81        .filter(|blob| !blob.is_empty());
82    // SAFETY: `data` is this build's own `get_data` output, and `fallback:
83    // true` has wgpu validate the header and fall back to an empty cache.
84    #[expect(unsafe_code)]
85    let cache = unsafe {
86        device.create_pipeline_cache(&wgpu::PipelineCacheDescriptor {
87            label: Some("cranpose pipeline disk cache"),
88            data,
89            fallback: true,
90        })
91    };
92    match (data, &file) {
93        (Some(data), _) => log::info!("[pipeline-cache] loaded {} B from disk", data.len()),
94        (None, Some(file))
95            if contents
96                .as_ref()
97                .is_none_or(|contents| contents.blob.is_none()) =>
98        {
99            log::info!(
100                "[pipeline-cache] cold: dropped {} B compiled from other shaders",
101                file.len()
102            );
103        }
104        (None, _) => log::info!("[pipeline-cache] cold (no blob on disk)"),
105    }
106    Loaded {
107        cache: Some(cache),
108        first_screen,
109    }
110}
111
112const FILE_LAYOUT: u32 = 3;
113
114/// Names what fills a file: its layout, the framework's WGSL sources, and
115/// this crate's version, which changes with each release of the shader
116/// rewrites, pipeline layouts and translator between those sources and the
117/// driver, and with every pipeline key's layout.
118///
119/// The driver keeps every pipeline of the blob it loads in the cache it saves,
120/// so a blob kept across shader changes only grows, and the driver holds all
121/// of it resident. A file under another key loads cold and is replaced.
122fn blob_key() -> [u8; 8] {
123    let mut hasher = FxHasher::default();
124    hasher.write_u32(FILE_LAYOUT);
125    hasher.write_u64(cranpose_ui_graphics::framework_shaders::SOURCES_KEY);
126    hasher.write(env!("CARGO_PKG_VERSION").as_bytes());
127    hasher.finish().to_le_bytes()
128}
129
130struct Contents<'a> {
131    first_screen: FirstScreenKeys<'a>,
132    blob: Option<&'a [u8]>,
133}
134
135#[derive(Clone)]
136struct FirstScreenKeys<'a>(&'a [u8]);
137
138impl Iterator for FirstScreenKeys<'_> {
139    type Item = u64;
140
141    fn next(&mut self) -> Option<u64> {
142        let (key, rest) = self.0.split_first_chunk::<8>()?;
143        self.0 = rest;
144        Some(u64::from_le_bytes(*key))
145    }
146
147    fn size_hint(&self) -> (usize, Option<usize>) {
148        (self.len(), Some(self.len()))
149    }
150}
151
152impl ExactSizeIterator for FirstScreenKeys<'_> {
153    fn len(&self) -> usize {
154        self.0.len() / 8
155    }
156}
157
158fn current_contents(file: &[u8]) -> Option<Contents<'_>> {
159    let (build, rest) = file.split_first_chunk::<8>()?;
160    let rest = rest.strip_prefix(crate::render::ShapePipelineKey::DISK_LAYOUT.as_slice())?;
161    let (count, rest) = rest.split_first_chunk::<4>()?;
162    let keys_len = usize::try_from(u32::from_le_bytes(*count))
163        .ok()?
164        .checked_mul(8)?;
165    let (keys, blob) = rest.split_at_checked(keys_len)?;
166    Some(Contents {
167        first_screen: FirstScreenKeys(keys),
168        blob: (*build == blob_key()).then_some(blob),
169    })
170}
171
172fn file_bytes(first_screen: impl ExactSizeIterator<Item = u64>, blob: &[u8]) -> Option<Vec<u8>> {
173    let count = u32::try_from(first_screen.len()).ok()?;
174    let mut bytes = Vec::with_capacity(20 + first_screen.len() * 8 + blob.len());
175    bytes.extend_from_slice(&blob_key());
176    bytes.extend_from_slice(&crate::render::ShapePipelineKey::DISK_LAYOUT);
177    bytes.extend_from_slice(&count.to_le_bytes());
178    for key in first_screen {
179        bytes.extend_from_slice(&key.to_le_bytes());
180    }
181    bytes.extend_from_slice(blob);
182    Some(bytes)
183}
184
185fn write_file(path: &Path, bytes: &[u8]) -> std::io::Result<()> {
186    let mut file = std::fs::File::create(path)?;
187    file.write_all(bytes)
188}
189
190pub(crate) fn persist(cache: Option<&wgpu::PipelineCache>, path: &Path) {
191    let started = Instant::now();
192    let blob = cache
193        .and_then(wgpu::PipelineCache::get_data)
194        .unwrap_or_default();
195    let existing = std::fs::read(path).ok();
196    let previous = existing.as_deref().and_then(current_contents);
197    let first_screen = first_screen_keys();
198    let key_count = first_screen.as_ref().map_or_else(
199        || {
200            previous
201                .as_ref()
202                .map_or(0, |contents| contents.first_screen.len())
203        },
204        Vec::len,
205    );
206    let bytes = match (first_screen.as_ref(), previous) {
207        (Some(keys), _) => file_bytes(keys.iter().copied(), &blob),
208        (None, Some(previous)) => file_bytes(previous.first_screen, &blob),
209        (None, None) => file_bytes(std::iter::empty(), &blob),
210    };
211    let Some(bytes) = bytes else {
212        return;
213    };
214    if existing.is_some_and(|existing| existing == bytes) {
215        return;
216    }
217    if let Some(parent) = path.parent()
218        && let Err(error) = std::fs::create_dir_all(parent)
219    {
220        log::warn!("[pipeline-cache] create_dir_all {parent:?}: {error}");
221        return;
222    }
223    let tmp = path.with_extension("tmp");
224    let written = write_file(&tmp, &bytes).and_then(|()| std::fs::rename(&tmp, path));
225    match written {
226        Ok(()) => log::info!(
227            "[pipeline-cache] persisted {} B and {} first-screen pipelines in {:.1} ms",
228            blob.len(),
229            key_count,
230            crate::render::instant_ms(started, Instant::now()),
231        ),
232        Err(error) => log::warn!("[pipeline-cache] write {path:?}: {error}"),
233    }
234}
235
236pub(crate) const FIRST_SCREEN_SPAN: Duration = Duration::from_secs(2);
237
238/// The shape pipelines, by their keys' bits, that this process's renderers
239/// drew their first screens with, in the order first drawn.
240static FIRST_SCREEN: Mutex<Vec<u64>> = Mutex::new(Vec::new());
241static FIRST_FRAME_DRAWN: AtomicBool = AtomicBool::new(false);
242
243/// Notes a shape pipeline a renderer drew its first screen with. The notes
244/// are written with the cache, for the next launch to build ahead of its
245/// first frame.
246pub(crate) fn note_first_screen_pipeline(key: u64) {
247    let mut keys = FIRST_SCREEN.lock().unwrap_or_else(PoisonError::into_inner);
248    if !keys.contains(&key) {
249        keys.push(key);
250        CHANGES.note_change();
251    }
252}
253
254pub(crate) fn note_frame_drawn() {
255    if !FIRST_FRAME_DRAWN.load(Ordering::Acquire) {
256        FIRST_FRAME_DRAWN.store(true, Ordering::Release);
257        CHANGES.note_change();
258    }
259}
260
261fn first_screen_keys() -> Option<Vec<u64>> {
262    FIRST_FRAME_DRAWN.load(Ordering::Acquire).then(|| {
263        FIRST_SCREEN
264            .lock()
265            .unwrap_or_else(PoisonError::into_inner)
266            .clone()
267    })
268}
269
270/// How long pipeline builds must pause before the cache is written. A burst
271/// of builds, such as a screen's first frames, is written once after its last
272/// build, and a session closed a moment after launch still keeps what the
273/// launch compiled, so the next launch skips those compiles.
274const PERSIST_QUIET: Duration = Duration::from_millis(500);
275
276static CHANGES: PersistSignal = PersistSignal::new();
277
278/// Counts a pipeline build and wakes the watchers waiting to write it back.
279pub(crate) fn note_pipeline_built() {
280    CHANGES.note_change();
281}
282
283struct PersistSignal {
284    revision: Mutex<u64>,
285    changed: Condvar,
286}
287
288/// Why a watcher woke, with the build count at that moment.
289#[derive(Debug, PartialEq, Eq)]
290enum Wake {
291    /// Builds past the written count paused for the quiet period.
292    Quiet(u64),
293    /// The watcher's owner let it go.
294    Stopped(u64),
295}
296
297impl PersistSignal {
298    const fn new() -> Self {
299        Self {
300            revision: Mutex::new(0),
301            changed: Condvar::new(),
302        }
303    }
304
305    fn lock(&self) -> MutexGuard<'_, u64> {
306        self.revision.lock().unwrap_or_else(PoisonError::into_inner)
307    }
308
309    fn note_change(&self) {
310        *self.lock() += 1;
311        self.changed.notify_all();
312    }
313
314    fn stop(&self, stopped: &AtomicBool) {
315        stopped.store(true, Ordering::Release);
316        // Holding the lock orders the flag against a watcher between its
317        // check and its wait, so the wake below cannot be lost.
318        let _held = self.lock();
319        self.changed.notify_all();
320    }
321
322    fn wait_for_quiet(&self, written: u64, quiet: Duration, stopped: &AtomicBool) -> Wake {
323        let is_stopped = || stopped.load(Ordering::Acquire);
324        let mut revision = self
325            .changed
326            .wait_while(self.lock(), |revision| {
327                *revision == written && !is_stopped()
328            })
329            .unwrap_or_else(PoisonError::into_inner);
330        loop {
331            if is_stopped() {
332                return Wake::Stopped(*revision);
333            }
334            let seen = *revision;
335            let (next, waited) = self
336                .changed
337                .wait_timeout_while(revision, quiet, |revision| {
338                    *revision == seen && !is_stopped()
339                })
340                .unwrap_or_else(PoisonError::into_inner);
341            revision = next;
342            if waited.timed_out() {
343                return Wake::Quiet(*revision);
344            }
345        }
346    }
347}
348
349/// Keeps the cache's watcher writing new pipelines to disk. Dropping it
350/// writes what is still unwritten and ends the watcher.
351pub(crate) struct PersistWatcher {
352    stopped: Arc<AtomicBool>,
353    thread: Option<std::thread::JoinHandle<()>>,
354}
355
356impl Drop for PersistWatcher {
357    fn drop(&mut self) {
358        CHANGES.stop(&self.stopped);
359        if let Some(thread) = self.thread.take()
360            && thread.join().is_err()
361        {
362            log::warn!("[pipeline-cache] persist thread panicked during shutdown");
363        }
364    }
365}
366
367pub(crate) fn spawn_persist_watcher(cache: Option<wgpu::PipelineCache>) -> Option<PersistWatcher> {
368    let path = file_path()?;
369    let stopped = Arc::new(AtomicBool::new(false));
370    let watcher_stopped = Arc::clone(&stopped);
371    let mut written = *CHANGES.lock();
372    let spawned = std::thread::Builder::new()
373        .name("cranpose-pl-cache".into())
374        .spawn(move || {
375            loop {
376                match CHANGES.wait_for_quiet(written, PERSIST_QUIET, &watcher_stopped) {
377                    Wake::Quiet(revision) => {
378                        persist(cache.as_ref(), &path);
379                        written = revision;
380                    }
381                    Wake::Stopped(revision) => {
382                        if revision != written {
383                            persist(cache.as_ref(), &path);
384                        }
385                        return;
386                    }
387                }
388            }
389        });
390    match spawned {
391        Ok(thread) => Some(PersistWatcher {
392            stopped,
393            thread: Some(thread),
394        }),
395        Err(error) => {
396            log::warn!("[pipeline-cache] persist thread failed to spawn: {error}");
397            None
398        }
399    }
400}
401
402#[cfg(test)]
403#[path = "tests/pipeline_disk_cache_tests.rs"]
404mod tests;
cranpose_render_wgpu/pipeline_disk_cache.rs

cranpose_render_wgpu/
pipeline_disk_cache.rs
