rusthound_ce/storage/
buffer.rs

1use std::error::Error;
2use std::fs::OpenOptions;
3use std::io::{BufReader, BufWriter, Seek, Write};
4use std::path::Path;
5
6pub use super::iter::BincodeIterator;
7
8const DEFAULT_BUFFER_SIZE: usize = 1000;
9
10pub trait Storage<T>
11where
12    Self: Sized,
13{
14    /// Returns a mutable reference to the internal buffer
15    fn buffer_mut(&mut self) -> &mut Vec<T>;
16
17    /// Flush the buffer to disk
18    fn flush(&mut self) -> Result<(), Box<dyn Error>>;
19
20    /// Add an item to the buffer,
21    /// Default implemetation calls [`Storage::flush`] if it reaches capacity
22    fn add(&mut self, item: T) -> Result<(), Box<dyn Error>> {
23        self.buffer_mut().push(item);
24
25        if self.buffer_mut().len() >= self.buffer_mut().capacity() {
26            self.flush()?;
27        }
28        Ok(())
29    }
30
31    /// Flush the buffer and consume `self`
32    fn finish(mut self) -> Result<(), Box<dyn Error>> {
33        self.flush()
34    }
35}
36
37impl<T> Storage<T> for Vec<T> {
38    fn buffer_mut(&mut self) -> &mut Vec<T> {
39        self
40    }
41
42    // no need to check capacity
43    fn add(&mut self, item: T) -> Result<(), Box<dyn Error>> {
44        self.push(item);
45        Ok(())
46    }
47
48    fn flush(&mut self) -> Result<(), Box<dyn Error>> {
49        Ok(())
50    }
51}
52
53pub struct BincodeObjectBuffer<T> {
54    /// `BufWriter` to a file handle that is opened for reading and writing.
55    writer: BufWriter<RWHandle>,
56
57    /// Buffer for storing objects to be written to disk
58    buffer: Vec<T>,
59
60    /// Intermediate buffer for encoding objects to bincode before writing to disk
61    encode_buffer: Vec<u8>,
62}
63
64impl<T> BincodeObjectBuffer<T> {
65    pub fn new(file_path: impl AsRef<Path>) -> Result<Self, Box<dyn Error>> {
66        Ok(BincodeObjectBuffer {
67            writer: BufWriter::new(RWHandle::open(file_path)?),
68            buffer: Vec::with_capacity(DEFAULT_BUFFER_SIZE),
69            encode_buffer: Vec::new(),
70        })
71    }
72
73    pub fn new_with_capacity(
74        file_path: impl AsRef<Path>,
75        capacity: usize,
76    ) -> Result<Self, Box<dyn Error>> {
77        Ok(BincodeObjectBuffer {
78            writer: BufWriter::new(RWHandle::open(file_path)?),
79            buffer: Vec::with_capacity(capacity),
80            encode_buffer: Vec::new(),
81        })
82    }
83}
84
85impl<T: bincode::Decode<()>> BincodeObjectBuffer<T> {
86    pub fn into_reader(
87        self,
88    ) -> Result<BincodeIterator<T, BufReader<std::fs::File>>, Box<dyn Error>> {
89        let mut inner = self.writer.into_inner()?;
90        inner.0.seek(std::io::SeekFrom::Start(0))?;
91        Ok(BincodeIterator::from_file(inner.0))
92    }
93}
94
95impl<T> Storage<T> for BincodeObjectBuffer<T>
96where
97    T: bincode::Encode,
98{
99    #[inline]
100    fn buffer_mut(&mut self) -> &mut Vec<T> {
101        &mut self.buffer
102    }
103
104    fn flush(&mut self) -> Result<(), Box<dyn Error>> {
105        for item in self.buffer.drain(..) {
106            self.encode_buffer.clear();
107            bincode::encode_into_std_write(
108                &item,
109                &mut self.encode_buffer,
110                bincode::config::standard(),
111            )?;
112
113            let len = self.encode_buffer.len() as u32;
114            self.writer.write_all(&len.to_le_bytes())?;
115            self.writer.write_all(&self.encode_buffer)?;
116        }
117
118        self.writer.flush()?;
119
120        Ok(())
121    }
122}
123
124/// Wrapper around a file handle. Used to indicate that the file is opened for reading and writing.
125///
126/// Will truncate the file if it already exists.
127#[derive(Debug)]
128struct RWHandle(std::fs::File);
129
130impl RWHandle {
131    pub fn open(file_path: impl AsRef<Path>) -> Result<Self, Box<dyn Error>> {
132        let file = OpenOptions::new()
133            .create(true)
134            .write(true)
135            .read(true) // read so we can read back the file later
136            .truncate(true)
137            .open(file_path)?;
138
139        Ok(RWHandle(file))
140    }
141}
142
143impl std::io::Write for RWHandle {
144    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
145        self.0.write(buf)
146    }
147
148    fn flush(&mut self) -> std::io::Result<()> {
149        self.0.flush()
150    }
151}
rusthound_ce/storage/buffer.rs

rusthound_ce/storage/
buffer.rs
