rotating_file/
lib.rs

1//! A thread-safe rotating file with customizable rotation behavior.
2//!
3//! ## Example
4//!
5//! ```
6//! use rotating_file::RotatingFile;
7//!
8//! let root_dir = "./target/tmp";
9//! let s = "The quick brown fox jumps over the lazy dog";
10//! let _ = std::fs::remove_dir_all(root_dir);
11//!
12//! // rotated by 1 kilobyte, compressed with gzip
13//! let rotating_file = RotatingFile::new(root_dir, Some(1), None, None, None, None, None);
14//! for _ in 0..24 {
15//!     rotating_file.writeln(s).unwrap();
16//! }
17//! rotating_file.close();
18//!
19//! assert_eq!(2, std::fs::read_dir(root_dir).unwrap().count());
20//! std::fs::remove_dir_all(root_dir).unwrap();
21//! ```
22use std::io::Write;
23use std::path::Path;
24use std::thread::JoinHandle;
25use std::time::{SystemTime, UNIX_EPOCH};
26use std::{ffi::OsString, fs, io::Error, sync::Mutex};
27use std::{io::BufWriter, sync::Arc};
28
29use chrono::{DateTime, NaiveDateTime, Utc};
30use flate2::write::GzEncoder;
31use log::*;
32
33#[derive(Copy, Clone)]
34pub enum Compression {
35    GZip,
36    Zip,
37}
38
39struct CurrentContext {
40    file: BufWriter<fs::File>,
41    file_path: OsString,
42    timestamp: u64,
43    total_written: usize,
44}
45
46/// A thread-safe rotating file with customizable rotation behavior.
47pub struct RotatingFile {
48    /// Root directory
49    root_dir: String,
50    /// Max size(in kilobytes) of the file after which it will rotate, 0 means unlimited
51    size: usize,
52    /// How often(in seconds) to rotate, 0 means unlimited
53    interval: u64,
54    /// Compression method, default to None
55    compression: Option<Compression>,
56
57    /// Format as used in chrono <https://docs.rs/chrono/latest/chrono/format/strftime/>, default to `%Y-%m-%d-%H-%M-%S`
58    date_format: String,
59    /// File name prefix, default to empty
60    prefix: String,
61    /// File name suffix, default to `.log`
62    suffix: String,
63
64    // current context
65    context: Mutex<CurrentContext>,
66    // compression threads
67    handles: Arc<Mutex<Vec<JoinHandle<Result<(), Error>>>>>,
68}
69
70unsafe impl Send for RotatingFile {}
71unsafe impl Sync for RotatingFile {}
72
73impl RotatingFile {
74    /// Creates a new RotatingFile.
75    ///
76    /// ## Arguments
77    ///
78    /// - `root_dir` The directory to store files.
79    /// - `size` Max size(in kilobytes) of the file after which it will rotate,
80    /// `None` and `0` mean unlimited.
81    /// - `interval` How often(in seconds) to rotate, 0 means unlimited.
82    /// - `compression` Available values are `GZip` and `Zip`, default to `None`
83    /// - `date_format` uses the syntax from chrono
84    /// <https://docs.rs/chrono/latest/chrono/format/strftime/>, default to `%Y-%m-%d-%H-%M-%S`
85    /// - `prefix` File name prefix, default to empty
86    /// - `suffix` File name suffix, default to `.log`
87    pub fn new(
88        root_dir: &str,
89        size: Option<usize>,
90        interval: Option<u64>,
91        compression: Option<Compression>,
92        date_format: Option<String>,
93        prefix: Option<String>,
94        suffix: Option<String>,
95    ) -> Self {
96        if let Err(e) = std::fs::create_dir_all(root_dir) {
97            error!("{}", e);
98        }
99
100        let interval = interval.unwrap_or(0);
101
102        let date_format = date_format.unwrap_or_else(|| "%Y-%m-%d-%H-%M-%S".to_string());
103        let prefix = prefix.unwrap_or_else(|| "".to_string());
104        let suffix = suffix.unwrap_or_else(|| ".log".to_string());
105
106        let context = Self::create_context(
107            interval,
108            root_dir,
109            date_format.as_str(),
110            prefix.as_str(),
111            suffix.as_str(),
112        );
113
114        RotatingFile {
115            root_dir: root_dir.to_string(),
116            size: size.unwrap_or(0),
117            interval,
118            compression,
119            date_format,
120            prefix,
121            suffix,
122            context: Mutex::new(context),
123            handles: Arc::new(Mutex::new(Vec::new())),
124        }
125    }
126
127    pub fn writeln(&self, s: &str) -> Result<(), Error> {
128        let mut guard = self.context.lock().unwrap();
129
130        let now = SystemTime::now()
131            .duration_since(UNIX_EPOCH)
132            .unwrap()
133            .as_secs();
134
135        if (self.size > 0 && guard.total_written + s.len() + 1 >= self.size * 1024)
136            || (self.interval > 0 && now >= (guard.timestamp + self.interval))
137        {
138            guard.file.flush()?;
139            guard.file.get_ref().sync_all()?;
140            let old_file = guard.file_path.clone();
141
142            // reset context
143            *guard = Self::create_context(
144                self.interval,
145                self.root_dir.as_str(),
146                self.date_format.as_str(),
147                self.prefix.as_str(),
148                self.suffix.as_str(),
149            );
150
151            // compress in a background thread
152            if let Some(c) = self.compression {
153                let handles_clone = self.handles.clone();
154                let handle = std::thread::spawn(move || Self::compress(old_file, c, handles_clone));
155                self.handles.lock().unwrap().push(handle);
156            }
157        }
158
159        if let Err(e) = writeln!(&mut guard.file, "{}", s) {
160            error!(
161                "Failed to write to file {}: {}",
162                guard.file_path.to_str().unwrap(),
163                e
164            );
165        } else {
166            guard.total_written += s.len() + 1;
167        }
168
169        Ok(())
170    }
171
172    pub fn close(&self) {
173        // wait for compression threads
174        let mut handles = self.handles.lock().unwrap();
175        for handle in handles.drain(..) {
176            if let Err(e) = handle.join().unwrap() {
177                error!("{}", e);
178            }
179        }
180        drop(handles);
181
182        let mut guard = self.context.lock().unwrap();
183        if let Err(e) = guard.file.flush() {
184            error!("{}", e);
185        }
186        if let Err(e) = guard.file.get_ref().sync_all() {
187            error!("{}", e);
188        }
189    }
190
191    fn create_context(
192        interval: u64,
193        root_dir: &str,
194        date_format: &str,
195        prefix: &str,
196        suffix: &str,
197    ) -> CurrentContext {
198        let now = SystemTime::now()
199            .duration_since(UNIX_EPOCH)
200            .unwrap()
201            .as_secs();
202        let timestamp = if interval > 0 {
203            now / interval * interval
204        } else {
205            now
206        };
207
208        let dt = DateTime::<Utc>::from_utc(NaiveDateTime::from_timestamp(timestamp as i64, 0), Utc);
209        let dt_str = dt.format(date_format).to_string();
210
211        let mut file_name = format!("{}{}{}", prefix, dt_str, suffix);
212        let mut index = 1;
213        while Path::new(root_dir).join(file_name.as_str()).exists()
214            || Path::new(root_dir).join(file_name.clone() + ".gz").exists()
215            || Path::new(root_dir)
216                .join(file_name.clone() + ".zip")
217                .exists()
218        {
219            file_name = format!("{}{}-{}{}", prefix, dt_str, index, suffix);
220            index += 1;
221        }
222
223        let file_path = Path::new(root_dir).join(file_name).into_os_string();
224
225        let file = fs::OpenOptions::new()
226            .append(true)
227            .create(true)
228            .open(file_path.as_os_str())
229            .unwrap();
230
231        CurrentContext {
232            file: BufWriter::new(file),
233            file_path,
234            timestamp,
235            total_written: 0,
236        }
237    }
238
239    fn compress(
240        file: OsString,
241        compress: Compression,
242        handles: Arc<Mutex<Vec<JoinHandle<Result<(), Error>>>>>,
243    ) -> Result<(), Error> {
244        let mut out_file_path = file.clone();
245        match compress {
246            Compression::GZip => out_file_path.push(".gz"),
247            Compression::Zip => out_file_path.push(".zip"),
248        }
249
250        let out_file = fs::OpenOptions::new()
251            .write(true)
252            .create(true)
253            .open(out_file_path.as_os_str())?;
254
255        let input_buf = fs::read(file.as_os_str())?;
256
257        match compress {
258            Compression::GZip => {
259                let mut encoder = GzEncoder::new(out_file, flate2::Compression::new(9));
260                encoder.write_all(&input_buf)?;
261                encoder.flush()?;
262            }
263            Compression::Zip => {
264                let file_name = Path::new(file.as_os_str())
265                    .file_name()
266                    .unwrap()
267                    .to_str()
268                    .unwrap();
269                let mut zip = zip::ZipWriter::new(out_file);
270                zip.start_file(file_name, zip::write::FileOptions::default())?;
271                zip.write_all(&input_buf)?;
272                zip.finish()?;
273            }
274        }
275
276        let ret = fs::remove_file(file.as_os_str());
277
278        // remove from the handles vector
279        if let Ok(ref mut guard) = handles.try_lock() {
280            let current_id = std::thread::current().id();
281            if let Some(pos) = guard.iter().position(|h| h.thread().id() == current_id) {
282                guard.remove(pos);
283            }
284        }
285
286        ret
287    }
288}
289
290#[cfg(test)]
291mod tests {
292    use chrono::{DateTime, Utc};
293    use once_cell::sync::Lazy;
294    use std::path::Path;
295    use std::time::Duration;
296    use std::time::SystemTime;
297
298    const TEXT: &'static str = "The quick brown fox jumps over the lazy dog";
299
300    #[test]
301    fn rotate_by_size() {
302        let root_dir = "./target/tmp1";
303        let _ = std::fs::remove_dir_all(root_dir);
304        let timestamp = current_timestamp_str();
305        let rotating_file =
306            super::RotatingFile::new(root_dir, Some(1), None, None, None, None, None);
307
308        for _ in 0..23 {
309            rotating_file.writeln(TEXT).unwrap();
310        }
311
312        rotating_file.close();
313
314        assert!(Path::new(root_dir)
315            .join(timestamp.clone() + ".log")
316            .exists());
317        assert!(!Path::new(root_dir)
318            .join(timestamp.clone() + "-1.log")
319            .exists());
320
321        std::fs::remove_dir_all(root_dir).unwrap();
322
323        let timestamp = current_timestamp_str();
324        let rotating_file =
325            super::RotatingFile::new(root_dir, Some(1), None, None, None, None, None);
326
327        for _ in 0..24 {
328            rotating_file.writeln(TEXT).unwrap();
329        }
330
331        rotating_file.close();
332
333        assert!(Path::new(root_dir)
334            .join(timestamp.clone() + ".log")
335            .exists());
336        assert!(Path::new(root_dir)
337            .join(timestamp.clone() + "-1.log")
338            .exists());
339        assert_eq!(
340            format!("{}\n", TEXT),
341            std::fs::read_to_string(Path::new(root_dir).join(timestamp + "-1.log")).unwrap()
342        );
343
344        std::fs::remove_dir_all(root_dir).unwrap();
345    }
346
347    #[test]
348    fn rotate_by_time() {
349        let root_dir = "./target/tmp2";
350        let _ = std::fs::remove_dir_all(root_dir);
351        let rotating_file =
352            super::RotatingFile::new(root_dir, None, Some(1), None, None, None, None);
353
354        let timestamp1 = current_timestamp_str();
355        rotating_file.writeln(TEXT).unwrap();
356
357        std::thread::sleep(Duration::from_secs(1));
358
359        let timestamp2 = current_timestamp_str();
360        rotating_file.writeln(TEXT).unwrap();
361
362        rotating_file.close();
363
364        assert!(Path::new(root_dir).join(timestamp1 + ".log").exists());
365        assert!(Path::new(root_dir).join(timestamp2 + ".log").exists());
366
367        std::fs::remove_dir_all(root_dir).unwrap();
368    }
369
370    #[test]
371    fn rotate_by_size_and_gzip() {
372        let root_dir = "./target/tmp3";
373        let _ = std::fs::remove_dir_all(root_dir);
374        let timestamp = current_timestamp_str();
375        let rotating_file = super::RotatingFile::new(
376            root_dir,
377            Some(1),
378            None,
379            Some(super::Compression::GZip),
380            None,
381            None,
382            None,
383        );
384
385        for _ in 0..24 {
386            rotating_file.writeln(TEXT).unwrap();
387        }
388
389        rotating_file.close();
390
391        assert!(Path::new(root_dir)
392            .join(timestamp.clone() + ".log.gz")
393            .exists());
394        assert!(Path::new(root_dir).join(timestamp + "-1.log").exists());
395
396        std::fs::remove_dir_all(root_dir).unwrap();
397    }
398
399    #[test]
400    fn rotate_by_size_and_zip() {
401        let root_dir = "./target/tmp4";
402        let _ = std::fs::remove_dir_all(root_dir);
403        let timestamp = current_timestamp_str();
404        let rotating_file = super::RotatingFile::new(
405            root_dir,
406            Some(1),
407            None,
408            Some(super::Compression::Zip),
409            None,
410            None,
411            None,
412        );
413
414        for _ in 0..24 {
415            rotating_file.writeln(TEXT).unwrap();
416        }
417
418        rotating_file.close();
419
420        assert!(Path::new(root_dir)
421            .join(timestamp.clone() + ".log.zip")
422            .exists());
423        assert!(Path::new(root_dir).join(timestamp + "-1.log").exists());
424
425        std::fs::remove_dir_all(root_dir).unwrap();
426    }
427
428    #[test]
429    fn rotate_by_time_and_gzip() {
430        let root_dir = "./target/tmp5";
431        let _ = std::fs::remove_dir_all(root_dir);
432        let rotating_file = super::RotatingFile::new(
433            root_dir,
434            None,
435            Some(1),
436            Some(super::Compression::GZip),
437            None,
438            None,
439            None,
440        );
441
442        let timestamp1 = current_timestamp_str();
443        rotating_file.writeln(TEXT).unwrap();
444
445        std::thread::sleep(Duration::from_secs(1));
446
447        let timestamp2 = current_timestamp_str();
448        rotating_file.writeln(TEXT).unwrap();
449
450        rotating_file.close();
451
452        assert!(Path::new(root_dir).join(timestamp1 + ".log.gz").exists());
453        assert!(Path::new(root_dir).join(timestamp2 + ".log").exists());
454
455        std::fs::remove_dir_all(root_dir).unwrap();
456    }
457
458    #[test]
459    fn rotate_by_time_and_zip() {
460        let root_dir = "./target/tmp6";
461        let _ = std::fs::remove_dir_all(root_dir);
462        let rotating_file = super::RotatingFile::new(
463            root_dir,
464            None,
465            Some(1),
466            Some(super::Compression::Zip),
467            None,
468            None,
469            None,
470        );
471
472        let timestamp1 = current_timestamp_str();
473        rotating_file.writeln(TEXT).unwrap();
474
475        std::thread::sleep(Duration::from_secs(1));
476
477        let timestamp2 = current_timestamp_str();
478        rotating_file.writeln(TEXT).unwrap();
479
480        rotating_file.close();
481
482        assert!(Path::new(root_dir).join(timestamp1 + ".log.zip").exists());
483        assert!(Path::new(root_dir).join(timestamp2 + ".log").exists());
484
485        std::fs::remove_dir_all(root_dir).unwrap();
486    }
487
488    #[test]
489    fn referred_in_two_threads() {
490        static ROOT_DIR: Lazy<&'static str> = Lazy::new(|| "./target/tmp7");
491        static ROTATING_FILE: Lazy<super::RotatingFile> = Lazy::new(|| {
492            super::RotatingFile::new(
493                *ROOT_DIR,
494                Some(1),
495                None,
496                Some(super::Compression::GZip),
497                None,
498                None,
499                None,
500            )
501        });
502        let _ = std::fs::remove_dir_all(*ROOT_DIR);
503
504        let timestamp = current_timestamp_str();
505        let handle1 = std::thread::spawn(move || {
506            for _ in 0..23 {
507                ROTATING_FILE.writeln(TEXT).unwrap();
508            }
509        });
510
511        let handle2 = std::thread::spawn(move || {
512            for _ in 0..23 {
513                ROTATING_FILE.writeln(TEXT).unwrap();
514            }
515        });
516
517        // trigger the third file creation
518        ROTATING_FILE.writeln(TEXT).unwrap();
519
520        let _ = handle1.join();
521        let _ = handle2.join();
522
523        ROTATING_FILE.close();
524
525        assert!(Path::new(*ROOT_DIR)
526            .join(timestamp.clone() + ".log.gz")
527            .exists());
528        assert!(Path::new(*ROOT_DIR)
529            .join(timestamp.clone() + "-1.log.gz")
530            .exists());
531
532        let third_file = Path::new(*ROOT_DIR).join(timestamp.clone() + "-2.log");
533        assert!(third_file.exists());
534        assert_eq!(
535            TEXT.len() + 1,
536            std::fs::metadata(third_file).unwrap().len() as usize
537        );
538
539        std::fs::remove_dir_all(*ROOT_DIR).unwrap();
540    }
541
542    fn current_timestamp_str() -> String {
543        let dt: DateTime<Utc> = SystemTime::now().into();
544        let dt_str = dt.format("%Y-%m-%d-%H-%M-%S").to_string();
545        dt_str
546    }
547}
rotating_file/lib.rs

rotating_file/
lib.rs
