use std::collections::HashMap;
use std::hash::Hash;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::OnceLock;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::{Duration, Instant, SystemTime};
use tokio::sync::{Mutex, OwnedMutexGuard, OwnedSemaphorePermit, Semaphore};
use tokio_util::sync::CancellationToken;
use crate::adapters::lance::{LanceStore, SearchHit};
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::{Crossencoder, FastembedCrossencoder, canonical_crossencoder_model};
const CHUNK_BUDGET_TOKENS: usize = 384;
pub(crate) const STALE_BACKUP_MAX_AGE: Duration = Duration::from_secs(30 * 24 * 60 * 60);
const MAINTENANCE_INTERVAL_SECS: u64 = 1800;
const MAINTENANCE_PRUNE_GRACE_SECS: u64 = 300;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum MaintenanceTick {
Continue,
Stop,
}
static PROCESS_START: OnceLock<Instant> = OnceLock::new();
fn monotonic_secs() -> u64 {
PROCESS_START.get_or_init(Instant::now).elapsed().as_secs()
}
fn is_idle(last_use_secs: u64, now_secs: u64, idle_secs: u64) -> bool {
now_secs.saturating_sub(last_use_secs) >= idle_secs
}
struct KeyedLockMap<K> {
inner: Mutex<HashMap<K, Arc<Mutex<()>>>>,
}
impl<K> Default for KeyedLockMap<K> {
fn default() -> Self {
Self {
inner: Mutex::new(HashMap::new()),
}
}
}
impl<K: Eq + Hash> KeyedLockMap<K> {
async fn lock<Q>(&self, key: &Q) -> OwnedMutexGuard<()>
where
Q: ToOwned<Owned = K> + ?Sized,
{
let mutex = {
let mut map = self.inner.lock().await;
map.entry(key.to_owned())
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone()
};
mutex.lock_owned().await
}
}
#[derive(Clone, PartialEq, Eq, Hash)]
struct ResourceKey {
ground_dir: PathBuf,
model: String,
quantized: bool,
enabled: bool,
}
impl ResourceKey {
fn from_config(cfg: &Config) -> Self {
Self {
ground_dir: expand_tilde(&cfg.storage.ground_dir),
model: cfg.embeddings.model.clone(),
quantized: cfg.embeddings.quantized,
enabled: cfg.embeddings.enabled,
}
}
}
pub struct RequestResources {
pub store: Arc<LanceStore>,
pub tokenizer: tokenizers::Tokenizer,
pub embeddings_enabled: bool,
pub ground_dir: PathBuf,
embedder: Arc<Mutex<Option<Embedder>>>,
model: String,
quantized: bool,
cache_dir: PathBuf,
last_activity_secs: Arc<AtomicU64>,
}
impl RequestResources {
pub async fn embedder(&self) -> anyhow::Result<EmbedderGuard> {
let mut guard = Arc::clone(&self.embedder).lock_owned().await;
if guard.is_none() {
let (model, quantized, cache_dir) =
(self.model.clone(), self.quantized, self.cache_dir.clone());
let embedder =
tokio::task::block_in_place(|| Embedder::try_new(&model, quantized, &cache_dir))
.map_err(|e| anyhow::anyhow!("init embedder ({model}): {e}"))?;
*guard = Some(embedder);
}
self.last_activity_secs
.store(monotonic_secs(), Ordering::Relaxed);
Ok(EmbedderGuard {
guard,
last_use_secs: Arc::clone(&self.last_activity_secs),
})
}
}
#[derive(Clone)]
pub struct DaemonState {
inner: Arc<DaemonStateInner>,
}
struct DaemonStateInner {
baseline: Config,
baseline_xdg_path: Option<PathBuf>,
baseline_resources: Arc<RequestResources>,
resources: Mutex<HashMap<ResourceKey, Arc<RequestResources>>>,
resource_build_locks: KeyedLockMap<ResourceKey>,
corpus_locks: KeyedLockMap<String>,
write_lane: Arc<Semaphore>,
crossencoders: Arc<Mutex<HashMap<String, FastembedCrossencoder>>>,
last_activity_secs: Arc<AtomicU64>,
active_connections: Arc<AtomicUsize>,
shutdown: CancellationToken,
}
pub struct MutationGuard {
_permit: OwnedSemaphorePermit,
_corpus: OwnedMutexGuard<()>,
}
pub struct ConnectionGuard {
active: Arc<AtomicUsize>,
}
impl Drop for ConnectionGuard {
fn drop(&mut self) {
self.active.fetch_sub(1, Ordering::SeqCst);
}
}
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()
));
}
};
if let Err(e) =
prune_stale_backups(&ground_dir, SystemTime::now(), STALE_BACKUP_MAX_AGE).await
{
tracing::warn!(
target: "hallouminate::daemon",
error = %e,
"failed to prune stale ground store backups",
);
}
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 shutdown = CancellationToken::new();
let embedder_arc = Arc::new(Mutex::new(embedder));
let crossencoders_arc = Arc::new(Mutex::new(crossencoders));
let last_activity = Arc::new(AtomicU64::new(monotonic_secs()));
let store = Arc::new(store);
let write_lane = Arc::new(Semaphore::new(1));
if cfg.embeddings.idle_evict_secs
!= crate::app::config::EmbeddingsConfig::default().idle_evict_secs
{
tracing::warn!(
target: "hallouminate::daemon",
idle_evict_secs = cfg.embeddings.idle_evict_secs,
"embeddings.idle_evict_secs is deprecated and does nothing; \
set [daemon].idle_exit_secs to control idle-exit instead",
);
}
let baseline_key = ResourceKey {
ground_dir: ground_dir.clone(),
model: cfg.embeddings.model.clone(),
quantized: cfg.embeddings.quantized,
enabled: cfg.embeddings.enabled,
};
let baseline_resources = Arc::new(RequestResources {
store,
tokenizer,
embeddings_enabled: cfg.embeddings.enabled,
ground_dir,
embedder: embedder_arc,
model: cfg.embeddings.model.clone(),
quantized: cfg.embeddings.quantized,
cache_dir,
last_activity_secs: Arc::clone(&last_activity),
});
let mut resources_map = HashMap::new();
resources_map.insert(baseline_key, Arc::clone(&baseline_resources));
let state = DaemonState {
inner: Arc::new(DaemonStateInner {
baseline: cfg,
baseline_xdg_path: xdg_path,
baseline_resources,
resources: Mutex::new(resources_map),
resource_build_locks: KeyedLockMap::default(),
corpus_locks: KeyedLockMap::default(),
write_lane,
crossencoders: crossencoders_arc,
last_activity_secs: last_activity,
active_connections: Arc::new(AtomicUsize::new(0)),
shutdown,
}),
};
{
let state = state.clone();
let cancel = state.shutdown_token().clone();
tokio::spawn(async move {
loop {
tokio::select! {
biased;
_ = cancel.cancelled() => break,
_ = tokio::time::sleep(Duration::from_secs(MAINTENANCE_INTERVAL_SECS)) => {
if state.run_maintenance_tick().await == MaintenanceTick::Stop {
break;
}
}
}
}
});
}
Ok(state)
}
async fn run_maintenance_tick(&self) -> MaintenanceTick {
let _conn = self.enter_connection();
let Ok(_permit) = self.inner.write_lane.acquire().await else {
return MaintenanceTick::Stop;
};
match self
.inner
.baseline_resources
.store
.maintain(lancedb::table::Duration::seconds(
MAINTENANCE_PRUNE_GRACE_SECS as i64,
))
.await
{
Ok(stats) => {
tracing::info!(
target: "hallouminate::lance",
fragments_removed = stats.compaction.as_ref().map(|c| c.fragments_removed),
fragments_added = stats.compaction.as_ref().map(|c| c.fragments_added),
old_versions_pruned = stats.prune.as_ref().map(|p| p.old_versions),
"periodic LanceDB maintenance completed",
);
}
Err(e) => {
tracing::warn!(
target: "hallouminate::lance",
error = %e,
"periodic LanceDB maintenance failed",
);
}
}
self.touch_activity();
MaintenanceTick::Continue
}
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.baseline_resources.store.clone()
}
pub fn ground_dir(&self) -> &std::path::Path {
&self.inner.baseline_resources.ground_dir
}
pub fn embeddings_enabled(&self) -> bool {
self.inner.baseline_resources.embeddings_enabled
}
pub async fn embedder(&self) -> anyhow::Result<EmbedderGuard> {
self.inner.baseline_resources.embedder().await
}
pub async fn resources_for(&self, cfg: &Config) -> anyhow::Result<Arc<RequestResources>> {
let key = ResourceKey::from_config(cfg);
if let Some(existing) = self.inner.resources.lock().await.get(&key) {
return Ok(Arc::clone(existing));
}
let _build = self.inner.resource_build_locks.lock(&key).await;
if let Some(existing) = self.inner.resources.lock().await.get(&key) {
return Ok(Arc::clone(existing));
}
let ground_dir = key.ground_dir.clone();
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 tokenizer = load_tokenizer(&cfg.embeddings.model)
.map_err(|e| anyhow::anyhow!("load tokenizer for {}: {e}", cfg.embeddings.model))?;
let store = LanceStore::open_or_create(
&ground_dir,
&cfg.embeddings.model,
cfg.embeddings.quantized,
cfg.embeddings.enabled,
)
.await
.map_err(|e| anyhow::anyhow!("open ground dir {}: {e}", ground_dir.display()))?;
let resources = Arc::new(RequestResources {
store: Arc::new(store),
tokenizer,
embeddings_enabled: cfg.embeddings.enabled,
ground_dir: ground_dir.clone(),
embedder: Arc::new(Mutex::new(None)),
model: cfg.embeddings.model.clone(),
quantized: cfg.embeddings.quantized,
cache_dir,
last_activity_secs: Arc::clone(&self.inner.last_activity_secs),
});
self.inner
.resources
.lock()
.await
.insert(key, Arc::clone(&resources));
Ok(resources)
}
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 = Arc::clone(&self.inner.crossencoders).lock_owned().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_activity_secs
.store(monotonic_secs(), Ordering::Relaxed);
Ok(Some(CrossencoderGuard {
guard,
key: canonical.to_string(),
last_use_secs: Arc::clone(&self.inner.last_activity_secs),
}))
}
pub fn touch_activity(&self) {
self.inner
.last_activity_secs
.store(monotonic_secs(), Ordering::Relaxed);
}
pub fn enter_connection(&self) -> ConnectionGuard {
self.inner.active_connections.fetch_add(1, Ordering::SeqCst);
ConnectionGuard {
active: Arc::clone(&self.inner.active_connections),
}
}
pub(crate) fn should_idle_exit(&self, idle_secs: u64) -> bool {
self.should_idle_exit_at(idle_secs, monotonic_secs())
}
fn should_idle_exit_at(&self, idle_secs: u64, now_secs: u64) -> bool {
if idle_secs == 0 {
return false;
}
if self.inner.active_connections.load(Ordering::SeqCst) != 0 {
return false;
}
is_idle(
self.inner.last_activity_secs.load(Ordering::Relaxed),
now_secs,
idle_secs,
)
}
pub(crate) fn secs_until_idle(&self, idle_exit_secs: u64) -> u64 {
self.secs_until_idle_at(idle_exit_secs, monotonic_secs())
}
fn secs_until_idle_at(&self, idle_exit_secs: u64, now_secs: u64) -> u64 {
let elapsed =
now_secs.saturating_sub(self.inner.last_activity_secs.load(Ordering::Relaxed));
idle_exit_secs.saturating_sub(elapsed)
}
#[cfg(test)]
pub(crate) fn last_activity_secs(&self) -> u64 {
self.inner.last_activity_secs.load(Ordering::Relaxed)
}
#[cfg(test)]
pub(crate) fn set_last_activity_secs_for_test(&self, secs: u64) {
self.inner.last_activity_secs.store(secs, Ordering::Relaxed);
}
pub fn make_registry(&self) -> HandlerRegistry {
HandlerRegistry::new(
self.inner.baseline_resources.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?;
let stamp_target = bak.clone();
tokio::task::spawn_blocking(move || {
std::fs::File::open(&stamp_target)?.set_modified(SystemTime::now())
})
.await??;
tracing::info!(
target: "hallouminate::daemon",
backup = %bak.display(),
"moved stale ground store aside; recoverable until pruned",
);
Ok(())
}
async fn prune_stale_backups(
ground_dir: &Path,
now: SystemTime,
max_age: Duration,
) -> anyhow::Result<()> {
let parent = ground_dir
.parent()
.filter(|p| !p.as_os_str().is_empty())
.unwrap_or_else(|| Path::new("."));
let prefix = format!(
"{}.bak-v",
ground_dir
.file_name()
.and_then(|s| s.to_str())
.unwrap_or("ground"),
);
let mut entries = match tokio::fs::read_dir(parent).await {
Ok(entries) => entries,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(e) => return Err(e.into()),
};
while let Some(entry) = entries.next_entry().await? {
let file_name = entry.file_name();
let Some(name) = file_name.to_str() else {
continue;
};
let Some(suffix) = name.strip_prefix(&prefix) else {
continue;
};
if suffix.is_empty() || !suffix.chars().all(|c| c.is_ascii_digit()) {
continue;
}
let metadata = match entry.metadata().await {
Ok(metadata) => metadata,
Err(e) => {
tracing::warn!(
target: "hallouminate::daemon",
entry = %name,
error = %e,
"skipping stale-backup entry: failed to read metadata",
);
continue;
}
};
if !metadata.is_dir() {
continue;
}
let modified = match metadata.modified() {
Ok(modified) => modified,
Err(e) => {
tracing::warn!(
target: "hallouminate::daemon",
entry = %name,
error = %e,
"skipping stale-backup entry: failed to read modified time",
);
continue;
}
};
let age = now.duration_since(modified).unwrap_or(Duration::ZERO);
if age < max_age {
continue;
}
let path = entry.path();
if let Err(e) = tokio::fs::remove_dir_all(&path).await {
tracing::warn!(
target: "hallouminate::daemon",
entry = %name,
error = %e,
"failed to remove stale ground store backup",
);
continue;
}
tracing::info!(
target: "hallouminate::daemon",
backup = %path.display(),
age_days = age.as_secs() / 86_400,
"pruned stale ground store backup",
);
}
Ok(())
}
pub struct EmbedderGuard {
guard: OwnedMutexGuard<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(monotonic_secs(), Ordering::Relaxed);
}
}
pub struct CrossencoderGuard {
guard: OwnedMutexGuard<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(monotonic_secs(), Ordering::Relaxed);
}
}
impl Crossencoder for CrossencoderGuard {
fn rerank(&mut self, query: &str, hits: &mut [SearchHit]) -> crate::domain::common::Result<()> {
(**self).rerank(query, hits)
}
}
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.baseline_resources.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 should_idle_exit_is_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.should_idle_exit_at(0, u64::MAX),
"idle_secs=0 disables idle-exit; must never fire",
);
}
#[tokio::test]
async fn should_idle_exit_is_false_when_recently_active() {
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_activity_secs();
assert!(
!state.should_idle_exit_at(300, last + 1),
"1 s elapsed < 300 s idle; must not exit",
);
}
#[tokio::test]
async fn should_idle_exit_is_true_when_idle_and_no_connections() {
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_activity_secs();
assert!(
state.should_idle_exit_at(300, last + 300),
"elapsed == idle_secs (>= threshold); must exit",
);
assert!(
state.should_idle_exit_at(300, last + 301),
"elapsed > idle_secs; must exit",
);
}
#[tokio::test]
async fn should_idle_exit_is_false_one_second_below_threshold() {
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_activity_secs();
assert!(
!state.should_idle_exit_at(300, last + 299),
"elapsed = idle_secs - 1 (< threshold); must not exit",
);
}
#[tokio::test]
async fn secs_until_idle_counts_down_to_the_deadline() {
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_activity_secs();
assert_eq!(
state.secs_until_idle_at(300, last),
300,
"no time elapsed since activity; the full window remains",
);
assert_eq!(
state.secs_until_idle_at(300, last + 100),
200,
"100 s elapsed of a 300 s window; 200 s remain",
);
assert_eq!(
state.secs_until_idle_at(300, last + 300),
0,
"window exactly elapsed; deadline reached",
);
assert_eq!(
state.secs_until_idle_at(300, last + 10_000),
0,
"well past the window; saturates to zero, never underflows",
);
}
#[tokio::test]
async fn active_connection_defers_idle_exit_even_when_clock_is_idle() {
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_activity_secs();
let guard = state.enter_connection();
assert!(
!state.should_idle_exit_at(300, last + 10_000),
"an active connection must defer idle-exit even when long idle",
);
drop(guard);
assert!(
state.should_idle_exit_at(300, last + 10_000),
"once the connection count returns to zero, idle-exit fires",
);
}
#[tokio::test]
async fn maintenance_tick_stamps_the_idle_clock_and_continues() {
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.set_last_activity_secs_for_test(1);
assert!(
state.should_idle_exit_at(300, 1000),
"sanity: stale clock with no connections is idle-eligible",
);
let tick = state.run_maintenance_tick().await;
assert_eq!(tick, MaintenanceTick::Continue);
assert!(
!state.should_idle_exit(300),
"a maintenance pass must stamp the activity clock so idle-exit \
does not fire immediately after it",
);
}
#[tokio::test]
async fn maintenance_tick_stops_when_the_write_lane_is_closed() {
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.write_lane().close();
let tick = state.run_maintenance_tick().await;
assert_eq!(tick, MaintenanceTick::Stop);
}
#[tokio::test]
async fn touch_activity_advances_the_idle_clock() {
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_activity_secs.store(1, Ordering::Relaxed);
assert!(
state.should_idle_exit_at(300, 1000),
"clock stale at 1 s; now=1000 is well past idle",
);
state.touch_activity();
assert!(
!state.should_idle_exit_at(300, state.last_activity_secs() + 1),
"touch_activity must reset the clock so a fresh now is not idle",
);
}
#[test]
fn idle_clock_is_monotonic_not_wall_clock() {
let wall_clock_secs = SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system clock after epoch")
.as_secs();
let monotonic = monotonic_secs();
assert!(
monotonic < wall_clock_secs / 2,
"idle clock must be process-relative (Instant-based), not wall-clock \
Unix seconds: monotonic={monotonic}, wall_clock={wall_clock_secs}",
);
}
#[tokio::test]
async fn embedder_guard_updates_last_use_on_drop() {
let last_use_secs = Arc::new(AtomicU64::new(1));
let guard = Arc::new(Mutex::new(None)).lock_owned().await;
let before_drop = monotonic_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 before_drop = monotonic_secs();
drop(CrossencoderGuard {
guard: Arc::new(Mutex::new(HashMap::new())).lock_owned().await,
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}",
);
}
#[tokio::test]
async fn prune_stale_backups_removes_dirs_at_or_past_max_age() {
let tmp = tempfile::tempdir().expect("tempdir");
let ground_dir = tmp.path().join("ground");
tokio::fs::create_dir_all(&ground_dir)
.await
.expect("create ground dir");
let stale = tmp.path().join("ground.bak-v1");
let unrelated = tmp.path().join("other-dir");
tokio::fs::create_dir_all(&stale)
.await
.expect("create stale backup");
tokio::fs::create_dir_all(&unrelated)
.await
.expect("create unrelated dir");
let max_age = Duration::from_secs(1);
let now = SystemTime::now() + Duration::from_secs(10);
prune_stale_backups(&ground_dir, now, max_age)
.await
.expect("prune");
assert!(!stale.exists(), "backup past max_age must be pruned");
assert!(
unrelated.exists(),
"dirs that don't match the backup prefix must be left alone",
);
}
#[tokio::test]
async fn prune_stale_backups_keeps_dirs_within_max_age() {
let tmp = tempfile::tempdir().expect("tempdir");
let ground_dir = tmp.path().join("ground");
tokio::fs::create_dir_all(&ground_dir)
.await
.expect("create ground dir");
let fresh = tmp.path().join("ground.bak-v2");
tokio::fs::create_dir_all(&fresh)
.await
.expect("create fresh backup");
prune_stale_backups(&ground_dir, SystemTime::now(), STALE_BACKUP_MAX_AGE)
.await
.expect("prune");
assert!(fresh.exists(), "backup younger than max_age must survive");
}
#[tokio::test]
async fn prune_stale_backups_keeps_non_numeric_suffix_even_if_stale() {
let tmp = tempfile::tempdir().expect("tempdir");
let ground_dir = tmp.path().join("ground");
tokio::fs::create_dir_all(&ground_dir)
.await
.expect("create ground dir");
let lookalike = tmp.path().join("ground.bak-vault");
tokio::fs::create_dir_all(&lookalike)
.await
.expect("create lookalike dir");
let max_age = Duration::from_secs(1);
let now = SystemTime::now() + Duration::from_secs(31 * 24 * 60 * 60);
prune_stale_backups(&ground_dir, now, max_age)
.await
.expect("prune");
assert!(
lookalike.exists(),
"non-numeric-suffix dir must survive even when older than max_age",
);
}
#[tokio::test]
async fn move_stale_store_stamps_backup_mtime_to_now() {
let tmp = tempfile::tempdir().expect("tempdir");
let ground_dir = tmp.path().join("ground");
tokio::fs::create_dir_all(&ground_dir)
.await
.expect("create ground dir");
let old_mtime = SystemTime::now() - Duration::from_secs(60 * 24 * 60 * 60);
let dir = ground_dir.clone();
tokio::task::spawn_blocking(move || std::fs::File::open(&dir)?.set_modified(old_mtime))
.await
.expect("join")
.expect("backdate ground dir mtime");
move_stale_store(&ground_dir, 1).await.expect("move");
let bak = tmp.path().join("ground.bak-v1");
prune_stale_backups(&ground_dir, SystemTime::now(), STALE_BACKUP_MAX_AGE)
.await
.expect("prune");
assert!(
bak.exists(),
"backup just moved aside must survive its own boot's prune, \
even though the source dir's mtime was 60d old",
);
}
#[tokio::test]
async fn resources_for_keys_on_ground_dir() {
let tmp_a = tempfile::tempdir().expect("tempdir a");
let tmp_b = tempfile::tempdir().expect("tempdir b");
let mut cfg_a = Config::default();
cfg_a.embeddings.enabled = false;
cfg_a.storage.ground_dir = tmp_a.path().to_string_lossy().into_owned();
let mut cfg_b = cfg_a.clone();
cfg_b.storage.ground_dir = tmp_b.path().to_string_lossy().into_owned();
let state = DaemonState::open(cfg_a.clone(), None)
.await
.expect("open daemon state");
let res_a1 = state
.resources_for(&cfg_a)
.await
.expect("resources_for cfg_a (first call)");
let res_a2 = state
.resources_for(&cfg_a)
.await
.expect("resources_for cfg_a (second call)");
assert!(
Arc::ptr_eq(&res_a1, &res_a2),
"same config must resolve the same cached RequestResources Arc, \
not rebuild or reopen a second store",
);
let res_b = state
.resources_for(&cfg_b)
.await
.expect("resources_for cfg_b");
assert!(
!Arc::ptr_eq(&res_a1, &res_b),
"a different ground_dir must key a distinct RequestResources entry",
);
assert_eq!(
res_a1.ground_dir,
tmp_a.path(),
"cfg_a's resources must be rooted at tmp_a's ground_dir",
);
assert_eq!(
res_b.ground_dir,
tmp_b.path(),
"cfg_b's resources must be rooted at tmp_b's ground_dir, not cfg_a's",
);
assert!(
tmp_b.path().join("meta.toml").exists(),
"resources_for must open (and initialize) a store at the new \
ground_dir on first use",
);
}
#[tokio::test]
async fn resources_for_builds_once_under_concurrent_same_key_calls() {
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.clone(), None)
.await
.expect("open daemon state");
let tmp_race = tempfile::tempdir().expect("tempdir race");
cfg.storage.ground_dir = tmp_race.path().to_string_lossy().into_owned();
let mut handles = Vec::new();
for _ in 0..16 {
let state = state.clone();
let cfg = cfg.clone();
handles.push(tokio::spawn(async move {
state.resources_for(&cfg).await.expect("resources_for")
}));
}
let mut resolved = Vec::new();
for h in handles {
resolved.push(h.await.expect("join resources_for task"));
}
let first = &resolved[0];
for (i, res) in resolved.iter().enumerate() {
assert!(
Arc::ptr_eq(first, res),
"racer {i} resolved a different RequestResources Arc — the \
store was opened more than once for one ground dir",
);
}
}
}