tracing_rolling/
lib.rs

1use std::{
2    fs::File,
3    io::{self, BufWriter, Write},
4    path::{Path, PathBuf},
5    sync::Arc,
6};
7
8use parking_lot::{Mutex, MutexGuard};
9use time::{
10    format_description::{parse_owned, Component, OwnedFormatItem},
11    Date, Duration, OffsetDateTime, Time, UtcOffset,
12};
13use tracing_subscriber::fmt::MakeWriter;
14
15pub trait Checker: Sized {
16    type W: Write;
17    fn should_update(&self) -> bool;
18    fn new_writer(&self) -> io::Result<Self::W>;
19    /// create a buffered writer with default size: 4K
20    fn buffered(self) -> Buffered<Self, Self::W> {
21        Buffered {
22            checker: self,
23            size: 4096,
24        }
25    }
26    /// create a buffered writer with specified buffer size
27    fn buffer_with(self, size: usize) -> Buffered<Self, Self::W> {
28        Buffered {
29            checker: self,
30            size,
31        }
32    }
33
34    fn build(self) -> io::Result<(Rolling<Self, Self::W>, Token<Self::W>)> {
35        let fd = Arc::new(Mutex::new(self.new_writer()?));
36        let t = Token(fd.clone());
37        let r = Rolling::new(self, fd);
38        Ok((r, t))
39    }
40}
41
42#[must_use = "must manual drop to ensure flush file when process exits"]
43pub struct Token<W: Write>(Arc<Mutex<W>>);
44
45impl<W: Write> Drop for Token<W> {
46    fn drop(&mut self) {
47        if let Err(e) = self.0.lock().flush() {
48            eprintln!("drop writer {e}");
49        }
50    }
51}
52
53pub struct Rolling<C: Checker<W = W>, W: Write> {
54    writer: Arc<Mutex<W>>,
55    checker: C,
56}
57
58impl<C: Checker<W = W>, W: Write> Rolling<C, W> {
59    pub fn new(checker: C, writer: Arc<Mutex<W>>) -> Self {
60        Self { writer, checker }
61    }
62
63    fn update_writer(&self) -> io::Result<()> {
64        {
65            let mut writer = self.writer.lock();
66            writer.flush()?;
67        }
68        let writer = self.checker.new_writer()?;
69        *self.writer.lock() = writer;
70        Ok(())
71    }
72}
73
74impl<'a, C: Checker<W = W>, W: Write + 'a> MakeWriter<'a> for Rolling<C, W> {
75    type Writer = RollingWriter<'a, W>;
76
77    fn make_writer(&'a self) -> Self::Writer {
78        if self.checker.should_update() {
79            if let Err(e) = self.update_writer() {
80                eprintln!("can not update log file {e}")
81            }
82        }
83        RollingWriter(self.writer.lock())
84    }
85}
86
87pub struct RollingWriter<'a, W: Write>(MutexGuard<'a, W>);
88
89impl<'a, W: Write> Write for RollingWriter<'a, W> {
90    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
91        self.0.write(buf)
92    }
93
94    fn flush(&mut self) -> std::io::Result<()> {
95        self.0.flush()
96    }
97}
98
99pub trait Period {
100    fn previous_dt(&self) -> Result<OffsetDateTime, String>;
101    fn now(&self) -> OffsetDateTime;
102    fn new_path(&self) -> String;
103    fn duration(&self) -> &Duration;
104}
105
106impl<P: Period> Checker for P {
107    type W = File;
108    fn should_update(&self) -> bool {
109        let file_dt = match self.previous_dt() {
110            Ok(v) => v,
111            Err(e) => {
112                eprintln!("parse previous file failed: {e}");
113                return false;
114            }
115        };
116        self.now() - file_dt >= *self.duration()
117    }
118
119    fn new_writer(&self) -> io::Result<File> {
120        let path = self.new_path();
121        let file = File::options().append(true).create(true).open(path)?;
122        Ok(file)
123    }
124}
125
126pub struct Minute {
127    offset: UtcOffset,
128    fmt: OwnedFormatItem,
129    active: Mutex<String>,
130}
131
132impl Minute {
133    pub const DURATION: Duration = Duration::MINUTE;
134
135    pub fn new(path: impl AsRef<Path>, offset: impl Into<Option<UtcOffset>>) -> Self {
136        let ext = path
137            .as_ref()
138            .extension()
139            .and_then(|ext| ext.to_str())
140            .unwrap_or_default();
141        let fmt = path
142            .as_ref()
143            .with_extension(format!("[year]-[month]-[day]-[hour]-[minute].{ext}"));
144        let fmt = parse_owned::<1>(&format!("{}", fmt.display())).unwrap();
145        Self {
146            offset: offset.into().unwrap_or(UtcOffset::UTC),
147            fmt,
148            active: Default::default(),
149        }
150    }
151}
152
153impl Period for Minute {
154    fn previous_dt(&self) -> Result<OffsetDateTime, String> {
155        let file = self.active.lock();
156        let date = Date::parse(&file, &self.fmt).map_err(|e| e.to_string())?;
157        let time = Time::parse(&file, &self.fmt).map_err(|e| e.to_string())?;
158        Ok(date.with_time(time).assume_offset(self.offset))
159    }
160
161    fn now(&self) -> OffsetDateTime {
162        OffsetDateTime::now_utc().to_offset(self.offset)
163    }
164
165    fn new_path(&self) -> String {
166        let now = self.now();
167        let file = now.format(&self.fmt).unwrap();
168        *self.active.lock() = file.clone();
169        file
170    }
171
172    fn duration(&self) -> &Duration {
173        &Self::DURATION
174    }
175}
176
177pub struct Hourly {
178    offset: UtcOffset,
179    fmt: OwnedFormatItem,
180    hour_regex: regex::Regex,
181    active: Mutex<String>,
182}
183
184impl Hourly {
185    pub const DURATION: Duration = Duration::HOUR;
186
187    pub fn new(path: impl AsRef<Path>, offset: impl Into<Option<UtcOffset>>) -> Self {
188        let ext = path
189            .as_ref()
190            .extension()
191            .and_then(|ext| ext.to_str())
192            .unwrap_or_default();
193        let fmt = path
194            .as_ref()
195            .with_extension(format!("[year]-[month]-[day]-[hour].{ext}"));
196        let hour_regex =
197            regex::Regex::new(&format!(r".*\d{{4}}-\d{{2}}-\d{{2}}-(\d{{2}})\.{ext}")).unwrap();
198        let fmt = parse_owned::<1>(&format!("{}", fmt.display())).unwrap();
199        Self {
200            offset: offset.into().unwrap_or(UtcOffset::UTC),
201            fmt,
202            active: Default::default(),
203            hour_regex,
204        }
205    }
206}
207
208impl Period for Hourly {
209    fn previous_dt(&self) -> Result<OffsetDateTime, String> {
210        let file = self.active.lock();
211        let date = Date::parse(&file, &self.fmt).map_err(|e| e.to_string())?;
212        let hour = self
213            .hour_regex
214            .captures(&file)
215            .and_then(|cap| cap.get(1))
216            .and_then(|m| m.as_str().parse::<u8>().ok())
217            .ok_or_else(|| format!("invalid hour component of {file}"))?;
218        let time = Time::from_hms(hour, 0, 0).unwrap();
219        Ok(date.with_time(time).assume_offset(self.offset))
220    }
221
222    fn now(&self) -> OffsetDateTime {
223        OffsetDateTime::now_utc().to_offset(self.offset)
224    }
225
226    fn new_path(&self) -> String {
227        let now = self.now();
228        let file = now.format(&self.fmt).unwrap();
229        *self.active.lock() = file.clone();
230        file
231    }
232
233    fn duration(&self) -> &Duration {
234        &Self::DURATION
235    }
236}
237
238pub struct Daily {
239    offset: UtcOffset,
240    fmt: OwnedFormatItem,
241    active: Mutex<String>,
242}
243
244impl Daily {
245    pub const DURATION: Duration = Duration::DAY;
246
247    fn ensure_year_month_day(fmt: &OwnedFormatItem) {
248        match fmt {
249            OwnedFormatItem::Compound(items) => {
250                let mut year = false;
251                let mut month = false;
252                let mut day = false;
253                for item in &items[..] {
254                    match item {
255                        OwnedFormatItem::Component(Component::Year(_)) => {
256                            year = !year;
257                        }
258                        OwnedFormatItem::Component(Component::Month(_)) => {
259                            month = !month;
260                        }
261                        OwnedFormatItem::Component(Component::Day(_)) => {
262                            day = !day;
263                        }
264                        _ => {}
265                    }
266                }
267                if !(year && month && day) {
268                    panic!("invalid daily format");
269                }
270            }
271            _ => panic!("expect compound format"),
272        }
273    }
274
275    /// **NOTE: if fmt is specified, it should be valid time format_description and contain
276    /// year, month, day**
277    ///
278    /// default fmt is `[year]-[month]-[day]`
279    pub fn new<S>(
280        path: impl AsRef<Path>,
281        fmt: impl Into<Option<S>>,
282        offset: impl Into<Option<UtcOffset>>,
283    ) -> Self
284    where
285        S: std::fmt::Display,
286    {
287        let ext = path
288            .as_ref()
289            .extension()
290            .and_then(|ext| ext.to_str())
291            .unwrap_or_default();
292
293        let file_stem = path
294            .as_ref()
295            .file_stem()
296            .and_then(|stem| stem.to_str())
297            .unwrap_or_default();
298
299        let file_name = fmt
300            .into()
301            .map(|f| format!("{file_stem}-{f}.{ext}"))
302            .unwrap_or_else(|| format!("{file_stem}-[year]-[month]-[day].{ext}"));
303        let fmt = parse_owned::<1>(&format!(
304            "{}",
305            path.as_ref().with_file_name(file_name).display()
306        ))
307        .unwrap();
308        Self::ensure_year_month_day(&fmt);
309        Self {
310            offset: offset.into().unwrap_or(UtcOffset::UTC),
311            fmt,
312            active: Default::default(),
313        }
314    }
315}
316
317impl Period for Daily {
318    fn previous_dt(&self) -> Result<OffsetDateTime, String> {
319        let file = self.active.lock();
320        let date = Date::parse(&file, &self.fmt).map_err(|e| e.to_string())?;
321        Ok(date
322            .with_time(time::macros::time!(0:0:0))
323            .assume_offset(self.offset))
324    }
325
326    fn now(&self) -> OffsetDateTime {
327        OffsetDateTime::now_utc().to_offset(self.offset)
328    }
329
330    fn new_path(&self) -> String {
331        let now = self.now();
332        let file = now.format(&self.fmt).unwrap();
333        *self.active.lock() = file.clone();
334        file
335    }
336
337    fn duration(&self) -> &Duration {
338        &Self::DURATION
339    }
340}
341
342pub struct Buffered<C: Checker<W = W>, W: Write> {
343    checker: C,
344    size: usize,
345}
346
347impl<C: Checker<W = W>, W: Write> Buffered<C, W> {
348    pub fn new(checker: C, size: usize) -> Self {
349        Self { checker, size }
350    }
351}
352
353impl<C: Checker<W = W>, W: Write> Checker for Buffered<C, W> {
354    type W = BufWriter<W>;
355    fn should_update(&self) -> bool {
356        self.checker.should_update()
357    }
358    fn new_writer(&self) -> io::Result<BufWriter<W>> {
359        Ok(BufWriter::with_capacity(
360            self.size,
361            self.checker.new_writer()?,
362        ))
363    }
364}
365
366/// construct a non rolling file
367pub struct ConstFile(PathBuf);
368
369impl Checker for ConstFile {
370    type W = File;
371
372    fn should_update(&self) -> bool {
373        false
374    }
375
376    fn new_writer(&self) -> io::Result<Self::W> {
377        File::options().append(true).create(true).open(&self.0)
378    }
379}
380
381impl ConstFile {
382    pub fn new(path: impl AsRef<Path>) -> Self {
383        Self(path.as_ref().to_path_buf())
384    }
385}
tracing_rolling/lib.rs

tracing_rolling/
lib.rs
