pikadick/logger/
delay_writer.rs

1use parking_lot::Mutex;
2use std::{
3    io::Write,
4    sync::Arc,
5};
6
7const DEFAULT_CAPACITY: usize = 1_000_000;
8
9/// The mut impl of a DelayWriter
10#[derive(Debug)]
11pub(crate) enum DelayWriterInner<W> {
12    /// The buffered data.
13    Buffer(Vec<u8>),
14
15    /// The file being written to.
16    Writer(W),
17}
18
19impl<W> DelayWriterInner<W> {
20    /// Make a new DelayWriterInner with an empty buffer
21    fn new() -> Self {
22        Self::Buffer(Vec::with_capacity(DEFAULT_CAPACITY))
23    }
24}
25
26impl<W> DelayWriterInner<W>
27where
28    W: Write,
29{
30    /// Try to init this DelayWriterInner with a file.
31    ///
32    /// # Error
33    /// Will return an error if this is already initalized.
34    fn init(&mut self, mut writer: W) -> Result<(), std::io::Error> {
35        let buffer = match self {
36            Self::Buffer(bytes) => bytes,
37            Self::Writer(_) => {
38                return Err(std::io::Error::other("already initalized"));
39            }
40        };
41
42        writer.write_all(buffer)?;
43
44        *self = Self::Writer(writer);
45
46        Ok(())
47    }
48}
49
50impl<W> Write for DelayWriterInner<W>
51where
52    W: Write,
53{
54    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
55        match self {
56            Self::Buffer(buffer) => buffer.write(buf),
57            Self::Writer(writer) => writer.write(buf),
58        }
59    }
60
61    fn flush(&mut self) -> std::io::Result<()> {
62        match self {
63            Self::Buffer(buffer) => buffer.flush(),
64            Self::Writer(writer) => writer.flush(),
65        }
66    }
67
68    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
69        match self {
70            Self::Buffer(buffer) => buffer.write_all(buf),
71            Self::Writer(writer) => writer.write_all(buf),
72        }
73    }
74}
75
76impl<W> DelayWriter<W> {
77    /// Create a new DelayWriter
78    pub fn new() -> Self {
79        Self(Arc::new(Mutex::new(DelayWriterInner::new())))
80    }
81}
82
83/// A writer that buffers data until it is assigned a file to write to.
84#[derive(Debug)]
85pub struct DelayWriter<W>(Arc<Mutex<DelayWriterInner<W>>>);
86
87impl<W> DelayWriter<W>
88where
89    W: Write,
90{
91    /// Try to init this DelayWriter
92    pub fn init(&self, writer: W) -> Result<(), std::io::Error> {
93        let mut lock = self.0.lock();
94        lock.init(writer)
95    }
96}
97
98impl<W> Default for DelayWriter<W> {
99    fn default() -> Self {
100        Self::new()
101    }
102}
103
104impl<W> Write for DelayWriter<W>
105where
106    W: Write,
107{
108    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
109        let mut lock = self.0.lock();
110
111        lock.write(buf)
112    }
113
114    fn flush(&mut self) -> std::io::Result<()> {
115        let mut lock = self.0.lock();
116
117        lock.flush()
118    }
119
120    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
121        let mut lock = self.0.lock();
122
123        lock.write_all(buf)
124    }
125}
126
127impl<W> Clone for DelayWriter<W> {
128    fn clone(&self) -> Self {
129        Self(self.0.clone())
130    }
131}