slate_kv/
file_counter.rs

1use crate::io_util::{read_exact_at, write_all_at};
2use hmac::{Hmac, KeyInit, Mac};
3use sha2::Sha256;
4use slate_kv_hal::{CounterKind, MonotonicCounter};
5use std::fs::File;
6#[cfg(target_os = "macos")]
7use std::os::unix::io::AsRawFd;
8
9#[cfg(target_os = "macos")]
10fn flush_durable(f: &File) -> std::io::Result<()> {
11    let rc = unsafe { libc::fcntl(f.as_raw_fd(), libc::F_FULLFSYNC) };
12    if rc == -1 {
13        Err(std::io::Error::last_os_error())
14    } else {
15        Ok(())
16    }
17}
18
19#[cfg(not(target_os = "macos"))]
20fn flush_durable(f: &File) -> std::io::Result<()> {
21    f.sync_data()
22}
23
24type HmacSha256 = Hmac<Sha256>;
25
26const SLOT_SIZE: usize = 40; // 8 bytes value + 32 bytes HMAC
27
28#[derive(Debug)]
29pub enum FileCounterError {
30    Io(std::io::Error),
31    Exhausted,
32    FormatError,
33}
34
35impl From<std::io::Error> for FileCounterError {
36    fn from(err: std::io::Error) -> Self {
37        FileCounterError::Io(err)
38    }
39}
40
41pub struct FileCounter {
42    file: File,
43    key: [u8; 32],
44    budget: u64,
45    val: u64,
46}
47
48impl FileCounter {
49    pub fn new(file: File, key: [u8; 32], budget: u64) -> Result<Self, FileCounterError> {
50        let meta = file.metadata()?;
51        if meta.len() != 80 {
52            file.set_len(80)?;
53            // Init both slots to 0
54            let mut c = Self {
55                file,
56                key,
57                budget,
58                val: 0,
59            };
60            c.write_slot(0, 0)?;
61            c.write_slot(1, 0)?;
62            return Ok(c);
63        }
64
65        let mut c = Self {
66            file,
67            key,
68            budget,
69            val: 0,
70        };
71        c.val = c.read_max()?;
72        Ok(c)
73    }
74
75    fn read_max(&self) -> Result<u64, FileCounterError> {
76        let v0 = self.read_slot(0);
77        let v1 = self.read_slot(1);
78        match (v0, v1) {
79            (Ok(a), Ok(b)) => Ok(a.max(b)),
80            (Ok(a), Err(_)) => Ok(a),
81            (Err(_), Ok(b)) => Ok(b),
82            (Err(e), Err(_)) => Err(e),
83        }
84    }
85
86    fn read_slot(&self, idx: u64) -> Result<u64, FileCounterError> {
87        let mut buf = [0u8; SLOT_SIZE];
88        read_exact_at(&self.file, &mut buf, idx * SLOT_SIZE as u64)?;
89
90        let mut val_bytes = [0u8; 8];
91        val_bytes.copy_from_slice(&buf[0..8]);
92        let val = u64::from_le_bytes(val_bytes);
93
94        let mut mac = HmacSha256::new_from_slice(&self.key).unwrap();
95        mac.update(&val_bytes);
96
97        if mac.verify_slice(&buf[8..40]).is_ok() {
98            Ok(val)
99        } else {
100            Err(FileCounterError::FormatError)
101        }
102    }
103
104    fn write_slot(&mut self, idx: u64, val: u64) -> Result<(), FileCounterError> {
105        let mut buf = [0u8; SLOT_SIZE];
106        let val_bytes = val.to_le_bytes();
107        buf[0..8].copy_from_slice(&val_bytes);
108
109        let mut mac = HmacSha256::new_from_slice(&self.key).unwrap();
110        mac.update(&val_bytes);
111        let tag = mac.finalize().into_bytes();
112        buf[8..40].copy_from_slice(&tag);
113
114        write_all_at(&self.file, &buf, idx * SLOT_SIZE as u64)?;
115        flush_durable(&self.file)?;
116        Ok(())
117    }
118}
119
120impl MonotonicCounter for FileCounter {
121    type Error = FileCounterError;
122
123    fn kind(&self) -> CounterKind {
124        CounterKind::BestEffort
125    }
126
127    fn read(&mut self) -> Result<u64, Self::Error> {
128        Ok(self.val)
129    }
130
131    fn increment(&mut self) -> Result<u64, Self::Error> {
132        if self.budget == 0 {
133            return Err(FileCounterError::Exhausted);
134        }
135
136        // Find which slot to overwrite: the older one
137        let v0 = self.read_slot(0).unwrap_or(0);
138        let v1 = self.read_slot(1).unwrap_or(0);
139
140        let target_idx = if v0 <= v1 { 0 } else { 1 };
141
142        self.val += 1;
143        self.budget -= 1;
144
145        self.write_slot(target_idx, self.val)?;
146        Ok(self.val)
147    }
148}
slate_kv/file_counter.rs

slate_kv/
file_counter.rs
