use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::fs::{self, File, OpenOptions};
use std::io::{self, Write};
use std::os::unix::fs::FileExt;
use std::path::{Path, PathBuf};
use std::sync::{Mutex, RwLock};
pub const TX_ID_BYTES: usize = 12;
pub const SECTOR_BYTES: u64 = 4_096;
pub const INLINE_LIMIT: usize = 262_144;
const PAYLOAD_ALPHABET: &[u8; 64] =
b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct TxId([u8; TX_ID_BYTES]);
impl TxId {
pub const fn from_bytes(bytes: [u8; TX_ID_BYTES]) -> Self {
Self(bytes)
}
pub const fn as_bytes(&self) -> &[u8; TX_ID_BYTES] {
&self.0
}
pub const fn into_bytes(self) -> [u8; TX_ID_BYTES] {
self.0
}
pub fn for_transaction(transaction: &[u8]) -> Self {
let digest = Sha256::digest(transaction);
let mut bytes = [0; TX_ID_BYTES];
bytes.copy_from_slice(&digest[..TX_ID_BYTES]);
Self(bytes)
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum PutOutcome {
Inserted(TxId),
Duplicate(TxId),
}
#[derive(Debug)]
pub enum StoreError {
Io(io::Error),
AlreadyExists,
InvalidStore,
StoreFull,
IdCollision(TxId),
OutcomeUnknown(TxId),
ReopenRequired,
CorruptTransaction(TxId),
}
impl std::fmt::Display for StoreError {
fn fmt(&self, output: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Io(error) => error.fmt(output),
Self::AlreadyExists => output.write_str("the store already exists"),
Self::InvalidStore => output.write_str("the store is invalid"),
Self::StoreFull => output.write_str("the transaction file is full"),
Self::IdCollision(id) => write!(output, "transaction identifier collision: {id:?}"),
Self::OutcomeUnknown(id) => write!(output, "transaction outcome is unknown: {id:?}"),
Self::ReopenRequired => output.write_str("the store must be reopened before writing"),
Self::CorruptTransaction(id) => write!(output, "transaction is corrupt: {id:?}"),
}
}
}
impl std::error::Error for StoreError {}
impl From<io::Error> for StoreError {
fn from(error: io::Error) -> Self {
Self::Io(error)
}
}
pub struct TransactionStore {
data: File,
payload: PathBuf,
locations: RwLock<HashMap<TxId, u32>>,
next_sector: Mutex<u64>,
publication: Mutex<Publication>,
}
struct Publication {
lookup: File,
reopen_required: bool,
}
impl TransactionStore {
pub fn create(root: &Path) -> Result<Self, StoreError> {
match fs::create_dir(root) {
Ok(()) => {}
Err(error) if error.kind() == io::ErrorKind::AlreadyExists => {
return Err(StoreError::AlreadyExists);
}
Err(error) => return Err(error.into()),
}
let payload = root.join("payload");
fs::create_dir(&payload)?;
for first in PAYLOAD_ALPHABET {
for second in PAYLOAD_ALPHABET {
fs::create_dir(payload.join(format!(
"{}{}",
char::from(*first),
char::from(*second)
)))?;
}
}
File::open(&payload)?.sync_all()?;
File::create(root.join("transactions.dat"))?.sync_all()?;
File::create(root.join("lookup.dat"))?.sync_all()?;
File::open(root)?.sync_all()?;
let parent = root
.parent()
.filter(|path| !path.as_os_str().is_empty())
.unwrap_or(Path::new("."));
File::open(parent)?.sync_all()?;
Self::open(root)
}
pub fn open(root: &Path) -> Result<Self, StoreError> {
let data_path = root.join("transactions.dat");
let lookup_path = root.join("lookup.dat");
let payload = root.join("payload");
if !root.is_dir() || !data_path.is_file() || !lookup_path.is_file() || !payload.is_dir() {
return Err(StoreError::InvalidStore);
}
let mut lookup_bytes = fs::read(&lookup_path)?;
let complete = lookup_bytes.len() / 16 * 16;
if complete != lookup_bytes.len() {
let lookup = OpenOptions::new().write(true).open(&lookup_path)?;
lookup.set_len(complete as u64)?;
lookup.sync_data()?;
lookup_bytes.truncate(complete);
}
let mut locations = HashMap::with_capacity(lookup_bytes.len() / 16);
for entry in lookup_bytes.chunks_exact(16) {
let mut id = [0; TX_ID_BYTES];
let mut sector = [0; 4];
id.copy_from_slice(&entry[..12]);
sector.copy_from_slice(&entry[12..]);
if locations
.insert(TxId::from_bytes(id), u32::from_le_bytes(sector))
.is_some()
{
return Err(StoreError::InvalidStore);
}
}
let data = OpenOptions::new().read(true).write(true).open(data_path)?;
let next_sector = data.metadata()?.len().div_ceil(SECTOR_BYTES);
if next_sector > u32::MAX as u64 + 1 {
return Err(StoreError::StoreFull);
}
let lookup = OpenOptions::new().append(true).open(lookup_path)?;
Ok(Self {
data,
payload,
locations: RwLock::new(locations),
next_sector: Mutex::new(next_sector),
publication: Mutex::new(Publication {
lookup,
reopen_required: false,
}),
})
}
pub fn put(&self, transaction: &[u8]) -> Result<PutOutcome, StoreError> {
let id = TxId::for_transaction(transaction);
if let Some(sector) = self.location(id) {
return self.compare(id, sector, transaction);
}
let external = transaction.len() > INLINE_LIMIT;
if external {
self.write_payload(id, transaction)?;
}
let sectors = if external {
1
} else {
transaction
.len()
.checked_add(9)
.ok_or(StoreError::StoreFull)?
.div_ceil(SECTOR_BYTES as usize) as u64
};
let sector = self.allocate(sectors)?;
let mut record = vec![0; sectors as usize * SECTOR_BYTES as usize];
record[0] = u8::from(external);
record[1..9].copy_from_slice(&(transaction.len() as u64).to_le_bytes());
if !external {
record[9..9 + transaction.len()].copy_from_slice(transaction);
}
self.data
.write_all_at(&record, sector as u64 * SECTOR_BYTES)?;
self.data.sync_data()?;
self.publish(id, sector, transaction)
}
pub fn contains(&self, id: TxId) -> bool {
self.location(id).is_some()
}
pub fn get(&self, id: TxId) -> Result<Option<Vec<u8>>, StoreError> {
let Some(sector) = self.location(id) else {
return Ok(None);
};
self.read_transaction(id, sector).map(Some)
}
fn location(&self, id: TxId) -> Option<u32> {
self.locations
.read()
.unwrap_or_else(|error| error.into_inner())
.get(&id)
.copied()
}
fn allocate(&self, sectors: u64) -> Result<u32, StoreError> {
let mut next = self
.next_sector
.lock()
.unwrap_or_else(|error| error.into_inner());
let end = next.checked_add(sectors).ok_or(StoreError::StoreFull)?;
if end > u32::MAX as u64 + 1 {
return Err(StoreError::StoreFull);
}
let sector = *next as u32;
*next = end;
Ok(sector)
}
fn write_payload(&self, id: TxId, transaction: &[u8]) -> Result<(), StoreError> {
let path = self.payload_path(id);
let shard = path.parent().ok_or(StoreError::InvalidStore)?.to_path_buf();
let file = OpenOptions::new()
.create(true)
.truncate(false)
.write(true)
.open(path)?;
file.write_all_at(transaction, 0)?;
file.set_len(transaction.len() as u64)?;
file.sync_all()?;
File::open(shard)?.sync_all()?;
Ok(())
}
fn payload_path(&self, id: TxId) -> PathBuf {
let encoded = URL_SAFE_NO_PAD.encode(id.as_bytes());
self.payload
.join(&encoded[..2])
.join(format!("{}.dat", &encoded[2..]))
}
fn publish(&self, id: TxId, sector: u32, transaction: &[u8]) -> Result<PutOutcome, StoreError> {
let mut publication = self
.publication
.lock()
.unwrap_or_else(|error| error.into_inner());
if publication.reopen_required {
return Err(StoreError::ReopenRequired);
}
if let Some(existing) = self.location(id) {
drop(publication);
return self.compare(id, existing, transaction);
}
let mut entry = [0; 16];
entry[..12].copy_from_slice(id.as_bytes());
entry[12..].copy_from_slice(§or.to_le_bytes());
if publication
.lookup
.write_all(&entry)
.and_then(|()| publication.lookup.sync_data())
.is_err()
{
publication.reopen_required = true;
return Err(StoreError::OutcomeUnknown(id));
}
self.locations
.write()
.unwrap_or_else(|error| error.into_inner())
.insert(id, sector);
Ok(PutOutcome::Inserted(id))
}
fn compare(&self, id: TxId, sector: u32, transaction: &[u8]) -> Result<PutOutcome, StoreError> {
if self.read_transaction(id, sector)? == transaction {
Ok(PutOutcome::Duplicate(id))
} else {
Err(StoreError::IdCollision(id))
}
}
fn read_transaction(&self, id: TxId, sector: u32) -> Result<Vec<u8>, StoreError> {
let offset = sector as u64 * SECTOR_BYTES;
let mut header = [0; 9];
self.read_data(&mut header, offset, id)?;
let length = u64::from_le_bytes(header[1..9].try_into().unwrap());
let bytes = match header[0] {
0 if length <= INLINE_LIMIT as u64 => {
let mut bytes = vec![0; length as usize];
self.read_data(&mut bytes, offset + 9, id)?;
bytes
}
1 if length > INLINE_LIMIT as u64 => {
let path = self.payload_path(id);
let metadata = fs::metadata(&path).map_err(|error| match error.kind() {
io::ErrorKind::NotFound => StoreError::CorruptTransaction(id),
_ => StoreError::Io(error),
})?;
if metadata.len() != length {
return Err(StoreError::CorruptTransaction(id));
}
let bytes = fs::read(path)?;
if bytes.len() as u64 != length {
return Err(StoreError::CorruptTransaction(id));
}
bytes
}
_ => return Err(StoreError::CorruptTransaction(id)),
};
if TxId::for_transaction(&bytes) != id {
return Err(StoreError::CorruptTransaction(id));
}
Ok(bytes)
}
fn read_data(&self, bytes: &mut [u8], offset: u64, id: TxId) -> Result<(), StoreError> {
self.data
.read_exact_at(bytes, offset)
.map_err(|error| match error.kind() {
io::ErrorKind::UnexpectedEof => StoreError::CorruptTransaction(id),
_ => StoreError::Io(error),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
fn root() -> PathBuf {
static NEXT: AtomicU64 = AtomicU64::new(0);
let path = std::env::temp_dir().join(format!(
"k1-store-{}-{}",
std::process::id(),
NEXT.fetch_add(1, Ordering::Relaxed)
));
let _ = fs::remove_dir_all(&path);
path
}
#[test]
fn lifecycle_boundaries_duplicates_corruption_and_tail() {
let root = root();
let store = TransactionStore::create(&root).unwrap();
assert_eq!(fs::read_dir(root.join("payload")).unwrap().count(), 4_096);
let values = [vec![], vec![3; INLINE_LIMIT], vec![7; INLINE_LIMIT + 1]];
let mut ids = Vec::new();
for value in &values {
let id = TxId::for_transaction(value);
assert_eq!(store.put(value).unwrap(), PutOutcome::Inserted(id));
assert_eq!(store.put(value).unwrap(), PutOutcome::Duplicate(id));
assert_eq!(store.get(id).unwrap().unwrap(), *value);
ids.push(id);
}
assert_eq!(fs::metadata(root.join("lookup.dat")).unwrap().len(), 48);
let sector = store.location(ids[0]).unwrap();
assert!(matches!(
store.compare(ids[0], sector, b"different"),
Err(StoreError::IdCollision(id)) if id == ids[0]
));
drop(store);
OpenOptions::new()
.append(true)
.open(root.join("lookup.dat"))
.unwrap()
.write_all(&[1, 2, 3])
.unwrap();
let store = TransactionStore::open(&root).unwrap();
assert_eq!(fs::metadata(root.join("lookup.dat")).unwrap().len(), 48);
for (id, value) in ids.into_iter().zip(values) {
assert!(store.contains(id));
assert_eq!(store.get(id).unwrap().unwrap(), value);
}
let id = TxId::for_transaction(b"intact");
assert_eq!(store.put(b"intact").unwrap(), PutOutcome::Inserted(id));
let sector = store.location(id).unwrap();
store
.data
.write_all_at(b"x", sector as u64 * SECTOR_BYTES + 9)
.unwrap();
assert!(matches!(
store.get(id),
Err(StoreError::CorruptTransaction(found)) if found == id
));
fs::remove_dir_all(root).unwrap();
}
#[test]
fn concurrent_writes_allocate_disjoint_sectors_and_publish_once() {
let root = root();
let store = Arc::new(TransactionStore::create(&root).unwrap());
let threads: Vec<_> = (0_u8..8)
.map(|byte| {
let store = Arc::clone(&store);
std::thread::spawn(move || {
let value = vec![byte; 8_193];
let outcome = store.put(&value).unwrap();
(value, outcome)
})
})
.collect();
for thread in threads {
let (value, outcome) = thread.join().unwrap();
let PutOutcome::Inserted(id) = outcome else {
panic!()
};
assert_eq!(store.get(id).unwrap().unwrap(), value);
}
let value = vec![11; INLINE_LIMIT + 1];
let id = TxId::for_transaction(&value);
let threads: Vec<_> = (0..2)
.map(|_| {
let store = Arc::clone(&store);
let value = value.clone();
std::thread::spawn(move || store.put(&value))
})
.collect();
let outcomes: Vec<_> = threads
.into_iter()
.map(|thread| thread.join().unwrap())
.collect();
assert_eq!(
outcomes
.iter()
.filter(|outcome| matches!(outcome, Ok(PutOutcome::Inserted(_))))
.count(),
1
);
assert!(outcomes.iter().all(|outcome| matches!(
outcome,
Ok(PutOutcome::Inserted(_)) | Ok(PutOutcome::Duplicate(_))
)));
assert_eq!(store.get(id).unwrap().unwrap(), value);
assert_eq!(store.put(&value).unwrap(), PutOutcome::Duplicate(id));
let sectors = store
.locations
.read()
.unwrap_or_else(|error| error.into_inner());
let mut locations: Vec<_> = sectors.values().copied().collect();
locations.sort_unstable();
locations.dedup();
assert_eq!(locations.len(), sectors.len());
fs::remove_dir_all(root).unwrap();
}
#[test]
fn capacity_is_bounded() {
let root = root();
let store = TransactionStore::create(&root).unwrap();
*store
.next_sector
.lock()
.unwrap_or_else(|error| error.into_inner()) = u32::MAX as u64 + 1;
assert!(matches!(store.put(b"full"), Err(StoreError::StoreFull)));
fs::remove_dir_all(root).unwrap();
}
}