use std::alloc::{alloc, dealloc, Layout};
use std::fs::OpenOptions;
use std::io::{self, Write};
use std::path::Path;
pub const DEFAULT_ALIGNMENT: usize = 4096;
pub const DEFAULT_DIRECT_IO_THRESHOLD: usize = 1024 * 1024;
#[derive(Debug, Clone)]
pub struct DirectIoConfig {
pub alignment: usize,
pub threshold: usize,
pub enabled: bool,
}
impl Default for DirectIoConfig {
fn default() -> Self {
Self {
alignment: DEFAULT_ALIGNMENT,
threshold: DEFAULT_DIRECT_IO_THRESHOLD,
enabled: true,
}
}
}
impl DirectIoConfig {
pub fn new() -> Self {
Self::default()
}
pub fn with_alignment(mut self, alignment: usize) -> Self {
assert!(alignment.is_power_of_two(), "Alignment must be power of 2");
self.alignment = alignment;
self
}
pub fn with_threshold(mut self, threshold: usize) -> Self {
self.threshold = threshold;
self
}
pub fn with_enabled(mut self, enabled: bool) -> Self {
self.enabled = enabled;
self
}
pub fn should_use_direct_io(&self, size: usize) -> bool {
self.enabled && size >= self.threshold
}
}
pub struct AlignedBuffer {
ptr: *mut u8,
capacity: usize,
len: usize,
alignment: usize,
}
impl AlignedBuffer {
pub fn new(capacity: usize, alignment: usize) -> io::Result<Self> {
assert!(alignment.is_power_of_two(), "Alignment must be power of 2");
let aligned_capacity = (capacity + alignment - 1) & !(alignment - 1);
let layout = Layout::from_size_align(aligned_capacity, alignment)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
let ptr = unsafe { alloc(layout) };
if ptr.is_null() {
return Err(io::Error::new(
io::ErrorKind::OutOfMemory,
"Failed to allocate aligned buffer",
));
}
Ok(Self {
ptr,
capacity: aligned_capacity,
len: 0,
alignment,
})
}
pub fn capacity(&self) -> usize {
self.capacity
}
pub fn len(&self) -> usize {
self.len
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
pub fn alignment(&self) -> usize {
self.alignment
}
pub fn as_slice(&self) -> &[u8] {
unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
}
pub fn as_mut_slice(&mut self) -> &mut [u8] {
unsafe { std::slice::from_raw_parts_mut(self.ptr, self.capacity) }
}
pub fn write(&mut self, data: &[u8]) -> io::Result<()> {
if self.len + data.len() > self.capacity {
return Err(io::Error::new(io::ErrorKind::WriteZero, "Buffer capacity exceeded"));
}
unsafe {
std::ptr::copy_nonoverlapping(data.as_ptr(), self.ptr.add(self.len), data.len());
}
self.len += data.len();
Ok(())
}
pub fn clear(&mut self) {
self.len = 0;
}
pub fn pad_to_alignment(&mut self) {
let remainder = self.len % self.alignment;
if remainder != 0 {
let padding = self.alignment - remainder;
if self.len + padding <= self.capacity {
unsafe {
std::ptr::write_bytes(self.ptr.add(self.len), 0, padding);
}
self.len += padding;
}
}
}
pub fn data_len(&self) -> usize {
self.len
}
}
impl Drop for AlignedBuffer {
fn drop(&mut self) {
if !self.ptr.is_null() {
unsafe {
let layout = Layout::from_size_align_unchecked(self.capacity, self.alignment);
dealloc(self.ptr, layout);
}
}
}
}
unsafe impl Send for AlignedBuffer {}
unsafe impl Sync for AlignedBuffer {}
pub struct DirectIoWriter {
config: DirectIoConfig,
}
impl DirectIoWriter {
pub fn new() -> Self {
Self {
config: DirectIoConfig::default(),
}
}
pub fn with_config(config: DirectIoConfig) -> Self {
Self { config }
}
pub fn write_to_file<P: AsRef<Path>>(&self, path: P, data: &[u8]) -> io::Result<()> {
if self.config.should_use_direct_io(data.len()) {
self.write_direct(path, data)
} else {
std::fs::write(path, data)
}
}
#[cfg(target_os = "linux")]
fn write_direct<P: AsRef<Path>>(&self, path: P, data: &[u8]) -> io::Result<()> {
use std::os::unix::fs::OpenOptionsExt;
let mut buffer = AlignedBuffer::new(data.len(), self.config.alignment)?;
buffer.write(data)?;
buffer.pad_to_alignment();
let mut file = OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.custom_flags(libc::O_DIRECT)
.open(path)?;
file.write_all(buffer.as_slice())?;
file.sync_all()?;
Ok(())
}
#[cfg(target_os = "windows")]
fn write_direct<P: AsRef<Path>>(&self, path: P, data: &[u8]) -> io::Result<()> {
use std::os::windows::fs::OpenOptionsExt;
let mut buffer = AlignedBuffer::new(data.len(), self.config.alignment)?;
buffer.write(data)?;
buffer.pad_to_alignment();
const FILE_FLAG_NO_BUFFERING: u32 = 0x20000000;
let mut file = OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.custom_flags(FILE_FLAG_NO_BUFFERING)
.open(path)?;
file.write_all(buffer.as_slice())?;
file.sync_all()?;
Ok(())
}
#[cfg(not(any(target_os = "linux", target_os = "windows")))]
fn write_direct<P: AsRef<Path>>(&self, path: P, data: &[u8]) -> io::Result<()> {
std::fs::write(path, data)
}
}
impl Default for DirectIoWriter {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_aligned_buffer_creation() {
let buffer = AlignedBuffer::new(1024, 512).unwrap();
assert_eq!(buffer.capacity(), 1024);
assert_eq!(buffer.len(), 0);
assert_eq!(buffer.alignment(), 512);
assert!(buffer.is_empty());
}
#[test]
fn test_aligned_buffer_write() {
let mut buffer = AlignedBuffer::new(1024, 512).unwrap();
let data = b"Hello, World!";
buffer.write(data).unwrap();
assert_eq!(buffer.len(), data.len());
assert_eq!(buffer.as_slice(), data);
}
#[test]
fn test_aligned_buffer_padding() {
let mut buffer = AlignedBuffer::new(1024, 512).unwrap();
buffer.write(b"Hello").unwrap();
assert_eq!(buffer.len(), 5);
buffer.pad_to_alignment();
assert_eq!(buffer.len(), 512); assert_eq!(buffer.as_slice()[0..5], b"Hello"[..]);
}
#[test]
fn test_config_threshold() {
let config = DirectIoConfig::new().with_threshold(1024).with_alignment(4096);
assert!(!config.should_use_direct_io(512));
assert!(config.should_use_direct_io(2048));
}
#[test]
fn test_config_disabled() {
let config = DirectIoConfig::new().with_enabled(false);
assert!(!config.should_use_direct_io(10_000_000));
}
#[test]
fn test_direct_io_writer() {
let writer = DirectIoWriter::new();
let temp_dir = tempfile::tempdir().unwrap();
let file_path = temp_dir.path().join("test.bin");
let data = vec![42u8; 1024];
writer.write_to_file(&file_path, &data).unwrap();
let read_data = std::fs::read(&file_path).unwrap();
assert!(read_data.len() >= data.len());
assert_eq!(&read_data[..data.len()], &data[..]);
}
#[test]
fn test_small_file_uses_regular_io() {
let config = DirectIoConfig::new().with_threshold(1024 * 1024);
let writer = DirectIoWriter::with_config(config);
let temp_dir = tempfile::tempdir().unwrap();
let file_path = temp_dir.path().join("small.bin");
let data = vec![42u8; 1024]; writer.write_to_file(&file_path, &data).unwrap();
let read_data = std::fs::read(&file_path).unwrap();
assert_eq!(read_data, data); }
#[test]
fn test_alignment_power_of_two() {
assert!(AlignedBuffer::new(1024, 512).is_ok());
assert!(AlignedBuffer::new(1024, 4096).is_ok());
}
#[test]
#[should_panic(expected = "Alignment must be power of 2")]
fn test_invalid_alignment() {
let _ = AlignedBuffer::new(1024, 1000); }
}