use crate::memory_core::community::KnowledgeGap;
use crate::memory_core::palace::{Palace, PalaceId};
use crate::memory_core::retrieval::PalaceHandle;
use crate::memory_core::store::concurrent_open::OpenIntent;
use crate::memory_core::store::palace_store::PalaceStore;
use crate::palace_alias::PalaceAliasStore;
use anyhow::{Context, Result};
use dashmap::DashMap;
use lru::LruCache;
use parking_lot::Mutex;
use std::num::NonZeroUsize;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
pub const MAX_OPEN_PALACES_ENV: &str = "TRUSTY_MEMORY_MAX_OPEN_PALACES";
pub fn max_open_palaces_from_env() -> usize {
std::env::var(MAX_OPEN_PALACES_ENV)
.ok()
.and_then(|v| v.trim().parse::<usize>().ok())
.filter(|&n| n >= 1)
.unwrap_or(DEFAULT_MAX_OPEN_PALACES)
}
pub const DEFAULT_MAX_OPEN_PALACES: usize = 64;
#[derive(Clone)]
pub struct PalaceRegistry {
handles: Arc<Mutex<LruCache<PalaceId, Arc<PalaceHandle>>>>,
gaps_cache: Arc<DashMap<PalaceId, Vec<KnowledgeGap>>>,
open_intent: OpenIntent,
open_locks: Arc<DashMap<PalaceId, Arc<Mutex<()>>>>,
}
impl Default for PalaceRegistry {
fn default() -> Self {
Self::with_max_open(DEFAULT_MAX_OPEN_PALACES)
}
}
impl PalaceRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn with_max_open(max_open_palaces: usize) -> Self {
let cap = NonZeroUsize::new(max_open_palaces.max(1)).expect("max(1) is always nonzero");
Self {
handles: Arc::new(Mutex::new(LruCache::new(cap))),
gaps_cache: Arc::new(DashMap::new()),
open_intent: OpenIntent::ReadOnlyClient,
open_locks: Arc::new(DashMap::new()),
}
}
pub fn from_env() -> Self {
Self::with_max_open(max_open_palaces_from_env())
}
#[must_use]
pub fn with_writer_intent(mut self) -> Self {
self.open_intent = OpenIntent::Writer;
self
}
#[must_use]
pub fn open_intent(&self) -> OpenIntent {
self.open_intent
}
pub fn register(&self, handle: PalaceHandle) {
let id = handle.id.clone();
let arc = Arc::new(handle);
let _evicted = {
let mut cache = self.handles.lock();
cache.put(id, arc)
};
}
pub fn register_arc(&self, handle: Arc<PalaceHandle>) {
let id = handle.id.clone();
let _evicted = {
let mut cache = self.handles.lock();
cache.put(id, handle)
};
}
pub fn get(&self, id: &PalaceId) -> Option<Arc<PalaceHandle>> {
let mut cache = self.handles.lock();
cache.get(id).cloned()
}
pub fn peek(&self, id: &PalaceId) -> Option<Arc<PalaceHandle>> {
let cache = self.handles.lock();
cache.peek(id).cloned()
}
pub fn list(&self) -> Vec<PalaceId> {
let cache = self.handles.lock();
cache.iter().map(|(k, _)| k.clone()).collect()
}
pub fn len(&self) -> usize {
self.handles.lock().len()
}
pub fn is_empty(&self) -> bool {
self.handles.lock().is_empty()
}
pub fn set_gaps(&self, palace_id: PalaceId, gaps: Vec<KnowledgeGap>) {
self.gaps_cache.insert(palace_id, gaps);
}
pub fn get_gaps(&self, palace_id: &PalaceId) -> Option<Vec<KnowledgeGap>> {
self.gaps_cache.get(palace_id).map(|r| r.value().clone())
}
pub fn clear_gaps(&self, palace_id: &PalaceId) {
self.gaps_cache.remove(palace_id);
}
pub fn remove(&self, palace_id: &PalaceId) {
let _evicted = {
let mut cache = self.handles.lock();
cache.pop(palace_id)
};
self.gaps_cache.remove(palace_id);
}
pub fn open_palace(&self, data_root: &Path, palace_id: &PalaceId) -> Result<Arc<PalaceHandle>> {
if let Some(h) = self.get(palace_id) {
return Ok(h);
}
let effective_id = Self::resolve_palace_alias(data_root, palace_id);
if let Some(h) = self.get(&effective_id) {
return Ok(h);
}
let open_lock = self
.open_locks
.entry(effective_id.clone())
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone();
let _open_guard = open_lock.lock();
if let Some(h) = self.get(&effective_id) {
return Ok(h);
}
let palace_dir = data_root.join(effective_id.as_str());
let palace = PalaceStore::load_palace(&palace_dir)
.with_context(|| format!("load palace metadata for {palace_id}"))?;
let handle = PalaceHandle::open_with_intent(&palace, self.open_intent)?;
self.register_arc(handle.clone());
Ok(handle)
}
fn resolve_palace_alias(data_root: &Path, palace_id: &PalaceId) -> PalaceId {
let exists = |id: &str| data_root.join(id).join("palace.json").exists();
if exists(palace_id.as_str()) {
return palace_id.clone();
}
match PalaceAliasStore::resolve_alias(data_root, palace_id.as_str()) {
Ok(Some(target)) if exists(&target) => PalaceId::new(target),
_ => palace_id.clone(),
}
}
pub fn create_palace(&self, data_root: &Path, mut palace: Palace) -> Result<Arc<PalaceHandle>> {
let palace_dir = data_root.join(palace.id.as_str());
palace.data_dir = palace_dir.clone();
std::fs::create_dir_all(&palace_dir)
.with_context(|| format!("create palace dir {}", palace_dir.display()))?;
PalaceStore::save_palace(&palace)
.with_context(|| format!("save palace metadata for {}", palace.id))?;
let handle = PalaceHandle::open_with_intent(&palace, self.open_intent)?;
self.register_arc(handle.clone());
Ok(handle)
}
pub fn list_palaces(data_root: &Path) -> Result<Vec<Palace>> {
PalaceStore::list_palaces(data_root)
.with_context(|| format!("list palaces under {}", data_root.display()))
}
pub fn open(data_root: &Path) -> Result<Self> {
std::fs::create_dir_all(data_root)
.with_context(|| format!("create registry root {}", data_root.display()))?;
let registry = Self::new();
let palaces = PalaceStore::list_palaces(data_root)
.with_context(|| format!("list palaces under {}", data_root.display()))?;
for palace in palaces {
match PalaceHandle::open_with_intent(&palace, registry.open_intent) {
Ok(handle) => registry.register_arc(handle),
Err(e) => {
tracing::warn!(palace = %palace.id, "skipping palace during registry open: {e:#}");
}
}
}
Ok(registry)
}
pub fn evict_idle(&self, threshold: Duration) -> usize {
let threshold_secs = threshold.as_secs();
if threshold_secs == 0 {
return 0;
}
let evicted: Vec<Arc<PalaceHandle>> = {
let mut cache = self.handles.lock();
let victims: Vec<PalaceId> = cache
.iter()
.filter(|(_, h)| Arc::strong_count(h) == 1 && h.idle_secs() >= threshold_secs)
.map(|(id, _)| id.clone())
.collect();
victims
.into_iter()
.filter_map(|id| cache.pop(&id))
.collect()
};
let count = evicted.len();
if count > 0 {
tracing::info!(
count,
idle_threshold_secs = threshold_secs,
"idle-evict: dropped {count} idle palace handle(s); redb remains the source of truth"
);
}
count
}
}
#[cfg(test)]
#[path = "registry_tests.rs"]
mod tests;