use std::path::Path;
use yo_common::{Addr, Code, Error, Result, Space};
use yo_file::{CreateOptions, LogFile, Yo};
use yo_format::{RecordHeader, RecordKind};
use yo_kv::cold::{Blocks, Store as Boxed};
use yo_record::{Durability, Log, LogConfig};
const SHARDS: u32 = 16;
pub struct Store {
yo: Yo,
}
impl Store {
pub fn create(path: &Path) -> Result<Store> {
let yo = Yo::create(
path,
&CreateOptions {
shard_count: SHARDS,
..CreateOptions::default()
},
)?;
Ok(Store { yo })
}
pub fn source(mut self) -> impl FnMut(usize) -> Option<Boxed> {
move |at| {
let sink = self.yo.log(u32::try_from(at).ok()?).ok()?;
let cfg = LogConfig {
shard: u32::try_from(at).ok()?,
durability: Durability::None,
..LogConfig::default()
};
let log = Log::new(cfg, sink).ok()?;
Some(Box::new(Chunks {
log,
staged: std::cell::UnsafeCell::new(Vec::new()),
}))
}
}
}
struct Chunks {
log: Log<LogFile>,
staged: std::cell::UnsafeCell<Vec<Box<[u8]>>>,
}
impl Blocks for Chunks {
fn put(&mut self, bytes: &[u8]) -> Result<Addr> {
let h = RecordHeader::new(RecordKind::CollectionChunk);
let at = self.log.append(&h, b"", bytes)?;
Ok(Addr::new(Space::Log, at.addr))
}
fn get(&self, at: Addr) -> Result<&[u8]> {
if at.space() != Some(Space::Log) {
return Err(Error::new(Code::Invalid, "that address is not in the log"));
}
let off = at.offset();
match self.log.read(off) {
Ok(r) => return Ok(r.value),
Err(e) if e.code() == Code::NotFound => {}
Err(e) => return Err(e),
}
let mut buf = Vec::new();
self.log.read_value_into(off, &mut buf)?;
let staged = unsafe { &mut *self.staged.get() };
staged.push(buf.into_boxed_slice());
let bytes: *const [u8] = &*staged[staged.len() - 1];
Ok(unsafe { &*bytes })
}
fn bytes(&self) -> u64 {
self.log.tail().saturating_sub(self.log.begin())
}
fn release(&mut self) {
self.staged.get_mut().clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
use yo_kv::{Keyspace, Str};
struct Tmp(std::path::PathBuf);
impl Drop for Tmp {
fn drop(&mut self) {
let _ = std::fs::remove_file(&self.0);
}
}
fn tmp(name: &str) -> Tmp {
let mut p = std::env::temp_dir();
p.push(format!("yo-store-{name}-{}.yo", std::process::id()));
let _ = std::fs::remove_file(&p);
Tmp(p)
}
#[test]
fn a_value_written_to_the_file_reads_back() {
let path = tmp("roundtrip");
let store = Store::create(&path.0).expect("a file");
let mut source = store.source();
let mut blocks = source(0).expect("database zero has a log");
let val = b"the quick brown fox jumps over the lazy dog".repeat(4);
let at = blocks.put(&val).expect("written");
assert_eq!(blocks.get(at).expect("read back"), &val[..]);
assert!(blocks.bytes() > 0, "the log knows it holds something");
}
#[test]
fn every_database_gets_a_log_of_its_own() {
let path = tmp("shards");
let store = Store::create(&path.0).expect("a file");
let mut source = store.source();
for at in 0..SHARDS as usize {
assert!(source(at).is_some(), "database {at} has no log");
}
assert!(
source(SHARDS as usize).is_none(),
"and there is not a seventeenth"
);
}
#[test]
fn a_keyspace_demotes_into_the_file_and_still_answers() {
let path = tmp("demote");
let store = Store::create(&path.0).expect("a file");
let mut source = store.source();
let mut k = Keyspace::new();
k.attach(source(0).expect("database zero has a log"));
let val = b"a value long enough to be worth moving out of memory".repeat(2);
for i in 0..200u32 {
k.set_plain(&i.to_le_bytes(), &val).expect("stored");
}
let swept = k.relieve(usize::MAX).expect("swept");
assert!(swept.moved > 0, "nothing was moved");
assert!(k.store_bytes().expect("attached") > 0, "the file is empty");
for i in 0..200u32 {
assert_eq!(
k.get(&i.to_le_bytes()).expect("read"),
Some(Str::Bytes(&val)),
"key {i} did not read back off the file"
);
}
assert_eq!(k.len(), 200, "a sweep moves values and not keys");
}
#[cfg_attr(miri, ignore = "128 MiB, and the volume is the claim")]
#[test]
fn chunks_past_the_resident_window_read_back_off_the_device() {
let path = tmp("window");
let store = Store::create(&path.0).expect("a file");
let mut source = store.source();
let mut blocks = source(0).expect("database zero has a log");
let piece = vec![b'q'; yo_kv::cold::CHUNK];
let mut addrs = Vec::new();
for i in 0..(4 * 32 * 1024 * 1024 / yo_kv::cold::CHUNK) {
let mut v = piece.clone();
v[..8].copy_from_slice(&(i as u64).to_le_bytes());
addrs.push((blocks.put(&v).expect("written"), i as u64));
}
let mut held = Vec::new();
for (at, i) in &addrs {
let got = blocks.get(*at).expect("read back");
assert_eq!(got.len(), yo_kv::cold::CHUNK, "chunk {i} came back short");
held.push((got, *i));
}
for (got, i) in held {
assert_eq!(
u64::from_le_bytes(got[..8].try_into().expect("eight bytes")),
i,
"chunk {i} came back as somebody else"
);
}
blocks.release();
}
#[test]
fn a_path_that_is_already_there_is_refused() {
let path = tmp("existing");
let first = Store::create(&path.0);
assert!(first.is_ok(), "{:?}", first.err());
assert!(
Store::create(&path.0).is_err(),
"the second start wrote over the first one's file"
);
}
}