use std::io;
use crc32fast::Hasher;
use super::{
BufferedFileWriter,
CompressionType,
Error,
IOError,
RecordType,
Result,
WritableFile,
BLOCK_SIZE,
HEADER_SIZE,
};
pub struct Writer {
dest: BufferedFileWriter,
block_offset: usize,
manual_flush: bool,
compression_type: CompressionType,
#[allow(dead_code)]
compressed_buffer: Option<Vec<u8>>,
}
impl Writer {
pub fn new(
dest: BufferedFileWriter,
manual_flush: bool,
compression_type: CompressionType,
block_offset: usize,
) -> Self {
let compressed_buffer = if compression_type != CompressionType::None {
Some(Vec::with_capacity(BLOCK_SIZE))
} else {
None
};
Self {
dest,
block_offset,
manual_flush,
compression_type,
compressed_buffer,
}
}
pub fn add_record(&mut self, slice: &[u8]) -> Result<()> {
let compressed;
let data_to_write = if self.compression_type == CompressionType::Lz4 {
compressed = lz4_flex::compress_prepend_size(slice);
&compressed[..]
} else {
slice
};
let mut ptr = data_to_write;
let mut begin = true;
while begin || !ptr.is_empty() {
self.maybe_switch_to_new_block()?;
let avail = BLOCK_SIZE - self.block_offset - HEADER_SIZE;
let fragment_length = ptr.len().min(avail);
let fragment = &ptr[..fragment_length];
let is_end = fragment_length == ptr.len();
let record_type = if begin && is_end {
RecordType::Full
} else if begin {
RecordType::First
} else if is_end {
RecordType::Last
} else {
RecordType::Middle
};
self.emit_physical_record(record_type, fragment)?;
ptr = &ptr[fragment_length..];
begin = false;
}
if !self.manual_flush {
self.write_buffer()?;
}
Ok(())
}
pub fn add_compression_type_record(&mut self) -> Result<()> {
if self.block_offset != 0 {
return Err(Error::IO(IOError::new(
io::ErrorKind::Other,
"Compression type record must be first",
)));
}
if self.compression_type == CompressionType::None {
return Ok(());
}
let data = [self.compression_type as u8];
self.emit_physical_record(RecordType::SetCompressionType, &data)?;
if !self.manual_flush {
self.write_buffer()?;
}
Ok(())
}
pub fn write_buffer(&mut self) -> Result<()> {
self.dest.flush() }
pub fn sync(&mut self) -> Result<()> {
self.dest.sync() }
pub fn close(&mut self) -> Result<()> {
self.sync()?;
self.dest.close()
}
fn maybe_switch_to_new_block(&mut self) -> Result<()> {
let leftover = BLOCK_SIZE - self.block_offset;
if leftover < HEADER_SIZE {
let padding = vec![0u8; leftover];
self.dest.append(&padding)?;
self.block_offset = 0;
}
Ok(())
}
fn emit_physical_record(&mut self, record_type: RecordType, data: &[u8]) -> Result<()> {
let length = data.len();
if length > 0xffff {
return Err(Error::IO(IOError::new(io::ErrorKind::InvalidInput, "Record too large")));
}
debug_assert!(
self.block_offset + HEADER_SIZE + length <= BLOCK_SIZE,
"Record exceeds block boundary: offset={}, header={}, data={}, block_size={}",
self.block_offset,
HEADER_SIZE,
length,
BLOCK_SIZE
);
let type_byte = record_type as u8;
let mut hasher = Hasher::new();
hasher.update(&[type_byte]); hasher.update(data); let crc = hasher.finalize();
let mut header = Vec::with_capacity(HEADER_SIZE);
header.extend_from_slice(&crc.to_be_bytes());
header.extend_from_slice(&(length as u16).to_be_bytes());
header.push(record_type as u8);
self.dest.append(&header)?;
self.dest.append(data)?;
self.block_offset += HEADER_SIZE + length;
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::fs::File;
use tempdir::TempDir;
use super::*;
#[test]
fn test_writer_basic() {
let temp_dir = TempDir::new("test").unwrap();
let file_path = temp_dir.path().join("test.wal");
let file = File::create(&file_path).unwrap();
let buffered_writer = BufferedFileWriter::new(file, BLOCK_SIZE);
let mut writer = Writer::new(buffered_writer, false, CompressionType::None, 0);
writer.add_record(b"Hello, World!").unwrap();
writer.close().unwrap();
let metadata = std::fs::metadata(&file_path).unwrap();
assert!(metadata.len() > 0);
}
#[test]
fn test_manual_flush() {
let temp_dir = TempDir::new("test").unwrap();
let file_path = temp_dir.path().join("test.wal");
let file = File::create(&file_path).unwrap();
let buffered_writer = BufferedFileWriter::new(file, BLOCK_SIZE);
let mut writer = Writer::new(buffered_writer, true, CompressionType::None, 0);
writer.add_record(b"Test").unwrap();
writer.write_buffer().unwrap();
writer.close().unwrap();
}
#[test]
fn test_fragmentation() {
let temp_dir = TempDir::new("test").unwrap();
let file_path = temp_dir.path().join("test.wal");
let file = File::create(&file_path).unwrap();
let buffered_writer = BufferedFileWriter::new(file, BLOCK_SIZE);
let mut writer = Writer::new(buffered_writer, false, CompressionType::None, 0);
let large_data = vec![b'A'; BLOCK_SIZE * 2];
writer.add_record(&large_data).unwrap();
writer.close().unwrap();
let metadata = std::fs::metadata(&file_path).unwrap();
assert!(metadata.len() > BLOCK_SIZE as u64 * 2);
}
}