Skip to content

Commit

Permalink
test: add unit tests for rolling file
Browse files Browse the repository at this point in the history
  • Loading branch information
1996fanrui committed Aug 10, 2024
1 parent 421de63 commit 7ec331f
Show file tree
Hide file tree
Showing 5 changed files with 252 additions and 13 deletions.
4 changes: 3 additions & 1 deletion .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,9 @@ jobs:
- uses: actions/checkout@v4
- uses: Swatinem/rust-cache@v2
- name: Run unit tests
run: cargo test -- --nocapture
run: |
cargo test -- --nocapture
cargo test --features rolling_file -- --nocapture
- name: Run examples
run: |
cargo run --example simple_stdio
Expand Down
4 changes: 4 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,10 @@ time = { version = "0.3", features = [
"macros",
], optional = true }

[dev-dependencies]
rand = "0.8.5"
tempfile = "3.3"

## Fastrace dependencies
[dependencies.fastrace]
optional = true
Expand Down
96 changes: 96 additions & 0 deletions src/append/rolling_file/clock.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
// Copyright 2024 tison <[email protected]>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt;
use time::OffsetDateTime;

/// A clock providing access to the current time.
pub trait Clock: fmt::Debug + Send {
fn now(&self) -> OffsetDateTime;
}

#[derive(Debug)]
pub struct DefaultClock;

impl Clock for DefaultClock {
fn now(&self) -> OffsetDateTime {
OffsetDateTime::now_utc()
}
}

/// The time could be reset.
#[derive(Debug)]
#[cfg(test)]
pub struct ManualClock {
fixed_time: OffsetDateTime,
}

#[cfg(test)]
impl Clock for ManualClock {
fn now(&self) -> OffsetDateTime {
self.fixed_time
}
}

#[cfg(test)]
impl ManualClock {
pub fn new(fixed_time: OffsetDateTime) -> ManualClock {
ManualClock { fixed_time }
}

pub fn set_now(&mut self, new_time: OffsetDateTime) {
self.fixed_time = new_time;
}
}

#[derive(Debug)]
pub enum StateClock {
DefaultClock(DefaultClock),
#[cfg(test)]
ManualClock(ManualClock),
}

impl StateClock {
pub fn now(&self) -> OffsetDateTime {
match self {
StateClock::DefaultClock(clock) => clock.now(),
#[cfg(test)]
StateClock::ManualClock(clock) => clock.now(),
}
}

#[cfg(test)]
pub fn set_now(&mut self, new_time: OffsetDateTime) {
if let StateClock::ManualClock(clock) = self {
clock.set_now(new_time);
}
}
}

#[cfg(test)]
mod tests {
use super::*;
use time::macros::datetime;

#[test]
fn test_manual_clock_adjusting() {
let mut clock = ManualClock {
fixed_time: datetime!(2023-01-01 12:00:00 UTC),
};
assert_eq!(clock.now(), datetime!(2023-01-01 12:00:00 UTC));

clock.set_now(datetime!(2024-01-01 12:00:00 UTC));
assert_eq!(clock.now(), datetime!(2024-01-01 12:00:00 UTC));
}
}
1 change: 1 addition & 0 deletions src/append/rolling_file/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ pub use rolling::RollingFileWriterBuilder;
pub use rotation::TimeRotation;

mod append;
mod clock;
mod non_blocking;
mod rolling;
mod rotation;
Expand Down
160 changes: 148 additions & 12 deletions src/append/rolling_file/rolling.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ use std::io::Write;
use std::path::Path;
use std::path::PathBuf;

use crate::append::rolling_file::clock::{DefaultClock, StateClock};
use crate::append::rolling_file::TimeRotation;
use anyhow::Context;
use parking_lot::RwLock;
Expand All @@ -39,15 +40,11 @@ impl RollingFileWriter {
pub fn builder() -> RollingFileWriterBuilder {
RollingFileWriterBuilder::new()
}

fn now(&self) -> OffsetDateTime {
OffsetDateTime::now_utc()
}
}

impl Write for RollingFileWriter {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let now = self.now();
let now = self.state.clock.now();
let writer = self.writer.get_mut();
if self.state.should_rollover_on_date(now) {
self.state.advance_date(now);
Expand Down Expand Up @@ -77,6 +74,7 @@ pub struct RollingFileWriterBuilder {
suffix: Option<String>,
max_size: usize,
max_files: Option<usize>,
clock: Option<StateClock>,
}

impl Default for RollingFileWriterBuilder {
Expand All @@ -94,6 +92,7 @@ impl RollingFileWriterBuilder {
suffix: None,
max_size: usize::MAX,
max_files: None,
clock: None,
}
}

Expand Down Expand Up @@ -138,18 +137,30 @@ impl RollingFileWriterBuilder {
self
}

#[cfg(test)]
fn clock(mut self, clock: StateClock) -> Self {
self.clock = Some(clock);
self
}

pub fn build(self, dir: impl AsRef<Path>) -> anyhow::Result<RollingFileWriter> {
let Self {
rotation,
prefix,
suffix,
max_size,
max_files,
clock,
} = self;
let directory = dir.as_ref().to_path_buf();
let now = OffsetDateTime::now_utc();
let (state, writer) = State::new(
now, rotation, directory, prefix, suffix, max_size, max_files,
rotation,
directory,
prefix,
suffix,
max_size,
max_files,
clock.unwrap_or(StateClock::DefaultClock(DefaultClock)),
)?;
Ok(RollingFileWriter { state, writer })
}
Expand All @@ -162,29 +173,29 @@ struct State {
log_filename_suffix: Option<String>,
date_format: Vec<format_description::FormatItem<'static>>,
rotation: TimeRotation,
current_date: OffsetDateTime,
current_count: usize,
current_filesize: usize,
next_date_timestamp: Option<usize>,
max_size: usize,
max_files: Option<usize>,
clock: StateClock,
}

impl State {
fn new(
now: OffsetDateTime,
rotation: TimeRotation,
dir: impl AsRef<Path>,
log_filename_prefix: Option<String>,
log_filename_suffix: Option<String>,
max_size: usize,
max_files: Option<usize>,
clock: StateClock,
) -> anyhow::Result<(Self, RwLock<File>)> {
let log_dir = dir.as_ref().to_path_buf();
let date_format = rotation.date_format();
let now = clock.now();
let next_date_timestamp = rotation.next_date_timestamp(&now);

let current_date = now;
let current_count = 0;
let current_filesize = 0;

Expand All @@ -193,13 +204,13 @@ impl State {
log_filename_prefix,
log_filename_suffix,
date_format,
current_date,
current_count,
current_filesize,
next_date_timestamp,
rotation,
max_size,
max_files,
clock,
};

let file = state.create_log_writer(now, 0)?;
Expand Down Expand Up @@ -331,9 +342,134 @@ impl State {
}

fn advance_date(&mut self, now: OffsetDateTime) {
self.current_date = now;
self.current_count = 0;
self.current_filesize = 0;
self.next_date_timestamp = self.rotation.next_date_timestamp(&now);
}
}
#[cfg(test)]
mod tests {
use crate::append::rolling_file::clock::{ManualClock, StateClock};
use crate::append::rolling_file::{RollingFileWriterBuilder, TimeRotation};
use rand::{distributions::Alphanumeric, Rng};
use std::cmp::min;
use std::fs;
use std::io::Write;
use std::ops::Add;
use tempfile::TempDir;
use time::macros::datetime;
use time::Duration;

#[test]
fn test_file_rolling_via_file_size() {
test_file_rolling_for_specific_file_size(3, 1000);
test_file_rolling_for_specific_file_size(3, 10000);
test_file_rolling_for_specific_file_size(10, 8888);
test_file_rolling_for_specific_file_size(10, 10000);
test_file_rolling_for_specific_file_size(20, 6666);
test_file_rolling_for_specific_file_size(20, 10000);
}
fn test_file_rolling_for_specific_file_size(max_files: usize, max_size: usize) {
let temp_dir = TempDir::new().expect("failed to create a temporary directory");

let mut writer = RollingFileWriterBuilder::new()
.rotation(TimeRotation::Never)
.filename_prefix("test_prefix")
.filename_suffix("log")
.max_log_files(max_files)
.max_file_size(max_size)
.build(&temp_dir)
.unwrap();

for i in 1..=(max_files * 2) {
let mut expected_file_size = 0;
while expected_file_size < max_size {
let rand_str = generate_random_string();
expected_file_size += rand_str.len();
assert_eq!(writer.write(rand_str.as_bytes()).unwrap(), rand_str.len());
assert_eq!(writer.state.current_filesize, expected_file_size);
}

writer.flush().unwrap();
assert_eq!(
fs::read_dir(&writer.state.log_dir).unwrap().count(),
min(i, max_files)
);
}
}

#[test]
fn test_file_rolling_via_time_rotation() {
test_file_rolling_for_specific_time_rotation(
TimeRotation::Minutely,
Duration::minutes(1),
Duration::seconds(1),
);
test_file_rolling_for_specific_time_rotation(
TimeRotation::Hourly,
Duration::hours(1),
Duration::minutes(1),
);
test_file_rolling_for_specific_time_rotation(
TimeRotation::Daily,
Duration::days(1),
Duration::hours(1),
);
}

fn test_file_rolling_for_specific_time_rotation(
rotation: TimeRotation,
rotation_duration: Duration,
write_interval: Duration,
) {
let temp_dir = TempDir::new().expect("failed to create a temporary directory");
let max_files = 10;

let start_time = datetime!(2024-08-10 00:00:00 +0);
let mut writer = RollingFileWriterBuilder::new()
.rotation(rotation)
.filename_prefix("test_prefix")
.filename_suffix("log")
.max_log_files(max_files)
.max_file_size(usize::MAX)
.clock(StateClock::ManualClock(ManualClock::new(start_time)))
.build(&temp_dir)
.unwrap();

let mut cur_time = start_time;

for i in 1..=(max_files * 2) {
let mut expected_file_size = 0;
let end_time = cur_time.add(rotation_duration);
while cur_time < end_time {
writer.state.clock.set_now(cur_time);

let rand_str = generate_random_string();
expected_file_size += rand_str.len();

assert_eq!(writer.write(rand_str.as_bytes()).unwrap(), rand_str.len());
assert_eq!(writer.state.current_filesize, expected_file_size);

cur_time = cur_time.add(write_interval);
}

writer.flush().unwrap();
assert_eq!(
fs::read_dir(&writer.state.log_dir).unwrap().count(),
min(i, max_files)
);
}
}

fn generate_random_string() -> String {
let mut rng = rand::thread_rng();
let len = rng.gen_range(1..=100);
let random_string: String = std::iter::repeat(())
.map(|()| rng.sample(Alphanumeric))
.map(char::from)
.take(len)
.collect();

random_string
}
}

0 comments on commit 7ec331f

Please sign in to comment.