1use crate::error::{Error, Result};
13use crate::format::base::{nflags, sflags};
14use crate::format::wal::{encode_frame, Op, OwnedSeg};
15use crate::key::{EdgeType, UrlKey};
16use crate::snapshot::Snapshot;
17use crate::store::{wal_path, CommitLock, Durability, Held, Seen, Store};
18use crate::ttl::Outcome;
19use compact_str::CompactString;
20use std::io::Write as _;
21use std::sync::atomic::Ordering;
22use std::sync::MutexGuard;
23
24pub const MAX_LABEL: usize = 2048;
26pub const MAX_SEGS: usize = 128;
28pub use crate::format::wal::MAX_EDGES;
30pub use crate::format::wal::{MAX_ATTRS, MAX_ATTR_LEN};
32pub const MAX_ATTR_BYTES: usize = crate::format::wal::MAX_FRAME - (1 + 8 + 1) - MAX_ATTRS * (1 + 4);
35pub const MAX_ATTR_NAME: usize = 255;
37
38#[derive(Clone, Copy, Debug)]
42pub struct DocRecord<'a> {
43 pub url: &'a str,
45 pub url_key: UrlKey,
47 pub content_hash: u64,
49 pub fetched_at_ms: u64,
51 pub title: Option<&'a str>,
53 pub snippet: Option<&'a str>,
55 pub etag: Option<&'a str>,
57 pub pinned: bool,
59}
60
61#[derive(Clone, Copy, Debug)]
63pub struct SegmentRecord<'a> {
64 pub label: &'a str,
66 pub byte_range: Option<(u32, u32)>,
68 pub depth: u8,
70 pub importance: u16,
72}
73
74#[derive(Clone, Copy, Debug)]
76pub struct Touch<'a> {
77 pub checked_at_ms: u64,
79 pub outcome: Outcome,
81 pub content_hash: Option<u64>,
83 pub etag: Option<&'a str>,
85}
86
87pub(crate) struct WriteState {
88 next_seq: u64,
89 staged: Vec<Op>,
90 pub(crate) ops_since_compact: usize,
91 pub(crate) wal_bytes: u64,
92 commits_since_sync: u32,
93 wal: Option<std::fs::File>,
94 pub(crate) disk_generation: u64,
99 pub(crate) seen: Seen,
101 pub(crate) lock_file: Option<std::fs::File>,
104}
105
106#[allow(clippy::missing_fields_in_debug)] impl std::fmt::Debug for WriteState {
108 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109 f.debug_struct("WriteState")
110 .field("next_seq", &self.next_seq)
111 .field("staged", &self.staged.len())
112 .finish()
113 }
114}
115
116impl WriteState {
117 pub(crate) fn new(
118 next_seq: u64,
119 replayed_ops: usize,
120 wal_bytes: u64,
121 disk_generation: u64,
122 seen: Seen,
123 ) -> Self {
124 Self {
125 next_seq,
126 staged: Vec::new(),
127 ops_since_compact: replayed_ops,
128 wal_bytes,
129 commits_since_sync: 0,
130 wal: None,
131 disk_generation,
132 seen,
133 lock_file: None,
134 }
135 }
136
137 pub(crate) fn drop_wal_handle(&mut self) {
139 self.wal = None;
140 }
141
142 pub(crate) fn adopt(&mut self, next_seq: u64, overlay_ops: usize, generation: u64) {
145 self.next_seq = self.next_seq.max(next_seq);
146 self.ops_since_compact = overlay_ops;
147 self.wal_bytes = self.seen.wal_pos;
148 self.disk_generation = self.disk_generation.max(generation);
149 }
150
151 pub(crate) fn release(&mut self, held: Held) {
153 let _ = held.0.unlock();
154 self.lock_file = Some(held.0);
155 }
156
157 pub(crate) fn reset_after_compact(&mut self, next_seq: u64) {
158 self.next_seq = next_seq;
159 self.ops_since_compact = 0;
160 self.wal_bytes = u64::try_from(crate::format::wal::HEADER_LEN).unwrap_or(64);
161 self.commits_since_sync = 0;
162 self.wal = None; }
164
165 pub(crate) fn staged_is_empty(&self) -> bool {
166 self.staged.is_empty()
167 }
168
169 pub(crate) fn last_committed_seq(&self) -> u64 {
170 self.next_seq.saturating_sub(1)
171 }
172}
173
174#[derive(Debug)]
179pub struct Writer<'s> {
180 store: &'s Store,
181 state: MutexGuard<'s, WriteState>,
182 snap: Option<Snapshot<'s>>,
196}
197
198fn check_label(name: &str, s: &str) -> Result<()> {
199 if s.len() > MAX_LABEL {
200 return Err(Error::format(format!("{name} label exceeds {MAX_LABEL} bytes")));
201 }
202 Ok(())
203}
204
205impl<'s> Writer<'s> {
206 pub(crate) fn new(store: &'s Store, state: MutexGuard<'s, WriteState>) -> Self {
207 Self { store, state, snap: None }
208 }
209
210 fn view(&mut self) -> &Snapshot<'s> {
212 if self.snap.is_none() {
213 self.snap = Some(self.store.snapshot());
214 }
215 self.snap.as_ref().expect("just filled")
216 }
217
218 pub fn upsert_node(&mut self, doc: &DocRecord<'_>) -> Result<()> {
220 check_label("url", doc.url)?;
221 if doc.url.is_empty() {
222 return Err(Error::format("url must not be empty"));
223 }
224 for (name, v) in [("title", doc.title), ("snippet", doc.snippet), ("etag", doc.etag)] {
225 if let Some(v) = v {
226 check_label(name, v)?;
227 }
228 }
229 let flags = if doc.pinned { nflags::PINNED } else { 0 };
230 self.state.staged.push(Op::UpsertNode {
231 key: doc.url_key,
232 content_hash: doc.content_hash,
233 fetched_at_ms: doc.fetched_at_ms,
234 flags,
235 url: doc.url.into(),
236 title: doc.title.map(Into::into),
237 snippet: doc.snippet.map(Into::into),
238 etag: doc.etag.map(Into::into),
239 });
240 Ok(())
241 }
242
243 pub fn set_segments(&mut self, key: UrlKey, segs: &[SegmentRecord<'_>]) -> Result<()> {
245 if segs.len() > MAX_SEGS {
246 return Err(Error::format(format!("more than {MAX_SEGS} segments")));
247 }
248 let mut owned = Vec::with_capacity(segs.len());
249 for s in segs {
250 check_label("segment", s.label)?;
251 let (importance, extra_flags) = (s.importance, 0);
252 let (byte_start, byte_len, range_flag) = match s.byte_range {
253 Some((start, len)) => (start, len, 0),
254 None => (0, 0, sflags::NO_RANGE),
255 };
256 owned.push(OwnedSeg {
257 label: s.label.into(),
258 byte_start,
259 byte_len,
260 depth: s.depth.max(1),
261 importance,
262 flags: range_flag | extra_flags,
263 });
264 }
265 self.state.staged.push(Op::SetSegments { key, segs: owned });
266 Ok(())
267 }
268
269 pub fn set_edges(&mut self, key: UrlKey, edges: &[(UrlKey, EdgeType, u16)]) {
273 self.stage_edges(key, edges.iter().map(|&(d, t, w)| (d, t, w, 0)));
274 }
275
276 pub fn set_edges_with_flags(&mut self, key: UrlKey, edges: &[(UrlKey, EdgeType, u16, u8)]) {
279 let mask = crate::format::base::eflags::EXACT;
280 self.stage_edges(key, edges.iter().map(|&(d, t, w, f)| (d, t, w, f & mask)));
281 }
282
283 fn stage_edges(
284 &mut self,
285 key: UrlKey,
286 edges: impl Iterator<Item = (UrlKey, EdgeType, u16, u8)>,
287 ) {
288 let mut seen = std::collections::HashSet::new();
289 let mut owned: Vec<(UrlKey, u8, u8, u16)> = Vec::new();
290 for (dst, etype, weight, mut flags) in edges {
291 if dst == key || !seen.insert(dst.0) {
292 continue;
293 }
294 if matches!(etype, EdgeType::Related) {
295 flags |= crate::format::base::eflags::INFERRED;
296 }
297 owned.push((dst, etype.as_tag(), flags, weight));
298 if owned.len() == MAX_EDGES {
299 break;
300 }
301 }
302 owned.sort_unstable_by_key(|(dst, _, _, _)| dst.0);
303 self.state.staged.push(Op::SetEdges { key, edges: owned });
304 }
305
306 pub fn set_attrs(&mut self, key: UrlKey, attrs: &[(u8, &str)]) -> Result<()> {
313 if attrs.len() > MAX_ATTRS {
314 return Err(Error::format("too many attributes"));
315 }
316 let mut owned: Vec<(u8, CompactString)> = Vec::with_capacity(attrs.len());
317 for &(k, v) in attrs.iter().rev() {
318 if v.len() > MAX_ATTR_LEN {
319 return Err(Error::format("attribute too long"));
320 }
321 if !owned.iter().any(|(o, _)| *o == k) {
322 owned.push((k, CompactString::from(v)));
323 }
324 }
325 if owned.iter().map(|(_, v)| v.len()).sum::<usize>() > MAX_ATTR_BYTES {
329 return Err(Error::format(format!(
330 "attributes exceed {MAX_ATTR_BYTES} bytes in total"
331 )));
332 }
333 owned.sort_unstable_by_key(|(k, _)| *k);
334 self.state.staged.push(Op::SetAttrs { key, attrs: owned });
335 Ok(())
336 }
337
338 pub fn set_named_attrs(&mut self, key: UrlKey, attrs: &[(&str, &str)]) -> Result<()> {
347 if attrs.len() > MAX_ATTRS {
348 return Err(Error::format(format!("more than {MAX_ATTRS} attributes")));
349 }
350 let mut named: Vec<(&str, &str)> = Vec::with_capacity(attrs.len());
351 for &(name, value) in attrs.iter().rev() {
352 if name.is_empty() || name.len() > MAX_ATTR_NAME || name.contains('\0') {
353 return Err(Error::format(format!(
354 "attribute name must be 1..={MAX_ATTR_NAME} bytes without NUL"
355 )));
356 }
357 if name.len() + 1 + value.len() > MAX_ATTR_LEN {
358 return Err(Error::format(format!(
359 "attribute {name:?} exceeds {MAX_ATTR_LEN} bytes"
360 )));
361 }
362 if !named.iter().any(|(n, _)| *n == name) {
363 named.push((name, value));
364 }
365 }
366 named.sort_unstable_by_key(|(n, _)| *n);
367 let mut owned: Vec<(u8, CompactString)> = Vec::with_capacity(named.len());
368 let mut total = 0usize;
369 for (k, (name, value)) in named.into_iter().enumerate() {
370 let mut rec = CompactString::with_capacity(name.len() + 1 + value.len());
371 rec.push_str(name);
372 rec.push('\0');
373 rec.push_str(value);
374 total += rec.len();
375 owned.push((k as u8, rec));
376 }
377 if total > MAX_ATTR_BYTES {
378 return Err(Error::format(format!(
379 "attributes exceed {MAX_ATTR_BYTES} bytes in total"
380 )));
381 }
382 self.state.staged.push(Op::SetAttrs { key, attrs: owned });
383 Ok(())
384 }
385
386 pub fn touch(&mut self, key: UrlKey, t: Touch<'_>) -> Result<bool> {
403 if let Some(e) = t.etag {
404 check_label("etag", e)?;
405 }
406 let staged = self.state.staged.iter().rev().find_map(|op| match op {
407 Op::UpsertNode { key: k, content_hash, etag, .. } if *k == key => {
408 Some((*content_hash, etag.clone()))
409 }
410 _ => None,
411 });
412 let (interval_s, checks, changes, last_change_ms, content_hash, etag) = {
413 let snap = self.view();
414 let Some(node) = snap.node(key) else { return Ok(false) };
415 let cfg = snap.ttl();
416 let next = cfg.next_interval(node.interval_s.max(1), t.outcome);
417 let checks = node.checks.saturating_add(1);
418 let changes = if t.outcome == Outcome::Changed {
419 node.changes.saturating_add(1)
420 } else {
421 node.changes
422 };
423 let last_change =
424 if t.outcome == Outcome::Changed { t.checked_at_ms } else { node.last_change_ms };
425 let (base_hash, base_etag) = match staged {
430 Some((h, e)) => (h, e),
431 None => (node.content_hash, node.etag.map(Into::into)),
432 };
433 let hash = t.content_hash.unwrap_or(base_hash);
434 let etag: Option<compact_str::CompactString> = t.etag.map(Into::into).or(base_etag);
435 (next, checks, changes, last_change, hash, etag)
436 };
437 self.state.staged.push(Op::Touch {
438 key,
439 checked_at_ms: t.checked_at_ms,
440 outcome: t.outcome.as_tag(),
441 content_hash: Some(content_hash),
442 etag,
443 interval_s,
444 checks,
445 changes,
446 last_change_ms,
447 tombstone: t.outcome == Outcome::Gone,
448 });
449 Ok(true)
450 }
451
452 pub fn set_pinned(&mut self, key: UrlKey, pinned: bool) -> Result<bool> {
456 let exists = self
457 .state
458 .staged
459 .iter()
460 .any(|op| matches!(op, Op::UpsertNode { key: k, .. } if *k == key))
461 || self.view().node(key).is_some();
462 if !exists {
463 return Ok(false);
464 }
465 self.state.staged.push(Op::SetPinned { key, pinned });
466 Ok(true)
467 }
468
469 pub fn remove(&mut self, key: UrlKey) -> Result<()> {
471 self.state.staged.push(Op::Remove { key });
472 Ok(())
473 }
474
475 pub fn commit(&mut self) -> Result<u64> {
482 self.snap = None;
485 if self.state.staged.is_empty() {
486 return Ok(self.state.next_seq.saturating_sub(1));
487 }
488 let held = self.store.acquire(&mut self.state, self.store.cfg.lock_wait)?;
489 let r = self.store.sync(&mut self.state, true).and_then(|_| self.append_held());
490 self.state.release(held);
491 r
492 }
493
494 fn append_held(&mut self) -> Result<u64> {
497 if self.state.staged.is_empty() {
498 return Ok(self.state.next_seq.saturating_sub(1));
499 }
500 let staged = std::mem::take(&mut self.state.staged);
501 let count = staged.len() as u64;
502 let first_seq = self.state.next_seq;
503 let mut buf = Vec::new();
504 for (i, op) in staged.iter().enumerate() {
505 buf.extend_from_slice(&encode_frame(first_seq + i as u64, op));
506 }
507 if self.state.wal.is_none() {
509 self.state.wal =
510 Some(std::fs::OpenOptions::new().append(true).open(wal_path(&self.store.dir))?);
511 }
512 let durability = self.store.cfg.durability;
513 let state = &mut *self.state;
514 let f = state.wal.as_mut().expect("opened above");
515 let sync = match durability {
516 Durability::Always => true,
517 Durability::Batch(n) => state.commits_since_sync + 1 >= n.max(1),
518 Durability::Never => false,
519 };
520 let appended = f.write_all(&buf).and_then(|()| f.flush()).and_then(|()| {
521 if !sync {
522 return Ok(());
523 }
524 #[cfg(test)]
525 crate::store::probe::bump(&crate::store::probe::FSYNCS, 1);
526 f.sync_all()
527 });
528 if let Err(e) = appended {
529 state.wal = None;
533 if let Ok(f) = std::fs::OpenOptions::new().write(true).open(wal_path(&self.store.dir)) {
534 let _ = f.set_len(state.seen.wal_pos);
535 let _ = f.sync_all();
536 }
537 return Err(e.into());
538 }
539 state.commits_since_sync = if sync { 0 } else { state.commits_since_sync + 1 };
540 state.seen.appended(buf.len() as u64);
541 let cur = self.store.current.load(Ordering::Acquire);
544 {
545 let guard = self.store.gens[cur]
546 .inner
547 .try_read()
548 .map_err(|_| Error::busy("current slot mid-install"))?;
549 let inner = guard.as_ref().ok_or_else(|| Error::corrupt("current slot empty"))?;
550 for (i, op) in staged.into_iter().enumerate() {
551 if !inner.arena.publish(first_seq + i as u64, op) {
552 return Err(Error::busy("overlay arena full"));
553 }
554 }
555 }
556 self.state.wal_bytes += buf.len() as u64;
557 self.state.next_seq += count;
558 self.state.ops_since_compact += count as usize;
559 let last = self.state.next_seq - 1;
560 if self.state.ops_since_compact >= self.store.cfg.compact_after_ops
561 || self.state.wal_bytes >= self.store.cfg.compact_after_wal_bytes
562 {
563 if self.store.compact_inline.load(Ordering::Relaxed) {
564 crate::compact::compact_locked(self.store, &mut self.state)?;
565 } else {
566 self.store.due.store(true, Ordering::Release);
568 }
569 }
570 Ok(last)
571 }
572}
573
574#[derive(Debug)]
578pub struct Txn<'t> {
579 w: Writer<'t>,
580}
581
582impl<'t> Txn<'t> {
583 #[must_use]
587 pub fn snapshot(&self) -> &Snapshot<'t> {
588 self.w.snap.as_ref().expect("a transaction holds its view for its whole life")
589 }
590
591 pub fn upsert_node(&mut self, doc: &DocRecord<'_>) -> Result<()> {
593 self.w.upsert_node(doc)
594 }
595
596 pub fn set_segments(&mut self, key: UrlKey, segs: &[SegmentRecord<'_>]) -> Result<()> {
598 self.w.set_segments(key, segs)
599 }
600
601 pub fn set_edges(&mut self, key: UrlKey, edges: &[(UrlKey, EdgeType, u16)]) {
603 self.w.set_edges(key, edges);
604 }
605
606 pub fn set_edges_with_flags(&mut self, key: UrlKey, edges: &[(UrlKey, EdgeType, u16, u8)]) {
608 self.w.set_edges_with_flags(key, edges);
609 }
610
611 pub fn set_attrs(&mut self, key: UrlKey, attrs: &[(u8, &str)]) -> Result<()> {
613 self.w.set_attrs(key, attrs)
614 }
615
616 pub fn set_named_attrs(&mut self, key: UrlKey, attrs: &[(&str, &str)]) -> Result<()> {
618 self.w.set_named_attrs(key, attrs)
619 }
620
621 pub fn touch(&mut self, key: UrlKey, t: Touch<'_>) -> Result<bool> {
623 self.w.touch(key, t)
624 }
625
626 pub fn set_pinned(&mut self, key: UrlKey, pinned: bool) -> Result<bool> {
628 self.w.set_pinned(key, pinned)
629 }
630
631 pub fn remove(&mut self, key: UrlKey) -> Result<()> {
633 self.w.remove(key)
634 }
635}
636
637pub(crate) fn transact<F, R>(
641 store: &Store,
642 lock: Option<&CommitLock>,
643 wait: std::time::Duration,
644 f: F,
645) -> Result<R>
646where
647 F: for<'t> FnOnce(&mut Txn<'t>) -> Result<R>,
648{
649 let mut txn = Txn { w: store.writer()? };
650 let held = match lock {
651 Some(_) => None,
652 None => Some(store.acquire(&mut txn.w.state, wait)?),
653 };
654 let r = run(store, &mut txn, f);
655 if let Some(h) = held {
656 txn.w.state.release(h);
657 }
658 r
659}
660
661fn run<'t, F, R>(store: &'t Store, txn: &mut Txn<'t>, f: F) -> Result<R>
662where
663 F: for<'x> FnOnce(&mut Txn<'x>) -> Result<R>,
664{
665 store.sync(&mut txn.w.state, true)?;
666 txn.w.snap = Some(store.snapshot());
667 let mark = txn.w.state.staged.len();
670 let out = f(txn);
671 txn.w.snap = None; match out {
673 Ok(r) => txn.w.append_held().map(|_| r),
674 Err(e) => {
675 txn.w.state.staged.truncate(mark);
676 Err(e)
677 }
678 }
679}