mentedb_server/
databases.rs

1//! Named databases: one MenteDB per tenant (a person, a team, an
2//! organization), each in its own directory with its own lock, indexes and
3//! files, so nothing is shared between them but the process.
4//!
5//! The legacy default database lives at the data directory root and backs the
6//! unprefixed `/v1/...` routes. Named databases live under `<data>/dbs/<name>`
7//! and back `/v1/db/{name}/...`. They open on first use, and the least recently
8//! used idle handle is flushed and closed once more than `max_open` are open.
9
10use std::collections::{HashMap, HashSet};
11use std::path::{Path, PathBuf};
12use std::sync::{Arc, Condvar, Mutex, MutexGuard};
13use std::time::Instant;
14
15use mentedb::{CognitiveConfig, MenteDb};
16use mentedb_core::error::MenteResult;
17use mentedb_embedding::provider::EmbeddingProvider;
18
19use crate::error::ApiError;
20
21/// Name of the database at the data directory root.
22pub const DEFAULT_DB: &str = "default";
23
24/// Most characters a database name may have.
25const MAX_NAME_LEN: usize = 128;
26
27/// Limits for named databases.
28#[derive(Clone, Debug)]
29pub struct DatabasesConfig {
30    /// Named databases kept open at once; idle ones beyond this are closed.
31    pub max_open: usize,
32    /// Create a named database on first use instead of returning 404.
33    pub auto_create: bool,
34}
35
36impl Default for DatabasesConfig {
37    fn default() -> Self {
38        Self {
39            max_open: 256,
40            auto_create: true,
41        }
42    }
43}
44
45/// One embedder shared by every open database.
46struct SharedEmbedder(Arc<dyn EmbeddingProvider>);
47
48impl EmbeddingProvider for SharedEmbedder {
49    fn embed(&self, text: &str) -> MenteResult<Vec<f32>> {
50        self.0.embed(text)
51    }
52    fn embed_batch(&self, texts: &[&str]) -> MenteResult<Vec<Vec<f32>>> {
53        self.0.embed_batch(texts)
54    }
55    fn dimensions(&self) -> usize {
56        self.0.dimensions()
57    }
58    fn model_name(&self) -> &str {
59        self.0.model_name()
60    }
61}
62
63struct OpenDb {
64    db: Arc<MenteDb>,
65    last_used: Instant,
66}
67
68/// Open handles plus the names whose old handle is still being closed. A
69/// name in `closing` must not be reopened until its close has released the
70/// directory lock, or the open fails with "locked by another process".
71#[derive(Default)]
72struct Registry {
73    open: HashMap<String, OpenDb>,
74    closing: HashSet<String>,
75}
76
77/// Registry of the default database and the named ones.
78pub struct Databases {
79    root: PathBuf,
80    default: Arc<MenteDb>,
81    embedder: Option<Arc<dyn EmbeddingProvider>>,
82    cognitive: CognitiveConfig,
83    config: DatabasesConfig,
84    registry: Mutex<Registry>,
85    /// Signalled whenever a name leaves `Registry::closing`.
86    closed: Condvar,
87}
88
89impl Databases {
90    /// Open the default database at `root` and prepare the registry.
91    pub fn open(
92        root: &Path,
93        embedder: Option<Arc<dyn EmbeddingProvider>>,
94        cognitive: CognitiveConfig,
95        config: DatabasesConfig,
96    ) -> MenteResult<Self> {
97        std::fs::create_dir_all(root)?;
98        let default = Arc::new(open_at(root, embedder.clone(), cognitive.clone())?);
99        Ok(Self {
100            root: root.to_path_buf(),
101            default,
102            embedder,
103            cognitive,
104            config,
105            registry: Mutex::new(Registry::default()),
106            closed: Condvar::new(),
107        })
108    }
109
110    /// The database at the data directory root.
111    pub fn default_db(&self) -> Arc<MenteDb> {
112        self.default.clone()
113    }
114
115    /// The data directory root (the default database's directory).
116    pub(crate) fn root(&self) -> &Path {
117        &self.root
118    }
119
120    /// The shared embedder, when one is configured.
121    pub fn embedder(&self) -> Option<Arc<dyn EmbeddingProvider>> {
122        self.embedder.clone()
123    }
124
125    fn lock(&self) -> MutexGuard<'_, Registry> {
126        self.registry.lock().unwrap_or_else(|e| e.into_inner())
127    }
128
129    /// Wait until no close of `name` is in flight. Returns the reacquired
130    /// guard.
131    fn wait_until_closed<'a>(
132        &'a self,
133        mut reg: MutexGuard<'a, Registry>,
134        name: &str,
135    ) -> MutexGuard<'a, Registry> {
136        while reg.closing.contains(name) {
137            reg = self.closed.wait(reg).unwrap_or_else(|e| e.into_inner());
138        }
139        reg
140    }
141
142    /// Close handles taken out of the registry, then release their names so
143    /// waiting opens can proceed. Each handle is dropped (releasing its
144    /// directory lock) before its name is released.
145    fn finish_closing(&self, handles: Vec<(String, Arc<MenteDb>)>) {
146        let mut names = Vec::with_capacity(handles.len());
147        for (name, handle) in handles {
148            if let Err(e) = handle.close() {
149                tracing::warn!(db = %name, error = %e, "closing idle database failed");
150            }
151            drop(handle);
152            names.push(name);
153        }
154        let mut reg = self.lock();
155        for name in names {
156            reg.closing.remove(&name);
157        }
158        drop(reg);
159        self.closed.notify_all();
160    }
161
162    fn dir(&self, name: &str) -> PathBuf {
163        self.root.join("dbs").join(name)
164    }
165
166    /// Whether a named database exists on disk.
167    pub fn exists(&self, name: &str) -> bool {
168        name == DEFAULT_DB || self.dir(name).join(".mentedb-db").exists()
169    }
170
171    /// Get a database by name, opening it if needed. Blocking: call from
172    /// `spawn_blocking`. `create` allows making a new database.
173    pub fn get(&self, name: &str, create: bool) -> Result<Arc<MenteDb>, ApiError> {
174        if name == DEFAULT_DB {
175            return Ok(self.default.clone());
176        }
177        validate_name(name)?;
178        let reg = self.lock();
179        let mut reg = self.wait_until_closed(reg, name);
180        if let Some(entry) = reg.open.get_mut(name) {
181            entry.last_used = Instant::now();
182            return Ok(entry.db.clone());
183        }
184        let dir = self.dir(name);
185        let marker = dir.join(".mentedb-db");
186        if !marker.exists() {
187            if !(create || self.config.auto_create) {
188                return Err(ApiError::NotFound(format!("database '{name}' not found")));
189            }
190            std::fs::create_dir_all(&dir)
191                .map_err(|e| ApiError::Internal(format!("create database '{name}': {e}")))?;
192            std::fs::write(&marker, name.as_bytes())
193                .map_err(|e| ApiError::Internal(format!("create database '{name}': {e}")))?;
194        }
195        // Opening holds the registry lock so one name is never opened twice;
196        // the engine's own file lock would refuse the second open anyway.
197        let db = Arc::new(
198            open_at(&dir, self.embedder.clone(), self.cognitive.clone())
199                .map_err(|e| ApiError::Internal(format!("open database '{name}': {e}")))?,
200        );
201        reg.open.insert(
202            name.to_string(),
203            OpenDb {
204                db: db.clone(),
205                last_used: Instant::now(),
206            },
207        );
208        let evicted = take_idle_overflow(&mut reg.open, self.config.max_open);
209        for (evicted_name, _) in &evicted {
210            reg.closing.insert(evicted_name.clone());
211        }
212        drop(reg);
213        if !evicted.is_empty() {
214            self.finish_closing(evicted);
215        }
216        Ok(db)
217    }
218
219    /// Named databases on disk, sorted.
220    pub fn list(&self) -> Vec<String> {
221        let mut names: Vec<String> = std::fs::read_dir(self.root.join("dbs"))
222            .into_iter()
223            .flatten()
224            .flatten()
225            .filter(|e| e.path().join(".mentedb-db").exists())
226            .filter_map(|e| e.file_name().into_string().ok())
227            .collect();
228        names.sort();
229        names
230    }
231
232    /// Close and permanently delete a named database. Blocking.
233    pub fn delete(&self, name: &str) -> Result<(), ApiError> {
234        if name == DEFAULT_DB {
235            return Err(ApiError::BadRequest(
236                "the default database cannot be deleted".into(),
237            ));
238        }
239        validate_name(name)?;
240        let reg = self.lock();
241        let mut reg = self.wait_until_closed(reg, name);
242        if let Some(entry) = reg.open.get(name)
243            && Arc::strong_count(&entry.db) > 1
244        {
245            return Err(ApiError::Conflict(format!(
246                "database '{name}' is in use; retry"
247            )));
248        }
249        // Closed and removed while holding the registry, so no open of this
250        // name can start until the directory is gone.
251        if let Some(entry) = reg.open.remove(name) {
252            let _ = entry.db.close();
253        }
254        let dir = self.dir(name);
255        if !dir.join(".mentedb-db").exists() {
256            return Err(ApiError::NotFound(format!("database '{name}' not found")));
257        }
258        std::fs::remove_dir_all(&dir)
259            .map_err(|e| ApiError::Internal(format!("delete database '{name}': {e}")))
260    }
261
262    /// Directory a database's files live in.
263    pub fn path_of(&self, name: &str) -> Result<PathBuf, ApiError> {
264        if name == DEFAULT_DB {
265            return Ok(self.root.clone());
266        }
267        validate_name(name)?;
268        Ok(self.dir(name))
269    }
270
271    /// Every currently open database, the default first.
272    pub fn open_handles(&self) -> Vec<(String, Arc<MenteDb>)> {
273        let mut all = vec![(DEFAULT_DB.to_string(), self.default.clone())];
274        let reg = self.lock();
275        all.extend(reg.open.iter().map(|(n, e)| (n.clone(), e.db.clone())));
276        all
277    }
278
279    /// Flush and close every open database. Used at shutdown.
280    pub fn close_all(&self) {
281        let handles: Vec<(String, Arc<MenteDb>)> = {
282            let mut reg = self.lock();
283            let drained: Vec<_> = reg.open.drain().map(|(n, e)| (n, e.db)).collect();
284            for (n, _) in &drained {
285                reg.closing.insert(n.clone());
286            }
287            drained
288        };
289        self.finish_closing(handles);
290        if let Err(e) = self.default.close() {
291            tracing::warn!(error = %e, "closing default database failed");
292        }
293    }
294}
295
296fn open_at(
297    dir: &Path,
298    embedder: Option<Arc<dyn EmbeddingProvider>>,
299    cognitive: CognitiveConfig,
300) -> MenteResult<MenteDb> {
301    match embedder {
302        Some(e) => {
303            MenteDb::open_with_embedder_and_config(dir, Box::new(SharedEmbedder(e)), cognitive)
304        }
305        None => MenteDb::open_with_config(dir, cognitive),
306    }
307}
308
309/// Remove the least recently used handles nobody else holds until at most
310/// `max_open` remain. Returned handles are closed by the caller off the lock.
311fn take_idle_overflow(
312    open: &mut HashMap<String, OpenDb>,
313    max_open: usize,
314) -> Vec<(String, Arc<MenteDb>)> {
315    let mut evicted = Vec::new();
316    while open.len() > max_open {
317        let victim = open
318            .iter()
319            .filter(|(_, e)| Arc::strong_count(&e.db) == 1)
320            .min_by_key(|(_, e)| e.last_used)
321            .map(|(n, _)| n.clone());
322        match victim {
323            Some(name) => {
324                if let Some(entry) = open.remove(&name) {
325                    evicted.push((name, entry.db));
326                }
327            }
328            // Every handle is in use; stay over the limit until one frees up.
329            None => break,
330        }
331    }
332    evicted
333}
334
335/// Database names are path segments: letters, digits, `-`, `_` and `.`, not
336/// starting with `.`, at most 128 characters.
337pub fn validate_name(name: &str) -> Result<(), ApiError> {
338    let ok = !name.is_empty()
339        && name.len() <= MAX_NAME_LEN
340        && !name.starts_with('.')
341        && name
342            .chars()
343            .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.'));
344    if ok {
345        Ok(())
346    } else {
347        Err(ApiError::BadRequest(format!(
348            "invalid database name '{name}': use letters, digits, '-', '_' or '.', at most \
349             {MAX_NAME_LEN} characters, not starting with '.'"
350        )))
351    }
352}
353
354/// Paths under a database that only read, so read-only tokens may POST them.
355const READ_ONLY_POSTS: &[&str] = &["/recall", "/search", "/search/hybrid", "/context", "/query"];
356
357/// The database a request targets: `{db}` in `/v1/db/{db}/...`, otherwise the
358/// default. Checks the caller's token scope and opens the database.
359pub struct Db {
360    pub name: String,
361    pub handle: Arc<MenteDb>,
362}
363
364impl std::ops::Deref for Db {
365    type Target = MenteDb;
366    fn deref(&self) -> &MenteDb {
367        &self.handle
368    }
369}
370
371impl axum::extract::FromRequestParts<Arc<crate::state::AppState>> for Db {
372    type Rejection = ApiError;
373
374    async fn from_request_parts(
375        parts: &mut axum::http::request::Parts,
376        state: &Arc<crate::state::AppState>,
377    ) -> Result<Self, ApiError> {
378        use axum::extract::RawPathParams;
379        let params = RawPathParams::from_request_parts(parts, state)
380            .await
381            .map_err(|e| ApiError::BadRequest(e.to_string()))?;
382        let name = params
383            .iter()
384            .find(|(k, _)| *k == "db")
385            .map(|(_, v)| v.to_string())
386            .unwrap_or_else(|| DEFAULT_DB.to_string());
387        let path = parts.uri.path();
388        let write = match parts.method {
389            axum::http::Method::GET | axum::http::Method::HEAD => false,
390            axum::http::Method::POST => !READ_ONLY_POSTS.iter().any(|p| path.ends_with(p)),
391            _ => true,
392        };
393        if let Some(agent) = parts.extensions.get::<crate::auth::AuthenticatedAgent>() {
394            agent.authorize_db(&name, write)?;
395        }
396        let dbs = state.databases.clone();
397        let lookup = name.clone();
398        let handle = tokio::task::spawn_blocking(move || dbs.get(&lookup, false))
399            .await
400            .map_err(|e| ApiError::Internal(format!("open database task failed: {e}")))??;
401        Ok(Self { name, handle })
402    }
403}
404
405#[cfg(test)]
406mod tests {
407    use super::*;
408
409    fn registry(dir: &Path, max_open: usize) -> Databases {
410        Databases::open(
411            dir,
412            None,
413            CognitiveConfig::default(),
414            DatabasesConfig {
415                max_open,
416                auto_create: true,
417            },
418        )
419        .unwrap()
420    }
421
422    #[test]
423    fn names_are_safe_path_segments() {
424        for ok in ["alice", "team-7", "org_1", "u.123"] {
425            assert!(validate_name(ok).is_ok(), "{ok}");
426        }
427        for bad in ["", "..", ".hidden", "a/b", "a\\b", "x y", &"a".repeat(129)] {
428            assert!(validate_name(bad).is_err(), "{bad}");
429        }
430    }
431
432    #[test]
433    fn named_databases_are_isolated_directories() {
434        let tmp = tempfile::tempdir().unwrap();
435        let dbs = registry(tmp.path(), 8);
436        let a = dbs.get("alice", false).unwrap();
437        let b = dbs.get("bob", false).unwrap();
438        assert!(!Arc::ptr_eq(&a, &b));
439        assert!(tmp.path().join("dbs/alice/.mentedb-db").exists());
440        assert!(tmp.path().join("dbs/bob/.mentedb-db").exists());
441        assert_eq!(dbs.list(), vec!["alice".to_string(), "bob".to_string()]);
442        assert!(
443            Arc::ptr_eq(&a, &dbs.get("alice", false).unwrap()),
444            "cached handle"
445        );
446    }
447
448    #[test]
449    fn idle_databases_beyond_the_limit_are_closed() {
450        let tmp = tempfile::tempdir().unwrap();
451        let dbs = registry(tmp.path(), 2);
452        drop(dbs.get("a", false).unwrap());
453        drop(dbs.get("b", false).unwrap());
454        drop(dbs.get("c", false).unwrap());
455        let open: Vec<String> = dbs.open_handles().into_iter().map(|(n, _)| n).collect();
456        assert_eq!(open.len(), 3, "default plus two named: {open:?}");
457        assert!(
458            !open.contains(&"a".to_string()),
459            "least recently used closed"
460        );
461        // A closed database reopens with its data intact.
462        assert!(dbs.get("a", false).is_ok());
463    }
464
465    #[test]
466    fn missing_database_is_404_without_auto_create() {
467        let tmp = tempfile::tempdir().unwrap();
468        let dbs = Databases::open(
469            tmp.path(),
470            None,
471            CognitiveConfig::default(),
472            DatabasesConfig {
473                max_open: 4,
474                auto_create: false,
475            },
476        )
477        .unwrap();
478        assert!(matches!(
479            dbs.get("nobody", false),
480            Err(ApiError::NotFound(_))
481        ));
482        assert!(dbs.get("nobody", true).is_ok(), "explicit create works");
483    }
484
485    /// More tenants than `max_open`, hit from many threads at once, so idle
486    /// handles are closed while other threads reopen the same names. Every
487    /// lookup must succeed; a reopen must never race the close of the old
488    /// handle and find the directory still locked.
489    #[test]
490    fn reopening_a_database_while_it_is_being_closed_waits_for_the_close() {
491        let tmp = tempfile::tempdir().unwrap();
492        let dbs = Arc::new(registry(tmp.path(), 2));
493        let names: Vec<String> = (0..8).map(|i| format!("t{i}")).collect();
494        let errors = Arc::new(std::sync::atomic::AtomicUsize::new(0));
495        let threads: Vec<_> = (0..16)
496            .map(|t| {
497                let dbs = dbs.clone();
498                let names = names.clone();
499                let errors = errors.clone();
500                std::thread::spawn(move || {
501                    for i in 0..40 {
502                        let name = &names[(t * 7 + i) % names.len()];
503                        match dbs.get(name, false) {
504                            Ok(db) => {
505                                let _ = db.memory_count();
506                            }
507                            Err(e) => {
508                                eprintln!("{name}: {e}");
509                                errors.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
510                            }
511                        }
512                    }
513                })
514            })
515            .collect();
516        for t in threads {
517            t.join().unwrap();
518        }
519        assert_eq!(errors.load(std::sync::atomic::Ordering::SeqCst), 0);
520    }
521
522    #[test]
523    fn delete_removes_the_directory() {
524        let tmp = tempfile::tempdir().unwrap();
525        let dbs = registry(tmp.path(), 4);
526        drop(dbs.get("gone", false).unwrap());
527        dbs.delete("gone").unwrap();
528        assert!(!tmp.path().join("dbs/gone").exists());
529        assert!(dbs.delete(DEFAULT_DB).is_err());
530    }
531}
mentedb_server/databases.rs

mentedb_server/
databases.rs
