1use crate::error::{Error, Result};
11use crate::format::base::{nflags, sflags};
12use crate::format::wal::{encode_frame, Op, OwnedSeg};
13use crate::key::{EdgeType, UrlKey};
14use crate::store::{wal_path, Durability, Store};
15use crate::ttl::Outcome;
16use compact_str::CompactString;
17use std::io::Write as _;
18use std::sync::atomic::Ordering;
19use std::sync::MutexGuard;
20
21pub const MAX_LABEL: usize = 2048;
23pub const MAX_SEGS: usize = 128;
25pub use crate::format::wal::MAX_EDGES;
27pub use crate::format::wal::{MAX_ATTRS, MAX_ATTR_LEN};
29pub const MAX_ATTR_BYTES: usize = crate::format::wal::MAX_FRAME - (1 + 8 + 1) - MAX_ATTRS * (1 + 4);
32pub const MAX_ATTR_NAME: usize = 255;
34
35#[derive(Clone, Copy, Debug)]
39pub struct DocRecord<'a> {
40 pub url: &'a str,
42 pub url_key: UrlKey,
44 pub content_hash: u64,
46 pub fetched_at_ms: u64,
48 pub title: Option<&'a str>,
50 pub snippet: Option<&'a str>,
52 pub etag: Option<&'a str>,
54 pub pinned: bool,
56}
57
58#[derive(Clone, Copy, Debug)]
60pub struct SegmentRecord<'a> {
61 pub label: &'a str,
63 pub byte_range: Option<(u32, u32)>,
65 pub depth: u8,
67 pub importance: u16,
69}
70
71#[derive(Clone, Copy, Debug)]
73pub struct Touch<'a> {
74 pub checked_at_ms: u64,
76 pub outcome: Outcome,
78 pub content_hash: Option<u64>,
80 pub etag: Option<&'a str>,
82}
83
84pub(crate) struct WriteState {
85 next_seq: u64,
86 staged: Vec<Op>,
87 pub(crate) ops_since_compact: usize,
88 pub(crate) wal_bytes: u64,
89 commits_since_sync: u32,
90 wal: Option<std::fs::File>,
91 pub(crate) disk_generation: u64,
96}
97
98#[allow(clippy::missing_fields_in_debug)] impl std::fmt::Debug for WriteState {
100 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101 f.debug_struct("WriteState")
102 .field("next_seq", &self.next_seq)
103 .field("staged", &self.staged.len())
104 .finish()
105 }
106}
107
108impl WriteState {
109 pub(crate) fn new(
110 next_seq: u64,
111 replayed_ops: usize,
112 wal_bytes: u64,
113 disk_generation: u64,
114 ) -> Self {
115 Self {
116 next_seq,
117 staged: Vec::new(),
118 ops_since_compact: replayed_ops,
119 wal_bytes,
120 commits_since_sync: 0,
121 wal: None,
122 disk_generation,
123 }
124 }
125
126 pub(crate) fn reset_after_compact(&mut self, next_seq: u64) {
127 self.next_seq = next_seq;
128 self.ops_since_compact = 0;
129 self.wal_bytes = u64::try_from(crate::format::wal::HEADER_LEN).unwrap_or(64);
130 self.commits_since_sync = 0;
131 self.wal = None; }
133
134 pub(crate) fn staged_is_empty(&self) -> bool {
135 self.staged.is_empty()
136 }
137
138 pub(crate) fn last_committed_seq(&self) -> u64 {
139 self.next_seq.saturating_sub(1)
140 }
141}
142
143#[derive(Debug)]
147pub struct Writer<'s> {
148 store: &'s Store,
149 state: MutexGuard<'s, WriteState>,
150 snap: Option<crate::snapshot::Snapshot<'s>>,
161}
162
163fn check_label(name: &str, s: &str) -> Result<()> {
164 if s.len() > MAX_LABEL {
165 return Err(Error::format(format!("{name} label exceeds {MAX_LABEL} bytes")));
166 }
167 Ok(())
168}
169
170impl<'s> Writer<'s> {
171 pub(crate) fn new(store: &'s Store, state: MutexGuard<'s, WriteState>) -> Self {
172 Self { store, state, snap: None }
173 }
174
175 fn view(&mut self) -> &crate::snapshot::Snapshot<'s> {
177 if self.snap.is_none() {
178 self.snap = Some(self.store.snapshot());
179 }
180 self.snap.as_ref().expect("just filled")
181 }
182
183 pub fn upsert_node(&mut self, doc: &DocRecord<'_>) -> Result<()> {
185 check_label("url", doc.url)?;
186 if doc.url.is_empty() {
187 return Err(Error::format("url must not be empty"));
188 }
189 for (name, v) in [("title", doc.title), ("snippet", doc.snippet), ("etag", doc.etag)] {
190 if let Some(v) = v {
191 check_label(name, v)?;
192 }
193 }
194 let flags = if doc.pinned { nflags::PINNED } else { 0 };
195 self.state.staged.push(Op::UpsertNode {
196 key: doc.url_key,
197 content_hash: doc.content_hash,
198 fetched_at_ms: doc.fetched_at_ms,
199 flags,
200 url: doc.url.into(),
201 title: doc.title.map(Into::into),
202 snippet: doc.snippet.map(Into::into),
203 etag: doc.etag.map(Into::into),
204 });
205 Ok(())
206 }
207
208 pub fn set_segments(&mut self, key: UrlKey, segs: &[SegmentRecord<'_>]) -> Result<()> {
210 if segs.len() > MAX_SEGS {
211 return Err(Error::format(format!("more than {MAX_SEGS} segments")));
212 }
213 let mut owned = Vec::with_capacity(segs.len());
214 for s in segs {
215 check_label("segment", s.label)?;
216 let (importance, extra_flags) = (s.importance, 0);
217 let (byte_start, byte_len, range_flag) = match s.byte_range {
218 Some((start, len)) => (start, len, 0),
219 None => (0, 0, sflags::NO_RANGE),
220 };
221 owned.push(OwnedSeg {
222 label: s.label.into(),
223 byte_start,
224 byte_len,
225 depth: s.depth.max(1),
226 importance,
227 flags: range_flag | extra_flags,
228 });
229 }
230 self.state.staged.push(Op::SetSegments { key, segs: owned });
231 Ok(())
232 }
233
234 pub fn set_edges(&mut self, key: UrlKey, edges: &[(UrlKey, EdgeType, u16)]) {
238 self.stage_edges(key, edges.iter().map(|&(d, t, w)| (d, t, w, 0)));
239 }
240
241 pub fn set_edges_with_flags(&mut self, key: UrlKey, edges: &[(UrlKey, EdgeType, u16, u8)]) {
244 let mask = crate::format::base::eflags::EXACT;
245 self.stage_edges(key, edges.iter().map(|&(d, t, w, f)| (d, t, w, f & mask)));
246 }
247
248 fn stage_edges(
249 &mut self,
250 key: UrlKey,
251 edges: impl Iterator<Item = (UrlKey, EdgeType, u16, u8)>,
252 ) {
253 let mut seen = std::collections::HashSet::new();
254 let mut owned: Vec<(UrlKey, u8, u8, u16)> = Vec::new();
255 for (dst, etype, weight, mut flags) in edges {
256 if dst == key || !seen.insert(dst.0) {
257 continue;
258 }
259 if matches!(etype, EdgeType::Related) {
260 flags |= crate::format::base::eflags::INFERRED;
261 }
262 owned.push((dst, etype.as_tag(), flags, weight));
263 if owned.len() == MAX_EDGES {
264 break;
265 }
266 }
267 owned.sort_unstable_by_key(|(dst, _, _, _)| dst.0);
268 self.state.staged.push(Op::SetEdges { key, edges: owned });
269 }
270
271 pub fn set_attrs(&mut self, key: UrlKey, attrs: &[(u8, &str)]) -> Result<()> {
278 if attrs.len() > MAX_ATTRS {
279 return Err(Error::format("too many attributes"));
280 }
281 let mut owned: Vec<(u8, CompactString)> = Vec::with_capacity(attrs.len());
282 for &(k, v) in attrs.iter().rev() {
283 if v.len() > MAX_ATTR_LEN {
284 return Err(Error::format("attribute too long"));
285 }
286 if !owned.iter().any(|(o, _)| *o == k) {
287 owned.push((k, CompactString::from(v)));
288 }
289 }
290 if owned.iter().map(|(_, v)| v.len()).sum::<usize>() > MAX_ATTR_BYTES {
294 return Err(Error::format(format!(
295 "attributes exceed {MAX_ATTR_BYTES} bytes in total"
296 )));
297 }
298 owned.sort_unstable_by_key(|(k, _)| *k);
299 self.state.staged.push(Op::SetAttrs { key, attrs: owned });
300 Ok(())
301 }
302
303 pub fn set_named_attrs(&mut self, key: UrlKey, attrs: &[(&str, &str)]) -> Result<()> {
312 if attrs.len() > MAX_ATTRS {
313 return Err(Error::format(format!("more than {MAX_ATTRS} attributes")));
314 }
315 let mut named: Vec<(&str, &str)> = Vec::with_capacity(attrs.len());
316 for &(name, value) in attrs.iter().rev() {
317 if name.is_empty() || name.len() > MAX_ATTR_NAME || name.contains('\0') {
318 return Err(Error::format(format!(
319 "attribute name must be 1..={MAX_ATTR_NAME} bytes without NUL"
320 )));
321 }
322 if name.len() + 1 + value.len() > MAX_ATTR_LEN {
323 return Err(Error::format(format!(
324 "attribute {name:?} exceeds {MAX_ATTR_LEN} bytes"
325 )));
326 }
327 if !named.iter().any(|(n, _)| *n == name) {
328 named.push((name, value));
329 }
330 }
331 named.sort_unstable_by_key(|(n, _)| *n);
332 let mut owned: Vec<(u8, CompactString)> = Vec::with_capacity(named.len());
333 let mut total = 0usize;
334 for (k, (name, value)) in named.into_iter().enumerate() {
335 let mut rec = CompactString::with_capacity(name.len() + 1 + value.len());
336 rec.push_str(name);
337 rec.push('\0');
338 rec.push_str(value);
339 total += rec.len();
340 owned.push((k as u8, rec));
341 }
342 if total > MAX_ATTR_BYTES {
343 return Err(Error::format(format!(
344 "attributes exceed {MAX_ATTR_BYTES} bytes in total"
345 )));
346 }
347 self.state.staged.push(Op::SetAttrs { key, attrs: owned });
348 Ok(())
349 }
350
351 pub fn touch(&mut self, key: UrlKey, t: Touch<'_>) -> Result<bool> {
368 if let Some(e) = t.etag {
369 check_label("etag", e)?;
370 }
371 let staged = self.state.staged.iter().rev().find_map(|op| match op {
372 Op::UpsertNode { key: k, content_hash, etag, .. } if *k == key => {
373 Some((*content_hash, etag.clone()))
374 }
375 _ => None,
376 });
377 let (interval_s, checks, changes, last_change_ms, content_hash, etag) = {
378 let snap = self.view();
379 let Some(node) = snap.node(key) else { return Ok(false) };
380 let cfg = snap.ttl();
381 let next = cfg.next_interval(node.interval_s.max(1), t.outcome);
382 let checks = node.checks.saturating_add(1);
383 let changes = if t.outcome == Outcome::Changed {
384 node.changes.saturating_add(1)
385 } else {
386 node.changes
387 };
388 let last_change =
389 if t.outcome == Outcome::Changed { t.checked_at_ms } else { node.last_change_ms };
390 let (base_hash, base_etag) = match staged {
395 Some((h, e)) => (h, e),
396 None => (node.content_hash, node.etag.map(Into::into)),
397 };
398 let hash = t.content_hash.unwrap_or(base_hash);
399 let etag: Option<compact_str::CompactString> = t.etag.map(Into::into).or(base_etag);
400 (next, checks, changes, last_change, hash, etag)
401 };
402 self.state.staged.push(Op::Touch {
403 key,
404 checked_at_ms: t.checked_at_ms,
405 outcome: t.outcome.as_tag(),
406 content_hash: Some(content_hash),
407 etag,
408 interval_s,
409 checks,
410 changes,
411 last_change_ms,
412 tombstone: t.outcome == Outcome::Gone,
413 });
414 Ok(true)
415 }
416
417 pub fn set_pinned(&mut self, key: UrlKey, pinned: bool) -> Result<bool> {
421 let exists = self
422 .state
423 .staged
424 .iter()
425 .any(|op| matches!(op, Op::UpsertNode { key: k, .. } if *k == key))
426 || self.view().node(key).is_some();
427 if !exists {
428 return Ok(false);
429 }
430 self.state.staged.push(Op::SetPinned { key, pinned });
431 Ok(true)
432 }
433
434 pub fn remove(&mut self, key: UrlKey) -> Result<()> {
436 self.state.staged.push(Op::Remove { key });
437 Ok(())
438 }
439
440 pub fn commit(&mut self) -> Result<u64> {
445 self.snap = None;
448 if self.state.staged.is_empty() {
449 return Ok(self.state.next_seq.saturating_sub(1));
450 }
451 let staged = std::mem::take(&mut self.state.staged);
452 let count = staged.len() as u64;
453 let first_seq = self.state.next_seq;
454 let mut buf = Vec::new();
455 for (i, op) in staged.iter().enumerate() {
456 buf.extend_from_slice(&encode_frame(first_seq + i as u64, op));
457 }
458 if self.state.wal.is_none() {
460 self.state.wal =
461 Some(std::fs::OpenOptions::new().append(true).open(wal_path(&self.store.dir))?);
462 }
463 {
464 let state = &mut *self.state;
465 let f = state.wal.as_mut().expect("opened above");
466 f.write_all(&buf)?;
467 f.flush()?;
468 let sync = match self.store.cfg.durability {
469 Durability::Always => true,
470 Durability::Batch(n) => {
471 state.commits_since_sync += 1;
472 state.commits_since_sync >= n.max(1)
473 }
474 Durability::Never => false,
475 };
476 if sync {
477 f.sync_all()?;
478 state.commits_since_sync = 0;
479 }
480 }
481 let cur = self.store.current.load(Ordering::Acquire);
484 {
485 let guard = self.store.gens[cur]
486 .inner
487 .try_read()
488 .map_err(|_| Error::busy("current slot mid-install"))?;
489 let inner = guard.as_ref().ok_or_else(|| Error::corrupt("current slot empty"))?;
490 for (i, op) in staged.into_iter().enumerate() {
491 if !inner.arena.publish(first_seq + i as u64, op) {
492 return Err(Error::busy("overlay arena full"));
493 }
494 }
495 }
496 self.state.wal_bytes += buf.len() as u64;
497 self.state.next_seq += count;
498 self.state.ops_since_compact += count as usize;
499 let last = self.state.next_seq - 1;
500 if self.state.ops_since_compact >= self.store.cfg.compact_after_ops
501 || self.state.wal_bytes >= self.store.cfg.compact_after_wal_bytes
502 {
503 if self.store.compact_inline.load(Ordering::Relaxed) {
504 crate::compact::compact_locked(self.store, &mut self.state)?;
505 } else {
506 self.store.due.store(true, Ordering::Release);
508 }
509 }
510 Ok(last)
511 }
512}