use std::fs::{self, File};
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use std::time::Instant;
use anyhow::{Context, Result};
use memmap2::MmapOptions;
use crate::alloc_counter;
use crate::generated::{v1_generated, v2_generated};
#[derive(Debug, Clone)]
pub struct ModuleCWriteConfig {
pub records: u64,
pub wal_path: PathBuf,
}
#[derive(Debug, Clone)]
pub struct ModuleCReadConfig {
pub wal_path: PathBuf,
pub expect_records: Option<u64>,
}
#[derive(Debug, Clone)]
pub struct ModuleCWriteStats {
pub records_written: u64,
pub elapsed_ms: f64,
pub file_size_bytes: u64,
pub wal_path: PathBuf,
}
impl ModuleCWriteStats {
pub fn to_json(&self) -> String {
format!(
"{{\"module\":\"C_WRITE\",\"records_written\":{},\"elapsed_ms\":{:.3},\"file_size_bytes\":{},\"wal_path\":\"{}\"}}",
self.records_written,
self.elapsed_ms,
self.file_size_bytes,
self.wal_path.display()
)
}
}
#[derive(Debug, Clone)]
pub struct ModuleCReadStats {
pub records_read: u64,
pub elapsed_ms: f64,
pub defaults_ok: bool,
pub zero_copy_ok: bool,
pub read_loop_allocations: u64,
pub wal_path: PathBuf,
}
impl ModuleCReadStats {
pub fn to_json(&self) -> String {
format!(
"{{\"module\":\"C_READ\",\"records_read\":{},\"elapsed_ms\":{:.3},\"defaults_ok\":{},\"zero_copy_ok\":{},\"read_loop_allocations\":{},\"wal_path\":\"{}\"}}",
self.records_read,
self.elapsed_ms,
self.defaults_ok,
self.zero_copy_ok,
self.read_loop_allocations,
self.wal_path.display()
)
}
}
pub fn run_writer(config: ModuleCWriteConfig) -> Result<ModuleCWriteStats> {
if config.records == 0 {
anyhow::bail!("records must be > 0");
}
if let Some(parent) = config.wal_path.parent() {
if !parent.as_os_str().is_empty() {
fs::create_dir_all(parent).with_context(|| {
format!("failed to create parent directory {}", parent.display())
})?;
}
}
let file = File::create(&config.wal_path)
.with_context(|| format!("failed to create WAL at {}", config.wal_path.display()))?;
let mut writer = BufWriter::with_capacity(8 * 1024 * 1024, file);
let mut builder = flatbuffers::FlatBufferBuilder::with_capacity(1024);
let bucket_names: Vec<String> = (0..1024).map(|idx| format!("bucket-{idx:04}")).collect();
let start = Instant::now();
for i in 0..config.records {
builder.reset();
let bucket_name = &bucket_names[(i as usize) % bucket_names.len()];
let bucket_name_offset = builder.create_string(bucket_name);
let args = v1_generated::RecordV1Args {
object_id: i + 1,
bucket_name: Some(bucket_name_offset),
offset: i * 4096,
};
let record = v1_generated::RecordV1::create(&mut builder, &args);
v1_generated::finish_size_prefixed_record_v1_buffer(&mut builder, record);
writer
.write_all(builder.finished_data())
.with_context(|| format!("failed writing record {}", i))?;
}
writer.flush().context("failed flushing WAL writer")?;
let elapsed_ms = start.elapsed().as_secs_f64() * 1_000.0;
let file_size_bytes = fs::metadata(&config.wal_path)
.with_context(|| format!("failed to stat WAL {}", config.wal_path.display()))?
.len();
Ok(ModuleCWriteStats {
records_written: config.records,
elapsed_ms,
file_size_bytes,
wal_path: config.wal_path,
})
}
pub fn run_reader(config: ModuleCReadConfig) -> Result<ModuleCReadStats> {
let file = File::open(&config.wal_path)
.with_context(|| format!("failed to open WAL {}", config.wal_path.display()))?;
let mapped = unsafe { MmapOptions::new().map(&file) }.context("failed to mmap WAL file")?;
if mapped.is_empty() {
anyhow::bail!("WAL is empty")
}
let start = Instant::now();
let alloc_before = alloc_counter::snapshot();
let mut cursor = 0usize;
let mut records_read = 0_u64;
let mut defaults_ok = true;
let mut zero_copy_ok = true;
let base_addr = mapped.as_ptr() as usize;
let end_addr = base_addr + mapped.len();
while cursor < mapped.len() {
if cursor + 4 > mapped.len() {
anyhow::bail!("truncated record size prefix at byte offset {}", cursor);
}
let size = u32::from_le_bytes(
mapped[cursor..cursor + 4]
.try_into()
.expect("prefix slice is always 4 bytes"),
) as usize;
let total_size = 4 + size;
if cursor + total_size > mapped.len() {
anyhow::bail!("record overflows WAL bounds at byte offset {}", cursor);
}
let record_buf = &mapped[cursor..cursor + total_size];
let record = flatbuffers::size_prefixed_root::<v2_generated::RecordV2<'_>>(record_buf)
.map_err(|err| {
anyhow::anyhow!(
"failed to parse size-prefixed record {} at offset {}: {}",
records_read,
cursor,
err
)
})?;
if record.retain_until_date() != 0 {
defaults_ok = false;
}
let bucket_name = record
.bucket_name()
.ok_or_else(|| anyhow::anyhow!("record {} missing bucket_name", records_read))?;
let ptr = bucket_name.as_ptr() as usize;
let ptr_end = ptr + bucket_name.len();
if ptr < base_addr || ptr_end > end_addr {
zero_copy_ok = false;
}
records_read += 1;
cursor += total_size;
}
if let Some(expect_records) = config.expect_records {
if expect_records != records_read {
anyhow::bail!(
"record count mismatch: expected {}, got {}",
expect_records,
records_read
);
}
}
let elapsed_ms = start.elapsed().as_secs_f64() * 1_000.0;
let read_loop_allocations = alloc_counter::allocations_since(alloc_before);
Ok(ModuleCReadStats {
records_read,
elapsed_ms,
defaults_ok,
zero_copy_ok,
read_loop_allocations,
wal_path: config.wal_path,
})
}
#[cfg(test)]
mod tests {
use std::path::PathBuf;
use std::time::{SystemTime, UNIX_EPOCH};
use super::*;
#[test]
fn round_trip_v1_writer_v2_reader() {
let wal_path = unique_wal_path("module-c-roundtrip");
let write_stats = run_writer(ModuleCWriteConfig {
records: 1_000,
wal_path: wal_path.clone(),
})
.expect("writer should succeed");
assert_eq!(write_stats.records_written, 1_000);
assert!(write_stats.file_size_bytes > 0);
let read_stats = run_reader(ModuleCReadConfig {
wal_path: wal_path.clone(),
expect_records: Some(1_000),
})
.expect("reader should succeed");
assert_eq!(read_stats.records_read, 1_000);
assert!(read_stats.defaults_ok);
assert!(read_stats.zero_copy_ok);
std::fs::remove_file(wal_path).expect("temporary WAL should be removable");
}
#[test]
#[ignore = "heavy million-record scan"]
fn million_record_round_trip() {
let wal_path = unique_wal_path("module-c-million");
run_writer(ModuleCWriteConfig {
records: 1_000_000,
wal_path: wal_path.clone(),
})
.expect("million-record writer should succeed");
let read_stats = run_reader(ModuleCReadConfig {
wal_path: wal_path.clone(),
expect_records: Some(1_000_000),
})
.expect("million-record reader should succeed");
assert!(read_stats.defaults_ok);
assert!(read_stats.zero_copy_ok);
std::fs::remove_file(wal_path).expect("temporary WAL should be removable");
}
fn unique_wal_path(prefix: &str) -> PathBuf {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("system clock should be after epoch")
.as_nanos();
std::env::temp_dir().join(format!("{}-{}-{}.bin", prefix, std::process::id(), nanos))
}
}