Skip to main content

pfx_live/
warm.rs

1use crate::frame::{OpaqueKey, opaque_pipeline};
2use std::collections::VecDeque;
3use std::sync::mpsc::Sender;
4use std::sync::{Arc, Condvar, Mutex};
5
6pub(crate) type Arrival = (OpaqueKey, Option<wgpu::RenderPipeline>);
7
8pub(crate) struct Job {
9    pub device: wgpu::Device,
10    pub layout: wgpu::PipelineLayout,
11    pub module: wgpu::ShaderModule,
12    pub cache: Option<wgpu::PipelineCache>,
13}
14
15#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
16pub struct Progress {
17    pub made: usize,
18    pub total: usize,
19    pub first_made: usize,
20    pub first_total: usize,
21    pub failed: usize,
22}
23
24impl Progress {
25    pub fn fraction(self) -> f32 {
26        if self.total == 0 {
27            1.0
28        } else {
29            self.made as f32 / self.total as f32
30        }
31    }
32}
33
34#[derive(Debug)]
35struct Shared {
36    progress: Mutex<Progress>,
37    changed: Condvar,
38}
39
40#[derive(Clone, Debug)]
41pub struct Warming {
42    shared: Arc<Shared>,
43}
44
45impl Warming {
46    pub fn progress(&self) -> Progress {
47        *self.shared.progress.lock().expect("warm progress")
48    }
49
50    pub fn first_ready(&self) -> bool {
51        let progress = self.progress();
52        progress.first_made >= progress.first_total
53    }
54
55    pub fn done(&self) -> bool {
56        let progress = self.progress();
57        progress.made >= progress.total
58    }
59
60    pub fn wait(&self) -> Progress {
61        let mut progress = self.shared.progress.lock().expect("warm progress");
62        while progress.made < progress.total {
63            progress = self.shared.changed.wait(progress).expect("warm progress");
64        }
65        *progress
66    }
67
68    fn finished(&self, first: bool, made: bool) {
69        let mut progress = self.shared.progress.lock().expect("warm progress");
70        progress.made += 1;
71        if first {
72            progress.first_made += 1;
73        }
74        if !made {
75            progress.failed += 1;
76        }
77        drop(progress);
78        self.shared.changed.notify_all();
79    }
80}
81
82struct Promise {
83    key: OpaqueKey,
84    first: bool,
85    sender: Sender<Arrival>,
86    warming: Warming,
87    kept: bool,
88}
89
90impl Promise {
91    fn keep(mut self, made: wgpu::RenderPipeline) -> bool {
92        self.kept = true;
93        self.warming.finished(self.first, true);
94        self.sender.send((self.key, Some(made))).is_ok()
95    }
96}
97
98impl Drop for Promise {
99    fn drop(&mut self) {
100        if !self.kept {
101            self.warming.finished(self.first, false);
102            let _ = self.sender.send((self.key, None));
103        }
104    }
105}
106
107pub fn default_workers() -> usize {
108    std::thread::available_parallelism()
109        .map_or(1, |threads| threads.get())
110        .saturating_sub(1)
111        .max(1)
112}
113
114pub(crate) fn spawn(
115    job: Job,
116    keys: Vec<OpaqueKey>,
117    first: usize,
118    workers: usize,
119    sender: Sender<Arrival>,
120) -> Warming {
121    let warming = Warming {
122        shared: Arc::new(Shared {
123            progress: Mutex::new(Progress {
124                total: keys.len(),
125                first_total: first,
126                ..Progress::default()
127            }),
128            changed: Condvar::new(),
129        }),
130    };
131    let count = workers.max(1).min(keys.len());
132    let queue: Arc<Mutex<VecDeque<(usize, OpaqueKey)>>> =
133        Arc::new(Mutex::new(keys.into_iter().enumerate().collect()));
134    let job = Arc::new(job);
135    let abandon = {
136        let sender = sender.clone();
137        let warming = warming.clone();
138        move |queue: &Mutex<VecDeque<(usize, OpaqueKey)>>| {
139            let left: Vec<(usize, OpaqueKey)> =
140                queue.lock().expect("warm queue").drain(..).collect();
141            for (index, key) in left {
142                drop(Promise {
143                    key,
144                    first: index < first,
145                    sender: sender.clone(),
146                    warming: warming.clone(),
147                    kept: false,
148                });
149            }
150        }
151    };
152    let mut spawned = 0;
153    for _ in 0..count {
154        let queue = Arc::clone(&queue);
155        let job = Arc::clone(&job);
156        let sender = sender.clone();
157        let warming = warming.clone();
158        let abandon = abandon.clone();
159        let worker = std::thread::Builder::new()
160            .name("pito warm".into())
161            .spawn(move || {
162                loop {
163                    let next = queue.lock().expect("warm queue").pop_front();
164                    let Some((index, key)) = next else { break };
165                    let promise = Promise {
166                        key,
167                        first: index < first,
168                        sender: sender.clone(),
169                        warming: warming.clone(),
170                        kept: false,
171                    };
172                    let made = opaque_pipeline(
173                        &job.device,
174                        &job.layout,
175                        &job.module,
176                        job.cache.as_ref(),
177                        key,
178                    );
179                    if !promise.keep(made) {
180                        abandon(&queue);
181                        break;
182                    }
183                }
184            });
185        if worker.is_err() {
186            break;
187        }
188        spawned += 1;
189    }
190    if spawned == 0 {
191        abandon(&queue);
192    }
193    warming
194}