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] = &[
356    "/recall",
357    "/search",
358    "/search/hybrid",
359    "/context",
360    "/query",
361    "/recall_for_action",
362];
363
364/// The database a request targets: `{db}` in `/v1/db/{db}/...`, otherwise the
365/// default. Checks the caller's token scope and opens the database.
366pub struct Db {
367    pub name: String,
368    pub handle: Arc<MenteDb>,
369}
370
371impl std::ops::Deref for Db {
372    type Target = MenteDb;
373    fn deref(&self) -> &MenteDb {
374        &self.handle
375    }
376}
377
378impl axum::extract::FromRequestParts<Arc<crate::state::AppState>> for Db {
379    type Rejection = ApiError;
380
381    async fn from_request_parts(
382        parts: &mut axum::http::request::Parts,
383        state: &Arc<crate::state::AppState>,
384    ) -> Result<Self, ApiError> {
385        use axum::extract::RawPathParams;
386        let params = RawPathParams::from_request_parts(parts, state)
387            .await
388            .map_err(|e| ApiError::BadRequest(e.to_string()))?;
389        let name = params
390            .iter()
391            .find(|(k, _)| *k == "db")
392            .map(|(_, v)| v.to_string())
393            .unwrap_or_else(|| DEFAULT_DB.to_string());
394        let path = parts.uri.path();
395        let write = match parts.method {
396            axum::http::Method::GET | axum::http::Method::HEAD => false,
397            axum::http::Method::POST => !READ_ONLY_POSTS.iter().any(|p| path.ends_with(p)),
398            _ => true,
399        };
400        if let Some(agent) = parts.extensions.get::<crate::auth::AuthenticatedAgent>() {
401            agent.authorize_db(&name, write)?;
402        }
403        let dbs = state.databases.clone();
404        let lookup = name.clone();
405        let handle = tokio::task::spawn_blocking(move || dbs.get(&lookup, false))
406            .await
407            .map_err(|e| ApiError::Internal(format!("open database task failed: {e}")))??;
408        Ok(Self { name, handle })
409    }
410}
411
412#[cfg(test)]
413mod tests {
414    use super::*;
415
416    fn registry(dir: &Path, max_open: usize) -> Databases {
417        Databases::open(
418            dir,
419            None,
420            CognitiveConfig::default(),
421            DatabasesConfig {
422                max_open,
423                auto_create: true,
424            },
425        )
426        .unwrap()
427    }
428
429    /// The allowed set is a stable contract: callers build hierarchical names
430    /// such as `tenant.alice.notes`, and nothing that could leave the data
431    /// directory is ever accepted.
432    #[test]
433    fn names_are_safe_path_segments() {
434        for ok in [
435            "alice",
436            "team-7",
437            "org_1",
438            "u.123",
439            "a.b.c",
440            "tenant.alice.notes",
441            "a..b",
442            &"a".repeat(128),
443        ] {
444            assert!(validate_name(ok).is_ok(), "{ok}");
445        }
446        for bad in [
447            "",
448            ".",
449            "..",
450            ".hidden",
451            "a/b",
452            "/abs",
453            "a/../b",
454            "../x",
455            "a\\b",
456            "x y",
457            "a\0b",
458            "caf\u{e9}",
459            &"a".repeat(129),
460        ] {
461            assert!(validate_name(bad).is_err(), "{bad}");
462        }
463    }
464
465    #[test]
466    fn dotted_names_open_as_one_directory() {
467        let tmp = tempfile::tempdir().unwrap();
468        let dbs = registry(tmp.path(), 4);
469        drop(dbs.get("tenant.alice.notes", false).unwrap());
470        assert!(
471            tmp.path()
472                .join("dbs/tenant.alice.notes/.mentedb-db")
473                .exists()
474        );
475        assert!(
476            !tmp.path().join("dbs/tenant").exists(),
477            "not nested directories"
478        );
479    }
480
481    #[test]
482    fn named_databases_are_isolated_directories() {
483        let tmp = tempfile::tempdir().unwrap();
484        let dbs = registry(tmp.path(), 8);
485        let a = dbs.get("alice", false).unwrap();
486        let b = dbs.get("bob", false).unwrap();
487        assert!(!Arc::ptr_eq(&a, &b));
488        assert!(tmp.path().join("dbs/alice/.mentedb-db").exists());
489        assert!(tmp.path().join("dbs/bob/.mentedb-db").exists());
490        assert_eq!(dbs.list(), vec!["alice".to_string(), "bob".to_string()]);
491        assert!(
492            Arc::ptr_eq(&a, &dbs.get("alice", false).unwrap()),
493            "cached handle"
494        );
495    }
496
497    #[test]
498    fn idle_databases_beyond_the_limit_are_closed() {
499        let tmp = tempfile::tempdir().unwrap();
500        let dbs = registry(tmp.path(), 2);
501        drop(dbs.get("a", false).unwrap());
502        drop(dbs.get("b", false).unwrap());
503        drop(dbs.get("c", false).unwrap());
504        let open: Vec<String> = dbs.open_handles().into_iter().map(|(n, _)| n).collect();
505        assert_eq!(open.len(), 3, "default plus two named: {open:?}");
506        assert!(
507            !open.contains(&"a".to_string()),
508            "least recently used closed"
509        );
510        // A closed database reopens with its data intact.
511        assert!(dbs.get("a", false).is_ok());
512    }
513
514    #[test]
515    fn missing_database_is_404_without_auto_create() {
516        let tmp = tempfile::tempdir().unwrap();
517        let dbs = Databases::open(
518            tmp.path(),
519            None,
520            CognitiveConfig::default(),
521            DatabasesConfig {
522                max_open: 4,
523                auto_create: false,
524            },
525        )
526        .unwrap();
527        assert!(matches!(
528            dbs.get("nobody", false),
529            Err(ApiError::NotFound(_))
530        ));
531        assert!(dbs.get("nobody", true).is_ok(), "explicit create works");
532    }
533
534    /// More tenants than `max_open`, hit from many threads at once, so idle
535    /// handles are closed while other threads reopen the same names. Every
536    /// lookup must succeed; a reopen must never race the close of the old
537    /// handle and find the directory still locked.
538    #[test]
539    fn reopening_a_database_while_it_is_being_closed_waits_for_the_close() {
540        let tmp = tempfile::tempdir().unwrap();
541        let dbs = Arc::new(registry(tmp.path(), 2));
542        let names: Vec<String> = (0..8).map(|i| format!("t{i}")).collect();
543        let errors = Arc::new(std::sync::atomic::AtomicUsize::new(0));
544        let threads: Vec<_> = (0..16)
545            .map(|t| {
546                let dbs = dbs.clone();
547                let names = names.clone();
548                let errors = errors.clone();
549                std::thread::spawn(move || {
550                    for i in 0..40 {
551                        let name = &names[(t * 7 + i) % names.len()];
552                        match dbs.get(name, false) {
553                            Ok(db) => {
554                                let _ = db.memory_count();
555                            }
556                            Err(e) => {
557                                eprintln!("{name}: {e}");
558                                errors.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
559                            }
560                        }
561                    }
562                })
563            })
564            .collect();
565        for t in threads {
566            t.join().unwrap();
567        }
568        assert_eq!(errors.load(std::sync::atomic::Ordering::SeqCst), 0);
569    }
570
571    #[test]
572    fn delete_removes_the_directory() {
573        let tmp = tempfile::tempdir().unwrap();
574        let dbs = registry(tmp.path(), 4);
575        drop(dbs.get("gone", false).unwrap());
576        dbs.delete("gone").unwrap();
577        assert!(!tmp.path().join("dbs/gone").exists());
578        assert!(dbs.delete(DEFAULT_DB).is_err());
579    }
580}
mentedb_server/databases.rs

mentedb_server/
databases.rs
