Skip to main content

clawft_kernel/cog_ingest/
vector_dir.rs

1//! A node's per-project vector stores (ADR-100 decision 5): in-memory HNSW
2//! indexes on the owner daemon, optionally backed by a durable log.
3
4use std::collections::{HashMap, HashSet};
5use std::sync::{Arc, Mutex};
6
7use super::store::{IngestStore, StoreDirectory, VectorBackendStore};
8use super::store_log::{OpenError, run_blocking};
9
10/// A node's stores for the projects it owns, each an HNSW index created on
11/// first use. Without [`with_persistence`](Self::with_persistence) the
12/// indexes are memory-only and a restart empties them; with it, each store
13/// keeps a capped, append-only log under the given directory (one file per
14/// project) and is rebuilt from it on first use after a restart. The index
15/// map is shared between [`views`](Self::view), so two views with different
16/// allow-lists (local cogs, remote forwarders) reach the same store for a
17/// project they both allow.
18#[derive(Clone)]
19pub struct VectorDirectory {
20    projects: HashSet<String>,
21    fallback: bool,
22    stores: Arc<Mutex<HashMap<Option<String>, Slot>>>,
23    persist: Option<(std::path::PathBuf, u64)>,
24}
25
26/// One project's store, opened on first use. The slot has its own lock, so
27/// replaying a big log blocks only callers of that project, never the
28/// directory; a failed open leaves it empty and the next call retries.
29type Slot = Arc<Mutex<Option<Arc<VectorBackendStore>>>>;
30
31/// File of the controller-fallback store inside the persistence dir. Not a
32/// valid project id, so it cannot collide with a project's file.
33const FALLBACK_LOG: &str = "_controller.vec";
34
35impl VectorDirectory {
36    /// Directory owning `projects`, and the controller fallback if `fallback`.
37    pub fn new(projects: impl IntoIterator<Item = String>, fallback: bool) -> Self {
38        Self {
39            projects: projects.into_iter().collect(),
40            fallback,
41            stores: Arc::default(),
42            persist: None,
43        }
44    }
45
46    /// Keep each store's vectors in `<dir>/<project id>.vec` (the controller
47    /// fallback in `<dir>/_controller.vec`), each capped at
48    /// [`DEFAULT_MAX_LOG_BYTES`](super::store_log::DEFAULT_MAX_LOG_BYTES).
49    /// Call before taking [`view`](Self::view)s.
50    pub fn with_persistence(self, dir: std::path::PathBuf) -> Self {
51        self.with_persistence_capped(dir, super::store_log::DEFAULT_MAX_LOG_BYTES)
52    }
53
54    /// [`with_persistence`](Self::with_persistence) with an explicit cap per file.
55    pub fn with_persistence_capped(mut self, dir: std::path::PathBuf, max_log_bytes: u64) -> Self {
56        self.persist = Some((dir, max_log_bytes));
57        self
58    }
59
60    /// A view over the same stores allowing only `projects` (and the
61    /// fallback if `fallback`).
62    pub fn view(&self, projects: impl IntoIterator<Item = String>, fallback: bool) -> Self {
63        Self {
64            projects: projects.into_iter().collect(),
65            fallback,
66            stores: self.stores.clone(),
67            persist: self.persist.clone(),
68        }
69    }
70
71    fn open_store(&self, project_id: Option<&str>) -> Option<Arc<VectorBackendStore>> {
72        let backend = || {
73            Arc::new(crate::vector_hnsw::HnswBackend::new(
74                crate::hnsw_service::HnswServiceConfig::default(),
75            ))
76        };
77        let Some((dir, cap)) = &self.persist else {
78            return Some(Arc::new(VectorBackendStore::new(backend())));
79        };
80        let file = match project_id {
81            Some(p) => format!("{p}.vec"),
82            None => FALLBACK_LOG.to_string(),
83        };
84        let path = dir.join(&file);
85        match VectorBackendStore::persistent(backend(), &path, *cap) {
86            Ok(s) => Some(Arc::new(s)),
87            Err(OpenError::Io(e)) => {
88                // Possibly transient and the file is not implicated: refuse for now, retry next call.
89                tracing::warn!(error = %e, path = %path.display(), "cog ingest vector log could not be opened");
90                None
91            }
92            Err(OpenError::Corrupt(e)) => {
93                // Never overwrite what could not be read: keep it aside, under a name that
94                // cannot clobber earlier evidence, and start this project's store empty.
95                let ts = std::time::SystemTime::now()
96                    .duration_since(std::time::UNIX_EPOCH)
97                    .map(|d| d.as_nanos())
98                    .unwrap_or_default();
99                let aside = dir.join(format!("{file}.unreadable.{ts}"));
100                tracing::warn!(error = %e, path = %path.display(), aside = %aside.display(), "cog ingest vector log unusable; moving it aside");
101                std::fs::rename(&path, &aside).ok()?;
102                VectorBackendStore::persistent(backend(), &path, *cap).ok().map(Arc::new)
103            }
104        }
105    }
106}
107
108impl StoreDirectory for VectorDirectory {
109    fn store_for(&self, project_id: Option<&str>) -> Option<Arc<dyn IngestStore>> {
110        match project_id {
111            Some(p) if !self.projects.contains(p) => return None,
112            None if !self.fallback => return None,
113            _ => {}
114        }
115        let slot = {
116            let mut g = self.stores.lock().ok()?;
117            g.entry(project_id.map(String::from)).or_default().clone()
118        };
119        // The directory lock is released; only this project's slot is held while the log replays.
120        let mut cell = slot.lock().ok()?;
121        if cell.is_none() {
122            *cell = run_blocking(|| self.open_store(project_id));
123        }
124        cell.clone().map(|s| s as Arc<dyn IngestStore>)
125    }
126}