plugmem_host/db.rs
1//! `Database`: the engine + its file + the maintenance policy behind
2//! one lock (§3).
3//!
4//! The orchestration model in one paragraph: a `Database` handle is
5//! `Clone + Send + Sync` (an `Arc` around an `RwLock`-guarded engine), so
6//! any number of threads or agents in one process share one file by
7//! cloning the handle — the read verbs (`recall`/`get`/`stats`/…) run
8//! concurrently under a shared guard, the write verbs serialize under an
9//! exclusive one; at microsecond engine calls neither is a bottleneck. A
10//! second *process* (or a second `Database` on the same path) is refused
11//! with [`HostError::Locked`] by the file lock. Different files are fully
12//! independent — open as many `Database`s as you have files.
13//!
14//! Everything expensive and external — computing embeddings over HTTP —
15//! happens **before** the lock is taken: while one agent waits for its
16//! embedding provider, others keep reading and writing.
17//!
18//! (Under the `counters` perf-gate feature the engine's instrumentation
19//! `Cell`s are not `Sync`, so the lock falls back to a `Mutex` and reads
20//! serialize — a single-threaded measurement build; the public API is
21//! unchanged. See `StateLock`.)
22//!
23//! ## Overlay write path
24//!
25//! Opening a database does **not** copy its snapshot into RAM. `open`
26//! memory-maps the snapshot file and the engine *borrows* the mapped pages
27//! (an overlay over the base), replaying the journal into a small owned
28//! overlay; a mutation lands its appends in an owned tail and copies only
29//! the pages it rewrites (per-page copy-on-write in `plugmem-arena`). So a
30//! multi-gigabyte database is opened and written to while resident only in
31//! the pages it actually touches — the SQLite model. A snapshot
32//! materializes the base + overlay into a fresh file and **re-maps** it, so
33//! the overlay collapses and a long write session stays bounded. A brand-new
34//! database has no file to map yet: it opens *owned* and empty, and switches
35//! to the mapped overlay at its first snapshot.
36
37use std::cell::RefCell;
38use std::collections::BTreeMap;
39use std::fs::File;
40use std::path::{Path, PathBuf};
41#[cfg(feature = "counters")]
42use std::sync::MutexGuard;
43use std::sync::{Arc, Mutex};
44#[cfg(not(feature = "counters"))]
45use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
46
47use memmap2::Mmap;
48use plugmem_core::{
49 Config, Error, FactFault, FactRecord, LinkInput, MaintainReport, MemStorage, Memory,
50 OpenReport, RecallQuery, RecallResult, RecallScratch, RememberInput, RememberOutcome, Stats,
51 Storage,
52};
53
54thread_local! {
55 /// Per-thread recall scratch. `recall` takes `&self` on the engine, so many
56 /// reader threads recall one [`Database`] at once; each reuses its own
57 /// scratch here (zero re-alloc after warm-up, no lock on the hot path).
58 static RECALL_SCRATCH: RefCell<RecallScratch> = RefCell::new(RecallScratch::new());
59}
60
61use crate::embedder::Embedder;
62use crate::error::HostError;
63use crate::readonly::ReadOnlyDatabase;
64use crate::storage::{FileScratch, FileStorage, FsyncPolicy};
65
66self_cell::self_cell!(
67 /// Owns the memory map and the overlay [`Memory`] that borrows it — the
68 /// read-write sibling of `readonly::MappedMemory`. `self_cell` keeps the
69 /// self-reference safe: the only `unsafe` on this path is the inherent
70 /// mmap call, not the borrow.
71 struct OverlayMap {
72 owner: Mmap,
73 #[covariant]
74 dependent: OverlayMemory,
75 }
76);
77
78/// The dependent type constructor `self_cell` reborrows per access.
79/// [`Memory`] is covariant in its lifetime (its byte pools are
80/// `Cow<'a, [u8]>`), so borrowing the map is sound.
81type OverlayMemory<'a> = Memory<'a>;
82
83/// The engine backing a live [`Database`]: either an owned in-RAM engine
84/// (a brand-new database with no snapshot file yet) or an overlay over a
85/// memory-mapped snapshot (the common case). Both are mutable; verbs reach
86/// the engine through [`Engine::with`] / [`Engine::read`], which unify the
87/// two lifetimes (`'static` vs the map's) behind one closure.
88enum Engine {
89 /// No snapshot file to map yet — owned and (initially) empty. Switches to
90 /// `Mapped` at the first snapshot, once the file exists. Boxed so the
91 /// common `Mapped` case does not carry the whole owned engine inline.
92 Owned(Box<Memory<'static>>),
93 /// Overlay over a memory-mapped snapshot: the base is borrowed, mutations
94 /// live in the overlay (owned tail + per-page copy-on-write).
95 Mapped(OverlayMap),
96}
97
98impl Engine {
99 /// Reads through an immutable borrow of the engine (owned or mapped).
100 fn read<R>(&self, f: impl for<'a> FnOnce(&Memory<'a>) -> R) -> R {
101 match self {
102 Engine::Owned(mem) => f(mem),
103 Engine::Mapped(map) => f(map.borrow_dependent()),
104 }
105 }
106
107 /// Mutates the engine and its store together (disjoint borrows). The
108 /// closure is higher-ranked over the engine's lifetime so one body serves
109 /// both the `'static` owned engine and the map-bound overlay.
110 fn with<R>(
111 &mut self,
112 store: &mut FileStorage,
113 f: impl for<'a> FnOnce(&mut Memory<'a>, &mut FileStorage) -> R,
114 ) -> R {
115 match self {
116 Engine::Owned(mem) => f(mem, store),
117 Engine::Mapped(map) => map.with_dependent_mut(|_owner, mem| f(mem, store)),
118 }
119 }
120}
121
122/// Opens the engine at `store`'s path: memory-maps the snapshot
123/// and borrows it as an overlay, replaying the journal. A missing snapshot
124/// file (a brand-new database) opens owned and empty — the file appears at the
125/// first snapshot. `store` must already hold the exclusive lock.
126fn open_engine(store: &mut FileStorage, cfg: &Config) -> Result<(Engine, OpenReport), HostError> {
127 let journal = store.read_journal()?;
128 let Some(genp) = store.current_snapshot_path()? else {
129 // No published generation yet. The database is owned until the first
130 // checkpoint publishes one — but a journal may already exist (mutations
131 // before any snapshot), so still replay it into the owned engine.
132 let (mem, report) = Memory::from_bytes(None, &journal, cfg.clone())?;
133 return Ok((Engine::Owned(Box::new(mem)), report));
134 };
135 let file = File::open(&genp).map_err(|e| HostError::io(&genp, e))?;
136 // SAFETY: mapping a file is inherently unsafe — a concurrent truncate or
137 // overwrite of the mapped file would fault the process (SIGBUS/exception)
138 // on the next page access. Our correctness argument: the
139 // generation file is **immutable** (a checkpoint publishes a new one and
140 // never rewrites this), and the `store` holds the exclusive writer lock, so
141 // nothing overwrites it under the map. A foreign `truncate`/`rm` under a
142 // live handle is out of contract — the same caveat as corrupting any
143 // database file under a running engine.
144 let map = unsafe { Mmap::map(&file) }.map_err(|e| HostError::io(&genp, e))?;
145 // The `File` handle is no longer needed: `Mmap` owns the mapping.
146 drop(file);
147 // Replay the journal into the overlay: no whole-arena clone, only the
148 // touched pages copy up. `self_cell` builds the engine borrowing the map;
149 // the replay report is captured out of the constructor closure.
150 let mut report = None;
151 let mapped = OverlayMap::try_new(map, |m| {
152 let (mem, rep) = Memory::from_bytes_overlay(&m[..], &journal, cfg.clone())?;
153 report = Some(rep);
154 Ok::<_, Error>(mem)
155 })?;
156 Ok((Engine::Mapped(mapped), report.unwrap_or_default()))
157}
158
159/// An owned view of one fact — [`Memory::get`] returns borrows that
160/// cannot cross the lock, so the database hands out copies.
161#[derive(Clone, Debug, PartialEq)]
162#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
163pub struct FactSnapshot {
164 /// The raw record (temporality, flags, references).
165 pub record: FactRecord,
166 /// The fact text.
167 pub text: String,
168 /// The fact's metadata as a sorted key→value map (empty when the fact
169 /// carries none). The engine stores it opaquely; this is the decoded view.
170 pub metadata: BTreeMap<String, String>,
171}
172
173/// One exported fact — the human-readable, id-free shape [`Database::export`]
174/// dumps and an importer re-`remember`s. Internal ids and
175/// `recorded_at` are the engine's bookkeeping and are *not* preserved across
176/// a round-trip; the knowledge itself (text, subject name, tags, validity
177/// start) is.
178#[derive(Clone, Debug, PartialEq)]
179#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
180pub struct ExportedFact {
181 /// The fact text.
182 pub text: String,
183 /// Subject entity name, if the fact had one.
184 pub entity: Option<String>,
185 /// Tag strings.
186 pub tags: Vec<String>,
187 /// Metadata as a sorted key→value map (empty when none) — preserved on
188 /// import.
189 pub metadata: BTreeMap<String, String>,
190 /// When the memory learned it (informational; not restorable on import).
191 pub recorded_at: u64,
192 /// Validity start — preserved on import.
193 pub valid_from: u64,
194}
195
196/// The outcome of a [`Database::recover`] salvage.
197#[derive(Clone, Copy, Debug, PartialEq, Eq)]
198#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
199pub struct RecoverReport {
200 /// Facts written to the destination (the survivors after the purge).
201 pub kept: usize,
202 /// Facts dropped because their stored text was not valid UTF-8.
203 pub dropped_text: usize,
204 /// Facts dropped because their vector slot was out of range or mismatched.
205 pub dropped_vector: usize,
206 /// Facts dropped because their metadata blob did not decode to a
207 /// well-formed key→value map.
208 pub dropped_metadata: usize,
209}
210
211/// Visits the currently-open facts (skipping closed revisions and tombstones),
212/// resolving each subject name and tag string, calling `f` once per fact. The
213/// streaming core of export: a caller that writes each fact out (CLI `export`)
214/// never materializes the whole dump, so a huge database exports without a RAM
215/// spike. Shared by the read-write and read-only handles.
216/// Decodes a fact's metadata into an owned, sorted key→value map (empty when
217/// the fact carries none). Shared by `get` (read-write and read-only) and
218/// `export`; the pairs come back from the engine in canonical order, so the
219/// resulting `BTreeMap` matches the raw core view key-for-key.
220pub(crate) fn metadata_map(mem: &Memory, id: plugmem_core::FactId) -> BTreeMap<String, String> {
221 let mut pairs = Vec::new();
222 mem.metadata_of(id, &mut pairs);
223 pairs
224 .into_iter()
225 .map(|(k, v)| (k.to_string(), v.to_string()))
226 .collect()
227}
228
229pub(crate) fn export_facts_each(mem: &Memory, mut f: impl FnMut(ExportedFact)) {
230 use plugmem_core::{EntityId, FactId, VALID_TO_OPEN};
231 let next = mem.stats().next_fact;
232 let mut terms = Vec::new();
233 for i in 0..next {
234 let id = FactId(i);
235 let Some(view) = mem.get(id) else {
236 continue; // unknown or tombstoned
237 };
238 if view.record.valid_to != VALID_TO_OPEN {
239 continue; // a closed revision — export the current state only
240 }
241 let entity = (view.record.entity != EntityId::NONE)
242 .then(|| mem.entity_name(view.record.entity))
243 .flatten()
244 .map(str::to_string);
245 terms.clear();
246 mem.tags_of(id, &mut terms);
247 let tags = terms.iter().map(|t| mem.term(*t).to_string()).collect();
248 f(ExportedFact {
249 text: view.text.to_string(),
250 entity,
251 tags,
252 metadata: metadata_map(mem, id),
253 recorded_at: view.record.recorded_at,
254 valid_from: view.record.valid_from,
255 });
256 }
257}
258
259/// Collects the currently-open facts into a `Vec` (the owning form of
260/// [`export_facts_each`]). Used where the whole dump is wanted in memory.
261pub(crate) fn export_facts(mem: &Memory) -> Vec<ExportedFact> {
262 let mut out = Vec::new();
263 export_facts_each(mem, |e| out.push(e));
264 out
265}
266
267/// Tuning knobs of a [`Database`]. Construct through
268/// [`Database::builder`].
269pub struct DatabaseBuilder {
270 cfg: Config,
271 fsync: FsyncPolicy,
272 snapshot_every_ops: u64,
273 snapshot_journal_bytes: u64,
274 maintain_every_forgets: Option<u64>,
275 embedder: Option<Box<dyn Embedder>>,
276}
277
278impl DatabaseBuilder {
279 /// Journal fsync policy (default: every operation).
280 pub fn fsync(mut self, policy: FsyncPolicy) -> Self {
281 self.fsync = policy;
282 self
283 }
284
285 /// Auto-snapshot after this many mutations (default 1024; `0`
286 /// disables the count trigger).
287 pub fn snapshot_every_ops(mut self, ops: u64) -> Self {
288 self.snapshot_every_ops = ops;
289 self
290 }
291
292 /// Auto-snapshot when the journal outgrows this many bytes (default
293 /// 4 MiB; `0` disables the size trigger).
294 pub fn snapshot_journal_bytes(mut self, bytes: u64) -> Self {
295 self.snapshot_journal_bytes = bytes;
296 self
297 }
298
299 /// Optional auto-`maintain` after this many forgets (default off —
300 /// maintenance is O(database) and the first pass beyond the HNSW
301 /// threshold pays the graph build).
302 pub fn maintain_every_forgets(mut self, forgets: u64) -> Self {
303 self.maintain_every_forgets = Some(forgets);
304 self
305 }
306
307 /// The embedding provider. When set (and its `dim() > 0`),
308 /// `remember` without a vector embeds the fact text and `recall`
309 /// with a text but no vector embeds the query — both outside the
310 /// database lock. `Config::dim` must equal the embedder's dimension.
311 pub fn embedder(mut self, embedder: Box<dyn Embedder>) -> Self {
312 self.embedder = Some(embedder);
313 self
314 }
315
316 /// Opens (or creates) the database at `path`.
317 ///
318 /// # Errors
319 ///
320 /// [`HostError::Locked`] when the file is owned elsewhere;
321 /// [`HostError::Engine`] for config/snapshot/journal problems
322 /// (including an embedder dimension that disagrees with
323 /// `Config::dim`); [`HostError::Io`] for filesystem failures.
324 pub fn open(self, path: impl Into<PathBuf>) -> Result<(Database, OpenReport), HostError> {
325 if let Some(embedder) = &self.embedder {
326 let dim = embedder.dim();
327 if dim != 0 && dim != self.cfg.dim {
328 return Err(HostError::Engine(Error::ConfigMismatch(
329 "embedder dimension must equal Config::dim",
330 )));
331 }
332 }
333 let mut store = FileStorage::open(path, self.fsync)?;
334 let (engine, report) = open_engine(&mut store, &self.cfg)?;
335 let db = Database {
336 inner: Arc::new(Inner {
337 state: StateLock::new(State {
338 engine,
339 store,
340 ops: 0,
341 forgets: 0,
342 }),
343 embedder: self.embedder.map(Mutex::new),
344 cfg: self.cfg,
345 snapshot_every_ops: self.snapshot_every_ops,
346 snapshot_journal_bytes: self.snapshot_journal_bytes,
347 maintain_every_forgets: self.maintain_every_forgets,
348 }),
349 };
350 Ok((db, report))
351 }
352}
353
354/// The engine lock. Normally an `RwLock` so read-only verbs run concurrently
355/// (the whole point of Variant 1). Under `counters`, `State` embeds the arena's
356/// non-`Sync` instrumentation `Cell`s, and `RwLock<T>` needs `T: Sync` to hand
357/// out shared guards — so there we fall back to a `Mutex`. `counters` is a
358/// single-threaded perf-gate build, so serialized readers cost nothing there,
359/// and the `Mutex` keeps `Database: Send + Sync` so every test still builds.
360#[cfg(not(feature = "counters"))]
361type StateLock = RwLock<State>;
362#[cfg(feature = "counters")]
363type StateLock = Mutex<State>;
364
365struct Inner {
366 state: StateLock,
367 embedder: Option<Mutex<Box<dyn Embedder>>>,
368 /// Kept to rebuild the overlay engine after a re-map on snapshot.
369 cfg: Config,
370 snapshot_every_ops: u64,
371 snapshot_journal_bytes: u64,
372 maintain_every_forgets: Option<u64>,
373}
374
375struct State {
376 engine: Engine,
377 store: FileStorage,
378 /// Mutations since the last snapshot.
379 ops: u64,
380 /// Forgets since the last maintain.
381 forgets: u64,
382}
383
384/// A clonable, thread-safe handle to one database file. See the module
385/// docs for the concurrency model.
386#[derive(Clone)]
387pub struct Database {
388 inner: Arc<Inner>,
389}
390
391impl Database {
392 /// Opens `path` with every knob at its default and no embedder.
393 pub fn open(path: impl Into<PathBuf>, cfg: Config) -> Result<(Self, OpenReport), HostError> {
394 Self::builder(cfg).open(path)
395 }
396
397 /// Opens `path` read-only over a memory-mapped snapshot:
398 /// the engine borrows the mapped pages instead of copying the file
399 /// into RAM, so a large read-mostly database residents only the pages
400 /// `recall`/`get` touch. Requires a checkpointed database (empty
401 /// journal) and takes a shared lock (N readers or one writer).
402 /// See [`ReadOnlyDatabase`].
403 ///
404 /// # Errors
405 ///
406 /// [`HostError::Locked`], [`HostError::NeedsCheckpoint`],
407 /// [`HostError::Io`], [`HostError::Engine`] — see
408 /// [`ReadOnlyDatabase::open`] semantics.
409 pub fn open_readonly(
410 path: impl Into<PathBuf>,
411 cfg: Config,
412 ) -> Result<ReadOnlyDatabase, HostError> {
413 ReadOnlyDatabase::open(path, cfg)
414 }
415
416 /// Starts a configured open (knobs).
417 pub fn builder(cfg: Config) -> DatabaseBuilder {
418 DatabaseBuilder {
419 cfg,
420 fsync: FsyncPolicy::default(),
421 snapshot_every_ops: 1024,
422 snapshot_journal_bytes: 4 * 1024 * 1024,
423 maintain_every_forgets: None,
424 embedder: None,
425 }
426 }
427
428 /// A shared (read) guard — for the read-only verbs (`recall`/`get`/
429 /// `stats`/`export`/`verify`). Many run at once; they exclude only writers.
430 /// (Under `counters` the lock is a `Mutex`, so reads serialize — see
431 /// [`StateLock`].) A panicked verb cannot leave the engine half-mutated
432 /// (check first, mutate last is the engine's own law), so a poisoned lock
433 /// is recoverable.
434 #[cfg(not(feature = "counters"))]
435 fn read(&self) -> RwLockReadGuard<'_, State> {
436 self.inner.state.read().unwrap_or_else(|e| e.into_inner())
437 }
438
439 /// An exclusive (write) guard — for the mutating verbs. Serializes writers
440 /// against each other and against every concurrent reader.
441 #[cfg(not(feature = "counters"))]
442 fn write(&self) -> RwLockWriteGuard<'_, State> {
443 self.inner.state.write().unwrap_or_else(|e| e.into_inner())
444 }
445
446 /// Under `counters` the engine lock is a `Mutex`: `read` and `write` both
447 /// take the one exclusive guard (readers serialize — acceptable for the
448 /// single-threaded perf-gate build). See [`StateLock`].
449 #[cfg(feature = "counters")]
450 fn read(&self) -> MutexGuard<'_, State> {
451 self.inner.state.lock().unwrap_or_else(|e| e.into_inner())
452 }
453
454 #[cfg(feature = "counters")]
455 fn write(&self) -> MutexGuard<'_, State> {
456 self.inner.state.lock().unwrap_or_else(|e| e.into_inner())
457 }
458
459 /// Embeds `text` outside the state lock, when an embedder is
460 /// configured. `None` = leave the input as it was.
461 fn embed_one(&self, text: &str) -> Result<Option<Vec<f32>>, HostError> {
462 let Some(embedder) = &self.inner.embedder else {
463 return Ok(None);
464 };
465 let mut embedder = embedder.lock().unwrap_or_else(|e| e.into_inner());
466 if embedder.dim() == 0 {
467 return Ok(None);
468 }
469 let mut vs = embedder.embed(&[text])?;
470 Ok(Some(vs.remove(0)))
471 }
472
473 /// Embeds a whole batch of texts in a **single** embedder call — outside the
474 /// lock, like [`embed_one`](Self::embed_one). `Ok(None)` when no embedder is
475 /// configured or `dim == 0`; otherwise a vector aligned one-to-one with
476 /// `texts` (the provider contract, checked by [`OpenAiCompatEmbedder`]). An
477 /// empty `texts` yields an empty vector without a round-trip. This is the one
478 /// HTTP that [`remember_many`](Self::remember_many) makes for a bulk write.
479 fn embed_many(&self, texts: &[&str]) -> Result<Option<Vec<Vec<f32>>>, HostError> {
480 let Some(embedder) = &self.inner.embedder else {
481 return Ok(None);
482 };
483 let mut embedder = embedder.lock().unwrap_or_else(|e| e.into_inner());
484 if embedder.dim() == 0 {
485 return Ok(None);
486 }
487 if texts.is_empty() {
488 return Ok(Some(Vec::new()));
489 }
490 Ok(Some(embedder.embed(texts)?))
491 }
492
493 /// Writes a full snapshot and re-maps the fresh file.
494 ///
495 /// Materializes the borrowed base + overlay into an owned buffer, drops
496 /// the current map, writes the buffer (tmp + fsync + rename) and clears
497 /// the journal, then maps the new file into a fresh overlay. The re-map
498 /// collapses the overlay so a long write session stays bounded, and
499 /// dropping the map **before** the rename keeps the write portable
500 /// (a mapped file cannot be renamed over on Windows).
501 fn resnapshot(&self, st: &mut State, now: u64) -> Result<(), HostError> {
502 // Stream the image straight to the tmp file — never a full-image Vec
503 // This reads through the live map, so it happens
504 // **before** the map is dropped.
505 {
506 let State { engine, store, .. } = &mut *st;
507 store.stage_snapshot(|sink| {
508 engine
509 .read(|mem| mem.write_snapshot_to(now, &mut *sink))
510 .map_err(HostError::from)
511 })?;
512 }
513 // Drop the current map before the rename: park a cheap empty engine.
514 // It is replaced by the fresh overlay below — or, if the commit fails,
515 // rebuilt from the intact on-disk snapshot + journal.
516 st.engine = Engine::Owned(Box::new(Memory::new(self.inner.cfg.clone())?));
517 let write = st
518 .store
519 .commit_snapshot()
520 .and_then(|()| st.store.clear_journal());
521 // Re-open regardless: on success the fresh file, on failure the
522 // untouched old file + journal (journal replay is idempotent, so a
523 // failed `clear_journal` does not corrupt state). Then surface the
524 // commit error, if any.
525 let (engine, _) = open_engine(&mut st.store, &self.inner.cfg)?;
526 st.engine = engine;
527 write
528 }
529
530 /// The post-mutation policy hook: counts the op, fires auto-maintain
531 /// and auto-snapshot inside the same critical section.
532 fn after_mutation(&self, st: &mut State, now: u64) -> Result<(), HostError> {
533 st.ops += 1;
534 if let Some(threshold) = self.inner.maintain_every_forgets
535 && st.forgets >= threshold
536 {
537 let State { engine, store, .. } = &mut *st;
538 engine.with(store, |mem, store| mem.maintain(store, now))?;
539 st.forgets = 0;
540 }
541 let by_ops = self.inner.snapshot_every_ops > 0 && st.ops >= self.inner.snapshot_every_ops;
542 let by_bytes = self.inner.snapshot_journal_bytes > 0
543 && st.store.journal_bytes() >= self.inner.snapshot_journal_bytes;
544 if by_ops || by_bytes {
545 self.resnapshot(st, now)?;
546 st.ops = 0;
547 }
548 Ok(())
549 }
550
551 /// Remembers a fact. Without an explicit vector and with an embedder
552 /// configured, the text is embedded first — outside the lock.
553 pub fn remember(&self, input: RememberInput<'_>) -> Result<RememberOutcome, HostError> {
554 let embedded = match input.vector {
555 Some(_) => None,
556 None => self.embed_one(input.text)?,
557 };
558 let input = RememberInput {
559 vector: embedded.as_deref().or(input.vector),
560 ..input
561 };
562 let mut st = self.write();
563 let State { engine, store, .. } = &mut *st;
564 let out = engine.with(store, |mem, store| mem.remember(store, input))?;
565 self.after_mutation(&mut st, input.now)?;
566 Ok(out)
567 }
568
569 /// Remembers a **batch** of facts in one shot — the bulk-write path (CLI
570 /// `import`). Equivalent to [`remember`](Self::remember) on each input in
571 /// order, but far cheaper for a batch: the texts that need embedding are
572 /// embedded together in **one** embedder round-trip (outside the lock), and
573 /// all facts are written under **one** write-guard with **one** post-mutation
574 /// policy pass — instead of N HTTP calls and N critical sections.
575 ///
576 /// Inputs that already carry a `vector` are not re-embedded. **Chunking is
577 /// the caller's job**: this writes the whole slice it is given, so a caller
578 /// that needs bounded memory / a bounded HTTP body passes fixed-size batches
579 /// (CLI `import` streams the file in `--batch`-sized slices).
580 ///
581 /// **Fail-fast:** the first engine error returns `Err`; the facts written
582 /// before it stay written (exactly as separate `remember`s — the journal
583 /// replay is idempotent, so a retried bulk load is safe). Returns one
584 /// [`RememberOutcome`] per input, in order.
585 pub fn remember_many(
586 &self,
587 inputs: Vec<RememberInput<'_>>,
588 ) -> Result<Vec<RememberOutcome>, HostError> {
589 if inputs.is_empty() {
590 return Ok(Vec::new());
591 }
592 // One embedder round-trip for every vector-less input's text, outside the
593 // lock. `to_embed` is the vector-less inputs in order, so its result maps
594 // back onto them by a running cursor below.
595 let to_embed: Vec<&str> = inputs
596 .iter()
597 .filter(|i| i.vector.is_none())
598 .map(|i| i.text)
599 .collect();
600 let embedded = if to_embed.is_empty() {
601 None
602 } else {
603 self.embed_many(&to_embed)?
604 };
605
606 let mut st = self.write();
607 // Batch mode: journal appends skip their per-record fsync; one
608 // `sync_journal` at the end makes the whole batch durable at once.
609 st.store.set_batch(true);
610 let mut out = Vec::with_capacity(inputs.len());
611 let mut cursor = 0usize; // into `embedded`, over vector-less inputs in order
612 let mut latest = 0u64;
613 let mut failed = None;
614 for input in inputs {
615 latest = latest.max(input.now);
616 let vector = if input.vector.is_some() {
617 input.vector
618 } else if let Some(embedded) = &embedded {
619 let v = embedded[cursor].as_slice();
620 cursor += 1;
621 Some(v)
622 } else {
623 None // no embedder — lexical/structural only, as single remember
624 };
625 let input = RememberInput { vector, ..input };
626 let State { engine, store, .. } = &mut *st;
627 match engine.with(store, |mem, store| mem.remember(store, input)) {
628 Ok(o) => out.push(o),
629 Err(e) => {
630 failed = Some(HostError::from(e));
631 break;
632 }
633 }
634 }
635 // Always leave batch mode and fsync — this is the batch's durability
636 // point. On fail-fast it makes the facts written before the error durable
637 // (they stay, exactly like separate remembers).
638 st.store.set_batch(false);
639 st.store.sync_journal()?;
640 if let Some(e) = failed {
641 return Err(e);
642 }
643 // One policy pass for the whole batch. The op counter advances by one per
644 // batch; the journal-bytes threshold still fires on a large batch, so a
645 // snapshot is not starved.
646 self.after_mutation(&mut st, latest)?;
647 Ok(out)
648 }
649
650 /// Runs a recall. With a text, no vector and an embedder configured,
651 /// the query text is embedded first — outside the lock.
652 pub fn recall(&self, q: RecallQuery<'_>) -> Result<RecallResult, HostError> {
653 let embedded = match (q.vector, q.text) {
654 (None, Some(text)) => self.embed_one(text)?,
655 _ => None,
656 };
657 let q = RecallQuery {
658 vector: embedded.as_deref().or(q.vector),
659 ..q
660 };
661 // A shared guard: concurrent recalls run in parallel. `recall_into`
662 // takes `&self` on the engine and a per-thread scratch, so there is no
663 // writer path and no cross-reader contention on the hot path.
664 let st = self.read();
665 RECALL_SCRATCH.with(|scratch| {
666 let mut scratch = scratch.borrow_mut();
667 let mut out = RecallResult::default();
668 st.engine
669 .read(|mem| mem.recall_into(q, &mut scratch, &mut out))?;
670 Ok(out)
671 })
672 }
673
674 /// Revises `target` (same auto-embedding rule as `remember`).
675 pub fn revise(
676 &self,
677 target: plugmem_core::FactId,
678 input: RememberInput<'_>,
679 ) -> Result<RememberOutcome, HostError> {
680 let embedded = match input.vector {
681 Some(_) => None,
682 None => self.embed_one(input.text)?,
683 };
684 let input = RememberInput {
685 vector: embedded.as_deref().or(input.vector),
686 ..input
687 };
688 let mut st = self.write();
689 let State { engine, store, .. } = &mut *st;
690 let out = engine.with(store, |mem, store| mem.revise(store, target, input))?;
691 self.after_mutation(&mut st, input.now)?;
692 Ok(out)
693 }
694
695 /// Tombstones a fact.
696 pub fn forget(&self, now: u64, id: plugmem_core::FactId) -> Result<bool, HostError> {
697 let mut st = self.write();
698 let State { engine, store, .. } = &mut *st;
699 let fresh = engine.with(store, |mem, store| mem.forget(store, now, id))?;
700 st.forgets += 1;
701 self.after_mutation(&mut st, now)?;
702 Ok(fresh)
703 }
704
705 /// Upserts a typed edge.
706 pub fn link(&self, input: LinkInput<'_>) -> Result<(), HostError> {
707 let mut st = self.write();
708 let State { engine, store, .. } = &mut *st;
709 engine.with(store, |mem, store| mem.link(store, input))?;
710 self.after_mutation(&mut st, input.now)?;
711 Ok(())
712 }
713
714 /// An owned copy of one fact, or `None` for unknown/tombstoned ids.
715 pub fn get(&self, id: plugmem_core::FactId) -> Option<FactSnapshot> {
716 self.read().engine.read(|mem| {
717 mem.get(id).map(|v| FactSnapshot {
718 record: v.record,
719 text: v.text.to_string(),
720 metadata: metadata_map(mem, id),
721 })
722 })
723 }
724
725 /// Engine size counters.
726 pub fn stats(&self) -> Stats {
727 self.read().engine.read(|mem| mem.stats())
728 }
729
730 /// Dumps the currently-open facts for a human-readable backup
731 /// See [`ExportedFact`]. Collects the whole set; for a large
732 /// database prefer [`export_each`](Self::export_each), which streams.
733 pub fn export(&self) -> Vec<ExportedFact> {
734 self.read().engine.read(export_facts)
735 }
736
737 /// Streams the currently-open facts, calling `f` once per fact under the
738 /// read guard — the whole dump is never materialized, so a huge database
739 /// exports without a RAM spike (CLI `export` writes each line straight out).
740 /// See [`ExportedFact`].
741 pub fn export_each(&self, f: impl FnMut(ExportedFact)) {
742 self.read().engine.read(|mem| export_facts_each(mem, f));
743 }
744
745 /// Runs a maintenance pass now (purge, compaction, HNSW build past the
746 /// threshold — for the cost model).
747 ///
748 /// **Disk-first** (milestone H): the compacted image is written by streaming
749 /// the two big pools (vectors, text) through temp files and then re-mapped,
750 /// so peak RAM tracks the record count (metadata + graph), not the image
751 /// size — a database larger than RAM can be maintained. It writes a fresh
752 /// snapshot and clears the journal (like a checkpoint). The optional
753 /// auto-maintain policy (`maintain_every_forgets`) still runs in RAM inline
754 /// — it is for databases that fit.
755 ///
756 /// The report's byte counts are the on-disk image size before and after.
757 pub fn maintain(&self, now: u64) -> Result<MaintainReport, HostError> {
758 let mut st = self.write();
759 // The image size is the current snapshot generation's, not the tiny
760 // manifest at the base path.
761 let snap_len = |store: &FileStorage| -> usize {
762 store
763 .current_snapshot_path()
764 .ok()
765 .flatten()
766 .and_then(|p| std::fs::metadata(&p).ok())
767 .map(|m| m.len() as usize)
768 .unwrap_or(0)
769 };
770 let bytes_before = snap_len(&st.store);
771 let text_tmp = tmp_sibling(st.store.path(), "mtext");
772 let vec_tmp = tmp_sibling(st.store.path(), "mvec");
773
774 // Stage a compacted snapshot, streaming the big pools through scratch;
775 // this reads through the live map, so it happens before the map is
776 // dropped (as in `resnapshot`).
777 let mut purged = 0usize;
778 {
779 let State { engine, store, .. } = &mut *st;
780 store.stage_snapshot(|sink| {
781 engine.read(|mem| {
782 let mut text_scratch = FileScratch::create(&text_tmp)?;
783 let mut vec_scratch = FileScratch::create(&vec_tmp)?;
784 purged = mem
785 .snapshot_disk_first(now, &mut text_scratch, &mut vec_scratch, &mut *sink)
786 .map_err(HostError::from)?;
787 Ok(())
788 })
789 })?;
790 }
791 // Drop the current map before the rename (park a cheap empty engine),
792 // commit, clear the journal, then re-map the compacted file — exactly
793 // the `resnapshot` dance, so the map is never renamed over on Windows.
794 st.engine = Engine::Owned(Box::new(Memory::new(self.inner.cfg.clone())?));
795 st.store
796 .commit_snapshot()
797 .and_then(|()| st.store.clear_journal())?;
798 let (engine, _) = open_engine(&mut st.store, &self.inner.cfg)?;
799 st.engine = engine;
800 st.forgets = 0;
801 st.ops = 0;
802 let bytes_after = snap_len(&st.store);
803 Ok(MaintainReport {
804 purged,
805 bytes_before,
806 bytes_after,
807 })
808 }
809
810 /// Writes a full snapshot and clears the journal now (re-mapping the
811 /// fresh file — see [`Database::resnapshot`]).
812 pub fn checkpoint(&self, now: u64) -> Result<(), HostError> {
813 let mut st = self.write();
814 self.resnapshot(&mut st, now)?;
815 st.ops = 0;
816 Ok(())
817 }
818
819 /// Runs the on-demand integrity check — the equivalent of
820 /// SQLite's `integrity_check`. An open validates only the metadata, so the
821 /// large byte pools stay non-resident on an mmap'd base; this sweeps them
822 /// (text UTF-8, vector self-consistency and the fact↔slot bijection) and
823 /// reports any latent corruption. Skipping it is safe — the accessors never
824 /// panic on bad bytes; `verify` only turns corruption into an explicit
825 /// error.
826 ///
827 /// # Errors
828 ///
829 /// [`HostError::Engine`] wrapping [`Error::Corrupt`](plugmem_core::Error)
830 /// for the first inconsistency found.
831 pub fn verify(&self) -> Result<(), HostError> {
832 Ok(self.read().engine.read(|mem| mem.verify())?)
833 }
834
835 /// Salvages a content-corrupt database (Tier 2): opens `src`,
836 /// drops the facts that fail the per-fact content checks (`verify`'s
837 /// predicate), compacts the survivors and their indexes, and writes a clean
838 /// image to `dst`. `src` on disk is left untouched — the evidence is
839 /// preserved.
840 ///
841 /// It is **disk-first** (milestone H): `src` is opened as an mmap overlay
842 /// (its pages are reclaimable) and the compacted image is written by
843 /// streaming the two big pools (vectors, text) through temp files, so peak
844 /// RAM tracks the record count (metadata + HNSW graph), not the image size.
845 /// A database far larger than RAM can be recovered, as long as its graph
846 /// fits.
847 ///
848 /// This handles *content* corruption (bad text bytes, a broken fact↔slot
849 /// vector bijection). *Structural* damage — a snapshot that will not parse
850 /// — is not salvageable here: `src` fails to open and recover returns the
851 /// engine's typed error; restore from a backup instead (Tier 0).
852 ///
853 /// # Errors
854 ///
855 /// [`HostError::Locked`] if `src` or `dst` is owned elsewhere;
856 /// [`HostError::Engine`] if `src` will not parse (structural corruption) or
857 /// `dst` equals `src`; [`HostError::Io`] for filesystem failures.
858 pub fn recover(
859 src: impl AsRef<Path>,
860 dst: impl AsRef<Path>,
861 cfg: Config,
862 now: u64,
863 ) -> Result<RecoverReport, HostError> {
864 let src = src.as_ref();
865 let dst = dst.as_ref();
866
867 // Lock the source exclusively for the salvage's whole life. We never
868 // write it — the lock only excludes a cooperating writer while we read.
869 let mut src_store = FileStorage::open(src, FsyncPolicy::OnSnapshot)?;
870 let src_base = src_store.path().to_path_buf();
871
872 // The destination must be a different file: recover preserves the source
873 // as evidence and writes the clean image elsewhere.
874 let same = dst == src_base
875 || matches!(
876 (std::fs::canonicalize(dst), std::fs::canonicalize(&src_base)),
877 (Ok(a), Ok(b)) if a == b
878 );
879 if same {
880 return Err(HostError::Engine(Error::Invalid(
881 "recover destination must differ from the source",
882 )));
883 }
884
885 // Open the source as an overlay: borrow the mmap base (reclaimable
886 // pages) and replay its journal into a small owned overlay — never an
887 // owned copy of the image. A structurally corrupt image fails here —
888 // that is Tier 0, not salvageable content corruption.
889 let journal = src_store.read_journal()?;
890 let Some(genp) = src_store.current_snapshot_path()? else {
891 return Err(HostError::Engine(Error::Corrupt(
892 "source database has no published snapshot to recover",
893 )));
894 };
895 let file = File::open(&genp).map_err(|e| HostError::io(&genp, e))?;
896 // SAFETY: as in `open_engine` — the generation file is immutable and
897 // `src_store` holds the exclusive lock, so nothing touches it under us.
898 let map = unsafe { Mmap::map(&file) }.map_err(|e| HostError::io(&genp, e))?;
899 drop(file);
900 let (mut mem, _report) = Memory::from_bytes_overlay(&map[..], &journal, cfg.clone())?;
901
902 // Drop each content-faulty fact into a throwaway store, so the source
903 // file is never written. The disk-first rebuild below then physically
904 // purges them and rebuilds clean indexes + HNSW from the survivors.
905 let mut scratch = MemStorage::new();
906 let mut dropped_text = 0usize;
907 let mut dropped_vector = 0usize;
908 let mut dropped_metadata = 0usize;
909 for (id, fault) in mem.faulty_facts() {
910 mem.forget(&mut scratch, now, id)?;
911 match fault {
912 FactFault::Text => dropped_text += 1,
913 FactFault::Vector => dropped_vector += 1,
914 FactFault::Metadata => dropped_metadata += 1,
915 }
916 }
917
918 // Write the compacted image to `dst`, streaming the big pools through
919 // temp scratch files (metadata + graph are the only things resident).
920 let mut dst_store = FileStorage::open(dst, FsyncPolicy::OnSnapshot)?;
921 let text_tmp = tmp_sibling(dst_store.path(), "rectext");
922 let vec_tmp = tmp_sibling(dst_store.path(), "recvec");
923 let mut purged = 0usize;
924 dst_store.stage_snapshot(|sink| {
925 let mut text_scratch = FileScratch::create(&text_tmp)?;
926 let mut vec_scratch = FileScratch::create(&vec_tmp)?;
927 purged = mem
928 .snapshot_disk_first(now, &mut text_scratch, &mut vec_scratch, &mut *sink)
929 .map_err(HostError::from)?;
930 Ok(())
931 })?;
932 dst_store.commit_snapshot()?;
933
934 let kept = mem.stats().facts.saturating_sub(purged);
935 Ok(RecoverReport {
936 kept,
937 dropped_text,
938 dropped_vector,
939 dropped_metadata,
940 })
941 }
942}
943
944/// A temp-file path beside `base` with the given tag (for disk-first scratch).
945fn tmp_sibling(base: &Path, tag: &str) -> PathBuf {
946 let mut p = base.as_os_str().to_os_string();
947 p.push(".");
948 p.push(tag);
949 p.push(".tmp");
950 PathBuf::from(p)
951}
952
953impl std::fmt::Debug for Database {
954 /// Summary only — the contents are the user's memory.
955 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
956 let stats = self.stats();
957 f.debug_struct("Database")
958 .field("facts", &stats.facts)
959 .field("entities", &stats.entities)
960 .finish()
961 }
962}