use std::num::NonZeroU64;
use std::os::unix::fs::PermissionsExt;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tokio::fs::{File, OpenOptions};
use tokio::io::AsyncWriteExt;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use super::Record;
use super::counters::SinkCounters;
const DRAIN_DEADLINE: Duration = Duration::from_secs(5);
pub(super) async fn run(
path: PathBuf,
max_bytes: NonZeroU64,
mut rx: mpsc::Receiver<Record>,
drain: CancellationToken,
counters: Arc<SinkCounters>,
) {
let mut writer = Writer::new(path, max_bytes);
loop {
tokio::select! {
received = rx.recv() => match received {
Some(record) => writer.append(record, &counters).await,
None => break,
},
() = drain.cancelled() => {
self::drain(&mut writer, &mut rx, &counters, DRAIN_DEADLINE).await;
break;
}
}
}
}
async fn drain(
writer: &mut Writer,
rx: &mut mpsc::Receiver<Record>,
counters: &SinkCounters,
deadline: Duration,
) {
let drained = tokio::time::timeout(deadline, async {
while let Ok(record) = rx.try_recv() {
writer.append(record, counters).await;
}
})
.await;
if drained.is_err() {
counters.lose(rx.len() as u64 + 1);
}
}
struct Writer {
path: PathBuf,
max_bytes: u64,
file: Option<File>,
written: u64,
warned: bool,
}
impl Writer {
fn new(path: PathBuf, max_bytes: NonZeroU64) -> Writer {
Writer {
path,
max_bytes: max_bytes.get(),
file: None,
written: 0,
warned: false,
}
}
async fn append(&mut self, record: Record, counters: &SinkCounters) {
let Ok(mut line) = serde_json::to_string(&record) else {
counters.lose(1);
return;
};
line.push('\n');
let cost = line.len() as u64;
if !self.open(counters).await {
return;
}
if self.written.saturating_add(cost) >= self.max_bytes {
self.roll_over().await;
if !self.open(counters).await {
return;
}
}
let file = self.file.as_mut().expect("the file was just opened");
let written = async {
file.write_all(line.as_bytes()).await?;
file.flush().await
}
.await;
match written {
Ok(()) => self.written = self.written.saturating_add(cost),
Err(err) => {
tracing::warn!(
path = %self.path.display(),
error = %err,
"a decision record could not be written to the log file"
);
counters.lose(1);
self.file = None;
}
}
}
async fn open(&mut self, counters: &SinkCounters) -> bool {
if self.file.is_some() {
return true;
}
match OpenOptions::new()
.create(true)
.append(true)
.mode(0o600)
.open(&self.path)
.await
{
Ok(file) => {
self.written = match file.metadata().await {
Ok(metadata) => {
self.warn_if_shared(metadata.permissions().mode());
metadata.len()
}
Err(_) => 0,
};
self.file = Some(file);
true
}
Err(err) => {
tracing::warn!(
path = %self.path.display(),
error = %err,
"the decision log file could not be opened"
);
counters.lose(1);
false
}
}
}
fn warn_if_shared(&mut self, mode: u32) {
if self.warned || mode & 0o077 == 0 {
return;
}
self.warned = true;
tracing::warn!(
path = %self.path.display(),
mode = %format_args!("{:o}", mode & 0o777),
"the decision log file is accessible to users other than its owner"
);
}
async fn roll_over(&mut self) {
self.file = None;
let mut rolled = self.path.clone().into_os_string();
rolled.push(".1");
if let Err(err) = tokio::fs::rename(&self.path, PathBuf::from(rolled)).await {
tracing::warn!(
path = %self.path.display(),
error = %err,
"the decision log file could not be rolled over"
);
}
self.written = 0;
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::delivery::Summary;
#[tokio::test]
async fn tp19_file_drain_deadline_counts_cut_off_records() {
let dir = tempfile::tempdir().expect("a temporary directory");
let path = dir.path().join("decisions.ndjson");
const QUEUED: u64 = 10_000;
let (tx, mut rx) = mpsc::channel(QUEUED as usize);
for _ in 0..QUEUED {
tx.try_send(Record::RequestSummary(Summary {
timestamp: String::new(),
requests: 0,
errors: 0,
bytes: 0,
mean_duration_micros: 0,
window_micros: 0,
dropped_file: 0,
dropped_siem: 0,
}))
.expect("the channel has room");
}
let mut writer = Writer::new(path.clone(), NonZeroU64::MAX);
let counters = SinkCounters::new();
drain(&mut writer, &mut rx, &counters, Duration::ZERO).await;
let written = std::fs::read_to_string(&path)
.map(|text| text.lines().count() as u64)
.unwrap_or(0);
assert!(
written < QUEUED,
"inconclusive: no cut-off happened, all {written} records were written"
);
let dropped = counters.total();
assert!(
(QUEUED..=QUEUED + 1).contains(&(written + dropped)),
"every queued record is either in the file or counted as dropped: \
{written} written, {dropped} dropped"
);
}
fn summary() -> Record {
Record::RequestSummary(Summary {
timestamp: String::new(),
requests: 0,
errors: 0,
bytes: 0,
mean_duration_micros: 0,
window_micros: 0,
dropped_file: 0,
dropped_siem: 0,
})
}
#[tokio::test]
async fn rl7_file_drain_deadline_counts_in_the_window() {
let dir = tempfile::tempdir().expect("a temporary directory");
let (tx, mut rx) = mpsc::channel(10_000);
for _ in 0..10_000 {
tx.try_send(summary()).expect("the channel has room");
}
let mut writer = Writer::new(dir.path().join("d.ndjson"), NonZeroU64::MAX);
let counters = SinkCounters::new();
drain(&mut writer, &mut rx, &counters, Duration::ZERO).await;
assert!(counters.take_window() > 0);
}
#[tokio::test]
async fn rl7_file_write_error_counts_in_the_window() {
let mut writer = Writer::new(PathBuf::from("/dev/full"), NonZeroU64::MAX);
let counters = SinkCounters::new();
writer.append(summary(), &counters).await;
assert_eq!(counters.take_window(), 1);
}
#[tokio::test]
async fn rl7_file_unopenable_counts_in_the_window() {
let dir = tempfile::tempdir().expect("a temporary directory");
let mut writer = Writer::new(dir.path().join("missing/d.ndjson"), NonZeroU64::MAX);
let counters = SinkCounters::new();
writer.append(summary(), &counters).await;
assert_eq!(counters.take_window(), 1);
}
}