Skip to main content

hotfix_store/
file.rs

1//! A file-based message store for persistence.
2
3use crate::MessageStore;
4use crate::error::{Result, StoreError};
5use chrono::{DateTime, Utc};
6use std::collections::HashMap;
7use std::fs::{File, OpenOptions};
8use std::io::{BufRead, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
9use std::path::{Path, PathBuf};
10
11/// Metadata for a stored message, tracking its position and size in the body file
12#[derive(Debug, Clone)]
13struct MessageDef {
14    offset: u64,
15    size: usize,
16}
17
18/// File-based message store implementation.
19///
20/// Uses multiple files for storage:
21/// - `.body`: Append-only file containing raw message data
22/// - `.header`: Index file mapping sequence numbers to message positions
23/// - `.seqnums`: Stores sender and target sequence numbers
24/// - `.session`: Stores session creation time
25pub struct FileStore {
26    base_path: PathBuf,
27    body_file: BufWriter<File>,
28    header_file: BufWriter<File>,
29    seqnums_file: File,
30    sender_seq_number: u64,
31    target_seq_number: u64,
32    creation_time: DateTime<Utc>,
33    message_index: HashMap<u64, MessageDef>,
34    current_body_offset: u64,
35}
36
37impl FileStore {
38    /// Creates a new file store in the specified directory with the given name.
39    ///
40    /// # Errors
41    ///
42    /// Returns `StoreError::Initialization` if the store files cannot be created or read.
43    pub fn new(directory: impl AsRef<Path>, name: &str) -> Result<Self> {
44        Self::new_inner(directory.as_ref(), name).map_err(|e| StoreError::Initialization(e.into()))
45    }
46
47    fn new_inner(directory: &Path, name: &str) -> std::io::Result<Self> {
48        let base_path = directory.join(name);
49        std::fs::create_dir_all(directory)?;
50
51        let body_path = base_path.with_extension("body");
52        let header_path = base_path.with_extension("header");
53        let seqnums_path = base_path.with_extension("seqnums");
54
55        let creation_time = Self::get_or_create_session_time(&base_path)?;
56        let (sender_seq_number, target_seq_number) = Self::read_initial_seqnums(&base_path)?;
57
58        // open or create body and header files
59        let body_file = OpenOptions::new()
60            .create(true)
61            .append(true)
62            .open(&body_path)?;
63        let current_body_offset = body_file.metadata()?.len();
64        let body_file = BufWriter::new(body_file);
65
66        let header_file = OpenOptions::new()
67            .create(true)
68            .append(true)
69            .open(&header_path)?;
70        let header_file = BufWriter::new(header_file);
71
72        let seqnums_file = OpenOptions::new()
73            .create(true)
74            .truncate(true)
75            .read(true)
76            .write(true)
77            .open(&seqnums_path)?;
78
79        // load existing message index from header file
80        let message_index = Self::load_message_index(&header_path)?;
81
82        Ok(Self {
83            base_path,
84            body_file,
85            header_file,
86            seqnums_file,
87            sender_seq_number,
88            target_seq_number,
89            creation_time,
90            message_index,
91            current_body_offset,
92        })
93    }
94
95    /// Retrieves the session creation time from the session file.
96    ///
97    /// It initialises the session file if it doesn't exist.
98    fn get_or_create_session_time(base_path: &Path) -> std::io::Result<DateTime<Utc>> {
99        let session_path = base_path.with_extension("session");
100        let session_time = if session_path.exists() {
101            let content = std::fs::read_to_string(&session_path)?;
102            content.trim().parse::<DateTime<Utc>>().map_err(|e| {
103                std::io::Error::new(
104                    std::io::ErrorKind::InvalidData,
105                    format!("failed to parse session time: {e}"),
106                )
107            })?
108        } else {
109            let now = Utc::now();
110            std::fs::write(&session_path, now.to_rfc3339())?;
111            now
112        };
113
114        Ok(session_time)
115    }
116
117    /// Retrieves the sequence numbers from the seqnums file.
118    ///
119    /// It defaults to `(0, 0)` if the file doesn't exist or if it's empty.
120    fn read_initial_seqnums(base_path: &Path) -> std::io::Result<(u64, u64)> {
121        let seqnums_path = base_path.with_extension("seqnums");
122        let (sender_seq_number, target_seq_number) = if seqnums_path.exists() {
123            let content = std::fs::read_to_string(&seqnums_path).map_err(|e| {
124                std::io::Error::new(
125                    std::io::ErrorKind::InvalidData,
126                    format!("failed to read seqnums file: {e}"),
127                )
128            })?;
129            if content.trim().is_empty() {
130                (0u64, 0u64)
131            } else {
132                Self::parse_seqnums(&content)?
133            }
134        } else {
135            (0u64, 0u64)
136        };
137
138        Ok((sender_seq_number, target_seq_number))
139    }
140
141    fn parse_seqnums(content: &str) -> std::io::Result<(u64, u64)> {
142        let parts: Vec<&str> = content.trim().split(':').map(|s| s.trim()).collect();
143        if parts.len() != 2 {
144            return Err(std::io::Error::new(
145                std::io::ErrorKind::InvalidData,
146                "invalid seqnums format",
147            ));
148        }
149        let sender = parts[0].parse::<u64>().map_err(|e| {
150            std::io::Error::new(
151                std::io::ErrorKind::InvalidData,
152                format!("failed to parse sender sequence number: {e}"),
153            )
154        })?;
155        let target = parts[1].parse::<u64>().map_err(|e| {
156            std::io::Error::new(
157                std::io::ErrorKind::InvalidData,
158                format!("failed to parse target sequence number: {e}"),
159            )
160        })?;
161        Ok((sender, target))
162    }
163
164    fn load_message_index(header_path: &Path) -> std::io::Result<HashMap<u64, MessageDef>> {
165        let mut index = HashMap::new();
166
167        if !header_path.exists() {
168            return Ok(index);
169        }
170
171        let file = File::open(header_path)?;
172        let reader = BufReader::new(file);
173
174        for line in reader.lines() {
175            let line = line?;
176            let parts: Vec<&str> = line.trim().split(',').collect();
177            if parts.len() != 3 {
178                continue;
179            }
180
181            if let (Ok(seq_num), Ok(offset), Ok(size)) = (
182                parts[0].parse::<u64>(),
183                parts[1].parse::<u64>(),
184                parts[2].parse::<usize>(),
185            ) {
186                index.insert(seq_num, MessageDef { offset, size });
187            }
188        }
189
190        Ok(index)
191    }
192
193    fn write_seqnums_with(&mut self, sender: u64, target: u64) -> std::io::Result<()> {
194        self.seqnums_file.seek(SeekFrom::Start(0))?;
195        self.seqnums_file.set_len(0)?;
196        write!(self.seqnums_file, "{:020} : {:020}", sender, target)?;
197        self.seqnums_file.flush()?;
198        Ok(())
199    }
200
201    fn write_message(&mut self, sequence_number: u64, message: &[u8]) -> std::io::Result<()> {
202        let msg_size = message.len();
203        let offset = self.current_body_offset;
204
205        // write the message itself
206        self.body_file.write_all(message)?;
207        self.body_file.flush()?;
208
209        // write the offset to the header file
210        writeln!(
211            self.header_file,
212            "{},{},{}",
213            sequence_number, offset, msg_size
214        )?;
215        self.header_file.flush()?;
216
217        self.message_index.insert(
218            sequence_number,
219            MessageDef {
220                offset,
221                size: msg_size,
222            },
223        );
224        self.current_body_offset += msg_size as u64;
225
226        Ok(())
227    }
228
229    fn perform_reset(&mut self) -> std::io::Result<()> {
230        self.body_file.flush()?;
231        self.header_file.flush()?;
232
233        // remove all files
234        let body_path = self.base_path.with_extension("body");
235        let header_path = self.base_path.with_extension("header");
236        let seqnums_path = self.base_path.with_extension("seqnums");
237        let session_path = self.base_path.with_extension("session");
238
239        if body_path.exists() {
240            std::fs::remove_file(&body_path)?;
241        }
242        if header_path.exists() {
243            std::fs::remove_file(&header_path)?;
244        }
245        if seqnums_path.exists() {
246            std::fs::remove_file(&seqnums_path)?;
247        }
248        if session_path.exists() {
249            std::fs::remove_file(&session_path)?;
250        }
251
252        // reset in-memory state
253        self.sender_seq_number = 0;
254        self.target_seq_number = 0;
255        self.creation_time = Utc::now();
256        self.message_index.clear();
257        self.current_body_offset = 0;
258
259        // recreate files
260        let now = Utc::now();
261        std::fs::write(&session_path, now.to_rfc3339())?;
262
263        let body_file = OpenOptions::new()
264            .create(true)
265            .append(true)
266            .open(&body_path)?;
267        self.body_file = BufWriter::new(body_file);
268
269        let header_file = OpenOptions::new()
270            .create(true)
271            .append(true)
272            .open(&header_path)?;
273        self.header_file = BufWriter::new(header_file);
274
275        self.seqnums_file = OpenOptions::new()
276            .create(true)
277            .truncate(true)
278            .read(true)
279            .write(true)
280            .open(&seqnums_path)?;
281
282        self.creation_time = now;
283
284        Ok(())
285    }
286
287    fn read_messages(&self, begin: usize, end: usize) -> std::io::Result<Vec<Vec<u8>>> {
288        let mut messages = Vec::with_capacity(end - begin + 1);
289
290        let body_path = self.base_path.with_extension("body");
291        let mut body_file = File::open(body_path)?;
292
293        for seq_num in begin..=end {
294            if let Some(msg_def) = self.message_index.get(&(seq_num as u64)) {
295                body_file.seek(SeekFrom::Start(msg_def.offset))?;
296
297                let mut buffer = vec![0u8; msg_def.size];
298                body_file.read_exact(&mut buffer)?;
299
300                messages.push(buffer);
301            }
302        }
303
304        Ok(messages)
305    }
306}
307
308#[async_trait::async_trait]
309impl MessageStore for FileStore {
310    async fn add(&mut self, sequence_number: u64, message: &[u8]) -> Result<()> {
311        self.write_message(sequence_number, message)
312            .map_err(|err| StoreError::PersistMessage {
313                sequence_number,
314                source: err.into(),
315            })
316    }
317
318    async fn get_slice(&self, begin: usize, end: usize) -> Result<Vec<Vec<u8>>> {
319        self.read_messages(begin, end)
320            .map_err(|e| StoreError::RetrieveMessages {
321                begin,
322                end,
323                source: e.into(),
324            })
325    }
326
327    fn next_sender_seq_number(&self) -> u64 {
328        self.sender_seq_number + 1
329    }
330
331    fn next_target_seq_number(&self) -> u64 {
332        self.target_seq_number + 1
333    }
334
335    async fn increment_sender_seq_number(&mut self) -> Result<()> {
336        let new_value = self.sender_seq_number + 1;
337        self.write_seqnums_with(new_value, self.target_seq_number)
338            .map_err(|e| StoreError::UpdateSequenceNumber(e.into()))?;
339        self.sender_seq_number = new_value;
340        Ok(())
341    }
342
343    async fn increment_target_seq_number(&mut self) -> Result<()> {
344        let new_value = self.target_seq_number + 1;
345        self.write_seqnums_with(self.sender_seq_number, new_value)
346            .map_err(|e| StoreError::UpdateSequenceNumber(e.into()))?;
347        self.target_seq_number = new_value;
348        Ok(())
349    }
350
351    async fn set_target_seq_number(&mut self, seq_number: u64) -> Result<()> {
352        self.write_seqnums_with(self.sender_seq_number, seq_number)
353            .map_err(|e| StoreError::UpdateSequenceNumber(e.into()))?;
354        self.target_seq_number = seq_number;
355        Ok(())
356    }
357
358    async fn reset(&mut self) -> Result<()> {
359        self.perform_reset()
360            .map_err(|e| StoreError::Reset(e.into()))
361    }
362
363    fn creation_time(&self) -> DateTime<Utc> {
364        self.creation_time
365    }
366}