xref: /cloud-hypervisor/serial_buffer/src/lib.rs (revision eea9bcea38e0c5649f444c829f3a4f9c22aa486c)
1 // Copyright © 2021 Intel Corporation
2 //
3 // SPDX-License-Identifier: Apache-2.0
4 //
5 
6 use std::{
7     collections::VecDeque,
8     io::Write,
9     sync::{
10         atomic::{AtomicBool, Ordering},
11         Arc,
12     },
13 };
14 
15 const MAX_BUFFER_SIZE: usize = 1 << 20;
16 
17 // Circular buffer implementation for serial output.
18 // Read from head; push to tail
19 pub struct SerialBuffer {
20     buffer: VecDeque<u8>,
21     out: Box<dyn Write + Send>,
22     write_out: Arc<AtomicBool>,
23 }
24 
25 impl SerialBuffer {
26     pub fn new(out: Box<dyn Write + Send>, write_out: Arc<AtomicBool>) -> Self {
27         Self {
28             buffer: VecDeque::new(),
29             out,
30             write_out,
31         }
32     }
33 
34     fn fill_buffer(&mut self, buf: &[u8]) {
35         if buf.len() >= MAX_BUFFER_SIZE {
36             let offset = buf.len() - MAX_BUFFER_SIZE;
37             self.buffer = VecDeque::from(buf[offset..].to_vec());
38             return;
39         }
40 
41         let num_allowed_bytes = MAX_BUFFER_SIZE - buf.len();
42         if self.buffer.len() > num_allowed_bytes {
43             let num_bytes_to_remove = self.buffer.len() - num_allowed_bytes;
44             self.buffer.drain(..num_bytes_to_remove);
45         }
46 
47         self.buffer.extend(buf);
48     }
49 }
50 
51 impl Write for SerialBuffer {
52     fn write(&mut self, buf: &[u8]) -> Result<usize, std::io::Error> {
53         // Simply fill the buffer if we're not allowed to write to the out
54         // device.
55         if !self.write_out.load(Ordering::Acquire) {
56             self.fill_buffer(buf);
57             return Ok(buf.len());
58         }
59 
60         // In case we're allowed to write to the out device, we flush the
61         // content of the buffer.
62         self.flush()?;
63 
64         // If after flushing the buffer, it's still not empty, that means
65         // only a subset of the bytes was written and we should fill the buffer
66         // with what's coming from the serial.
67         if !self.buffer.is_empty() {
68             self.fill_buffer(buf);
69             return Ok(buf.len());
70         }
71 
72         // We reach this point if we're allowed to write to the out device
73         // and we know there's nothing left in the buffer.
74         let mut offset = 0;
75         loop {
76             match self.out.write(&buf[offset..]) {
77                 Ok(written_bytes) => {
78                     if written_bytes < buf.len() - offset {
79                         offset += written_bytes;
80                         continue;
81                     }
82                 }
83                 Err(e) => {
84                     if !matches!(e.kind(), std::io::ErrorKind::WouldBlock) {
85                         return Err(e);
86                     }
87                     self.fill_buffer(&buf[offset..]);
88                 }
89             }
90             break;
91         }
92 
93         // Make sure we flush anything that might have been written to the
94         // out device.
95         self.out.flush()?;
96 
97         Ok(buf.len())
98     }
99 
100     // This function flushes the content of the buffer to the out device if
101     // it is allowed to, otherwise this is a no-op.
102     fn flush(&mut self) -> Result<(), std::io::Error> {
103         if !self.write_out.load(Ordering::Acquire) {
104             return Ok(());
105         }
106 
107         while let Some(byte) = self.buffer.pop_front() {
108             if self.out.write_all(&[byte]).is_err() {
109                 self.buffer.push_front(byte);
110                 break;
111             }
112         }
113         self.out.flush()
114     }
115 }
116