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}