use crate::{Event, Offset, StorageError, StorageResult};
use derive_more::{Deref, DerefMut, From, Into};
use memmap2::{MmapMut, MmapOptions};
use serde::{Deserialize, Serialize};
use std::{
collections::HashMap,
fs::{File, OpenOptions},
path::{Path, PathBuf},
sync::Arc,
};
use tokio::sync::RwLock;
use uuid::Uuid;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Deref, DerefMut, From, Into)]
pub struct SegmentId(pub Uuid);
impl SegmentId {
pub fn new() -> Self {
Self(Uuid::new_v4())
}
pub fn from_uuid(uuid: Uuid) -> Self {
Self(uuid)
}
}
impl Default for SegmentId {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Display for SegmentId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SegmentHeader {
pub segment_id: SegmentId,
pub created_at: u64,
pub size: u64,
pub write_offset: u64,
pub entry_count: u64,
pub header_crc: u32,
}
impl SegmentHeader {
pub fn new(segment_id: SegmentId, size: u64) -> Self {
let created_at = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let mut header = Self {
segment_id,
created_at,
size,
write_offset: 256, entry_count: 0,
header_crc: 0,
};
header.header_crc = header.calculate_crc();
header
}
fn calculate_crc(&self) -> u32 {
let mut temp_header = self.clone();
temp_header.header_crc = 0;
let serialized = bincode::serialize(&temp_header).unwrap_or_default();
crc32fast::hash(&serialized)
}
pub fn verify_integrity(&self) -> StorageResult<()> {
let expected_crc = self.calculate_crc();
if self.header_crc != expected_crc {
return Err(StorageError::CrcMismatch {
offset: 0,
expected: expected_crc,
actual: self.header_crc,
});
}
Ok(())
}
}
#[derive(Debug)]
pub struct Segment {
id: SegmentId,
file_path: PathBuf,
mmap: MmapMut,
header: SegmentHeader,
read_only: bool,
}
impl Segment {
pub fn create<P: AsRef<Path>>(file_path: P, size: u64) -> StorageResult<Self> {
let file_path = file_path.as_ref().to_path_buf();
let segment_id = SegmentId::new();
let file = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.open(&file_path)?;
file.set_len(size)?;
let mut mmap = unsafe { MmapOptions::new().map_mut(&file)? };
let header = SegmentHeader::new(segment_id, size);
let header_bytes = bincode::serialize(&header)?;
if header_bytes.len() > mmap.len() {
return Err(StorageError::configuration("Header size exceeds segment size"));
}
let header_size = 256;
if header_bytes.len() > header_size {
return Err(StorageError::configuration("Header size exceeds reserved space"));
}
mmap[..header_bytes.len()].copy_from_slice(&header_bytes);
for i in header_bytes.len()..header_size {
mmap[i] = 0;
}
Ok(Self {
id: segment_id,
file_path,
mmap,
header,
read_only: false,
})
}
pub fn open<P: AsRef<Path>>(file_path: P, read_only: bool) -> StorageResult<Self> {
let file_path = file_path.as_ref().to_path_buf();
let file = if read_only {
File::open(&file_path)?
} else {
OpenOptions::new()
.read(true)
.write(true)
.open(&file_path)?
};
let mmap = if read_only {
unsafe { MmapOptions::new().map(&file)?.make_mut()? }
} else {
unsafe { MmapOptions::new().map_mut(&file)? }
};
let header_size = 256;
if mmap.len() < header_size {
return Err(StorageError::invalid_format("Segment file too small for header"));
}
let mut actual_header_end = 0;
for i in 0..header_size {
if mmap[i] != 0 {
actual_header_end = i + 1;
}
}
if actual_header_end == 0 {
return Err(StorageError::invalid_format("No header data found"));
}
for end in (actual_header_end..header_size).rev() {
if let Ok(header) = bincode::deserialize::<SegmentHeader>(&mmap[..end]) {
if header.verify_integrity().is_ok() {
return Ok(Self {
id: header.segment_id,
file_path,
mmap,
header,
read_only,
});
}
}
}
return Err(StorageError::invalid_format("Failed to deserialize segment header"));
}
pub fn append_event(&mut self, event: &Event) -> StorageResult<Offset> {
if self.read_only {
return Err(StorageError::internal("Cannot write to read-only segment"));
}
let serialized_event = bincode::serialize(event)?;
let data_size = serialized_event.len() as u32;
let entry_size = std::mem::size_of::<u32>() + serialized_event.len() + std::mem::size_of::<u32>();
if self.header.write_offset + entry_size as u64 > self.header.size {
return Err(StorageError::InsufficientSpace {
required: entry_size as u64,
available: self.header.size - self.header.write_offset,
});
}
let write_offset = self.header.write_offset as usize;
let offset = Offset::new(self.header.entry_count);
let crc = crc32fast::hash(&serialized_event);
let size_bytes = data_size.to_le_bytes();
let size_end = write_offset + size_bytes.len();
self.mmap[write_offset..size_end].copy_from_slice(&size_bytes);
let data_end = size_end + serialized_event.len();
self.mmap[size_end..data_end].copy_from_slice(&serialized_event);
let crc_bytes = crc.to_le_bytes();
let crc_end = data_end + crc_bytes.len();
self.mmap[data_end..crc_end].copy_from_slice(&crc_bytes);
self.header.write_offset += entry_size as u64;
self.header.entry_count += 1;
self.header.header_crc = self.header.calculate_crc();
let header_bytes = bincode::serialize(&self.header)?;
let header_size = 256;
if header_bytes.len() > header_size {
return Err(StorageError::configuration("Header size exceeds reserved space"));
}
self.mmap[..header_bytes.len()].copy_from_slice(&header_bytes);
for i in header_bytes.len()..header_size {
self.mmap[i] = 0;
}
Ok(offset)
}
pub fn read_event(&self, offset: Offset) -> StorageResult<Option<Event>> {
if offset.0 >= self.header.entry_count {
return Ok(None);
}
let mut current_position = 256; let mut current_entry_index = 0u64;
while current_entry_index <= offset.0 && current_position < self.header.write_offset as usize {
if current_position + std::mem::size_of::<u32>() > self.mmap.len() {
break;
}
let size_bytes = &self.mmap[current_position..current_position + std::mem::size_of::<u32>()];
let data_size = u32::from_le_bytes([size_bytes[0], size_bytes[1], size_bytes[2], size_bytes[3]]) as usize;
let data_start = current_position + std::mem::size_of::<u32>();
let data_end = data_start + data_size;
let crc_end = data_end + std::mem::size_of::<u32>();
if crc_end > self.mmap.len() {
return Err(StorageError::invalid_format("Entry extends beyond segment"));
}
if current_entry_index == offset.0 {
let event_data = &self.mmap[data_start..data_end];
let stored_crc_bytes = &self.mmap[data_end..crc_end];
let stored_crc = u32::from_le_bytes([
stored_crc_bytes[0], stored_crc_bytes[1],
stored_crc_bytes[2], stored_crc_bytes[3]
]);
let calculated_crc = crc32fast::hash(event_data);
if stored_crc != calculated_crc {
return Err(StorageError::CrcMismatch {
offset: current_position as u64,
expected: calculated_crc,
actual: stored_crc,
});
}
match bincode::deserialize::<Event>(event_data) {
Ok(event) => return Ok(Some(event)),
Err(_) => {
return Err(StorageError::invalid_format("Failed to deserialize event data"));
}
}
}
current_position = crc_end;
current_entry_index += 1;
}
Ok(None)
}
pub fn read_events_range(&self, start_offset: Offset, max_events: usize) -> StorageResult<Vec<Event>> {
let mut events = Vec::with_capacity(max_events.min(self.header.entry_count as usize));
let mut current_position = 256; let mut current_entry_index = 0u64;
let end_offset = (start_offset.0 + max_events as u64).min(self.header.entry_count);
while current_entry_index < start_offset.0 && current_position < self.header.write_offset as usize {
if current_position + std::mem::size_of::<u32>() > self.mmap.len() {
break;
}
let size_bytes = &self.mmap[current_position..current_position + std::mem::size_of::<u32>()];
let data_size = u32::from_le_bytes([size_bytes[0], size_bytes[1], size_bytes[2], size_bytes[3]]) as usize;
let entry_size = std::mem::size_of::<u32>() + data_size + std::mem::size_of::<u32>(); current_position += entry_size;
current_entry_index += 1;
}
while current_entry_index < end_offset && current_position < self.header.write_offset as usize && events.len() < max_events {
if current_position + std::mem::size_of::<u32>() > self.mmap.len() {
break;
}
let size_bytes = &self.mmap[current_position..current_position + std::mem::size_of::<u32>()];
let data_size = u32::from_le_bytes([size_bytes[0], size_bytes[1], size_bytes[2], size_bytes[3]]) as usize;
let data_start = current_position + std::mem::size_of::<u32>();
let data_end = data_start + data_size;
let crc_end = data_end + std::mem::size_of::<u32>();
if crc_end > self.mmap.len() {
return Err(StorageError::invalid_format("Entry extends beyond segment"));
}
let event_data = &self.mmap[data_start..data_end];
let stored_crc_bytes = &self.mmap[data_end..crc_end];
let stored_crc = u32::from_le_bytes([
stored_crc_bytes[0], stored_crc_bytes[1],
stored_crc_bytes[2], stored_crc_bytes[3]
]);
let calculated_crc = crc32fast::hash(event_data);
if stored_crc != calculated_crc {
return Err(StorageError::CrcMismatch {
offset: current_position as u64,
expected: calculated_crc,
actual: stored_crc,
});
}
match bincode::deserialize::<Event>(event_data) {
Ok(event) => {
events.push(event);
current_position = crc_end;
current_entry_index += 1;
}
Err(_) => {
return Err(StorageError::invalid_format("Failed to deserialize event data during read"));
}
}
}
Ok(events)
}
pub fn flush(&mut self) -> StorageResult<()> {
if !self.read_only {
self.mmap.flush()?;
}
Ok(())
}
pub fn stats(&self) -> SegmentStats {
SegmentStats {
id: self.id,
file_path: self.file_path.clone(),
total_size: self.header.size,
used_size: self.header.write_offset,
entry_count: self.header.entry_count,
read_only: self.read_only,
}
}
pub fn append_batch_data(&mut self, batch_data: &[u8], event_sizes: &[usize]) -> StorageResult<Offset> {
if self.read_only {
return Err(StorageError::invalid_format("Cannot write to read-only segment"));
}
let start_entry_index = self.header.entry_count;
let current_position = self.header.write_offset as usize;
let required_space = batch_data.len() + (event_sizes.len() * 8); if current_position + required_space > self.mmap.len() {
return Err(StorageError::InsufficientSpace {
required: required_space as u64,
available: (self.mmap.len() - current_position) as u64,
});
}
let mut write_position = current_position;
let mut data_offset = 0;
for &event_size in event_sizes {
let size_bytes = (event_size as u32).to_le_bytes();
self.mmap[write_position..write_position + 4].copy_from_slice(&size_bytes);
write_position += 4;
let event_data = &batch_data[data_offset..data_offset + event_size];
self.mmap[write_position..write_position + event_size].copy_from_slice(event_data);
write_position += event_size;
let crc = crc32fast::hash(event_data);
let crc_bytes = crc.to_le_bytes();
self.mmap[write_position..write_position + 4].copy_from_slice(&crc_bytes);
write_position += 4;
data_offset += event_size;
}
self.header.write_offset = write_position as u64;
self.header.entry_count += event_sizes.len() as u64;
let header_bytes = bincode::serialize(&self.header)?;
let header_size = 256;
if header_bytes.len() > header_size {
return Err(StorageError::configuration("Header size exceeds reserved space"));
}
self.mmap[..header_bytes.len()].copy_from_slice(&header_bytes);
for i in header_bytes.len()..header_size {
self.mmap[i] = 0;
}
tracing::debug!(
"Batch written: {} events, {} bytes, start_offset={}",
event_sizes.len(),
batch_data.len(),
start_entry_index
);
Ok(Offset::new(start_entry_index))
}
}
#[derive(Debug, Clone)]
pub struct SegmentStats {
pub id: SegmentId,
pub file_path: PathBuf,
pub total_size: u64,
pub used_size: u64,
pub entry_count: u64,
pub read_only: bool,
}
#[derive(Debug)]
pub struct SegmentManager {
data_dir: PathBuf,
segments: RwLock<HashMap<SegmentId, Arc<RwLock<Segment>>>>,
current_segment: RwLock<Option<SegmentId>>,
segment_size: u64,
}
impl SegmentManager {
pub async fn new<P: AsRef<Path>>(data_dir: P) -> StorageResult<Self> {
let data_dir = data_dir.as_ref().to_path_buf();
let segment_size = 1024 * 1024 * 1024;
let manager = Self {
data_dir,
segments: RwLock::new(HashMap::new()),
current_segment: RwLock::new(None),
segment_size,
};
manager.load_existing_segments().await?;
Ok(manager)
}
async fn load_existing_segments(&self) -> StorageResult<()> {
if !self.data_dir.exists() {
tokio::fs::create_dir_all(&self.data_dir).await?;
return Ok(());
}
let mut segments = self.segments.write().await;
let mut entries = tokio::fs::read_dir(&self.data_dir).await?;
while let Some(entry) = entries.next_entry().await? {
let path = entry.path();
if path.extension().and_then(|s| s.to_str()) == Some("seg") {
match Segment::open(&path, false) {
Ok(segment) => {
let segment_id = segment.id;
segments.insert(segment_id, Arc::new(RwLock::new(segment)));
}
Err(e) => {
tracing::warn!("Failed to load segment {:?}: {}", path, e);
}
}
}
}
Ok(())
}
pub async fn create_segment(&self) -> StorageResult<SegmentId> {
let segment_id = SegmentId::new();
let file_path = self.data_dir.join(format!("{}.seg", segment_id));
let segment = Segment::create(file_path, self.segment_size)?;
let segment_id = segment.id;
let mut segments = self.segments.write().await;
segments.insert(segment_id, Arc::new(RwLock::new(segment)));
Ok(segment_id)
}
pub async fn get_segment(&self, segment_id: SegmentId) -> Option<Arc<RwLock<Segment>>> {
let segments = self.segments.read().await;
segments.get(&segment_id).cloned()
}
pub async fn get_all_stats(&self) -> Vec<SegmentStats> {
let segments = self.segments.read().await;
let mut stats = Vec::new();
for segment_arc in segments.values() {
if let Ok(segment) = segment_arc.try_read() {
stats.push(segment.stats());
}
}
stats
}
pub async fn list_segments(&self) -> Vec<SegmentId> {
let segments = self.segments.read().await;
segments.keys().copied().collect()
}
pub async fn get_segments_stats(&self) -> Vec<SegmentStats> {
let segments = self.segments.read().await;
let mut stats = Vec::new();
for segment_arc in segments.values() {
let segment = segment_arc.read().await;
stats.push(segment.stats());
}
stats
}
pub async fn process_segments_parallel<T, F, Fut>(
&self,
processor: F,
) -> StorageResult<Vec<T>>
where
T: Send + 'static,
F: Fn(SegmentId, Arc<RwLock<Segment>>) -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = StorageResult<T>> + Send,
{
let segments = self.segments.read().await;
let processor = Arc::new(processor);
let mut handles = Vec::new();
for (&segment_id, segment_arc) in segments.iter() {
let processor = Arc::clone(&processor);
let segment_arc = Arc::clone(segment_arc);
let handle = tokio::spawn(async move {
processor(segment_id, segment_arc).await
});
handles.push(handle);
}
let mut results = Vec::new();
for handle in handles {
match handle.await {
Ok(result) => results.push(result?),
Err(e) => return Err(StorageError::internal(format!("Segment processing failed: {}", e))),
}
}
Ok(results)
}
}