Skip to main content

diffs_rs/input/
worker.rs

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
19/// Published files waiting for the screen are bounded by count and index bytes. A full
20/// queue pauses ingestion until the screen drains it; it never delays keyboard input.
21/// The count also bounds the screen's per-turn drain.
22const QUEUED_FILES: usize = 1024;
23const QUEUED_FILE_BYTES: usize = 8 << 20;
24/// Prepared sources retain both complete sides, so only two may wait.
25const 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    /// The worker was closed; a waiting producer must stop.
34    closed: bool,
35    /// The producer exited; nothing more will arrive.
36    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            // An oversized file still enters an empty queue; it is accounted for in peak memory.
65            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            // At most one completion and one error are ever published.
72            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/// Single-producer handoff from the input thread to the screen. The producer blocks on a
100/// condition variable when the queue is full; the screen takes everything queued at once.
101#[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        // Release queued files and sources outside the lock.
151        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
193/// Marks the producer finished before the exit notification, so the screen observes disconnection.
194struct 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    /// Take the oldest reply. `Disconnected` means the worker exited and nothing remains.
301    pub fn try_reply(&self) -> std::result::Result<Result<Message>, mpsc::TryRecvError> {
302        self.replies.try_reply()
303    }
304
305    /// Take every queued reply in order, releasing a producer that waits for space.
306    /// `Disconnected` means the worker exited and nothing remains.
307    pub fn try_replies(
308        &self,
309    ) -> std::result::Result<VecDeque<Result<Message>>, mpsc::TryRecvError> {
310        self.replies.try_replies()
311    }
312
313    /// Cancel and join. Input errors arrive through replies; this reports cleanup failure or panic.
314    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}