clawft_kernel/cog_ingest/
vector_dir.rs1use 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#[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
26type Slot = Arc<Mutex<Option<Arc<VectorBackendStore>>>>;
30
31const FALLBACK_LOG: &str = "_controller.vec";
34
35impl VectorDirectory {
36 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 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 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 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 tracing::warn!(error = %e, path = %path.display(), "cog ingest vector log could not be opened");
90 None
91 }
92 Err(OpenError::Corrupt(e)) => {
93 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 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}