kcode_k1_transaction_store/
lib.rs

1use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
2use sha2::{Digest, Sha256};
3use std::collections::HashMap;
4use std::fs::{self, File, OpenOptions};
5use std::io::{self, Write};
6use std::os::unix::fs::FileExt;
7use std::path::{Path, PathBuf};
8use std::sync::{Mutex, RwLock};
9
10pub const TX_ID_BYTES: usize = 12;
11pub const SECTOR_BYTES: u64 = 4_096;
12pub const INLINE_LIMIT: usize = 262_144;
13const PAYLOAD_ALPHABET: &[u8; 64] =
14    b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
15
16#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
17pub struct TxId([u8; TX_ID_BYTES]);
18
19impl TxId {
20    pub const fn from_bytes(bytes: [u8; TX_ID_BYTES]) -> Self {
21        Self(bytes)
22    }
23
24    pub const fn as_bytes(&self) -> &[u8; TX_ID_BYTES] {
25        &self.0
26    }
27
28    pub const fn into_bytes(self) -> [u8; TX_ID_BYTES] {
29        self.0
30    }
31
32    pub fn for_transaction(transaction: &[u8]) -> Self {
33        let digest = Sha256::digest(transaction);
34        let mut bytes = [0; TX_ID_BYTES];
35        bytes.copy_from_slice(&digest[..TX_ID_BYTES]);
36        Self(bytes)
37    }
38}
39
40#[derive(Clone, Copy, Debug, Eq, PartialEq)]
41pub enum PutOutcome {
42    Inserted(TxId),
43    Duplicate(TxId),
44}
45
46#[derive(Debug)]
47pub enum StoreError {
48    Io(io::Error),
49    AlreadyExists,
50    InvalidStore,
51    StoreFull,
52    IdCollision(TxId),
53    OutcomeUnknown(TxId),
54    ReopenRequired,
55    CorruptTransaction(TxId),
56}
57
58impl std::fmt::Display for StoreError {
59    fn fmt(&self, output: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60        match self {
61            Self::Io(error) => error.fmt(output),
62            Self::AlreadyExists => output.write_str("the store already exists"),
63            Self::InvalidStore => output.write_str("the store is invalid"),
64            Self::StoreFull => output.write_str("the transaction file is full"),
65            Self::IdCollision(id) => write!(output, "transaction identifier collision: {id:?}"),
66            Self::OutcomeUnknown(id) => write!(output, "transaction outcome is unknown: {id:?}"),
67            Self::ReopenRequired => output.write_str("the store must be reopened before writing"),
68            Self::CorruptTransaction(id) => write!(output, "transaction is corrupt: {id:?}"),
69        }
70    }
71}
72
73impl std::error::Error for StoreError {}
74
75impl From<io::Error> for StoreError {
76    fn from(error: io::Error) -> Self {
77        Self::Io(error)
78    }
79}
80
81pub struct TransactionStore {
82    data: File,
83    payload: PathBuf,
84    locations: RwLock<HashMap<TxId, u32>>,
85    next_sector: Mutex<u64>,
86    publication: Mutex<Publication>,
87}
88
89struct Publication {
90    lookup: File,
91    reopen_required: bool,
92}
93
94impl TransactionStore {
95    pub fn create(root: &Path) -> Result<Self, StoreError> {
96        match fs::create_dir(root) {
97            Ok(()) => {}
98            Err(error) if error.kind() == io::ErrorKind::AlreadyExists => {
99                return Err(StoreError::AlreadyExists);
100            }
101            Err(error) => return Err(error.into()),
102        }
103        let payload = root.join("payload");
104        fs::create_dir(&payload)?;
105        for first in PAYLOAD_ALPHABET {
106            for second in PAYLOAD_ALPHABET {
107                fs::create_dir(payload.join(format!(
108                    "{}{}",
109                    char::from(*first),
110                    char::from(*second)
111                )))?;
112            }
113        }
114        File::open(&payload)?.sync_all()?;
115        File::create(root.join("transactions.dat"))?.sync_all()?;
116        File::create(root.join("lookup.dat"))?.sync_all()?;
117        File::open(root)?.sync_all()?;
118        let parent = root
119            .parent()
120            .filter(|path| !path.as_os_str().is_empty())
121            .unwrap_or(Path::new("."));
122        File::open(parent)?.sync_all()?;
123        Self::open(root)
124    }
125
126    pub fn open(root: &Path) -> Result<Self, StoreError> {
127        let data_path = root.join("transactions.dat");
128        let lookup_path = root.join("lookup.dat");
129        let payload = root.join("payload");
130        if !root.is_dir() || !data_path.is_file() || !lookup_path.is_file() || !payload.is_dir() {
131            return Err(StoreError::InvalidStore);
132        }
133        let mut lookup_bytes = fs::read(&lookup_path)?;
134        let complete = lookup_bytes.len() / 16 * 16;
135        if complete != lookup_bytes.len() {
136            let lookup = OpenOptions::new().write(true).open(&lookup_path)?;
137            lookup.set_len(complete as u64)?;
138            lookup.sync_data()?;
139            lookup_bytes.truncate(complete);
140        }
141        let mut locations = HashMap::with_capacity(lookup_bytes.len() / 16);
142        for entry in lookup_bytes.chunks_exact(16) {
143            let mut id = [0; TX_ID_BYTES];
144            let mut sector = [0; 4];
145            id.copy_from_slice(&entry[..12]);
146            sector.copy_from_slice(&entry[12..]);
147            if locations
148                .insert(TxId::from_bytes(id), u32::from_le_bytes(sector))
149                .is_some()
150            {
151                return Err(StoreError::InvalidStore);
152            }
153        }
154        let data = OpenOptions::new().read(true).write(true).open(data_path)?;
155        let next_sector = data.metadata()?.len().div_ceil(SECTOR_BYTES);
156        if next_sector > u32::MAX as u64 + 1 {
157            return Err(StoreError::StoreFull);
158        }
159        let lookup = OpenOptions::new().append(true).open(lookup_path)?;
160        Ok(Self {
161            data,
162            payload,
163            locations: RwLock::new(locations),
164            next_sector: Mutex::new(next_sector),
165            publication: Mutex::new(Publication {
166                lookup,
167                reopen_required: false,
168            }),
169        })
170    }
171
172    pub fn put(&self, transaction: &[u8]) -> Result<PutOutcome, StoreError> {
173        let id = TxId::for_transaction(transaction);
174        if let Some(sector) = self.location(id) {
175            return self.compare(id, sector, transaction);
176        }
177        let external = transaction.len() > INLINE_LIMIT;
178        if external {
179            self.write_payload(id, transaction)?;
180        }
181        let sectors = if external {
182            1
183        } else {
184            transaction
185                .len()
186                .checked_add(9)
187                .ok_or(StoreError::StoreFull)?
188                .div_ceil(SECTOR_BYTES as usize) as u64
189        };
190        let sector = self.allocate(sectors)?;
191        let mut record = vec![0; sectors as usize * SECTOR_BYTES as usize];
192        record[0] = u8::from(external);
193        record[1..9].copy_from_slice(&(transaction.len() as u64).to_le_bytes());
194        if !external {
195            record[9..9 + transaction.len()].copy_from_slice(transaction);
196        }
197        self.data
198            .write_all_at(&record, sector as u64 * SECTOR_BYTES)?;
199        self.data.sync_data()?;
200        self.publish(id, sector, transaction)
201    }
202
203    pub fn contains(&self, id: TxId) -> bool {
204        self.location(id).is_some()
205    }
206
207    pub fn get(&self, id: TxId) -> Result<Option<Vec<u8>>, StoreError> {
208        let Some(sector) = self.location(id) else {
209            return Ok(None);
210        };
211        self.read_transaction(id, sector).map(Some)
212    }
213
214    fn location(&self, id: TxId) -> Option<u32> {
215        self.locations
216            .read()
217            .unwrap_or_else(|error| error.into_inner())
218            .get(&id)
219            .copied()
220    }
221
222    fn allocate(&self, sectors: u64) -> Result<u32, StoreError> {
223        let mut next = self
224            .next_sector
225            .lock()
226            .unwrap_or_else(|error| error.into_inner());
227        let end = next.checked_add(sectors).ok_or(StoreError::StoreFull)?;
228        if end > u32::MAX as u64 + 1 {
229            return Err(StoreError::StoreFull);
230        }
231        let sector = *next as u32;
232        *next = end;
233        Ok(sector)
234    }
235
236    fn write_payload(&self, id: TxId, transaction: &[u8]) -> Result<(), StoreError> {
237        let path = self.payload_path(id);
238        let shard = path.parent().ok_or(StoreError::InvalidStore)?.to_path_buf();
239        let file = OpenOptions::new()
240            .create(true)
241            .truncate(false)
242            .write(true)
243            .open(path)?;
244        file.write_all_at(transaction, 0)?;
245        file.set_len(transaction.len() as u64)?;
246        file.sync_all()?;
247        File::open(shard)?.sync_all()?;
248        Ok(())
249    }
250
251    fn payload_path(&self, id: TxId) -> PathBuf {
252        let encoded = URL_SAFE_NO_PAD.encode(id.as_bytes());
253        self.payload
254            .join(&encoded[..2])
255            .join(format!("{}.dat", &encoded[2..]))
256    }
257
258    fn publish(&self, id: TxId, sector: u32, transaction: &[u8]) -> Result<PutOutcome, StoreError> {
259        let mut publication = self
260            .publication
261            .lock()
262            .unwrap_or_else(|error| error.into_inner());
263        if publication.reopen_required {
264            return Err(StoreError::ReopenRequired);
265        }
266        if let Some(existing) = self.location(id) {
267            drop(publication);
268            return self.compare(id, existing, transaction);
269        }
270        let mut entry = [0; 16];
271        entry[..12].copy_from_slice(id.as_bytes());
272        entry[12..].copy_from_slice(&sector.to_le_bytes());
273        if publication
274            .lookup
275            .write_all(&entry)
276            .and_then(|()| publication.lookup.sync_data())
277            .is_err()
278        {
279            publication.reopen_required = true;
280            return Err(StoreError::OutcomeUnknown(id));
281        }
282        self.locations
283            .write()
284            .unwrap_or_else(|error| error.into_inner())
285            .insert(id, sector);
286        Ok(PutOutcome::Inserted(id))
287    }
288
289    fn compare(&self, id: TxId, sector: u32, transaction: &[u8]) -> Result<PutOutcome, StoreError> {
290        if self.read_transaction(id, sector)? == transaction {
291            Ok(PutOutcome::Duplicate(id))
292        } else {
293            Err(StoreError::IdCollision(id))
294        }
295    }
296
297    fn read_transaction(&self, id: TxId, sector: u32) -> Result<Vec<u8>, StoreError> {
298        let offset = sector as u64 * SECTOR_BYTES;
299        let mut header = [0; 9];
300        self.read_data(&mut header, offset, id)?;
301        let length = u64::from_le_bytes(header[1..9].try_into().unwrap());
302        let bytes = match header[0] {
303            0 if length <= INLINE_LIMIT as u64 => {
304                let mut bytes = vec![0; length as usize];
305                self.read_data(&mut bytes, offset + 9, id)?;
306                bytes
307            }
308            1 if length > INLINE_LIMIT as u64 => {
309                let path = self.payload_path(id);
310                let metadata = fs::metadata(&path).map_err(|error| match error.kind() {
311                    io::ErrorKind::NotFound => StoreError::CorruptTransaction(id),
312                    _ => StoreError::Io(error),
313                })?;
314                if metadata.len() != length {
315                    return Err(StoreError::CorruptTransaction(id));
316                }
317                let bytes = fs::read(path)?;
318                if bytes.len() as u64 != length {
319                    return Err(StoreError::CorruptTransaction(id));
320                }
321                bytes
322            }
323            _ => return Err(StoreError::CorruptTransaction(id)),
324        };
325        if TxId::for_transaction(&bytes) != id {
326            return Err(StoreError::CorruptTransaction(id));
327        }
328        Ok(bytes)
329    }
330
331    fn read_data(&self, bytes: &mut [u8], offset: u64, id: TxId) -> Result<(), StoreError> {
332        self.data
333            .read_exact_at(bytes, offset)
334            .map_err(|error| match error.kind() {
335                io::ErrorKind::UnexpectedEof => StoreError::CorruptTransaction(id),
336                _ => StoreError::Io(error),
337            })
338    }
339}
340
341#[cfg(test)]
342mod tests {
343    use super::*;
344    use std::sync::Arc;
345    use std::sync::atomic::{AtomicU64, Ordering};
346
347    fn root() -> PathBuf {
348        static NEXT: AtomicU64 = AtomicU64::new(0);
349        let path = std::env::temp_dir().join(format!(
350            "k1-store-{}-{}",
351            std::process::id(),
352            NEXT.fetch_add(1, Ordering::Relaxed)
353        ));
354        let _ = fs::remove_dir_all(&path);
355        path
356    }
357
358    #[test]
359    fn lifecycle_boundaries_duplicates_corruption_and_tail() {
360        let root = root();
361        let store = TransactionStore::create(&root).unwrap();
362        assert_eq!(fs::read_dir(root.join("payload")).unwrap().count(), 4_096);
363        let values = [vec![], vec![3; INLINE_LIMIT], vec![7; INLINE_LIMIT + 1]];
364        let mut ids = Vec::new();
365        for value in &values {
366            let id = TxId::for_transaction(value);
367            assert_eq!(store.put(value).unwrap(), PutOutcome::Inserted(id));
368            assert_eq!(store.put(value).unwrap(), PutOutcome::Duplicate(id));
369            assert_eq!(store.get(id).unwrap().unwrap(), *value);
370            ids.push(id);
371        }
372        assert_eq!(fs::metadata(root.join("lookup.dat")).unwrap().len(), 48);
373        let sector = store.location(ids[0]).unwrap();
374        assert!(matches!(
375            store.compare(ids[0], sector, b"different"),
376            Err(StoreError::IdCollision(id)) if id == ids[0]
377        ));
378        drop(store);
379        OpenOptions::new()
380            .append(true)
381            .open(root.join("lookup.dat"))
382            .unwrap()
383            .write_all(&[1, 2, 3])
384            .unwrap();
385        let store = TransactionStore::open(&root).unwrap();
386        assert_eq!(fs::metadata(root.join("lookup.dat")).unwrap().len(), 48);
387        for (id, value) in ids.into_iter().zip(values) {
388            assert!(store.contains(id));
389            assert_eq!(store.get(id).unwrap().unwrap(), value);
390        }
391        let id = TxId::for_transaction(b"intact");
392        assert_eq!(store.put(b"intact").unwrap(), PutOutcome::Inserted(id));
393        let sector = store.location(id).unwrap();
394        store
395            .data
396            .write_all_at(b"x", sector as u64 * SECTOR_BYTES + 9)
397            .unwrap();
398        assert!(matches!(
399            store.get(id),
400            Err(StoreError::CorruptTransaction(found)) if found == id
401        ));
402        fs::remove_dir_all(root).unwrap();
403    }
404
405    #[test]
406    fn concurrent_writes_allocate_disjoint_sectors_and_publish_once() {
407        let root = root();
408        let store = Arc::new(TransactionStore::create(&root).unwrap());
409        let threads: Vec<_> = (0_u8..8)
410            .map(|byte| {
411                let store = Arc::clone(&store);
412                std::thread::spawn(move || {
413                    let value = vec![byte; 8_193];
414                    let outcome = store.put(&value).unwrap();
415                    (value, outcome)
416                })
417            })
418            .collect();
419        for thread in threads {
420            let (value, outcome) = thread.join().unwrap();
421            let PutOutcome::Inserted(id) = outcome else {
422                panic!()
423            };
424            assert_eq!(store.get(id).unwrap().unwrap(), value);
425        }
426        let value = vec![11; INLINE_LIMIT + 1];
427        let id = TxId::for_transaction(&value);
428        let threads: Vec<_> = (0..2)
429            .map(|_| {
430                let store = Arc::clone(&store);
431                let value = value.clone();
432                std::thread::spawn(move || store.put(&value))
433            })
434            .collect();
435        let outcomes: Vec<_> = threads
436            .into_iter()
437            .map(|thread| thread.join().unwrap())
438            .collect();
439        assert_eq!(
440            outcomes
441                .iter()
442                .filter(|outcome| matches!(outcome, Ok(PutOutcome::Inserted(_))))
443                .count(),
444            1
445        );
446        assert!(outcomes.iter().all(|outcome| matches!(
447            outcome,
448            Ok(PutOutcome::Inserted(_)) | Ok(PutOutcome::Duplicate(_))
449        )));
450        assert_eq!(store.get(id).unwrap().unwrap(), value);
451        assert_eq!(store.put(&value).unwrap(), PutOutcome::Duplicate(id));
452        let sectors = store
453            .locations
454            .read()
455            .unwrap_or_else(|error| error.into_inner());
456        let mut locations: Vec<_> = sectors.values().copied().collect();
457        locations.sort_unstable();
458        locations.dedup();
459        assert_eq!(locations.len(), sectors.len());
460        fs::remove_dir_all(root).unwrap();
461    }
462
463    #[test]
464    fn capacity_is_bounded() {
465        let root = root();
466        let store = TransactionStore::create(&root).unwrap();
467        *store
468            .next_sector
469            .lock()
470            .unwrap_or_else(|error| error.into_inner()) = u32::MAX as u64 + 1;
471        assert!(matches!(store.put(b"full"), Err(StoreError::StoreFull)));
472        fs::remove_dir_all(root).unwrap();
473    }
474}
kcode_k1_transaction_store/lib.rs

kcode_k1_transaction_store/
lib.rs
