use crate::compact;
use crate::error::{Error, Result};
use crate::format::{base, wal};
use crate::mmap::MappedFile;
use crate::overlay::Arena;
use crate::snapshot::Snapshot;
use crate::ttl::TtlConfig;
use crate::writer::{WriteState, Writer};
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Mutex, RwLock};
const GEN_SLOTS: usize = 4;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Durability {
Always,
Batch(u32),
Never,
}
#[derive(Clone, Copy, Debug)]
pub struct Config {
pub durability: Durability,
pub compact_after_ops: usize,
pub compact_after_wal_bytes: u64,
pub ttl: TtlConfig,
}
impl Default for Config {
fn default() -> Self {
Self {
durability: Durability::Always,
compact_after_ops: 4096,
compact_after_wal_bytes: 4 << 20,
ttl: TtlConfig::default(),
}
}
}
#[derive(Debug)]
pub(crate) struct GenInner {
map: MappedFile,
pub(crate) dir: base::BaseDir,
pub(crate) arena: Arena,
}
impl GenInner {
pub(crate) fn bytes(&self) -> &[u8] {
self.map.bytes()
}
pub(crate) fn from_parts(map: MappedFile, dir: base::BaseDir) -> Self {
Self { map, dir, arena: Arena::new() }
}
}
#[derive(Debug)]
pub(crate) struct GenSlot {
pub(crate) seq: AtomicU64,
pub(crate) inner: RwLock<Option<GenInner>>,
}
#[derive(Debug)]
pub struct Store {
pub(crate) dir: PathBuf,
pub(crate) cfg: Config,
read_only: bool,
_lock: Option<LockFile>,
pub(crate) current: AtomicUsize,
pub(crate) gens: [GenSlot; GEN_SLOTS],
pub(crate) write: Mutex<WriteState>,
pub(crate) compact_inline: AtomicBool,
pub(crate) due: AtomicBool,
reload: Mutex<Reload>,
}
#[derive(Clone, Copy, Debug, Default)]
struct Reload {
stamp: Option<Stamp>,
wal_pos: u64,
}
#[derive(Debug)]
struct LockFile {
_file: fs::File,
}
impl LockFile {
fn acquire(path: &Path) -> Result<Self> {
match crate::fsio::try_lock(path)? {
Some(file) => Ok(Self { _file: file }),
None => Err(Error::Locked),
}
}
}
const RACE_RETRIES: usize = 3;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct Stamp {
base: FileId,
wal_len: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct FileId {
len: u64,
modified: Option<std::time::SystemTime>,
#[cfg(unix)]
ino: (u64, u64),
#[cfg(not(unix))]
created: Option<std::time::SystemTime>,
}
impl FileId {
fn of(md: &fs::Metadata) -> Self {
#[cfg(unix)]
use std::os::unix::fs::MetadataExt as _;
Self {
len: md.len(),
modified: md.modified().ok(),
#[cfg(unix)]
ino: (md.dev(), md.ino()),
#[cfg(not(unix))]
created: md.created().ok(),
}
}
}
impl Stamp {
fn read(dir: &Path) -> Option<Self> {
let base = FileId::of(&fs::metadata(base_path(dir)).ok()?);
let wal_len = match fs::metadata(wal_path(dir)) {
Ok(md) => md.len(),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => 0,
Err(_) => return None,
};
Some(Self { base, wal_len })
}
}
pub const BASE_FILE: &str = "graph.base";
pub const WAL_FILE: &str = "graph.wal";
pub const LOCK_FILE: &str = "graph.lock";
pub(crate) fn base_path(dir: &Path) -> PathBuf {
dir.join(BASE_FILE)
}
pub(crate) fn wal_path(dir: &Path) -> PathBuf {
dir.join(WAL_FILE)
}
fn lock_path(dir: &Path) -> PathBuf {
dir.join(LOCK_FILE)
}
use crate::fsio::now_ms;
pub(crate) fn atomic_write(path: &Path, bytes: &[u8]) -> Result<()> {
crate::fsio::write_atomic(path, &[bytes]).map_err(Into::into)
}
impl Store {
pub fn create(dir: impl AsRef<Path>, cfg: Config) -> Result<Self> {
let dir = dir.as_ref();
fs::create_dir_all(dir)?;
if base_path(dir).exists() {
return Err(Error::format("store already exists; use open"));
}
let empty = compact::render_empty_base(1);
atomic_write(&base_path(dir), &empty)?;
atomic_write(&wal_path(dir), &wal::encode_header(1, now_ms()))?;
Self::open_with(dir, cfg, false)
}
pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
Self::open_with(dir.as_ref(), Config::default(), false)
}
pub fn open_cfg(dir: impl AsRef<Path>, cfg: Config) -> Result<Self> {
Self::open_with(dir.as_ref(), cfg, false)
}
pub fn open_read_only(dir: impl AsRef<Path>) -> Result<Self> {
Self::open_with(dir.as_ref(), Config::default(), true)
}
fn open_with(dir: &Path, cfg: Config, read_only: bool) -> Result<Self> {
let lock = if read_only { None } else { Some(LockFile::acquire(&lock_path(dir))?) };
let stamp = Stamp::read(dir);
let built = load_generation(dir, read_only)?;
Ok(Self::assemble(dir.to_path_buf(), cfg, read_only, lock, stamp, built))
}
pub fn open_image(base: bytes::Bytes, wal: &[u8]) -> Result<Self> {
let built = build_generation(MappedFile::from_bytes(base), wal)?;
if built.log == LogState::Newer {
return Err(newer_image());
}
Ok(Self::assemble(PathBuf::new(), Config::default(), true, None, None, built))
}
pub fn reload_image(&self, base: Option<bytes::Bytes>, wal: &[u8]) -> Result<bool> {
if !self.read_only {
return Ok(false);
}
let _one_at_a_time = self.reload.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
match base {
Some(b) => {
self.install_or_apply(MappedFile::from_bytes(b), wal)?.ok_or_else(newer_image)
}
None => match self.apply_wal_bytes(wal)? {
Some(changed) => Ok(changed),
None => {
Err(Error::format("log image belongs to another generation; pass its base"))
}
},
}
}
#[must_use]
pub fn is_image(&self) -> bool {
self.dir.as_os_str().is_empty()
}
fn assemble(
dir: PathBuf,
cfg: Config,
read_only: bool,
lock: Option<LockFile>,
stamp: Option<Stamp>,
built: Built,
) -> Self {
let generation = built.inner.dir.header.generation;
let reload = Reload { stamp, wal_pos: built.good_len as u64 };
let store = Self {
dir,
cfg,
read_only,
_lock: lock,
current: AtomicUsize::new(0),
gens: std::array::from_fn(|_| GenSlot {
seq: AtomicU64::new(0),
inner: RwLock::new(None),
}),
write: Mutex::new(WriteState::new(
built.next_seq,
built.replayed,
built.good_len as u64,
generation,
)),
compact_inline: AtomicBool::new(true),
due: AtomicBool::new(false),
reload: Mutex::new(reload),
};
*store.gens[0].inner.write().expect("fresh lock") = Some(built.inner);
store.gens[0].seq.store(generation, Ordering::Release);
store
}
#[must_use]
pub fn len(&self) -> usize {
self.snapshot().len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[must_use]
pub fn dir(&self) -> &Path {
&self.dir
}
#[must_use]
pub fn generation(&self) -> u64 {
self.gens[self.current.load(Ordering::Acquire)].seq.load(Ordering::Acquire)
}
pub fn set_compact_inline(&self, inline: bool) {
self.compact_inline.store(inline, Ordering::Relaxed);
}
#[must_use]
pub fn compaction_due(&self) -> bool {
self.due.load(Ordering::Acquire)
}
#[must_use]
pub fn snapshot(&self) -> Snapshot<'_> {
loop {
let cur = self.current.load(Ordering::Acquire);
if let Ok(guard) = self.gens[cur].inner.try_read() {
if guard.is_some() && self.current.load(Ordering::Acquire) == cur {
return Snapshot::new(self, guard);
}
}
std::hint::spin_loop();
}
}
pub fn writer(&self) -> Result<Writer<'_>> {
if self.read_only {
return Err(Error::ReadOnly);
}
let state = self.write.lock().map_err(|_| Error::corrupt("writer mutex poisoned"))?;
Ok(Writer::new(self, state))
}
pub fn compact(&self) -> Result<compact::CompactStats> {
if self.read_only {
return Err(Error::ReadOnly);
}
let mut state = self.write.lock().map_err(|_| Error::corrupt("writer mutex poisoned"))?;
compact::compact_locked(self, &mut state)
}
pub fn reload_if_stale(&self) -> Result<bool> {
if self.is_image() || !self.read_only {
return Ok(false); }
let now = Stamp::read(&self.dir);
let mut st = self.reload.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
if let (Some(n), Some(l)) = (now, st.stamp) {
if n == l {
return Ok(false);
}
if n.base == l.base && n.wal_len >= st.wal_pos && st.wal_pos >= wal::HEADER_LEN as u64 {
let (changed, pos) = self.reload_wal_tail(st.wal_pos)?;
*st = Reload { stamp: now, wal_pos: pos };
return Ok(changed);
}
}
let (changed, pos) = self.reload_full()?;
*st = Reload { stamp: now, wal_pos: pos };
Ok(changed)
}
fn reload_wal_tail(&self, pos: u64) -> Result<(bool, u64)> {
use std::io::{Read as _, Seek as _, SeekFrom};
let mut tail = Vec::new();
match fs::File::open(wal_path(&self.dir)) {
Ok(mut f) => {
f.seek(SeekFrom::Start(pos))?;
f.read_to_end(&mut tail)?;
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok((false, pos)),
Err(e) => return Err(e.into()),
}
let (ops, used) = wal::replay_frames(&tail)?;
Ok((self.publish_new(ops)?, pos + used as u64))
}
fn publish_new(&self, ops: Vec<(u64, wal::Op)>) -> Result<bool> {
let cur = self.current.load(Ordering::Acquire);
let guard = self.gens[cur].inner.try_read().map_err(|_| Error::busy("slot busy"))?;
let Some(g) = guard.as_ref() else { return Ok(false) };
let seen = g.dir.header.wal_applied_seq.max(g.arena.last_seq());
let mut changed = false;
for (seq, op) in ops.into_iter().filter(|&(s, _)| s > seen) {
if !g.arena.publish(seq, op) {
return Err(Error::busy("overlay arena full"));
}
changed = true;
}
Ok(changed)
}
fn apply_wal_bytes(&self, raw: &[u8]) -> Result<Option<bool>> {
if raw.is_empty() {
return Ok(Some(false));
}
let replayed = wal::replay(raw)?;
if replayed.base_generation != self.generation() {
return Ok(None);
}
self.publish_new(replayed.ops).map(Some)
}
fn reload_full(&self) -> Result<(bool, u64)> {
for _ in 0..RACE_RETRIES {
let map = MappedFile::open(&base_path(&self.dir))?;
let raw = match fs::read(wal_path(&self.dir)) {
Ok(raw) => raw,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Vec::new(),
Err(e) => return Err(e.into()),
};
if let Some(changed) = self.install_or_apply(map, &raw)? {
return Ok((changed, wal::committed_len(&raw) as u64));
}
}
Err(Error::busy("compaction in progress; retry"))
}
fn install_or_apply(&self, map: MappedFile, raw: &[u8]) -> Result<Option<bool>> {
let cur_gen = self.generation();
let new_gen = base::Header::parse(map.bytes())?.generation;
if new_gen == cur_gen {
return match self.apply_wal_bytes(raw)? {
Some(changed) => Ok(Some(changed)),
None if wal::replay(raw)?.base_generation > new_gen => Ok(None),
None => Ok(Some(false)),
};
}
let built = build_generation(map, raw)?;
if built.log == LogState::Newer {
return Ok(None);
}
let cur = self.current.load(Ordering::Acquire);
for (i, slot) in self.gens.iter().enumerate() {
if i == cur {
continue;
}
if let Ok(mut w) = slot.inner.try_write() {
*w = Some(built.inner);
drop(w);
slot.seq.store(new_gen, Ordering::Release);
self.current.store(i, Ordering::Release);
return Ok(Some(true));
}
}
Err(Error::busy("all generation slots pinned by snapshots"))
}
}
fn newer_image() -> Error {
Error::busy("log image is newer than its base; fetch the base again")
}
struct Built {
inner: GenInner,
next_seq: u64,
replayed: usize,
good_len: usize,
log: LogState,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum LogState {
Empty,
Stale,
Current { torn: bool },
Newer,
}
fn build_generation(map: MappedFile, raw: &[u8]) -> Result<Built> {
let bdir = base::BaseDir::parse(map.bytes())?;
let header = bdir.header;
let inner = GenInner { map, dir: bdir, arena: Arena::new() };
if raw.is_empty() {
return Ok(Built {
inner,
next_seq: header.wal_applied_seq + 1,
replayed: 0,
good_len: 0,
log: LogState::Empty,
});
}
let replayed = wal::replay(raw)?;
if replayed.base_generation > header.generation {
return Ok(Built {
inner,
next_seq: header.wal_applied_seq + 1,
replayed: 0,
good_len: 0,
log: LogState::Newer,
});
}
if replayed.base_generation < header.generation {
return Ok(Built {
inner,
next_seq: header.wal_applied_seq + 1,
replayed: 0,
good_len: wal::HEADER_LEN,
log: LogState::Stale,
});
}
let torn = replayed.good_len < raw.len();
let next_seq = replayed
.ops
.last()
.map_or(header.wal_applied_seq + 1, |&(s, _)| s.max(header.wal_applied_seq) + 1);
let mut count = 0;
for (seq, op) in replayed.ops.into_iter().filter(|&(s, _)| s > header.wal_applied_seq) {
if !inner.arena.publish(seq, op) {
return Err(Error::busy("overlay arena full during replay"));
}
count += 1;
}
Ok(Built {
inner,
next_seq,
replayed: count,
good_len: replayed.good_len,
log: LogState::Current { torn },
})
}
fn load_generation(dir: &Path, read_only: bool) -> Result<Built> {
let wal_file = wal_path(dir);
let mut attempt = 0;
let mut built = loop {
let map = MappedFile::open(&base_path(dir))?;
let raw = match fs::read(&wal_file) {
Ok(raw) => raw,
Err(e) if e.kind() == std::io::ErrorKind::NotFound && read_only => Vec::new(),
Err(e) => return Err(e.into()),
};
let built = build_generation(map, &raw)?;
if built.log != LogState::Newer {
break built;
}
attempt += 1;
if !read_only {
return Err(Error::corrupt("wal is newer than base"));
}
if attempt == RACE_RETRIES {
return Err(Error::busy("compaction in progress; retry"));
}
};
if !read_only {
match built.log {
LogState::Stale | LogState::Empty => {
let generation = built.inner.dir.header.generation;
atomic_write(&wal_file, &wal::encode_header(generation, now_ms()))?;
built.good_len = wal::HEADER_LEN;
}
LogState::Current { torn: true } => {
let f = fs::OpenOptions::new().write(true).open(&wal_file)?;
f.set_len(built.good_len as u64)?;
f.sync_all()?;
}
LogState::Current { torn: false } | LogState::Newer => {}
}
}
Ok(built)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::writer::DocRecord;
use crate::UrlKey;
fn doc(url: &str, hash: u64) -> DocRecord<'_> {
DocRecord {
url,
url_key: UrlKey::of(url),
content_hash: hash,
fetched_at_ms: 1,
title: None,
snippet: None,
etag: None,
pinned: false,
}
}
fn write(store: &Store, url: &str, hash: u64) {
let mut w = store.writer().unwrap();
w.upsert_node(&doc(url, hash)).unwrap();
w.commit().unwrap();
}
#[test]
fn a_snapshot_does_not_observe_later_commits() {
let dir = tempfile::tempdir().unwrap();
let store = Store::create(dir.path(), Config::default()).unwrap();
write(&store, "https://x.dev/a", 1);
let before = store.snapshot();
assert_eq!(before.len(), 1);
write(&store, "https://x.dev/b", 2);
assert_eq!(before.len(), 1, "an existing snapshot must not grow");
assert_eq!(store.snapshot().len(), 2, "a fresh snapshot sees the commit");
}
#[test]
fn compaction_leaves_an_existing_snapshot_on_its_own_generation() {
let dir = tempfile::tempdir().unwrap();
let store = Store::create(dir.path(), Config::default()).unwrap();
write(&store, "https://x.dev/a", 1);
store.compact().unwrap();
let pinned = store.snapshot();
let pinned_gen = pinned.generation();
write(&store, "https://x.dev/b", 2);
store.compact().unwrap();
assert_eq!(pinned.generation(), pinned_gen, "a pinned generation must not move");
assert_eq!(pinned.len(), 1, "the old generation still reads its own state");
let fresh = store.snapshot();
assert!(fresh.generation() > pinned_gen, "the store advanced");
assert_eq!(fresh.len(), 2);
}
#[test]
fn each_compaction_advances_the_generation_exactly_once() {
let dir = tempfile::tempdir().unwrap();
let store = Store::create(dir.path(), Config::default()).unwrap();
write(&store, "https://x.dev/a", 1);
let mut prev = store.snapshot().generation();
for _ in 0..6 {
store.compact().unwrap();
let now = store.snapshot().generation();
assert_eq!(now, prev + 1, "a compaction must advance the generation by one");
prev = now;
}
}
#[test]
fn exhausting_the_generation_ring_is_a_retryable_refusal() {
let dir = tempfile::tempdir().unwrap();
let store = Store::create(dir.path(), Config::default()).unwrap();
write(&store, "https://x.dev/a", 1);
let mut pins = Vec::new();
for round in 0..GEN_SLOTS - 1 {
pins.push(store.snapshot());
store.compact().unwrap_or_else(|e| panic!("round {round} should still fit: {e}"));
}
pins.push(store.snapshot());
assert!(
matches!(store.compact(), Err(Error::Busy { .. })),
"a full ring must refuse with Busy, not stall or corrupt the swap"
);
drop(pins);
store.compact().expect("compaction must succeed once a slot frees up");
assert_eq!(store.snapshot().len(), 1);
}
#[test]
fn a_read_only_handle_reloads_a_newer_generation() {
let dir = tempfile::tempdir().unwrap();
let store = Store::create(dir.path(), Config::default()).unwrap();
write(&store, "https://x.dev/a", 1);
store.compact().unwrap();
let reader = Store::open_read_only(dir.path()).unwrap();
assert_eq!(reader.snapshot().len(), 1);
assert!(!reader.reload_if_stale().unwrap(), "nothing has changed yet");
write(&store, "https://x.dev/b", 2);
store.compact().unwrap();
assert!(reader.reload_if_stale().unwrap(), "a new generation is visible");
assert_eq!(reader.snapshot().len(), 2);
}
}