use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::sync::{Mutex, MutexGuard, OwnedMutexGuard, OwnedSemaphorePermit, Semaphore};
use tokio_util::sync::CancellationToken;
use crate::adapters::lance::LanceStore;
use crate::app::config::Config;
use crate::domain::common::{HallouminateError, expand_tilde};
use crate::domain::corpus::{load_tokenizer, missing_roots};
use crate::domain::embeddings::{EmbedBatch, Embedder};
use crate::domain::indexer::HandlerRegistry;
use crate::domain::indexer::index::index_corpus;
use crate::domain::search::{FastembedCrossencoder, canonical_crossencoder_model};
const CHUNK_BUDGET_TOKENS: usize = 384;
fn unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
fn is_idle(last_use_secs: u64, now_secs: u64, idle_secs: u64) -> bool {
now_secs.saturating_sub(last_use_secs) >= idle_secs
}
#[derive(Default)]
struct CorpusLockMap {
inner: Mutex<HashMap<String, Arc<Mutex<()>>>>,
}
impl CorpusLockMap {
async fn lock(&self, corpus: &str) -> OwnedMutexGuard<()> {
let mutex = {
let mut map = self.inner.lock().await;
map.entry(corpus.to_string())
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone()
};
mutex.lock_owned().await
}
}
#[derive(Clone)]
pub struct DaemonState {
inner: Arc<DaemonStateInner>,
}
struct DaemonStateInner {
baseline: Config,
baseline_xdg_path: Option<PathBuf>,
store: Arc<LanceStore>,
ground_dir: PathBuf,
embeddings_enabled: bool,
corpus_locks: CorpusLockMap,
write_lane: Arc<Semaphore>,
embedder: Arc<Mutex<Option<Embedder>>>,
crossencoders: Arc<Mutex<HashMap<String, FastembedCrossencoder>>>,
last_embed_use_secs: Arc<AtomicU64>,
tokenizer: tokenizers::Tokenizer,
shutdown: CancellationToken,
}
pub struct MutationGuard {
_permit: OwnedSemaphorePermit,
_corpus: OwnedMutexGuard<()>,
}
impl DaemonState {
pub async fn open(cfg: Config, xdg_path: Option<PathBuf>) -> anyhow::Result<Self> {
let ground_dir = expand_tilde(&cfg.storage.ground_dir);
if let Some(parent) = ground_dir.parent()
&& !parent.as_os_str().is_empty()
{
tokio::fs::create_dir_all(parent)
.await
.map_err(|e| anyhow::anyhow!("create ground dir parent: {e}"))?;
}
let cache_dir = expand_tilde(&cfg.embeddings.cache_dir);
let mut embedder: Option<Embedder> = if cfg.embeddings.enabled {
match Embedder::try_new(&cfg.embeddings.model, cfg.embeddings.quantized, &cache_dir) {
Ok(e) => Some(e),
Err(e) => {
tracing::warn!(
target: "hallouminate::daemon",
model = %cfg.embeddings.model,
error = %e,
"embedder unavailable at startup; will retry on first embedding request",
);
None
}
}
} else {
None
};
let tokenizer = load_tokenizer(&cfg.embeddings.model)
.map_err(|e| anyhow::anyhow!("load tokenizer for {}: {e}", cfg.embeddings.model))?;
let store = match LanceStore::open_or_create(
&ground_dir,
&cfg.embeddings.model,
cfg.embeddings.quantized,
cfg.embeddings.enabled,
)
.await
{
Ok(s) => s,
Err(HallouminateError::StoreSchemaStale {
found, expected, ..
}) => {
tracing::warn!(
target: "hallouminate::daemon",
%found,
%expected,
"ground store schema v{found} < expected v{expected}; rebuilding from source",
);
move_stale_store(&ground_dir, found).await?;
let rebuild_result: anyhow::Result<LanceStore> = async {
let fresh = LanceStore::open_or_create(
&ground_dir,
&cfg.embeddings.model,
cfg.embeddings.quantized,
cfg.embeddings.enabled,
)
.await
.map_err(|e| {
anyhow::anyhow!(
"rebuild: open fresh ground dir {}: {e}",
ground_dir.display()
)
})?;
let registry = HandlerRegistry::new(tokenizer.clone(), CHUNK_BUDGET_TOKENS);
for corpus in cfg
.effective_corpora()
.map_err(|e| anyhow::anyhow!("rebuild: list corpora: {e}"))?
{
let missing = missing_roots(&corpus);
if !missing.is_empty() {
tracing::warn!(
target: "hallouminate::daemon",
corpus = %corpus.name,
"rebuild: corpus root missing; skipped",
);
continue;
}
let emb: Option<&mut dyn EmbedBatch> =
embedder.as_mut().map(|e| e as &mut dyn EmbedBatch);
let stats = index_corpus(&corpus, &fresh, emb, ®istry)
.await
.map_err(|e| anyhow::anyhow!("rebuild: index {}: {e}", corpus.name))?;
tracing::info!(
target: "hallouminate::daemon",
corpus = %corpus.name,
files = stats.files_upserted,
chunks = stats.chunks_inserted,
"rebuild: reindexed",
);
}
Ok(fresh)
}
.await;
match rebuild_result {
Ok(fresh) => fresh,
Err(e) => {
if ground_dir.exists() {
let _ = tokio::fs::remove_dir_all(&ground_dir).await;
tracing::warn!(
target: "hallouminate::daemon",
"rebuild failed; removed partial ground dir so next boot retries. \
Backup preserved at {}.bak-v{found}",
ground_dir.display(),
);
}
return Err(e);
}
}
}
Err(e) => {
return Err(anyhow::anyhow!(
"open ground dir {}: {e}",
ground_dir.display()
));
}
};
let mut crossencoders: HashMap<String, FastembedCrossencoder> = HashMap::new();
if let Some(model) = cfg.search.crossencoder.as_deref() {
match canonical_crossencoder_model(model)
.map_err(anyhow::Error::from)
.and_then(|canonical| {
FastembedCrossencoder::try_new(canonical, &cache_dir)
.map(|c| (canonical, c))
.map_err(anyhow::Error::from)
}) {
Ok((canonical, c)) => {
crossencoders.insert(canonical.to_string(), c);
}
Err(e) => {
tracing::warn!(
target: "hallouminate::daemon",
model = %model,
error = %e,
"crossencoder unavailable at startup; ground will skip rerank until reload",
);
}
}
}
let idle_evict_secs = cfg.embeddings.idle_evict_secs;
let shutdown = CancellationToken::new();
let embedder_arc = Arc::new(Mutex::new(embedder));
let crossencoders_arc = Arc::new(Mutex::new(crossencoders));
let last_embed_use = Arc::new(AtomicU64::new(unix_secs()));
if idle_evict_secs > 0 {
let embedder_ref = Arc::clone(&embedder_arc);
let crossencoders_ref = Arc::clone(&crossencoders_arc);
let last_use_ref = Arc::clone(&last_embed_use);
let cancel = shutdown.clone();
tracing::debug!(
target: "hallouminate::daemon",
idle_evict_secs,
"idle evict task spawned",
);
tokio::spawn(async move {
loop {
tokio::select! {
biased;
_ = cancel.cancelled() => break,
_ = tokio::time::sleep(Duration::from_secs(idle_evict_secs)) => {
if is_idle(last_use_ref.load(Ordering::Relaxed), unix_secs(), idle_evict_secs) {
let mut emb = embedder_ref.lock().await;
let mut xenc = crossencoders_ref.lock().await;
if is_idle(last_use_ref.load(Ordering::Relaxed), unix_secs(), idle_evict_secs) {
let was_live = emb.is_some();
*emb = None;
let had_xenc = !xenc.is_empty();
xenc.clear();
if was_live || had_xenc {
tracing::info!(
target: "hallouminate::daemon",
idle_secs = idle_evict_secs,
"embedder evicted after idle; ORT BFCArena memory released",
);
}
}
}
}
}
}
});
}
Ok(DaemonState {
inner: Arc::new(DaemonStateInner {
embeddings_enabled: cfg.embeddings.enabled,
baseline: cfg,
baseline_xdg_path: xdg_path,
store: Arc::new(store),
ground_dir,
corpus_locks: CorpusLockMap::default(),
write_lane: Arc::new(Semaphore::new(1)),
embedder: embedder_arc,
crossencoders: crossencoders_arc,
last_embed_use_secs: last_embed_use,
tokenizer,
shutdown,
}),
})
}
pub fn shutdown_token(&self) -> &CancellationToken {
&self.inner.shutdown
}
pub fn baseline_xdg_path(&self) -> Option<&Path> {
self.inner.baseline_xdg_path.as_deref()
}
pub fn baseline(&self) -> &Config {
&self.inner.baseline
}
pub fn store(&self) -> Arc<LanceStore> {
self.inner.store.clone()
}
pub fn ground_dir(&self) -> &std::path::Path {
&self.inner.ground_dir
}
pub fn embeddings_enabled(&self) -> bool {
self.inner.embeddings_enabled
}
pub async fn embedder(&self) -> anyhow::Result<EmbedderGuard<'_>> {
let mut guard = self.inner.embedder.lock().await;
if guard.is_none() {
let cache_dir = expand_tilde(&self.inner.baseline.embeddings.cache_dir);
let embedder = Embedder::try_new(
&self.inner.baseline.embeddings.model,
self.inner.baseline.embeddings.quantized,
&cache_dir,
)
.map_err(|e| {
anyhow::anyhow!(
"init embedder ({}): {e}",
self.inner.baseline.embeddings.model
)
})?;
*guard = Some(embedder);
}
self.inner
.last_embed_use_secs
.store(unix_secs(), Ordering::Relaxed);
Ok(EmbedderGuard {
guard,
last_use_secs: Arc::clone(&self.inner.last_embed_use_secs),
})
}
pub async fn crossencoder(
&self,
model_name: Option<&str>,
) -> anyhow::Result<Option<CrossencoderGuard<'_>>> {
let Some(model_name) = model_name else {
return Ok(None);
};
let canonical = canonical_crossencoder_model(model_name)?;
let mut guard = self.inner.crossencoders.lock().await;
if !guard.contains_key(canonical) {
let cache_dir = expand_tilde(&self.inner.baseline.embeddings.cache_dir);
let model = FastembedCrossencoder::try_new(canonical, &cache_dir)
.map_err(|e| anyhow::anyhow!("init crossencoder ({canonical}): {e}"))?;
guard.insert(canonical.to_string(), model);
}
self.inner
.last_embed_use_secs
.store(unix_secs(), Ordering::Relaxed);
Ok(Some(CrossencoderGuard {
guard,
key: canonical.to_string(),
last_use_secs: Arc::clone(&self.inner.last_embed_use_secs),
}))
}
#[cfg(test)]
pub(crate) async fn evict_if_idle_at(&self, idle_secs: u64, now_secs: u64) -> bool {
if idle_secs == 0 {
return false;
}
let last = self.inner.last_embed_use_secs.load(Ordering::Relaxed);
if !is_idle(last, now_secs, idle_secs) {
return false;
}
let mut emb = self.inner.embedder.lock().await;
let mut xenc = self.inner.crossencoders.lock().await;
if !is_idle(
self.inner.last_embed_use_secs.load(Ordering::Relaxed),
now_secs,
idle_secs,
) {
return false;
}
let was_live = emb.is_some();
*emb = None;
let had_xenc = !xenc.is_empty();
xenc.clear();
if was_live || had_xenc {
tracing::info!(
target: "hallouminate::daemon",
idle_secs,
"embedder evicted after idle; ORT BFCArena memory released",
);
}
true
}
#[cfg(test)]
pub(crate) fn last_embed_use_secs(&self) -> u64 {
self.inner.last_embed_use_secs.load(Ordering::Relaxed)
}
pub fn make_registry(&self) -> HandlerRegistry {
HandlerRegistry::new(self.inner.tokenizer.clone(), CHUNK_BUDGET_TOKENS)
}
pub async fn lock_corpus(&self, corpus: &str) -> OwnedMutexGuard<()> {
self.inner.corpus_locks.lock(corpus).await
}
pub fn write_lane(&self) -> Arc<Semaphore> {
self.inner.write_lane.clone()
}
pub async fn acquire_mutation_guard(
&self,
corpus: &str,
) -> Result<MutationGuard, &'static str> {
let corpus = self.lock_corpus(corpus).await;
let permit = self
.write_lane()
.acquire_owned()
.await
.map_err(|_| "write lane closed")?;
Ok(MutationGuard {
_permit: permit,
_corpus: corpus,
})
}
}
async fn move_stale_store(ground_dir: &Path, found_version: u32) -> anyhow::Result<()> {
let bak = ground_dir.with_file_name(format!(
"{}.bak-v{found_version}",
ground_dir
.file_name()
.and_then(|s| s.to_str())
.unwrap_or("ground"),
));
if bak.exists() {
tokio::fs::remove_dir_all(&bak).await?;
}
tokio::fs::rename(ground_dir, &bak).await?;
tracing::info!(
target: "hallouminate::daemon",
backup = %bak.display(),
"moved stale ground store aside; recoverable until pruned",
);
Ok(())
}
pub struct EmbedderGuard<'a> {
guard: MutexGuard<'a, Option<Embedder>>,
last_use_secs: Arc<AtomicU64>,
}
impl std::ops::Deref for EmbedderGuard<'_> {
type Target = Embedder;
fn deref(&self) -> &Embedder {
self.guard.as_ref().expect("embedder loaded")
}
}
impl std::ops::DerefMut for EmbedderGuard<'_> {
fn deref_mut(&mut self) -> &mut Embedder {
self.guard.as_mut().expect("embedder loaded")
}
}
impl Drop for EmbedderGuard<'_> {
fn drop(&mut self) {
self.last_use_secs.store(unix_secs(), Ordering::Relaxed);
}
}
pub struct CrossencoderGuard<'a> {
guard: MutexGuard<'a, HashMap<String, FastembedCrossencoder>>,
key: String,
last_use_secs: Arc<AtomicU64>,
}
impl std::ops::Deref for CrossencoderGuard<'_> {
type Target = FastembedCrossencoder;
fn deref(&self) -> &FastembedCrossencoder {
self.guard.get(&self.key).expect("crossencoder loaded")
}
}
impl std::ops::DerefMut for CrossencoderGuard<'_> {
fn deref_mut(&mut self) -> &mut FastembedCrossencoder {
self.guard.get_mut(&self.key).expect("crossencoder loaded")
}
}
impl Drop for CrossencoderGuard<'_> {
fn drop(&mut self) {
self.last_use_secs.store(unix_secs(), Ordering::Relaxed);
}
}
impl std::fmt::Debug for DaemonState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DaemonState")
.field("ground_dir", &self.inner.ground_dir)
.field("model", &self.inner.baseline.embeddings.model)
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn baseline_returns_the_configured_config() {
let tmp = tempfile::tempdir().expect("tempdir");
let mut cfg = Config::default();
cfg.storage.ground_dir = tmp.path().to_string_lossy().into_owned();
let expected_model = cfg.embeddings.model.clone();
let state = DaemonState::open(cfg, None)
.await
.expect("open daemon state");
assert_eq!(state.baseline().embeddings.model, expected_model);
}
#[tokio::test]
async fn evict_if_idle_at_returns_false_when_disabled() {
let tmp = tempfile::tempdir().expect("tempdir");
let mut cfg = Config::default();
cfg.embeddings.enabled = false;
cfg.storage.ground_dir = tmp.path().to_string_lossy().into_owned();
let state = DaemonState::open(cfg, None).await.expect("open");
assert!(
!state.evict_if_idle_at(0, u64::MAX).await,
"idle_secs=0 means disabled; must not evict",
);
}
#[tokio::test]
async fn evict_if_idle_at_returns_false_when_recently_used() {
let tmp = tempfile::tempdir().expect("tempdir");
let mut cfg = Config::default();
cfg.embeddings.enabled = false;
cfg.storage.ground_dir = tmp.path().to_string_lossy().into_owned();
let state = DaemonState::open(cfg, None).await.expect("open");
let last = state.last_embed_use_secs();
assert!(
!state.evict_if_idle_at(300, last + 1).await,
"1 s elapsed < 300 s idle; must not evict",
);
}
#[tokio::test]
async fn evict_if_idle_at_returns_true_when_idle_threshold_exceeded() {
let tmp = tempfile::tempdir().expect("tempdir");
let mut cfg = Config::default();
cfg.embeddings.enabled = false;
cfg.storage.ground_dir = tmp.path().to_string_lossy().into_owned();
let state = DaemonState::open(cfg, None).await.expect("open");
let last = state.last_embed_use_secs();
assert!(
state.evict_if_idle_at(300, last + 301).await,
"301 s elapsed >= 300 s idle; must evict",
);
}
#[tokio::test]
async fn evict_if_idle_at_exact_boundary_evicts() {
let tmp = tempfile::tempdir().expect("tempdir");
let mut cfg = Config::default();
cfg.embeddings.enabled = false;
cfg.storage.ground_dir = tmp.path().to_string_lossy().into_owned();
let state = DaemonState::open(cfg, None).await.expect("open");
let last = state.last_embed_use_secs();
assert!(
state.evict_if_idle_at(300, last + 300).await,
"elapsed == idle_secs (>= threshold); must evict",
);
}
#[tokio::test]
async fn evict_if_idle_at_one_second_below_threshold_does_not_evict() {
let tmp = tempfile::tempdir().expect("tempdir");
let mut cfg = Config::default();
cfg.embeddings.enabled = false;
cfg.storage.ground_dir = tmp.path().to_string_lossy().into_owned();
let state = DaemonState::open(cfg, None).await.expect("open");
let last = state.last_embed_use_secs();
assert!(
!state.evict_if_idle_at(300, last + 299).await,
"elapsed = idle_secs - 1 (< threshold); must not evict",
);
}
#[tokio::test]
async fn evict_if_idle_at_rechecks_after_waiting_for_embedder_lock() {
let tmp = tempfile::tempdir().expect("tempdir");
let mut cfg = Config::default();
cfg.embeddings.enabled = false;
cfg.storage.ground_dir = tmp.path().to_string_lossy().into_owned();
let state = DaemonState::open(cfg, None).await.expect("open");
state.inner.last_embed_use_secs.store(1, Ordering::Relaxed);
let held = state.inner.embedder.lock().await;
let task_state = state.clone();
let eviction = tokio::spawn(async move { task_state.evict_if_idle_at(300, 1000).await });
state
.inner
.last_embed_use_secs
.store(999, Ordering::Relaxed);
drop(held);
assert!(
!eviction.await.expect("eviction task joins"),
"eviction must re-check idleness after waiting for an active guard",
);
}
#[tokio::test]
async fn embedder_guard_updates_last_use_on_drop() {
let last_use_secs = Arc::new(AtomicU64::new(1));
let guard = Mutex::new(None);
let guard = guard.lock().await;
let before_drop = unix_secs();
drop(EmbedderGuard {
guard,
last_use_secs: Arc::clone(&last_use_secs),
});
let observed = last_use_secs.load(Ordering::Relaxed);
assert!(
observed >= before_drop,
"drop should stamp embedder use at or after guard lifetime start: observed {observed}, before {before_drop}",
);
}
#[tokio::test]
async fn crossencoder_guard_updates_last_use_on_drop() {
let last_use_secs = Arc::new(AtomicU64::new(1));
let guard = Mutex::new(HashMap::new());
let guard = guard.lock().await;
let before_drop = unix_secs();
drop(CrossencoderGuard {
guard,
key: String::new(),
last_use_secs: Arc::clone(&last_use_secs),
});
let observed = last_use_secs.load(Ordering::Relaxed);
assert!(
observed >= before_drop,
"drop should stamp crossencoder use at or after guard lifetime start: observed {observed}, before {before_drop}",
);
}
}