use std::cell::RefCell;
use std::fs::File;
use std::path::{Path, PathBuf};
#[cfg(feature = "counters")]
use std::sync::Mutex;
use memmap2::Mmap;
use plugmem_core::snapshot::{DEFAULT_SCRUB_BUDGET, ScrubCursor, ScrubProgress, Snapshot};
use plugmem_core::{Config, FactId, Memory, RecallQuery, RecallResult, RecallScratch, Stats};
thread_local! {
static RECALL_SCRATCH: RefCell<RecallScratch> = RefCell::new(RecallScratch::new());
}
use crate::db::FactSnapshot;
use crate::error::HostError;
use crate::storage::{pin_current_generation, read_manifest};
self_cell::self_cell!(
struct MappedMemory {
owner: Mmap,
#[covariant]
dependent: BorrowedMemory,
}
);
type BorrowedMemory<'a> = Memory<'a>;
pub struct ReadOnlyDatabase {
#[cfg(not(feature = "counters"))]
mapped: MappedMemory,
#[cfg(feature = "counters")]
mapped: Mutex<MappedMemory>,
_pin: File,
path: PathBuf,
generation: u64,
cfg: Config,
}
impl ReadOnlyDatabase {
pub(crate) fn open(path: impl Into<PathBuf>, cfg: Config) -> Result<Self, HostError> {
let base: PathBuf = path.into();
let Some((pin, genp, generation)) = pin_current_generation(&base)? else {
return Err(HostError::NeedsCheckpoint { path: base });
};
let map = unsafe { Mmap::map(&pin) }.map_err(|e| HostError::io(&genp, e))?;
let mapped = MappedMemory::try_new(map, |map| {
Memory::from_bytes_borrowed(&map[..], &[], cfg.clone())
})?;
Ok(Self {
#[cfg(not(feature = "counters"))]
mapped,
#[cfg(feature = "counters")]
mapped: Mutex::new(mapped),
_pin: pin,
path: base,
generation,
cfg,
})
}
#[cfg(not(feature = "counters"))]
fn with_mem<R>(&self, f: impl FnOnce(&Memory<'_>) -> R) -> R {
f(self.mapped.borrow_dependent())
}
#[cfg(feature = "counters")]
fn with_mem<R>(&self, f: impl FnOnce(&Memory<'_>) -> R) -> R {
let guard = self.mapped.lock().unwrap_or_else(|e| e.into_inner());
f(guard.borrow_dependent())
}
pub fn recall(&self, q: RecallQuery<'_>) -> Result<RecallResult, HostError> {
self.with_mem(|mem| {
RECALL_SCRATCH.with(|scratch| {
let mut scratch = scratch.borrow_mut();
let mut out = RecallResult::default();
mem.recall_into(q, &mut scratch, &mut out)?;
Ok(out)
})
})
}
pub fn get(&self, id: FactId) -> Option<FactSnapshot> {
self.with_mem(|mem| {
mem.get(id).map(|v| FactSnapshot {
record: v.record,
text: v.text.to_string(),
metadata: crate::db::metadata_map(mem, id),
})
})
}
pub fn stats(&self) -> Stats {
self.with_mem(|mem| mem.stats())
}
pub fn verify(&self) -> Result<(), HostError> {
Ok(self.with_mem(|mem| mem.verify())?)
}
pub fn scrub(&self) -> Result<Scrub, HostError> {
self.scrub_with_budget(DEFAULT_SCRUB_BUDGET)
}
pub fn scrub_with_budget(&self, budget: usize) -> Result<Scrub, HostError> {
Scrub::open(&self.path, budget)
}
pub fn export(&self) -> Vec<crate::db::ExportedFact> {
self.with_mem(crate::db::export_facts)
}
pub fn export_each(&self, f: impl FnMut(crate::db::ExportedFact)) {
self.with_mem(|mem| crate::db::export_facts_each(mem, f));
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn generation(&self) -> u64 {
self.generation
}
pub fn refresh(&mut self) -> Result<bool, HostError> {
match read_manifest(&self.path)? {
Some(latest) if latest > self.generation => {}
_ => return Ok(false),
}
let Some((pin, genp, generation)) = pin_current_generation(&self.path)? else {
return Ok(false);
};
if generation <= self.generation {
return Ok(false);
}
let map = unsafe { Mmap::map(&pin) }.map_err(|e| HostError::io(&genp, e))?;
let cfg = self.cfg.clone();
let mapped =
MappedMemory::try_new(map, |map| Memory::from_bytes_borrowed(&map[..], &[], cfg))?;
#[cfg(not(feature = "counters"))]
{
self.mapped = mapped;
}
#[cfg(feature = "counters")]
{
self.mapped = Mutex::new(mapped);
}
self._pin = pin;
self.generation = generation;
Ok(true)
}
}
impl std::fmt::Debug for ReadOnlyDatabase {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let stats = self.stats();
f.debug_struct("ReadOnlyDatabase")
.field("path", &self.path)
.field("facts", &stats.facts)
.field("entities", &stats.entities)
.finish()
}
}
self_cell::self_cell!(
struct MappedScrub {
owner: Mmap,
#[covariant]
dependent: BorrowedScrub,
}
);
type BorrowedScrub<'a> = ScrubCursor<'a>;
pub struct Scrub {
mapped: MappedScrub,
_pin: File,
}
impl Scrub {
fn open(base: &Path, budget: usize) -> Result<Self, HostError> {
let Some((pin, genp, _generation)) = pin_current_generation(base)? else {
return Err(HostError::NeedsCheckpoint {
path: base.to_path_buf(),
});
};
let map = unsafe { Mmap::map(&pin) }.map_err(|e| HostError::io(&genp, e))?;
let mapped = MappedScrub::try_new(map, |map| {
Snapshot::parse(&map[..])
.map(|snap| snap.scrub_with_budget(budget))
.map_err(HostError::from)
})?;
Ok(Self { mapped, _pin: pin })
}
}
impl Iterator for Scrub {
type Item = Result<ScrubProgress, HostError>;
fn next(&mut self) -> Option<Self::Item> {
self.mapped
.with_dependent_mut(|_map, cur| cur.next())
.map(|step| step.map_err(HostError::from))
}
}
impl std::fmt::Debug for Scrub {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Scrub").finish_non_exhaustive()
}
}