mod registry;
use std::fmt;
use std::path::{Path, PathBuf};
use std::sync::{Mutex, MutexGuard};
use crate::{Database, HostError};
pub use registry::{
ARCHIVED_TAG, DbEntry, Description, ENTRY_TAG, ReindexReport, SELF_ENTITY, WorkspaceIssue,
};
pub const MAX_DB_NAME: usize = 64;
const DB_DIR: &str = "db";
const REGISTRY_FILE: &str = "registry.plugmem";
const DB_EXT: &str = "plugmem";
const RESERVED_DEVICE_NAMES: &[&str] = &[
"con", "prn", "aux", "nul", "com1", "com2", "com3", "com4", "com5", "com6", "com7", "com8",
"com9", "lpt1", "lpt2", "lpt3", "lpt4", "lpt5", "lpt6", "lpt7", "lpt8", "lpt9",
];
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum NameProblem {
Empty,
TooLong,
LeadingChar,
Character,
ReservedDevice,
}
impl fmt::Display for NameProblem {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Empty => f.write_str("it is empty"),
Self::TooLong => write!(f, "it is longer than {MAX_DB_NAME} bytes"),
Self::LeadingChar => f.write_str("it must start with a lowercase letter or a digit"),
Self::Character => {
f.write_str("it may hold only lowercase letters, digits, '-' and '_'")
}
Self::ReservedDevice => f.write_str(
"it is a Windows device name, which would open a device rather than a file there",
),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct DbName(String);
impl DbName {
pub fn parse(s: &str) -> Result<Self, WorkspaceError> {
let bad = |why| {
Err(WorkspaceError::BadName {
name: s.to_string(),
why,
})
};
let Some(&first) = s.as_bytes().first() else {
return bad(NameProblem::Empty);
};
if !first.is_ascii_lowercase() && !first.is_ascii_digit() {
return bad(NameProblem::LeadingChar);
}
if !s.bytes().all(is_name_byte) {
return bad(NameProblem::Character);
}
if s.len() > MAX_DB_NAME {
return bad(NameProblem::TooLong);
}
if RESERVED_DEVICE_NAMES.contains(&s) {
return bad(NameProblem::ReservedDevice);
}
Ok(DbName(s.to_string()))
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for DbName {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.0)
}
}
fn is_name_byte(b: u8) -> bool {
b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'-' || b == b'_'
}
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum WorkspaceError {
#[error("{name:?} is not a usable database name: {why}")]
BadName {
name: String,
why: NameProblem,
},
#[error("no database named {name} in this workspace (looked for {})", path.display())]
NoSuchDatabase {
name: DbName,
path: PathBuf,
},
#[error(
"database {name} is in use by another process; it is released once that process closes it (a pooled handle does so after its idle timeout)"
)]
Busy {
name: DbName,
},
#[error("i/o on {}: {source}", path.display())]
Io {
path: PathBuf,
#[source]
source: std::io::Error,
},
#[error(transparent)]
Host(#[from] HostError),
}
impl WorkspaceError {
pub(crate) fn io(path: &Path, source: std::io::Error) -> Self {
Self::Io {
path: path.to_path_buf(),
source,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WorkspaceLayout {
root: PathBuf,
}
impl WorkspaceLayout {
pub fn new(root: impl Into<PathBuf>) -> Self {
Self { root: root.into() }
}
pub fn root(&self) -> &Path {
&self.root
}
pub fn db_dir(&self) -> PathBuf {
self.root.join(DB_DIR)
}
pub fn path_of(&self, name: &DbName) -> PathBuf {
self.db_dir().join(format!("{}.{DB_EXT}", name.0))
}
pub fn registry_path(&self) -> PathBuf {
self.root.join(REGISTRY_FILE)
}
pub fn exists(&self, name: &DbName) -> bool {
crate::storage::database_exists(&self.path_of(name))
}
pub fn list(&self) -> Result<Vec<DbName>, WorkspaceError> {
let dir = self.db_dir();
let entries = match std::fs::read_dir(&dir) {
Ok(entries) => entries,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
Err(e) => return Err(WorkspaceError::io(&dir, e)),
};
let mut names = Vec::new();
for entry in entries {
let entry = entry.map_err(|e| WorkspaceError::io(&dir, e))?;
let file_name = entry.file_name();
let Some(candidate) = file_name.to_str().and_then(|n| n.split('.').next()) else {
continue;
};
if let Ok(name) = DbName::parse(candidate)
&& !names.contains(&name)
&& self.exists(&name)
{
names.push(name);
}
}
names.sort_unstable();
Ok(names)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct WorkspaceLimits {
pub max_open: usize,
pub idle_timeout_ms: u64,
}
pub const DEFAULT_MAX_OPEN: usize = 16;
pub const DEFAULT_IDLE_TIMEOUT_MS: u64 = 60_000;
const FILES_PER_OPEN_DATABASE: usize = 4;
const ASSUMED_FD_LIMIT: usize = 1024;
const RESERVED_FDS: usize = 64;
pub const MAX_OPEN_CEILING: usize = (ASSUMED_FD_LIMIT - RESERVED_FDS) / FILES_PER_OPEN_DATABASE;
const _: () = {
assert!(MAX_OPEN_CEILING > DEFAULT_MAX_OPEN);
};
impl Default for WorkspaceLimits {
fn default() -> Self {
Self {
max_open: DEFAULT_MAX_OPEN,
idle_timeout_ms: DEFAULT_IDLE_TIMEOUT_MS,
}
}
}
impl WorkspaceLimits {
pub fn ceiling(&self) -> usize {
self.max_open.clamp(1, MAX_OPEN_CEILING)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum IfMissing {
Create,
Fail,
}
pub type Opener = Box<dyn Fn(&Path) -> Result<Database, HostError> + Send + Sync>;
struct Pooled {
name: DbName,
db: Database,
last_used_ms: u64,
}
pub struct Workspace {
layout: WorkspaceLayout,
open: Opener,
limits: WorkspaceLimits,
pool: Mutex<Vec<Pooled>>,
registry: Mutex<Option<Database>>,
}
impl Workspace {
pub fn new(layout: WorkspaceLayout, open: Opener, limits: WorkspaceLimits) -> Self {
Self {
layout,
open,
limits,
pool: Mutex::new(Vec::new()),
registry: Mutex::new(None),
}
}
pub fn registry(&self) -> Result<Database, WorkspaceError> {
let mut slot = self.registry.lock().unwrap_or_else(|e| e.into_inner());
if let Some(db) = slot.as_ref() {
return Ok(db.clone());
}
let root = self.layout.root();
std::fs::create_dir_all(root).map_err(|e| WorkspaceError::io(root, e))?;
let db = (self.open)(&self.layout.registry_path())?;
*slot = Some(db.clone());
Ok(db)
}
pub fn close_registry(&self) -> bool {
self.registry
.lock()
.unwrap_or_else(|e| e.into_inner())
.take()
.is_some()
}
pub fn layout(&self) -> &WorkspaceLayout {
&self.layout
}
pub fn limits(&self) -> WorkspaceLimits {
self.limits
}
pub fn open_count(&self) -> usize {
self.pooled().len()
}
pub fn get(
&self,
name: &DbName,
now_ms: u64,
missing: IfMissing,
) -> Result<Database, WorkspaceError> {
let mut pool = self.pooled();
if let Some(slot) = pool.iter_mut().find(|p| &p.name == name) {
slot.last_used_ms = now_ms;
return Ok(slot.db.clone());
}
let path = self.layout.path_of(name);
if !self.layout.exists(name) {
if missing == IfMissing::Fail {
return Err(WorkspaceError::NoSuchDatabase {
name: name.clone(),
path,
});
}
let dir = self.layout.db_dir();
std::fs::create_dir_all(&dir).map_err(|e| WorkspaceError::io(&dir, e))?;
}
let ceiling = self.limits.ceiling();
while pool.len() >= ceiling {
let lru = Self::least_recently_used(&pool);
pool.remove(lru);
}
let db = (self.open)(&path).map_err(|e| match e {
HostError::Locked { .. } => WorkspaceError::Busy { name: name.clone() },
other => WorkspaceError::Host(other),
})?;
pool.push(Pooled {
name: name.clone(),
db: db.clone(),
last_used_ms: now_ms,
});
Ok(db)
}
pub fn close_idle(&self, now_ms: u64) -> usize {
let timeout = self.limits.idle_timeout_ms;
if timeout == 0 {
return 0;
}
let mut pool = self.pooled();
let before = pool.len();
pool.retain(|p| now_ms.saturating_sub(p.last_used_ms) < timeout);
before - pool.len()
}
pub fn close_all(&self) -> usize {
let mut pool = self.pooled();
std::mem::take(&mut *pool).len()
}
fn pooled(&self) -> MutexGuard<'_, Vec<Pooled>> {
self.pool.lock().unwrap_or_else(|e| e.into_inner())
}
fn least_recently_used(pool: &[Pooled]) -> usize {
let mut oldest = 0;
for (i, p) in pool.iter().enumerate() {
if p.last_used_ms < pool[oldest].last_used_ms {
oldest = i;
}
}
oldest
}
}
impl fmt::Debug for Workspace {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Workspace")
.field("root", &self.layout.root())
.field("limits", &self.limits)
.field("open", &self.open_count())
.finish()
}
}
#[cfg(test)]
pub(crate) mod testkit {
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use super::{DbName, Opener, Workspace, WorkspaceLayout, WorkspaceLimits};
use crate::Database;
pub(crate) struct TempDir(pub PathBuf);
impl TempDir {
pub(crate) fn new(tag: &str) -> Self {
let dir = std::env::temp_dir().join(format!(
"plugmem-workspace-{tag}-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
std::fs::create_dir_all(&dir).unwrap();
TempDir(dir)
}
}
impl Drop for TempDir {
fn drop(&mut self) {
let _ = std::fs::remove_dir_all(&self.0);
}
}
pub(crate) fn workspace(
tmp: &TempDir,
limits: WorkspaceLimits,
) -> (Workspace, Arc<AtomicUsize>) {
let opens = Arc::new(AtomicUsize::new(0));
let counted = Arc::clone(&opens);
let open: Opener = Box::new(move |path: &std::path::Path| {
counted.fetch_add(1, Ordering::SeqCst);
Ok(Database::open(path, crate::Config::default())?.0)
});
(
Workspace::new(WorkspaceLayout::new(&tmp.0), open, limits),
opens,
)
}
pub(crate) fn name(s: &str) -> DbName {
DbName::parse(s).unwrap()
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::Ordering;
use super::testkit::{TempDir, name, workspace};
use super::*;
fn problem(s: &str) -> NameProblem {
match DbName::parse(s) {
Err(WorkspaceError::BadName { why, .. }) => why,
other => panic!("expected {s:?} to be refused, got {other:?}"),
}
}
#[test]
fn a_name_admits_only_the_safe_alphabet() {
for ok in [
"a",
"0",
"chat-42",
"common",
"x_y-9",
&"a".repeat(MAX_DB_NAME),
] {
assert_eq!(DbName::parse(ok).unwrap().as_str(), ok);
}
assert_eq!(problem(""), NameProblem::Empty);
assert_eq!(problem(&"a".repeat(MAX_DB_NAME + 1)), NameProblem::TooLong);
for bad in ["-x", "_x", ".x", "..", "/x", "Ab", "чат"] {
assert_eq!(problem(bad), NameProblem::LeadingChar, "{bad:?}");
}
for bad in ["a/b", "a\\b", "a.b", "a b", "aB", "a:b", "aчат", "a\0b"] {
assert_eq!(problem(bad), NameProblem::Character, "{bad:?}");
}
for bad in ["con", "nul", "prn", "aux", "com1", "lpt9"] {
assert_eq!(problem(bad), NameProblem::ReservedDevice, "{bad:?}");
}
for ok in ["console", "con1", "com0", "com10", "nula"] {
assert!(DbName::parse(ok).is_ok(), "{ok:?}");
}
}
#[test]
fn a_name_prints_as_itself() {
assert_eq!(DbName::parse("chat-42").unwrap().to_string(), "chat-42");
assert_eq!(NameProblem::Empty.to_string(), "it is empty");
assert_eq!(
NameProblem::TooLong.to_string(),
format!("it is longer than {MAX_DB_NAME} bytes")
);
assert!(NameProblem::LeadingChar.to_string().contains("start with"));
assert!(NameProblem::Character.to_string().contains("lowercase"));
assert!(
NameProblem::ReservedDevice
.to_string()
.contains("Windows device name")
);
}
#[test]
fn the_layout_puts_the_registry_out_of_reach_of_names() {
let layout = WorkspaceLayout::new("/ws");
let name = DbName::parse("chat-42").unwrap();
assert_eq!(layout.root(), Path::new("/ws"));
assert_eq!(layout.db_dir(), Path::new("/ws/db"));
assert_eq!(layout.path_of(&name), Path::new("/ws/db/chat-42.plugmem"));
assert_eq!(layout.registry_path(), Path::new("/ws/registry.plugmem"));
let lookalike = DbName::parse("registry").unwrap();
assert_ne!(layout.path_of(&lookalike), layout.registry_path());
}
#[test]
fn listing_reads_the_directory_and_ignores_what_is_not_a_database() {
let tmp = TempDir::new("list");
let layout = WorkspaceLayout::new(&tmp.0);
assert!(layout.list().unwrap().is_empty());
std::fs::create_dir_all(layout.db_dir()).unwrap();
for file in [
"chat-42.plugmem",
"common.plugmem",
"chat-42.plugmem.lock",
"chat-42.plugmem.journal",
"chat-42.plugmem.snap.3",
"notes.txt",
"Chat-43.plugmem",
] {
std::fs::write(layout.db_dir().join(file), b"").unwrap();
}
let names: Vec<String> = layout
.list()
.unwrap()
.iter()
.map(DbName::to_string)
.collect();
assert_eq!(names, ["chat-42", "common"]);
assert!(layout.exists(&DbName::parse("chat-42").unwrap()));
assert!(!layout.exists(&DbName::parse("nope").unwrap()));
}
#[test]
fn a_database_exists_before_its_first_checkpoint() {
let tmp = TempDir::new("list-uncheckpointed");
let layout = WorkspaceLayout::new(&tmp.0);
let fresh = DbName::parse("fresh").unwrap();
std::fs::create_dir_all(layout.db_dir()).unwrap();
let (db, _) = Database::open(layout.path_of(&fresh), crate::Config::default()).unwrap();
db.remember(crate::RememberInput::text(1_000, "not yet checkpointed"))
.unwrap();
assert!(!layout.path_of(&fresh).exists());
assert!(layout.exists(&fresh));
assert_eq!(layout.list().unwrap(), [fresh]);
}
#[test]
fn an_unreadable_directory_is_an_error_not_an_empty_workspace() {
let tmp = TempDir::new("list-io");
let layout = WorkspaceLayout::new(&tmp.0);
std::fs::write(layout.db_dir(), b"not a directory").unwrap();
assert!(matches!(layout.list(), Err(WorkspaceError::Io { .. })));
}
#[test]
fn every_failure_names_what_the_caller_typed() {
let busy = WorkspaceError::Busy {
name: DbName::parse("chat-42").unwrap(),
};
assert!(busy.to_string().contains("chat-42"));
let host = WorkspaceError::from(HostError::Embed("no".into()));
assert!(matches!(host, WorkspaceError::Host(HostError::Embed(_))));
let io = WorkspaceError::io(
Path::new("/ws"),
std::io::Error::new(std::io::ErrorKind::PermissionDenied, "denied"),
);
assert!(io.to_string().contains("/ws"));
let missing = WorkspaceError::NoSuchDatabase {
name: DbName::parse("gone").unwrap(),
path: PathBuf::from("/ws/db/gone.plugmem"),
};
assert!(missing.to_string().contains("gone"));
}
#[test]
fn a_pooled_database_is_reused_and_a_missing_one_is_created_only_on_request() {
let tmp = TempDir::new("pool-reuse");
let (ws, opens) = workspace(&tmp, WorkspaceLimits::default());
let chat = name("chat-42");
let missed = ws.get(&chat, 1_000, IfMissing::Fail).unwrap_err();
assert!(
matches!(&missed, WorkspaceError::NoSuchDatabase { name, .. } if name == &chat),
"{missed}"
);
assert_eq!(opens.load(Ordering::SeqCst), 0);
assert!(!ws.layout().db_dir().exists());
let db = ws.get(&chat, 1_000, IfMissing::Create).unwrap();
db.remember(crate::RememberInput::text(1_000, "prefers tokio"))
.unwrap();
assert!(ws.layout().exists(&chat));
assert_eq!(ws.open_count(), 1);
let again = ws.get(&chat, 2_000, IfMissing::Fail).unwrap();
assert_eq!(again.stats().facts, 1);
assert_eq!(opens.load(Ordering::SeqCst), 1);
assert!(format!("{ws:?}").contains("open: 1"));
}
#[test]
fn databases_in_one_workspace_do_not_see_each_other() {
let tmp = TempDir::new("isolation");
let (ws, _) = workspace(&tmp, WorkspaceLimits::default());
for (db, text) in [
("chat-42", "the sky is blue"),
("chat-43", "the sky is red"),
] {
ws.get(&name(db), 1_000, IfMissing::Create)
.unwrap()
.remember(crate::RememberInput::text(1_000, text))
.unwrap();
}
for (db, expected) in [("chat-42", "blue"), ("chat-43", "red")] {
let out = ws
.get(&name(db), 2_000, IfMissing::Fail)
.unwrap()
.recall(crate::RecallQuery::text(2_000, "sky"))
.unwrap();
assert_eq!(out.facts.len(), 1, "{db}");
assert!(out.rendered.contains(expected), "{db}: {}", out.rendered);
}
}
#[test]
fn the_ceiling_evicts_the_least_recently_used() {
let tmp = TempDir::new("pool-evict");
let (ws, opens) = workspace(
&tmp,
WorkspaceLimits {
max_open: 2,
..WorkspaceLimits::default()
},
);
ws.get(&name("a"), 1_000, IfMissing::Create).unwrap();
ws.get(&name("b"), 2_000, IfMissing::Create).unwrap();
ws.get(&name("a"), 3_000, IfMissing::Fail).unwrap();
ws.get(&name("c"), 4_000, IfMissing::Create).unwrap();
assert_eq!(ws.open_count(), 2);
assert_eq!(opens.load(Ordering::SeqCst), 3);
ws.get(&name("a"), 5_000, IfMissing::Fail).unwrap();
assert_eq!(opens.load(Ordering::SeqCst), 3);
ws.get(&name("b"), 6_000, IfMissing::Fail).unwrap();
assert_eq!(opens.load(Ordering::SeqCst), 4);
}
#[test]
fn a_ceiling_of_zero_still_serves_one_database() {
let tmp = TempDir::new("pool-zero");
let (ws, opens) = workspace(
&tmp,
WorkspaceLimits {
max_open: 0,
idle_timeout_ms: 0,
},
);
ws.get(&name("a"), 1_000, IfMissing::Create).unwrap();
ws.get(&name("b"), 2_000, IfMissing::Create).unwrap();
assert_eq!(ws.open_count(), 1);
assert_eq!(opens.load(Ordering::SeqCst), 2);
assert_eq!(ws.close_idle(u64::MAX), 0);
assert_eq!(ws.open_count(), 1);
}
#[test]
fn an_idle_database_is_closed_and_its_lock_released() {
let tmp = TempDir::new("pool-idle");
let (ws, _) = workspace(
&tmp,
WorkspaceLimits {
max_open: 8,
idle_timeout_ms: 1_000,
},
);
let chat = name("chat-42");
let path = ws.layout().path_of(&chat);
drop(ws.get(&chat, 1_000, IfMissing::Create).unwrap());
assert_eq!(ws.close_idle(1_500), 0);
assert!(matches!(
Database::open(&path, crate::Config::default()),
Err(HostError::Locked { .. })
));
assert_eq!(ws.close_idle(500), 0);
assert_eq!(ws.close_idle(2_000), 1);
assert_eq!(ws.open_count(), 0);
assert!(Database::open(&path, crate::Config::default()).is_ok());
}
#[test]
fn a_handle_held_by_a_caller_outlives_its_pool_entry() {
let tmp = TempDir::new("pool-outlive");
let (ws, _) = workspace(&tmp, WorkspaceLimits::default());
let chat = name("chat-42");
let held = ws.get(&chat, 1_000, IfMissing::Create).unwrap();
let path = ws.layout().path_of(&chat);
assert_eq!(ws.close_all(), 1);
assert_eq!(ws.open_count(), 0);
held.remember(crate::RememberInput::text(2_000, "still mine"))
.unwrap();
assert!(matches!(
Database::open(&path, crate::Config::default()),
Err(HostError::Locked { .. })
));
drop(held);
assert!(Database::open(&path, crate::Config::default()).is_ok());
}
#[test]
fn a_database_held_by_another_process_is_reported_by_name() {
let tmp = TempDir::new("pool-busy");
let (ws, _) = workspace(&tmp, WorkspaceLimits::default());
let chat = name("chat-42");
let path = ws.layout().path_of(&chat);
std::fs::create_dir_all(ws.layout().db_dir()).unwrap();
let outsider = Database::open(&path, crate::Config::default()).unwrap().0;
let e = ws.get(&chat, 1_000, IfMissing::Create).unwrap_err();
assert!(matches!(&e, WorkspaceError::Busy { name } if name == &chat));
assert!(e.to_string().contains("chat-42"), "{e}");
drop(outsider);
assert!(ws.get(&chat, 2_000, IfMissing::Fail).is_ok());
}
#[test]
fn an_open_that_fails_for_another_reason_keeps_its_own_error() {
let tmp = TempDir::new("pool-open-err");
let open: Opener = Box::new(|_| Err(HostError::Embed("no provider".into())));
let ws = Workspace::new(
WorkspaceLayout::new(&tmp.0),
open,
WorkspaceLimits::default(),
);
let e = ws.get(&name("a"), 1_000, IfMissing::Create).unwrap_err();
assert!(
matches!(e, WorkspaceError::Host(HostError::Embed(_))),
"{e}"
);
}
#[test]
fn a_directory_that_cannot_be_created_is_an_error_not_a_panic() {
let tmp = TempDir::new("pool-mkdir");
std::fs::write(tmp.0.join(DB_DIR), b"in the way").unwrap();
let (ws, _) = workspace(&tmp, WorkspaceLimits::default());
assert!(matches!(
ws.get(&name("a"), 1_000, IfMissing::Create),
Err(WorkspaceError::Io { .. })
));
}
proptest::proptest! {
#[test]
fn a_name_that_parses_can_only_resolve_inside_the_workspace(s in ".*") {
let Ok(name) = DbName::parse(&s) else { return Ok(()) };
let layout = WorkspaceLayout::new("/ws");
let path = layout.path_of(&name);
let rest: Vec<_> = path
.strip_prefix(layout.db_dir())
.expect("resolved outside the workspace")
.components()
.collect();
let expected = format!("{s}.{DB_EXT}");
proptest::prop_assert_eq!(rest.len(), 1);
proptest::prop_assert_eq!(
path.file_name().and_then(|n| n.to_str()),
Some(expected.as_str())
);
}
}
}