/src/pipe.rs

https://github.com/phiresky/ripgrep-all · Rust · 196 lines · 139 code · 29 blank · 28 comment · 14 complexity · a9aa9f8c1b9bb5dda8b4339fa7834312 MD5 · raw file

  1. // https://github.com/arcnmx/pipe-rs/blob/master/src/lib.rs
  2. // extended to support sending io errors
  3. #![deny(missing_docs)]
  4. #![doc(html_root_url = "https://docs.rs/pipe/0.3.0")]
  5. #![cfg_attr(feature = "unstable-doc-cfg", feature(doc_cfg))]
  6. //! Synchronous in-memory pipe
  7. //!
  8. //! ## Example
  9. //!
  10. //! ```
  11. //! use std::thread::spawn;
  12. //! use std::io::{Read, Write};
  13. //!
  14. //! let (mut read, mut write) = ripgrep_all::pipe::pipe();
  15. //!
  16. //! let message = "Hello, world!";
  17. //! spawn(move || write.write_all(message.as_bytes()).unwrap());
  18. //!
  19. //! let mut s = String::new();
  20. //! read.read_to_string(&mut s).unwrap();
  21. //!
  22. //! assert_eq!(&s, message);
  23. //! ```
  24. use crossbeam_channel::{Receiver, Sender};
  25. use std::cmp::min;
  26. use std::io::{self, BufRead, Read, Result, Write};
  27. /// The `Read` end of a pipe (see `pipe()`)
  28. pub struct PipeReader {
  29. receiver: Receiver<Result<Vec<u8>>>,
  30. buffer: Vec<u8>,
  31. position: usize,
  32. }
  33. /// The `Write` end of a pipe (see `pipe()`)
  34. #[derive(Clone)]
  35. pub struct PipeWriter {
  36. sender: Sender<Result<Vec<u8>>>,
  37. }
  38. /// Creates a synchronous memory pipe
  39. pub fn pipe() -> (PipeReader, PipeWriter) {
  40. let (sender, receiver) = crossbeam_channel::bounded(0);
  41. (
  42. PipeReader {
  43. receiver,
  44. buffer: Vec::new(),
  45. position: 0,
  46. },
  47. PipeWriter { sender },
  48. )
  49. }
  50. impl PipeWriter {
  51. /// Extracts the inner `SyncSender` from the writer
  52. pub fn into_inner(self) -> Sender<Result<Vec<u8>>> {
  53. self.sender
  54. }
  55. /// Write any error into the pipe, will be handled as an IO error
  56. pub fn write_err(&self, e: std::io::Error) -> Result<()> {
  57. self.sender
  58. .send(Err(e))
  59. .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "pipe reader has been dropped"))
  60. }
  61. }
  62. impl PipeReader {
  63. /// Extracts the inner `Receiver` from the writer, and any pending buffered data
  64. pub fn into_inner(mut self) -> (Receiver<Result<Vec<u8>>>, Vec<u8>) {
  65. self.buffer.drain(..self.position);
  66. (self.receiver, self.buffer)
  67. }
  68. }
  69. impl BufRead for PipeReader {
  70. fn fill_buf(&mut self) -> io::Result<&[u8]> {
  71. while self.position >= self.buffer.len() {
  72. match self.receiver.recv() {
  73. // The only existing error is EOF
  74. Err(_) => break,
  75. Ok(Err(e)) => Err(e)?,
  76. Ok(Ok(data)) => {
  77. self.buffer = data;
  78. self.position = 0;
  79. }
  80. }
  81. }
  82. Ok(&self.buffer[self.position..])
  83. }
  84. fn consume(&mut self, amt: usize) {
  85. debug_assert!(self.buffer.len() - self.position >= amt);
  86. self.position += amt
  87. }
  88. }
  89. impl Read for PipeReader {
  90. fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
  91. if buf.is_empty() {
  92. return Ok(0);
  93. }
  94. let internal = self.fill_buf()?;
  95. let len = min(buf.len(), internal.len());
  96. if len > 0 {
  97. buf[..len].copy_from_slice(&internal[..len]);
  98. self.consume(len);
  99. }
  100. Ok(len)
  101. }
  102. }
  103. impl Write for PipeWriter {
  104. fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
  105. let data = buf.to_vec();
  106. self.sender
  107. .send(Ok(data))
  108. .map(|_| buf.len())
  109. .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "pipe reader has been dropped"))
  110. }
  111. fn flush(&mut self) -> io::Result<()> {
  112. Ok(())
  113. }
  114. }
  115. #[cfg(test)]
  116. mod tests {
  117. use super::*;
  118. use std::io::{Read, Write};
  119. use std::thread::spawn;
  120. #[test]
  121. fn pipe_reader() {
  122. let i = b"hello there";
  123. let mut o = Vec::with_capacity(i.len());
  124. let (mut r, mut w) = pipe();
  125. let guard = spawn(move || {
  126. w.write_all(&i[..5]).unwrap();
  127. w.write_all(&i[5..]).unwrap();
  128. drop(w);
  129. });
  130. r.read_to_end(&mut o).unwrap();
  131. assert_eq!(i, &o[..]);
  132. guard.join().unwrap();
  133. }
  134. #[test]
  135. fn pipe_writer_fail() {
  136. let i = b"hi";
  137. let (r, mut w) = pipe();
  138. let guard = spawn(move || {
  139. drop(r);
  140. });
  141. assert!(w.write_all(i).is_err());
  142. guard.join().unwrap();
  143. }
  144. #[test]
  145. fn small_reads() {
  146. let block_cnt = 20;
  147. const BLOCK: usize = 20;
  148. let (mut r, mut w) = pipe();
  149. let guard = spawn(move || {
  150. for _ in 0..block_cnt {
  151. let data = &[0; BLOCK];
  152. w.write_all(data).unwrap();
  153. }
  154. });
  155. let mut buff = [0; BLOCK / 2];
  156. let mut read = 0;
  157. while let Ok(size) = r.read(&mut buff) {
  158. // 0 means EOF
  159. if size == 0 {
  160. break;
  161. }
  162. read += size;
  163. }
  164. assert_eq!(block_cnt * BLOCK, read);
  165. guard.join().unwrap();
  166. }
  167. }