code.octet-stream.net Git - m17rt/blob - m17app/src/util/out_buffer.rs

   1 //! Buffer between `read()` calls
   2
   3 use std::{
   4     io::{self, ErrorKind, Read},
   5     sync::{mpsc::Receiver, Arc, Mutex},
   6 };
   7
   8 #[derive(Clone)]
   9 struct PartialOut {
  10     output: Arc<[u8]>,
  11     idx: usize,
  12 }
  13
  14 /// Buffer binary chunks from an MPSC receiver, feeding arbitrary chunks to `read()` calls.
  15 ///
  16 /// Can be cloned, but should only be read from once at a time.
  17 #[derive(Clone)]
  18 pub struct OutBuffer {
  19     rx: Arc<Mutex<Receiver<Arc<[u8]>>>>,
  20     partial_out: Arc<Mutex<Option<PartialOut>>>,
  21 }
  22
  23 impl OutBuffer {
  24     pub fn new(rx: Receiver<Arc<[u8]>>) -> Self {
  25         Self {
  26             rx: Arc::new(Mutex::new(rx)),
  27             partial_out: Arc::new(Mutex::new(None)),
  28         }
  29     }
  30 }
  31
  32 impl Read for OutBuffer {
  33     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
  34         {
  35             let mut partial_out = self.partial_out.lock().unwrap();
  36             if let Some(partial) = partial_out.as_mut() {
  37                 let remaining = partial.output.len() - partial.idx;
  38                 let to_write = remaining.min(buf.len());
  39                 buf[0..to_write]
  40                     .copy_from_slice(&partial.output[partial.idx..(partial.idx + to_write)]);
  41                 if to_write == remaining {
  42                     *partial_out = None;
  43                 } else {
  44                     partial.idx += to_write;
  45                 }
  46                 return Ok(to_write);
  47             }
  48         }
  49         let output = {
  50             let rx = self.rx.lock().unwrap();
  51             rx.recv()
  52                 .map_err(|s| io::Error::new(ErrorKind::Other, format!("{:?}", s)))?
  53         };
  54         let to_write = output.len().min(buf.len());
  55         buf[0..to_write].copy_from_slice(&output[0..to_write]);
  56         if to_write != output.len() {
  57             *self.partial_out.lock().unwrap() = Some(PartialOut {
  58                 output,
  59                 idx: to_write,
  60             })
  61         }
  62         Ok(to_write)
  63     }
  64 }