use crate::error::{Error, Result};
use crate::format::base::nflags;
use crate::format::wal::{self, OpView, RecordBuf, TouchRecord};
use crate::key::{EdgeType, UrlKey};
use crate::snapshot::Snapshot;
use crate::store::{CommitLock, Durability, Held, Seen, Store};
use crate::ttl::Outcome;
use std::io::Write as _;
use std::sync::atomic::Ordering;
use std::sync::MutexGuard;
pub use crate::format::wal::MAX_EDGES;
pub use crate::format::wal::{MAX_ATTRS, MAX_ATTR_LEN};
pub use crate::format::wal::{MAX_LABEL, MAX_SEGS};
pub const MAX_ATTR_BYTES: usize = wal::MAX_OP - (wal::OP_HEADER_LEN + 8) - MAX_ATTRS * 8;
pub const MAX_ATTR_NAME: usize = 255;
#[derive(Clone, Copy, Debug)]
pub struct DocRecord<'a> {
pub url: &'a str,
pub url_key: UrlKey,
pub content_hash: u64,
pub fetched_at_ms: u64,
pub title: Option<&'a str>,
pub snippet: Option<&'a str>,
pub etag: Option<&'a str>,
pub pinned: bool,
}
#[derive(Clone, Copy, Debug)]
pub struct SegmentRecord<'a> {
pub label: &'a str,
pub byte_range: Option<(u32, u32)>,
pub depth: u8,
pub importance: u16,
}
#[derive(Clone, Copy, Debug)]
pub struct Touch<'a> {
pub checked_at_ms: u64,
pub outcome: Outcome,
pub content_hash: Option<u64>,
pub etag: Option<&'a str>,
}
pub(crate) struct WriteState {
next_seq: u64,
staged: RecordBuf,
pub(crate) ops_since_compact: usize,
pub(crate) wal_bytes: u64,
commits_since_sync: u32,
wal: Option<std::fs::File>,
pub(crate) disk_generation: u64,
pub(crate) seen: Seen,
pub(crate) lock_file: Option<std::fs::File>,
pub(crate) wait_file: Option<std::fs::File>,
}
#[allow(clippy::missing_fields_in_debug)] impl std::fmt::Debug for WriteState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WriteState")
.field("next_seq", &self.next_seq)
.field("staged", &self.staged.ops())
.finish()
}
}
impl WriteState {
pub(crate) fn new(
next_seq: u64,
replayed_ops: usize,
wal_bytes: u64,
disk_generation: u64,
seen: Seen,
) -> Self {
Self {
next_seq,
staged: RecordBuf::new(),
ops_since_compact: replayed_ops,
wal_bytes,
commits_since_sync: 0,
wal: None,
disk_generation,
seen,
lock_file: None,
wait_file: None,
}
}
pub(crate) fn drop_wal_handle(&mut self) {
self.wal = None;
}
pub(crate) fn adopt(&mut self, next_seq: u64, overlay_ops: usize, generation: u64) {
self.next_seq = self.next_seq.max(next_seq);
self.ops_since_compact = overlay_ops;
self.wal_bytes = self.seen.wal_pos;
self.disk_generation = self.disk_generation.max(generation);
}
pub(crate) fn release(&mut self, held: Held) {
let _ = held.0.unlock();
self.lock_file = Some(held.0);
}
pub(crate) fn reset_after_compact(&mut self, next_seq: u64) {
self.next_seq = next_seq;
self.ops_since_compact = 0;
self.wal_bytes = u64::try_from(crate::format::wal::HEADER_LEN).unwrap_or(64);
self.commits_since_sync = 0;
self.wal = None; }
pub(crate) fn staged_is_empty(&self) -> bool {
self.staged.is_empty()
}
pub(crate) fn last_committed_seq(&self) -> u64 {
self.next_seq.saturating_sub(1)
}
}
#[derive(Debug)]
pub struct Writer<'s> {
store: &'s Store,
state: MutexGuard<'s, WriteState>,
snap: Option<Snapshot<'s>>,
}
fn check_label(name: &str, s: &str) -> Result<()> {
if s.len() > MAX_LABEL {
return Err(Error::format(format!("{name} label exceeds {MAX_LABEL} bytes")));
}
Ok(())
}
impl<'s> Writer<'s> {
pub(crate) fn new(store: &'s Store, state: MutexGuard<'s, WriteState>) -> Self {
Self { store, state, snap: None }
}
fn view(&mut self) -> &Snapshot<'s> {
if self.snap.is_none() {
self.snap = Some(self.store.snapshot());
}
self.snap.as_ref().expect("just filled")
}
pub fn upsert_node(&mut self, doc: &DocRecord<'_>) -> Result<()> {
check_label("url", doc.url)?;
if doc.url.is_empty() {
return Err(Error::format("url must not be empty"));
}
for (name, v) in [("title", doc.title), ("snippet", doc.snippet), ("etag", doc.etag)] {
if let Some(v) = v {
check_label(name, v)?;
}
}
let flags = if doc.pinned { nflags::PINNED } else { 0 };
self.state.staged.push_upsert(doc, flags)
}
pub fn set_segments(&mut self, key: UrlKey, segs: &[SegmentRecord<'_>]) -> Result<()> {
self.state.staged.push_segments(key, segs)
}
pub fn set_edges(&mut self, key: UrlKey, edges: &[(UrlKey, EdgeType, u16)]) {
self.stage_edges(key, edges.iter().map(|&(d, t, w)| (d, t, w, 0)));
}
pub fn set_edges_with_flags(&mut self, key: UrlKey, edges: &[(UrlKey, EdgeType, u16, u8)]) {
let mask = crate::format::base::eflags::EXACT;
self.stage_edges(key, edges.iter().map(|&(d, t, w, f)| (d, t, w, f & mask)));
}
fn stage_edges(
&mut self,
key: UrlKey,
edges: impl Iterator<Item = (UrlKey, EdgeType, u16, u8)>,
) {
let mut seen = std::collections::HashSet::new();
let mut owned: Vec<wal::Edge> = Vec::new();
for (dst, etype, weight, mut flags) in edges {
if dst == key || !seen.insert(dst.0) {
continue;
}
if matches!(etype, EdgeType::Related) {
flags |= crate::format::base::eflags::INFERRED;
}
owned.push((dst, etype.as_tag(), flags, weight));
if owned.len() == MAX_EDGES {
break;
}
}
owned.sort_unstable_by_key(|(dst, _, _, _)| dst.0);
let staged = self.state.staged.push_edges(key, &owned);
debug_assert!(staged.is_ok(), "normalized edges refused");
}
pub fn set_attrs(&mut self, key: UrlKey, attrs: &[(u8, &str)]) -> Result<()> {
if attrs.len() > MAX_ATTRS {
return Err(Error::format("too many attributes"));
}
let mut owned: Vec<(u8, [&str; 3])> = Vec::with_capacity(attrs.len());
for &(k, v) in attrs.iter().rev() {
if v.len() > MAX_ATTR_LEN {
return Err(Error::format("attribute too long"));
}
if !owned.iter().any(|(o, _)| *o == k) {
owned.push((k, [v, "", ""]));
}
}
if owned.iter().map(|(_, v)| v[0].len()).sum::<usize>() > MAX_ATTR_BYTES {
return Err(Error::format(format!(
"attributes exceed {MAX_ATTR_BYTES} bytes in total"
)));
}
owned.sort_unstable_by_key(|(k, _)| *k);
self.state.staged.push_attr_parts(key, &owned)
}
pub fn set_named_attrs(&mut self, key: UrlKey, attrs: &[(&str, &str)]) -> Result<()> {
if attrs.len() > MAX_ATTRS {
return Err(Error::format(format!("more than {MAX_ATTRS} attributes")));
}
let mut named: Vec<(&str, &str)> = Vec::with_capacity(attrs.len());
for &(name, value) in attrs.iter().rev() {
if name.is_empty() || name.len() > MAX_ATTR_NAME || name.contains('\0') {
return Err(Error::format(format!(
"attribute name must be 1..={MAX_ATTR_NAME} bytes without NUL"
)));
}
if name.len() + 1 + value.len() > MAX_ATTR_LEN {
return Err(Error::format(format!(
"attribute {name:?} exceeds {MAX_ATTR_LEN} bytes"
)));
}
if !named.iter().any(|(n, _)| *n == name) {
named.push((name, value));
}
}
named.sort_unstable_by_key(|(n, _)| *n);
let total: usize = named.iter().map(|(n, v)| n.len() + 1 + v.len()).sum();
if total > MAX_ATTR_BYTES {
return Err(Error::format(format!(
"attributes exceed {MAX_ATTR_BYTES} bytes in total"
)));
}
let parts: Vec<(u8, [&str; 3])> =
named.iter().enumerate().map(|(k, &(n, v))| (k as u8, [n, "\0", v])).collect();
self.state.staged.push_attr_parts(key, &parts)
}
pub fn touch(&mut self, key: UrlKey, t: Touch<'_>) -> Result<bool> {
if let Some(e) = t.etag {
check_label("etag", e)?;
}
let mut etag_buf = [0u8; MAX_LABEL];
let mut staged: Option<(u64, Option<usize>)> = None;
for op in self.state.staged.op_views() {
if let OpView::Upsert(u) = op {
if u.key() == key {
let etag = u.etag().map(|e| {
etag_buf[..e.len()].copy_from_slice(e.as_bytes());
e.len()
});
staged = Some((u.content_hash(), etag));
}
}
}
let staged_etag = |len: Option<usize>| {
len.map(|n| std::str::from_utf8(&etag_buf[..n]).unwrap_or_default())
};
let snap = self.snap.get_or_insert_with(|| self.store.snapshot());
let Some(node) = snap.node(key) else { return Ok(false) };
let cfg = snap.ttl();
let interval_s = cfg.next_interval(node.interval_s.max(1), t.outcome);
let changed = t.outcome == Outcome::Changed;
let (base_hash, base_etag) = match staged {
Some((h, e)) => (h, staged_etag(e)),
None => (node.content_hash, node.etag),
};
let rec = TouchRecord {
checked_at_ms: t.checked_at_ms,
outcome: t.outcome.as_tag(),
content_hash: Some(t.content_hash.unwrap_or(base_hash)),
etag: t.etag.or(base_etag),
interval_s,
checks: node.checks.saturating_add(1),
changes: if changed { node.changes.saturating_add(1) } else { node.changes },
last_change_ms: if changed { t.checked_at_ms } else { node.last_change_ms },
tombstone: t.outcome == Outcome::Gone,
};
self.state.staged.push_touch(key, &rec)?;
Ok(true)
}
pub fn set_pinned(&mut self, key: UrlKey, pinned: bool) -> Result<bool> {
let exists = self
.state
.staged
.op_views()
.any(|op| matches!(op, OpView::Upsert(u) if u.key() == key))
|| self.view().node(key).is_some();
if !exists {
return Ok(false);
}
self.state.staged.push_pin(key, pinned);
Ok(true)
}
pub fn remove(&mut self, key: UrlKey) -> Result<()> {
self.state.staged.push_remove(key);
Ok(())
}
pub fn discard(&mut self) {
self.state.staged = RecordBuf::new();
}
pub fn commit(&mut self) -> Result<u64> {
self.snap = None;
if self.state.staged.is_empty() {
return Ok(self.state.next_seq.saturating_sub(1));
}
let held = self.store.acquire(&mut self.state, self.store.cfg.lock_wait)?;
let r = self.store.sync(&mut self.state, true).and_then(|_| self.append_held());
self.state.release(held);
r
}
fn append_held(&mut self) -> Result<u64> {
if self.state.staged.is_empty() {
return Ok(self.state.next_seq.saturating_sub(1));
}
let count = self.state.staged.ops();
let first_seq = self.state.next_seq;
self.state.staged.seal(first_seq)?;
let cur = self.store.current.load(Ordering::Acquire);
let guard = self.store.gens[cur]
.inner
.try_read()
.map_err(|_| Error::busy("current slot mid-install"))?;
let inner = guard.as_ref().ok_or_else(|| Error::corrupt("current slot empty"))?;
if !inner.log.reserve(count as usize, self.state.staged.len()) {
return Err(Error::busy("overlay arena full"));
}
if self.state.wal.is_none() {
self.state.wal =
Some(std::fs::OpenOptions::new().append(true).open(&self.store.wal_file)?);
}
let durability = self.store.cfg.durability;
let state = &mut *self.state;
let Some(f) = state.wal.as_mut() else { return Err(Error::corrupt("log handle missing")) };
let sync = match durability {
Durability::Always => true,
Durability::Batch(n) => state.commits_since_sync + 1 >= n.max(1),
Durability::Never => false,
};
let record = state.staged.as_bytes();
let appended = f.write_all(record).and_then(|()| f.flush()).and_then(|()| {
#[cfg(test)]
crate::store::probe::bump(&crate::store::probe::WRITES, 1);
if !sync {
return Ok(());
}
#[cfg(test)]
crate::store::probe::bump(&crate::store::probe::FSYNCS, 1);
f.sync_all()
});
#[cfg(test)]
let (appended, stuck) = match appended {
Ok(()) if crate::store::probe::fail_append() => (
Err(std::io::Error::other("injected append failure")),
crate::store::probe::stuck(),
),
other => (other, false),
};
#[cfg(not(test))]
let stuck = false;
if let Err(e) = appended {
state.wal = None;
if !stuck {
if let Ok(f) = std::fs::OpenOptions::new().write(true).open(&self.store.wal_file) {
let _ = f.set_len(state.seen.wal_pos);
let _ = f.sync_all();
}
}
state.seen.stamp = None;
return Err(e.into());
}
let len = record.len() as u64;
state.commits_since_sync = if sync { 0 } else { state.commits_since_sync + 1 };
state.seen.appended(len);
state.wal_bytes += len;
state.next_seq += u64::from(count);
state.ops_since_compact += count as usize;
let record = std::mem::take(&mut state.staged).into_vec();
crate::overlay::publish_records(&inner.log, record, 0, first_seq - 1);
drop(guard);
let last = self.state.next_seq - 1;
if self.state.ops_since_compact >= self.store.cfg.compact_after_ops
|| self.state.wal_bytes >= self.store.cfg.compact_after_wal_bytes
{
if !self.store.compact_inline.load(Ordering::Relaxed)
|| crate::compact::compact_locked(self.store, &mut self.state).is_err()
{
self.store.due.store(true, Ordering::Release);
}
}
Ok(last)
}
}
#[derive(Debug)]
pub struct Txn<'t> {
w: Writer<'t>,
}
impl<'t> Txn<'t> {
#[must_use]
pub fn snapshot(&self) -> &Snapshot<'t> {
self.w.snap.as_ref().expect("a transaction holds its view for its whole life")
}
pub fn upsert_node(&mut self, doc: &DocRecord<'_>) -> Result<()> {
self.w.upsert_node(doc)
}
pub fn set_segments(&mut self, key: UrlKey, segs: &[SegmentRecord<'_>]) -> Result<()> {
self.w.set_segments(key, segs)
}
pub fn set_edges(&mut self, key: UrlKey, edges: &[(UrlKey, EdgeType, u16)]) {
self.w.set_edges(key, edges);
}
pub fn set_edges_with_flags(&mut self, key: UrlKey, edges: &[(UrlKey, EdgeType, u16, u8)]) {
self.w.set_edges_with_flags(key, edges);
}
pub fn set_attrs(&mut self, key: UrlKey, attrs: &[(u8, &str)]) -> Result<()> {
self.w.set_attrs(key, attrs)
}
pub fn set_named_attrs(&mut self, key: UrlKey, attrs: &[(&str, &str)]) -> Result<()> {
self.w.set_named_attrs(key, attrs)
}
pub fn touch(&mut self, key: UrlKey, t: Touch<'_>) -> Result<bool> {
self.w.touch(key, t)
}
pub fn set_pinned(&mut self, key: UrlKey, pinned: bool) -> Result<bool> {
self.w.set_pinned(key, pinned)
}
pub fn remove(&mut self, key: UrlKey) -> Result<()> {
self.w.remove(key)
}
}
pub(crate) fn transact<F, R>(
store: &Store,
lock: Option<&CommitLock>,
wait: std::time::Duration,
f: F,
) -> Result<R>
where
F: for<'t> FnOnce(&mut Txn<'t>) -> Result<R>,
{
let mut txn = Txn { w: store.writer()? };
let held = match lock {
Some(_) => None,
None => Some(store.acquire(&mut txn.w.state, wait)?),
};
let r = run(store, &mut txn, f);
if let Some(h) = held {
txn.w.state.release(h);
}
r
}
fn run<'t, F, R>(store: &'t Store, txn: &mut Txn<'t>, f: F) -> Result<R>
where
F: for<'x> FnOnce(&mut Txn<'x>) -> Result<R>,
{
store.sync(&mut txn.w.state, true)?;
txn.w.snap = Some(store.snapshot());
let mark = txn.w.state.staged.mark();
let out = f(txn);
txn.w.snap = None; match out.and_then(|r| txn.w.append_held().map(|_| r)) {
Ok(r) => Ok(r),
Err(e) => {
txn.w.state.staged.truncate(mark);
Err(e)
}
}
}