1use super::*;
2use std::{
3 collections::VecDeque,
4 sync::{
5 Condvar, Mutex, MutexGuard,
6 atomic::Ordering,
7 mpsc::{self, SyncSender},
8 },
9 thread::{self, JoinHandle},
10 time::Duration,
11};
12
13pub enum Message {
14 File(Arc<FileIndex>),
15 Complete { files: usize },
16 Source { file: usize, state: SourceState },
17}
18
19const QUEUED_FILES: usize = 1024;
23const QUEUED_FILE_BYTES: usize = 8 << 20;
24const QUEUED_SOURCES: usize = 2;
26
27#[derive(Default)]
28struct Queue {
29 replies: VecDeque<Result<Message>>,
30 files: usize,
31 file_bytes: usize,
32 sources: usize,
33 closed: bool,
35 finished: bool,
37}
38
39#[derive(Clone, Copy)]
40enum Cost {
41 File(usize),
42 Source,
43 Final,
44}
45
46impl Cost {
47 fn of(reply: &Result<Message>) -> Self {
48 match reply {
49 Ok(Message::File(file)) => Self::File(
50 size_of::<FileIndex>()
51 + file.source_index_bytes()
52 + file.projection_bytes()
53 + file.metadata_bytes(),
54 ),
55 Ok(Message::Source { .. }) => Self::Source,
56 Ok(Message::Complete { .. }) | Err(_) => Self::Final,
57 }
58 }
59}
60
61impl Queue {
62 fn admits(&self, cost: Cost) -> bool {
63 match cost {
64 Cost::File(bytes) => {
66 self.files == 0
67 || (self.files < QUEUED_FILES
68 && self.file_bytes.saturating_add(bytes) <= QUEUED_FILE_BYTES)
69 }
70 Cost::Source => self.sources < QUEUED_SOURCES,
71 Cost::Final => true,
73 }
74 }
75
76 fn add(&mut self, cost: Cost) {
77 match cost {
78 Cost::File(bytes) => {
79 self.files += 1;
80 self.file_bytes = self.file_bytes.saturating_add(bytes);
81 }
82 Cost::Source => self.sources += 1,
83 Cost::Final => {}
84 }
85 }
86
87 fn remove(&mut self, cost: Cost) {
88 match cost {
89 Cost::File(bytes) => {
90 self.files = self.files.saturating_sub(1);
91 self.file_bytes = self.file_bytes.saturating_sub(bytes);
92 }
93 Cost::Source => self.sources = self.sources.saturating_sub(1),
94 Cost::Final => {}
95 }
96 }
97}
98
99#[derive(Default)]
102struct Mailbox {
103 queue: Mutex<Queue>,
104 space: Condvar,
105}
106
107impl Mailbox {
108 fn lock(&self) -> Result<MutexGuard<'_, Queue>> {
109 self.queue
110 .lock()
111 .map_err(|_| anyhow::anyhow!("The input queue lock failed"))
112 }
113
114 fn deliver(&self, reply: Result<Message>, stop: &Cancellation) -> Result<()> {
115 let cost = Cost::of(&reply);
116 let mut queue = self.lock()?;
117 loop {
118 if queue.closed {
119 return Err(Cancelled.into());
120 }
121 stop.check()?;
122 if queue.admits(cost) {
123 break;
124 }
125 queue = self
126 .space
127 .wait(queue)
128 .map_err(|_| anyhow::anyhow!("The input queue lock failed"))?;
129 }
130 queue.add(cost);
131 queue.replies.push_back(reply);
132 Ok(())
133 }
134
135 fn finish(&self) {
136 if let Ok(mut queue) = self.queue.lock() {
137 queue.finished = true;
138 }
139 }
140
141 fn close(&self) {
142 let replies = self.queue.lock().map(|mut queue| {
143 queue.closed = true;
144 queue.files = 0;
145 queue.file_bytes = 0;
146 queue.sources = 0;
147 std::mem::take(&mut queue.replies)
148 });
149 self.space.notify_all();
150 drop(replies);
152 }
153
154 fn try_reply(&self) -> std::result::Result<Result<Message>, mpsc::TryRecvError> {
155 let mut queue = self
156 .queue
157 .lock()
158 .map_err(|_| mpsc::TryRecvError::Disconnected)?;
159 match queue.replies.pop_front() {
160 Some(reply) => {
161 queue.remove(Cost::of(&reply));
162 drop(queue);
163 self.space.notify_one();
164 Ok(reply)
165 }
166 None if queue.finished => Err(mpsc::TryRecvError::Disconnected),
167 None => Err(mpsc::TryRecvError::Empty),
168 }
169 }
170
171 fn try_replies(&self) -> std::result::Result<VecDeque<Result<Message>>, mpsc::TryRecvError> {
172 let mut queue = self
173 .queue
174 .lock()
175 .map_err(|_| mpsc::TryRecvError::Disconnected)?;
176 if queue.replies.is_empty() {
177 return Err(if queue.finished {
178 mpsc::TryRecvError::Disconnected
179 } else {
180 mpsc::TryRecvError::Empty
181 });
182 }
183 queue.files = 0;
184 queue.file_bytes = 0;
185 queue.sources = 0;
186 let replies = std::mem::take(&mut queue.replies);
187 drop(queue);
188 self.space.notify_one();
189 Ok(replies)
190 }
191}
192
193struct Finish(Arc<Mailbox>);
195
196impl Drop for Finish {
197 fn drop(&mut self) {
198 self.0.finish();
199 }
200}
201
202pub struct Worker {
203 cancel: Cancellation,
204 destination: Arc<AtomicUsize>,
205 wake: SyncSender<()>,
206 replies: Arc<Mailbox>,
207 join: Option<JoinHandle<Result<()>>>,
208}
209
210enum Input {
211 Local(LocalInput),
212 Patch(PatchInput),
213}
214
215impl Worker {
216 pub fn start(input: LocalInput, notify: SyncSender<()>) -> Self {
217 Self::spawn(Input::Local(input), notify)
218 }
219
220 pub fn start_patch(input: PatchInput, notify: SyncSender<()>) -> Self {
221 Self::spawn(Input::Patch(input), notify)
222 }
223
224 fn spawn(input: Input, notify: SyncSender<()>) -> Self {
225 let cancel = Cancellation::default();
226 let destination = Arc::new(AtomicUsize::new(usize::MAX));
227 let (wake, requests) = mpsc::sync_channel(1);
228 let replies = Arc::new(Mailbox::default());
229 let mailbox = replies.clone();
230 let stop = cancel.clone();
231 let wanted = destination.clone();
232 let join = thread::spawn(move || {
233 let _notify_exit = crate::wake::NotifyOnExit(notify.clone());
234 let _finish = Finish(mailbox.clone());
235 let deliver = |value| -> Result<()> {
236 mailbox.deliver(value, &stop)?;
237 let _ = notify.try_send(());
238 Ok(())
239 };
240 let result = (|| -> Result<()> {
241 let input = match input {
242 Input::Local(input) => input,
243 Input::Patch(input) => {
244 let files =
245 patch::read(input, &stop, |file| deliver(Ok(Message::File(file))))?;
246 return deliver(Ok(Message::Complete { files }));
247 }
248 };
249 let mut review =
250 Review::open(input, &stop, |file| deliver(Ok(Message::File(file))))?;
251 deliver(Ok(Message::Complete {
252 files: review.doc.files.len(),
253 }))?;
254 loop {
255 stop.check()?;
256 match requests.recv_timeout(Duration::from_millis(50)) {
257 Ok(()) => {}
258 Err(mpsc::RecvTimeoutError::Timeout) => continue,
259 Err(mpsc::RecvTimeoutError::Disconnected) => return Ok(()),
260 }
261 let id = wanted.load(Ordering::Acquire);
262 if id >= review.doc.files.len() {
263 continue;
264 }
265 let request = Cancellation {
266 stopped: stop.stopped.clone(),
267 destination: Some((wanted.clone(), id)),
268 };
269 match review.prepare(id, &request) {
270 Ok(state) => deliver(Ok(Message::Source { file: id, state }))?,
271 Err(error) if error.is::<Cancelled>() => {}
272 Err(error) => return Err(error),
273 }
274 }
275 })();
276 if let Err(error) = &result
277 && !error.is::<Cancelled>()
278 {
279 let _ = deliver(Err(anyhow::anyhow!("{error:#}")));
280 }
281 match result {
282 Err(error) if error.is::<patch::RestoreFailure>() => Err(error),
283 _ => Ok(()),
284 }
285 });
286 Self {
287 cancel,
288 destination,
289 wake,
290 replies,
291 join: Some(join),
292 }
293 }
294
295 pub fn request(&self, file: usize) {
296 if self.destination.swap(file, Ordering::AcqRel) != file {
297 let _ = self.wake.try_send(());
298 }
299 }
300 pub fn try_reply(&self) -> std::result::Result<Result<Message>, mpsc::TryRecvError> {
302 self.replies.try_reply()
303 }
304
305 pub fn try_replies(
308 &self,
309 ) -> std::result::Result<VecDeque<Result<Message>>, mpsc::TryRecvError> {
310 self.replies.try_replies()
311 }
312
313 pub fn close(&mut self) -> Result<()> {
315 self.cancel.stop();
316 self.replies.close();
317 let _ = self.wake.try_send(());
318 let result = self.join.take().map_or(Ok(()), |join| {
319 join.join()
320 .map_err(|_| anyhow::anyhow!("The source worker panicked"))
321 .and_then(|result| match result {
322 Err(error) if error.is::<Cancelled>() => Ok(()),
323 result => result,
324 })
325 });
326 self.replies.close();
327 result
328 }
329}
330
331impl Drop for Worker {
332 fn drop(&mut self) {
333 let _ = self.close();
334 }
335}
336
337#[cfg(test)]
338mod tests {
339 use super::*;
340
341 fn source(file: usize) -> Result<Message> {
342 Ok(Message::Source {
343 file,
344 state: SourceState::Pending,
345 })
346 }
347
348 #[test]
349 fn queued_files_are_bounded_by_count_and_bytes_but_one_always_fits() {
350 let mut queue = Queue::default();
351 assert!(queue.admits(Cost::File(QUEUED_FILE_BYTES * 2)));
352 for _ in 0..QUEUED_FILES {
353 assert!(queue.admits(Cost::File(1)));
354 queue.add(Cost::File(1));
355 }
356 assert!(!queue.admits(Cost::File(1)));
357 assert!(queue.admits(Cost::Final));
358 queue.remove(Cost::File(1));
359 assert!(queue.admits(Cost::File(1)));
360
361 let mut queue = Queue::default();
362 queue.add(Cost::File(QUEUED_FILE_BYTES - 10));
363 assert!(queue.admits(Cost::File(10)));
364 assert!(!queue.admits(Cost::File(11)));
365 assert!(queue.admits(Cost::Source));
366 }
367
368 #[test]
369 fn a_full_queue_blocks_the_producer_until_the_screen_drains_everything() {
370 let mailbox = Arc::new(Mailbox::default());
371 let stop = Cancellation::default();
372 for file in 0..QUEUED_SOURCES {
373 mailbox.deliver(source(file), &stop).unwrap();
374 }
375 let (done, delivered) = mpsc::sync_channel(1);
376 let producer = {
377 let mailbox = mailbox.clone();
378 let stop = stop.clone();
379 thread::spawn(move || {
380 let result = mailbox.deliver(source(QUEUED_SOURCES), &stop);
381 done.send(()).unwrap();
382 result
383 })
384 };
385 assert_eq!(
386 delivered.recv_timeout(Duration::from_millis(50)),
387 Err(mpsc::RecvTimeoutError::Timeout)
388 );
389 let replies = mailbox.try_replies().unwrap();
390 assert_eq!(replies.len(), QUEUED_SOURCES);
391 delivered.recv_timeout(Duration::from_secs(10)).unwrap();
392 producer.join().unwrap().unwrap();
393 let replies = mailbox.try_replies().unwrap();
394 assert!(matches!(
395 replies.front(),
396 Some(Ok(Message::Source { file, .. })) if *file == QUEUED_SOURCES
397 ));
398 assert!(matches!(
399 mailbox.try_replies(),
400 Err(mpsc::TryRecvError::Empty)
401 ));
402 mailbox.finish();
403 assert!(matches!(
404 mailbox.try_reply(),
405 Err(mpsc::TryRecvError::Disconnected)
406 ));
407 }
408
409 #[test]
410 fn closing_releases_a_blocked_producer_and_drops_queued_replies() {
411 let mailbox = Arc::new(Mailbox::default());
412 let stop = Cancellation::default();
413 for file in 0..QUEUED_SOURCES {
414 mailbox.deliver(source(file), &stop).unwrap();
415 }
416 let producer = {
417 let mailbox = mailbox.clone();
418 let stop = stop.clone();
419 thread::spawn(move || mailbox.deliver(source(QUEUED_SOURCES), &stop))
420 };
421 thread::sleep(Duration::from_millis(20));
422 stop.stop();
423 mailbox.close();
424 let error = producer.join().unwrap().unwrap_err();
425 assert!(error.is::<Cancelled>());
426 assert!(mailbox.lock().unwrap().replies.is_empty());
427 assert!(
428 mailbox
429 .deliver(source(0), &Cancellation::default())
430 .unwrap_err()
431 .is::<Cancelled>()
432 );
433 }
434}