use std::path::PathBuf;
use anyhow::Result;
#[cfg(target_os = "linux")]
const ALIGNMENT: usize = 4096;
#[derive(Debug, Clone)]
pub struct ModuleDConfig {
pub target_path: PathBuf,
pub total_bytes: u64,
pub group_bytes: usize,
pub request_bytes: usize,
pub producers: usize,
pub seed: u64,
pub allow_file_fallback: bool,
pub require_io_uring: bool,
}
#[derive(Debug, Clone)]
pub struct ModuleDStats {
pub bytes_written: u64,
pub commits: u64,
pub elapsed_ms: f64,
pub throughput_mb_s: f64,
pub io_wait_pct: f64,
pub mode: String,
pub alignment_violations: u64,
pub write_errors: u64,
pub target_path: PathBuf,
}
impl ModuleDStats {
pub fn to_json(&self) -> String {
format!(
"{{\"module\":\"D\",\"bytes_written\":{},\"commits\":{},\"elapsed_ms\":{:.3},\"throughput_mb_s\":{:.3},\"io_wait_pct\":{:.3},\"mode\":\"{}\",\"alignment_violations\":{},\"write_errors\":{},\"target_path\":\"{}\"}}",
self.bytes_written,
self.commits,
self.elapsed_ms,
self.throughput_mb_s,
self.io_wait_pct,
self.mode,
self.alignment_violations,
self.write_errors,
self.target_path.display()
)
}
}
#[cfg(target_os = "linux")]
mod linux {
use std::fs;
use std::fs::OpenOptions as StdOpenOptions;
use std::io;
use std::io::ErrorKind;
use std::os::unix::fs::FileExt;
use std::os::unix::fs::OpenOptionsExt;
use std::path::{Path, PathBuf};
use std::sync::mpsc;
use std::time::Instant;
use anyhow::{Context, Result};
use nix::libc;
use tokio::sync::mpsc as tokio_mpsc;
use tokio_uring::buf::BoundedBuf;
use tokio_uring::fs::{File, OpenOptions as TokioOpenOptions};
use crate::module_d::{ModuleDConfig, ModuleDStats, ALIGNMENT};
#[derive(Debug, Clone, Copy)]
struct CpuSample {
total: u64,
iowait: u64,
}
struct ProducerChunk {
payload: Vec<u8>,
}
pub(super) fn run(config: ModuleDConfig) -> Result<ModuleDStats> {
validate_config(&config)?;
if config.require_io_uring && !io_uring_probe() {
anyhow::bail!("io_uring is required but unavailable in this environment");
}
let async_config = config.clone();
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
tokio_uring::start(async move { run_async(async_config).await })
})) {
Ok(Ok(stats)) => Ok(stats),
Ok(Err(err)) => {
if should_fallback_to_sync(&err.to_string()) {
if config.require_io_uring {
return Err(err).context("io_uring required; refusing sync fallback");
}
eprintln!(
"module-d: io_uring unavailable ({}); retrying with sync direct-io fallback",
err
);
run_sync(config)
} else {
Err(err)
}
}
Err(payload) => {
let panic_message = panic_payload_to_string(&payload);
if should_fallback_to_sync(&panic_message) {
if config.require_io_uring {
anyhow::bail!(
"io_uring required; runtime init failed with panic: {}",
panic_message
);
}
eprintln!(
"module-d: io_uring unavailable ({}); retrying with sync direct-io fallback",
panic_message
);
run_sync(config)
} else {
std::panic::resume_unwind(payload);
}
}
}
}
async fn run_async(config: ModuleDConfig) -> Result<ModuleDStats> {
let cpu_before = read_cpu_sample().context("failed reading pre-run CPU sample")?;
let start = Instant::now();
let (file, mode, resolved_target) = open_target(&config).await.with_context(|| {
format!(
"failed opening target path {}",
config.target_path.display()
)
})?;
let request_bytes_u64 = config.request_bytes as u64;
let total_requests = config.total_bytes.div_ceil(request_bytes_u64);
let channel_capacity = config.producers.saturating_mul(2).max(4);
let (tx, mut rx) = tokio_mpsc::channel::<ProducerChunk>(channel_capacity);
for producer_id in 0..config.producers {
let tx = tx.clone();
let producer_count = config.producers as u64;
let total_bytes = config.total_bytes;
let request_bytes = config.request_bytes;
let seed = config.seed;
tokio_uring::spawn(async move {
let mut seq = producer_id as u64;
while seq < total_requests {
let offset = seq * request_bytes as u64;
let remaining = total_bytes.saturating_sub(offset);
if remaining == 0 {
break;
}
let chunk_bytes = remaining.min(request_bytes as u64) as usize;
let mut payload = vec![0_u8; chunk_bytes];
fill_payload(&mut payload, seed, producer_id as u64, seq);
if tx.send(ProducerChunk { payload }).await.is_err() {
break;
}
seq += producer_count;
}
});
}
drop(tx);
let mut pending = Vec::<u8>::with_capacity(config.group_bytes + config.request_bytes);
let mut write_buf = vec![0_u8; config.group_bytes + ALIGNMENT];
let write_buf_align_start = aligned_offset(write_buf.as_ptr() as usize, ALIGNMENT);
let mut file_offset = 0_u64;
let mut logical_bytes_written = 0_u64;
let mut commits = 0_u64;
let mut alignment_violations = 0_u64;
let mut write_errors = 0_u64;
while let Some(chunk) = rx.recv().await {
let mut idx = 0usize;
while idx < chunk.payload.len() {
let remaining_in_group = config.group_bytes - pending.len();
let take = remaining_in_group.min(chunk.payload.len() - idx);
pending.extend_from_slice(&chunk.payload[idx..idx + take]);
idx += take;
if pending.len() == config.group_bytes {
flush_pending(
&file,
&mut pending,
&mut write_buf,
write_buf_align_start,
&mut file_offset,
&mut logical_bytes_written,
&mut commits,
&mut alignment_violations,
&mut write_errors,
false,
)
.await?;
}
}
}
if !pending.is_empty() {
flush_pending(
&file,
&mut pending,
&mut write_buf,
write_buf_align_start,
&mut file_offset,
&mut logical_bytes_written,
&mut commits,
&mut alignment_violations,
&mut write_errors,
true,
)
.await?;
}
file.sync_data()
.await
.context("module-d sync_data failed")?;
file.close().await.context("module-d close failed")?;
if logical_bytes_written != config.total_bytes {
anyhow::bail!(
"module-d byte mismatch: expected {}, wrote {}",
config.total_bytes,
logical_bytes_written
);
}
let elapsed = start.elapsed();
let elapsed_ms = elapsed.as_secs_f64() * 1_000.0;
let throughput_mb_s =
(logical_bytes_written as f64 / (1024.0 * 1024.0)) / elapsed.as_secs_f64();
let cpu_after = read_cpu_sample().context("failed reading post-run CPU sample")?;
let io_wait_pct = compute_iowait_pct(cpu_before, cpu_after);
Ok(ModuleDStats {
bytes_written: logical_bytes_written,
commits,
elapsed_ms,
throughput_mb_s,
io_wait_pct,
mode,
alignment_violations,
write_errors,
target_path: resolved_target,
})
}
fn run_sync(config: ModuleDConfig) -> Result<ModuleDStats> {
let cpu_before = read_cpu_sample().context("failed reading pre-run CPU sample")?;
let start = Instant::now();
let (file, mode, resolved_target) = open_target_sync(&config).with_context(|| {
format!(
"failed opening target path {}",
config.target_path.display()
)
})?;
let request_bytes_u64 = config.request_bytes as u64;
let total_requests = config.total_bytes.div_ceil(request_bytes_u64);
let channel_capacity = config.producers.saturating_mul(2).max(4);
let (tx, rx) = mpsc::sync_channel::<ProducerChunk>(channel_capacity);
let mut producer_handles = Vec::with_capacity(config.producers);
for producer_id in 0..config.producers {
let tx = tx.clone();
let producer_count = config.producers as u64;
let total_bytes = config.total_bytes;
let request_bytes = config.request_bytes;
let seed = config.seed;
producer_handles.push(std::thread::spawn(move || {
let mut seq = producer_id as u64;
while seq < total_requests {
let offset = seq * request_bytes as u64;
let remaining = total_bytes.saturating_sub(offset);
if remaining == 0 {
break;
}
let chunk_bytes = remaining.min(request_bytes as u64) as usize;
let mut payload = vec![0_u8; chunk_bytes];
fill_payload(&mut payload, seed, producer_id as u64, seq);
if tx.send(ProducerChunk { payload }).is_err() {
break;
}
seq += producer_count;
}
}));
}
drop(tx);
let mut pending = Vec::<u8>::with_capacity(config.group_bytes + config.request_bytes);
let mut write_buf = vec![0_u8; config.group_bytes + ALIGNMENT];
let write_buf_align_start = aligned_offset(write_buf.as_ptr() as usize, ALIGNMENT);
let mut file_offset = 0_u64;
let mut logical_bytes_written = 0_u64;
let mut commits = 0_u64;
let mut alignment_violations = 0_u64;
let mut write_errors = 0_u64;
for chunk in rx {
let mut idx = 0usize;
while idx < chunk.payload.len() {
let remaining_in_group = config.group_bytes - pending.len();
let take = remaining_in_group.min(chunk.payload.len() - idx);
pending.extend_from_slice(&chunk.payload[idx..idx + take]);
idx += take;
if pending.len() == config.group_bytes {
flush_pending_sync(
&file,
&mut pending,
&mut write_buf,
write_buf_align_start,
&mut file_offset,
&mut logical_bytes_written,
&mut commits,
&mut alignment_violations,
&mut write_errors,
false,
)?;
}
}
}
if !pending.is_empty() {
flush_pending_sync(
&file,
&mut pending,
&mut write_buf,
write_buf_align_start,
&mut file_offset,
&mut logical_bytes_written,
&mut commits,
&mut alignment_violations,
&mut write_errors,
true,
)?;
}
for handle in producer_handles {
if handle.join().is_err() {
anyhow::bail!("producer thread panicked in module-d sync fallback");
}
}
file.sync_data()
.context("module-d sync_data failed (sync fallback)")?;
if logical_bytes_written != config.total_bytes {
anyhow::bail!(
"module-d byte mismatch: expected {}, wrote {}",
config.total_bytes,
logical_bytes_written
);
}
let elapsed = start.elapsed();
let elapsed_ms = elapsed.as_secs_f64() * 1_000.0;
let throughput_mb_s =
(logical_bytes_written as f64 / (1024.0 * 1024.0)) / elapsed.as_secs_f64();
let cpu_after = read_cpu_sample().context("failed reading post-run CPU sample")?;
let io_wait_pct = compute_iowait_pct(cpu_before, cpu_after);
Ok(ModuleDStats {
bytes_written: logical_bytes_written,
commits,
elapsed_ms,
throughput_mb_s,
io_wait_pct,
mode,
alignment_violations,
write_errors,
target_path: resolved_target,
})
}
fn flush_pending_sync(
file: &std::fs::File,
pending: &mut Vec<u8>,
write_buf: &mut Vec<u8>,
write_buf_align_start: usize,
file_offset: &mut u64,
logical_bytes_written: &mut u64,
commits: &mut u64,
alignment_violations: &mut u64,
write_errors: &mut u64,
pad_tail: bool,
) -> Result<()> {
let logical_len = pending.len();
let physical_len = if pad_tail {
align_up(logical_len, ALIGNMENT)
} else {
logical_len
};
if physical_len == 0 {
pending.clear();
return Ok(());
}
let ptr = write_buf.as_ptr() as usize + write_buf_align_start;
if ptr % ALIGNMENT != 0
|| physical_len % ALIGNMENT != 0
|| *file_offset % ALIGNMENT as u64 != 0
{
*alignment_violations += 1;
anyhow::bail!(
"unaligned write detected: ptr_mod={}, physical_len_mod={}, offset_mod={}",
ptr % ALIGNMENT,
physical_len % ALIGNMENT,
*file_offset % ALIGNMENT as u64
);
}
let needed = write_buf_align_start + physical_len;
if write_buf.len() < needed {
write_buf.resize(needed, 0);
}
write_buf[write_buf_align_start..write_buf_align_start + logical_len]
.copy_from_slice(pending.as_slice());
if physical_len > logical_len {
write_buf[write_buf_align_start + logical_len..write_buf_align_start + physical_len]
.fill(0);
}
let aligned = &write_buf[write_buf_align_start..write_buf_align_start + physical_len];
if let Err(err) = write_all_at(file, aligned, *file_offset) {
*write_errors += 1;
return Err(err).context("module-d write_at failed (sync fallback)");
}
*file_offset += physical_len as u64;
*logical_bytes_written += logical_len as u64;
*commits += 1;
pending.clear();
Ok(())
}
fn write_all_at(file: &std::fs::File, mut buf: &[u8], mut offset: u64) -> io::Result<()> {
while !buf.is_empty() {
let written = file.write_at(buf, offset)?;
if written == 0 {
return Err(io::Error::new(
ErrorKind::WriteZero,
"write_at returned 0 bytes",
));
}
buf = &buf[written..];
offset = offset.saturating_add(written as u64);
}
Ok(())
}
async fn flush_pending(
file: &File,
pending: &mut Vec<u8>,
write_buf: &mut Vec<u8>,
write_buf_align_start: usize,
file_offset: &mut u64,
logical_bytes_written: &mut u64,
commits: &mut u64,
alignment_violations: &mut u64,
write_errors: &mut u64,
pad_tail: bool,
) -> Result<()> {
let logical_len = pending.len();
let physical_len = if pad_tail {
align_up(logical_len, ALIGNMENT)
} else {
logical_len
};
if physical_len == 0 {
pending.clear();
return Ok(());
}
let ptr = write_buf.as_ptr() as usize + write_buf_align_start;
if ptr % ALIGNMENT != 0
|| physical_len % ALIGNMENT != 0
|| *file_offset % ALIGNMENT as u64 != 0
{
*alignment_violations += 1;
anyhow::bail!(
"unaligned write detected: ptr_mod={}, physical_len_mod={}, offset_mod={}",
ptr % ALIGNMENT,
physical_len % ALIGNMENT,
*file_offset % ALIGNMENT as u64
);
}
let needed = write_buf_align_start + physical_len;
if write_buf.len() < needed {
write_buf.resize(needed, 0);
}
write_buf[write_buf_align_start..write_buf_align_start + logical_len]
.copy_from_slice(pending.as_slice());
if physical_len > logical_len {
write_buf[write_buf_align_start + logical_len..write_buf_align_start + physical_len]
.fill(0);
}
let mut owned = std::mem::take(write_buf);
let slice = owned.slice(write_buf_align_start..write_buf_align_start + physical_len);
let (result, returned) = file.write_all_at(slice, *file_offset).await;
owned = returned.into_inner();
*write_buf = owned;
if let Err(err) = result {
*write_errors += 1;
return Err(err).context("module-d write_all_at failed");
}
*file_offset += physical_len as u64;
*logical_bytes_written += logical_len as u64;
*commits += 1;
pending.clear();
Ok(())
}
async fn open_target(config: &ModuleDConfig) -> Result<(File, String, PathBuf)> {
match open_direct(&config.target_path, false).await {
Ok(file) => Ok((file, "block".to_string(), config.target_path.clone())),
Err(primary_err) => {
if !config.allow_file_fallback {
return Err(primary_err).with_context(|| {
format!(
"opening {} as raw block target failed and fallback disabled",
config.target_path.display()
)
});
}
let fallback_path = fallback_path_for(&config.target_path);
if let Some(parent) = fallback_path.parent() {
if !parent.as_os_str().is_empty() {
fs::create_dir_all(parent).with_context(|| {
format!("failed to create fallback parent {}", parent.display())
})?;
}
}
let file = open_direct(&fallback_path, true).await.with_context(|| {
format!(
"failed opening fallback sparse file target {}",
fallback_path.display()
)
})?;
Ok((file, "file-fallback".to_string(), fallback_path))
}
}
}
fn open_target_sync(config: &ModuleDConfig) -> Result<(std::fs::File, String, PathBuf)> {
match open_direct_sync(&config.target_path, false) {
Ok(file) => Ok((file, "block".to_string(), config.target_path.clone())),
Err(primary_err) => {
if !config.allow_file_fallback {
return Err(primary_err).with_context(|| {
format!(
"opening {} as raw block target failed and fallback disabled",
config.target_path.display()
)
});
}
let fallback_path = fallback_path_for(&config.target_path);
if let Some(parent) = fallback_path.parent() {
if !parent.as_os_str().is_empty() {
fs::create_dir_all(parent).with_context(|| {
format!("failed to create fallback parent {}", parent.display())
})?;
}
}
let file = open_direct_sync(&fallback_path, true).with_context(|| {
format!(
"failed opening fallback sparse file target {}",
fallback_path.display()
)
})?;
Ok((file, "file-fallback".to_string(), fallback_path))
}
}
}
async fn open_direct(path: &Path, create_file: bool) -> Result<File> {
let mut opts = TokioOpenOptions::new();
opts.write(true);
if create_file {
opts.create(true).truncate(true).mode(0o644);
}
opts.custom_flags(libc::O_DIRECT | libc::O_DSYNC);
opts.open(path)
.await
.with_context(|| format!("open_direct failed for {}", path.display()))
}
fn open_direct_sync(path: &Path, create_file: bool) -> Result<std::fs::File> {
let mut opts = StdOpenOptions::new();
opts.write(true);
if create_file {
opts.create(true).truncate(true).mode(0o644);
}
opts.custom_flags(libc::O_DIRECT | libc::O_DSYNC);
opts.open(path)
.with_context(|| format!("open_direct failed for {}", path.display()))
}
fn fallback_path_for(target: &Path) -> PathBuf {
if target.starts_with("/dev") {
PathBuf::from("/tmp/tracer-bullet-module-d-direct.bin")
} else {
target.with_extension("direct.bin")
}
}
fn should_fallback_to_sync(message: &str) -> bool {
message.contains("Operation not permitted")
|| message.contains("io_uring")
|| message.contains("tokio-uring")
}
pub(super) fn io_uring_available() -> bool {
io_uring_probe()
}
fn io_uring_probe() -> bool {
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
tokio_uring::start(async { Ok::<(), anyhow::Error>(()) })
}))
.is_ok()
}
fn panic_payload_to_string(payload: &Box<dyn std::any::Any + Send>) -> String {
if let Some(text) = payload.downcast_ref::<&str>() {
return (*text).to_string();
}
if let Some(text) = payload.downcast_ref::<String>() {
return text.clone();
}
"unknown panic payload".to_string()
}
fn fill_payload(buffer: &mut [u8], seed: u64, producer_id: u64, sequence: u64) {
let mut state = seed
^ producer_id.wrapping_mul(0x9E37_79B9_7F4A_7C15)
^ sequence.wrapping_mul(0xBF58_476D_1CE4_E5B9);
for byte in buffer.iter_mut() {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
*byte = (state >> 24) as u8;
}
}
fn read_cpu_sample() -> Result<CpuSample> {
let stat = fs::read_to_string("/proc/stat").context("failed to read /proc/stat")?;
let line = stat
.lines()
.next()
.context("/proc/stat did not contain cpu header")?;
let mut fields = line.split_whitespace();
let cpu_tag = fields.next().context("missing cpu tag in /proc/stat")?;
if cpu_tag != "cpu" {
anyhow::bail!("unexpected cpu tag in /proc/stat: {}", cpu_tag);
}
let mut values = Vec::with_capacity(8);
for field in fields.take(8) {
values.push(
field
.parse::<u64>()
.with_context(|| format!("failed parsing /proc/stat field: {}", field))?,
);
}
if values.len() < 5 {
anyhow::bail!("/proc/stat cpu line missing expected counters");
}
let total = values.iter().copied().sum::<u64>();
let iowait = values[4];
Ok(CpuSample { total, iowait })
}
fn compute_iowait_pct(before: CpuSample, after: CpuSample) -> f64 {
let total_delta = after.total.saturating_sub(before.total);
if total_delta == 0 {
return 0.0;
}
let iowait_delta = after.iowait.saturating_sub(before.iowait);
(iowait_delta as f64 / total_delta as f64) * 100.0
}
fn validate_config(config: &ModuleDConfig) -> Result<()> {
if config.total_bytes == 0 {
anyhow::bail!("total_bytes must be > 0");
}
if config.group_bytes == 0 || config.group_bytes % ALIGNMENT != 0 {
anyhow::bail!("group_bytes must be > 0 and aligned to {} bytes", ALIGNMENT);
}
if config.request_bytes == 0 || config.request_bytes % ALIGNMENT != 0 {
anyhow::bail!(
"request_bytes must be > 0 and aligned to {} bytes",
ALIGNMENT
);
}
if config.group_bytes < config.request_bytes {
anyhow::bail!("group_bytes must be >= request_bytes");
}
if config.producers == 0 {
anyhow::bail!("producers must be > 0");
}
if config.total_bytes % ALIGNMENT as u64 != 0 {
anyhow::bail!(
"total_bytes must be aligned to {} bytes for O_DIRECT",
ALIGNMENT
);
}
Ok(())
}
fn align_up(value: usize, align: usize) -> usize {
if value % align == 0 {
value
} else {
value + (align - (value % align))
}
}
fn aligned_offset(ptr: usize, align: usize) -> usize {
(align - (ptr % align)) % align
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn alignment_math_is_correct() {
assert_eq!(align_up(4096, 4096), 4096);
assert_eq!(align_up(4097, 4096), 8192);
assert_eq!(align_up(8191, 4096), 8192);
}
#[test]
fn validates_direct_io_constraints() {
let err = validate_config(&ModuleDConfig {
target_path: PathBuf::from("/tmp/x"),
total_bytes: 123,
group_bytes: 16 * 1024 * 1024,
request_bytes: 256 * 1024,
producers: 1,
seed: 1,
allow_file_fallback: true,
require_io_uring: false,
})
.expect_err("unaligned total_bytes should fail");
assert!(err.to_string().contains("total_bytes"));
}
#[test]
fn aligned_offset_returns_expected_values() {
assert_eq!(aligned_offset(0, 4096), 0);
assert_eq!(aligned_offset(1, 4096), 4095);
assert_eq!(aligned_offset(4095, 4096), 1);
assert_eq!(aligned_offset(4096, 4096), 0);
}
}
}
#[cfg(target_os = "linux")]
pub fn run(config: ModuleDConfig) -> Result<ModuleDStats> {
linux::run(config)
}
#[cfg(target_os = "linux")]
pub fn io_uring_available() -> bool {
linux::io_uring_available()
}
#[cfg(not(target_os = "linux"))]
pub fn run(_config: ModuleDConfig) -> Result<ModuleDStats> {
anyhow::bail!("module-d requires Linux (io_uring + O_DIRECT)")
}
#[cfg(not(target_os = "linux"))]
pub fn io_uring_available() -> bool {
false
}