use crate::error::{Error, Result};
use crate::format::base::{nflags, sflags};
use crate::format::wal::{encode_frame, Op, OwnedSeg};
use crate::key::{EdgeType, UrlKey};
use crate::store::{wal_path, Durability, Store};
use crate::ttl::Outcome;
use compact_str::CompactString;
use std::io::Write as _;
use std::sync::atomic::Ordering;
use std::sync::MutexGuard;
pub const MAX_LABEL: usize = 2048;
pub const MAX_SEGS: usize = 128;
pub use crate::format::wal::MAX_EDGES;
pub use crate::format::wal::{MAX_ATTRS, MAX_ATTR_LEN};
pub const MAX_ATTR_BYTES: usize = crate::format::wal::MAX_FRAME - (1 + 8 + 1) - MAX_ATTRS * (1 + 4);
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: Vec<Op>,
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,
}
#[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.len())
.finish()
}
}
impl WriteState {
pub(crate) fn new(
next_seq: u64,
replayed_ops: usize,
wal_bytes: u64,
disk_generation: u64,
) -> Self {
Self {
next_seq,
staged: Vec::new(),
ops_since_compact: replayed_ops,
wal_bytes,
commits_since_sync: 0,
wal: None,
disk_generation,
}
}
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<crate::snapshot::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) -> &crate::snapshot::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(Op::UpsertNode {
key: doc.url_key,
content_hash: doc.content_hash,
fetched_at_ms: doc.fetched_at_ms,
flags,
url: doc.url.into(),
title: doc.title.map(Into::into),
snippet: doc.snippet.map(Into::into),
etag: doc.etag.map(Into::into),
});
Ok(())
}
pub fn set_segments(&mut self, key: UrlKey, segs: &[SegmentRecord<'_>]) -> Result<()> {
if segs.len() > MAX_SEGS {
return Err(Error::format(format!("more than {MAX_SEGS} segments")));
}
let mut owned = Vec::with_capacity(segs.len());
for s in segs {
check_label("segment", s.label)?;
let (importance, extra_flags) = (s.importance, 0);
let (byte_start, byte_len, range_flag) = match s.byte_range {
Some((start, len)) => (start, len, 0),
None => (0, 0, sflags::NO_RANGE),
};
owned.push(OwnedSeg {
label: s.label.into(),
byte_start,
byte_len,
depth: s.depth.max(1),
importance,
flags: range_flag | extra_flags,
});
}
self.state.staged.push(Op::SetSegments { key, segs: owned });
Ok(())
}
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<(UrlKey, u8, u8, u16)> = 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);
self.state.staged.push(Op::SetEdges { key, edges: owned });
}
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, CompactString)> = 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, CompactString::from(v)));
}
}
if owned.iter().map(|(_, v)| v.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(Op::SetAttrs { key, attrs: owned });
Ok(())
}
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 mut owned: Vec<(u8, CompactString)> = Vec::with_capacity(named.len());
let mut total = 0usize;
for (k, (name, value)) in named.into_iter().enumerate() {
let mut rec = CompactString::with_capacity(name.len() + 1 + value.len());
rec.push_str(name);
rec.push('\0');
rec.push_str(value);
total += rec.len();
owned.push((k as u8, rec));
}
if total > MAX_ATTR_BYTES {
return Err(Error::format(format!(
"attributes exceed {MAX_ATTR_BYTES} bytes in total"
)));
}
self.state.staged.push(Op::SetAttrs { key, attrs: owned });
Ok(())
}
pub fn touch(&mut self, key: UrlKey, t: Touch<'_>) -> Result<bool> {
if let Some(e) = t.etag {
check_label("etag", e)?;
}
let staged = self.state.staged.iter().rev().find_map(|op| match op {
Op::UpsertNode { key: k, content_hash, etag, .. } if *k == key => {
Some((*content_hash, etag.clone()))
}
_ => None,
});
let (interval_s, checks, changes, last_change_ms, content_hash, etag) = {
let snap = self.view();
let Some(node) = snap.node(key) else { return Ok(false) };
let cfg = snap.ttl();
let next = cfg.next_interval(node.interval_s.max(1), t.outcome);
let checks = node.checks.saturating_add(1);
let changes = if t.outcome == Outcome::Changed {
node.changes.saturating_add(1)
} else {
node.changes
};
let last_change =
if t.outcome == Outcome::Changed { t.checked_at_ms } else { node.last_change_ms };
let (base_hash, base_etag) = match staged {
Some((h, e)) => (h, e),
None => (node.content_hash, node.etag.map(Into::into)),
};
let hash = t.content_hash.unwrap_or(base_hash);
let etag: Option<compact_str::CompactString> = t.etag.map(Into::into).or(base_etag);
(next, checks, changes, last_change, hash, etag)
};
self.state.staged.push(Op::Touch {
key,
checked_at_ms: t.checked_at_ms,
outcome: t.outcome.as_tag(),
content_hash: Some(content_hash),
etag,
interval_s,
checks,
changes,
last_change_ms,
tombstone: t.outcome == Outcome::Gone,
});
Ok(true)
}
pub fn set_pinned(&mut self, key: UrlKey, pinned: bool) -> Result<bool> {
let exists = self
.state
.staged
.iter()
.any(|op| matches!(op, Op::UpsertNode { key: k, .. } if *k == key))
|| self.view().node(key).is_some();
if !exists {
return Ok(false);
}
self.state.staged.push(Op::SetPinned { key, pinned });
Ok(true)
}
pub fn remove(&mut self, key: UrlKey) -> Result<()> {
self.state.staged.push(Op::Remove { key });
Ok(())
}
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 staged = std::mem::take(&mut self.state.staged);
let count = staged.len() as u64;
let first_seq = self.state.next_seq;
let mut buf = Vec::new();
for (i, op) in staged.iter().enumerate() {
buf.extend_from_slice(&encode_frame(first_seq + i as u64, op));
}
if self.state.wal.is_none() {
self.state.wal =
Some(std::fs::OpenOptions::new().append(true).open(wal_path(&self.store.dir))?);
}
{
let state = &mut *self.state;
let f = state.wal.as_mut().expect("opened above");
f.write_all(&buf)?;
f.flush()?;
let sync = match self.store.cfg.durability {
Durability::Always => true,
Durability::Batch(n) => {
state.commits_since_sync += 1;
state.commits_since_sync >= n.max(1)
}
Durability::Never => false,
};
if sync {
f.sync_all()?;
state.commits_since_sync = 0;
}
}
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"))?;
for (i, op) in staged.into_iter().enumerate() {
if !inner.arena.publish(first_seq + i as u64, op) {
return Err(Error::busy("overlay arena full"));
}
}
}
self.state.wal_bytes += buf.len() as u64;
self.state.next_seq += count;
self.state.ops_since_compact += count as usize;
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)?;
} else {
self.store.due.store(true, Ordering::Release);
}
}
Ok(last)
}
}