use crate::MessageStore;
use crate::error::{Result, StoreError};
use chrono::{DateTime, Utc};
use std::collections::HashMap;
use std::fs::{File, OpenOptions};
use std::io::{BufRead, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
#[derive(Debug, Clone)]
struct MessageDef {
offset: u64,
size: usize,
}
pub struct FileStore {
base_path: PathBuf,
body_file: BufWriter<File>,
header_file: BufWriter<File>,
seqnums_file: File,
sender_seq_number: u64,
target_seq_number: u64,
creation_time: DateTime<Utc>,
message_index: HashMap<u64, MessageDef>,
current_body_offset: u64,
}
impl FileStore {
pub fn new(directory: impl AsRef<Path>, name: &str) -> Result<Self> {
Self::new_inner(directory.as_ref(), name).map_err(|e| StoreError::Initialization(e.into()))
}
fn new_inner(directory: &Path, name: &str) -> std::io::Result<Self> {
let base_path = directory.join(name);
std::fs::create_dir_all(directory)?;
let body_path = base_path.with_extension("body");
let header_path = base_path.with_extension("header");
let seqnums_path = base_path.with_extension("seqnums");
let creation_time = Self::get_or_create_session_time(&base_path)?;
let (sender_seq_number, target_seq_number) = Self::read_initial_seqnums(&base_path)?;
let body_file = OpenOptions::new()
.create(true)
.append(true)
.open(&body_path)?;
let current_body_offset = body_file.metadata()?.len();
let body_file = BufWriter::new(body_file);
let header_file = OpenOptions::new()
.create(true)
.append(true)
.open(&header_path)?;
let header_file = BufWriter::new(header_file);
let seqnums_file = OpenOptions::new()
.create(true)
.truncate(true)
.read(true)
.write(true)
.open(&seqnums_path)?;
let message_index = Self::load_message_index(&header_path)?;
Ok(Self {
base_path,
body_file,
header_file,
seqnums_file,
sender_seq_number,
target_seq_number,
creation_time,
message_index,
current_body_offset,
})
}
fn get_or_create_session_time(base_path: &Path) -> std::io::Result<DateTime<Utc>> {
let session_path = base_path.with_extension("session");
let session_time = if session_path.exists() {
let content = std::fs::read_to_string(&session_path)?;
content.trim().parse::<DateTime<Utc>>().map_err(|e| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("failed to parse session time: {e}"),
)
})?
} else {
let now = Utc::now();
std::fs::write(&session_path, now.to_rfc3339())?;
now
};
Ok(session_time)
}
fn read_initial_seqnums(base_path: &Path) -> std::io::Result<(u64, u64)> {
let seqnums_path = base_path.with_extension("seqnums");
let (sender_seq_number, target_seq_number) = if seqnums_path.exists() {
let content = std::fs::read_to_string(&seqnums_path).map_err(|e| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("failed to read seqnums file: {e}"),
)
})?;
if content.trim().is_empty() {
(0u64, 0u64)
} else {
Self::parse_seqnums(&content)?
}
} else {
(0u64, 0u64)
};
Ok((sender_seq_number, target_seq_number))
}
fn parse_seqnums(content: &str) -> std::io::Result<(u64, u64)> {
let parts: Vec<&str> = content.trim().split(':').map(|s| s.trim()).collect();
if parts.len() != 2 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"invalid seqnums format",
));
}
let sender = parts[0].parse::<u64>().map_err(|e| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("failed to parse sender sequence number: {e}"),
)
})?;
let target = parts[1].parse::<u64>().map_err(|e| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("failed to parse target sequence number: {e}"),
)
})?;
Ok((sender, target))
}
fn load_message_index(header_path: &Path) -> std::io::Result<HashMap<u64, MessageDef>> {
let mut index = HashMap::new();
if !header_path.exists() {
return Ok(index);
}
let file = File::open(header_path)?;
let reader = BufReader::new(file);
for line in reader.lines() {
let line = line?;
let parts: Vec<&str> = line.trim().split(',').collect();
if parts.len() != 3 {
continue;
}
if let (Ok(seq_num), Ok(offset), Ok(size)) = (
parts[0].parse::<u64>(),
parts[1].parse::<u64>(),
parts[2].parse::<usize>(),
) {
index.insert(seq_num, MessageDef { offset, size });
}
}
Ok(index)
}
fn write_seqnums_with(&mut self, sender: u64, target: u64) -> std::io::Result<()> {
self.seqnums_file.seek(SeekFrom::Start(0))?;
self.seqnums_file.set_len(0)?;
write!(self.seqnums_file, "{:020} : {:020}", sender, target)?;
self.seqnums_file.flush()?;
Ok(())
}
fn write_message(&mut self, sequence_number: u64, message: &[u8]) -> std::io::Result<()> {
let msg_size = message.len();
let offset = self.current_body_offset;
self.body_file.write_all(message)?;
self.body_file.flush()?;
writeln!(
self.header_file,
"{},{},{}",
sequence_number, offset, msg_size
)?;
self.header_file.flush()?;
self.message_index.insert(
sequence_number,
MessageDef {
offset,
size: msg_size,
},
);
self.current_body_offset += msg_size as u64;
Ok(())
}
fn perform_reset(&mut self) -> std::io::Result<()> {
self.body_file.flush()?;
self.header_file.flush()?;
let body_path = self.base_path.with_extension("body");
let header_path = self.base_path.with_extension("header");
let seqnums_path = self.base_path.with_extension("seqnums");
let session_path = self.base_path.with_extension("session");
if body_path.exists() {
std::fs::remove_file(&body_path)?;
}
if header_path.exists() {
std::fs::remove_file(&header_path)?;
}
if seqnums_path.exists() {
std::fs::remove_file(&seqnums_path)?;
}
if session_path.exists() {
std::fs::remove_file(&session_path)?;
}
self.sender_seq_number = 0;
self.target_seq_number = 0;
self.creation_time = Utc::now();
self.message_index.clear();
self.current_body_offset = 0;
let now = Utc::now();
std::fs::write(&session_path, now.to_rfc3339())?;
let body_file = OpenOptions::new()
.create(true)
.append(true)
.open(&body_path)?;
self.body_file = BufWriter::new(body_file);
let header_file = OpenOptions::new()
.create(true)
.append(true)
.open(&header_path)?;
self.header_file = BufWriter::new(header_file);
self.seqnums_file = OpenOptions::new()
.create(true)
.truncate(true)
.read(true)
.write(true)
.open(&seqnums_path)?;
self.creation_time = now;
Ok(())
}
fn read_messages(&self, begin: usize, end: usize) -> std::io::Result<Vec<Vec<u8>>> {
let mut messages = Vec::with_capacity(end - begin + 1);
let body_path = self.base_path.with_extension("body");
let mut body_file = File::open(body_path)?;
for seq_num in begin..=end {
if let Some(msg_def) = self.message_index.get(&(seq_num as u64)) {
body_file.seek(SeekFrom::Start(msg_def.offset))?;
let mut buffer = vec![0u8; msg_def.size];
body_file.read_exact(&mut buffer)?;
messages.push(buffer);
}
}
Ok(messages)
}
}
#[async_trait::async_trait]
impl MessageStore for FileStore {
async fn add(&mut self, sequence_number: u64, message: &[u8]) -> Result<()> {
self.write_message(sequence_number, message)
.map_err(|err| StoreError::PersistMessage {
sequence_number,
source: err.into(),
})
}
async fn get_slice(&self, begin: usize, end: usize) -> Result<Vec<Vec<u8>>> {
self.read_messages(begin, end)
.map_err(|e| StoreError::RetrieveMessages {
begin,
end,
source: e.into(),
})
}
fn next_sender_seq_number(&self) -> u64 {
self.sender_seq_number + 1
}
fn next_target_seq_number(&self) -> u64 {
self.target_seq_number + 1
}
async fn increment_sender_seq_number(&mut self) -> Result<()> {
let new_value = self.sender_seq_number + 1;
self.write_seqnums_with(new_value, self.target_seq_number)
.map_err(|e| StoreError::UpdateSequenceNumber(e.into()))?;
self.sender_seq_number = new_value;
Ok(())
}
async fn increment_target_seq_number(&mut self) -> Result<()> {
let new_value = self.target_seq_number + 1;
self.write_seqnums_with(self.sender_seq_number, new_value)
.map_err(|e| StoreError::UpdateSequenceNumber(e.into()))?;
self.target_seq_number = new_value;
Ok(())
}
async fn set_target_seq_number(&mut self, seq_number: u64) -> Result<()> {
self.write_seqnums_with(self.sender_seq_number, seq_number)
.map_err(|e| StoreError::UpdateSequenceNumber(e.into()))?;
self.target_seq_number = seq_number;
Ok(())
}
async fn reset(&mut self) -> Result<()> {
self.perform_reset()
.map_err(|e| StoreError::Reset(e.into()))
}
fn creation_time(&self) -> DateTime<Utc> {
self.creation_time
}
}