nodedb_fts/backend/
memory.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! In-memory FTS backend for Lite and WASM deployments.
4//!
5//! All data lives in HashMaps behind `RefCell` for interior mutability,
6//! matching the `&self` trait signature. Rebuilt from documents on cold
7//! start — acceptable for edge-scale datasets.
8//!
9//! Keys are fully structural tuples `(database_id, tid, collection, …)` —
10//! database and tenant isolation never depends on lexical-prefix ordering.
11
12use std::cell::RefCell;
13use std::collections::HashMap;
14use std::fmt;
15
16use nodedb_types::Surrogate;
17
18use crate::backend::FtsBackend;
19use crate::posting::Posting;
20
21/// In-memory backend error (infallible in practice, but trait requires it).
22#[derive(Debug)]
23pub struct MemoryError(String);
24
25impl fmt::Display for MemoryError {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        write!(f, "memory backend: {}", self.0)
28    }
29}
30
31type QuadKey = (u64, u64, String, String);
32type DocLenKey = (u64, u64, String, Surrogate);
33type TripleKey = (u64, u64, String);
34
35/// In-memory FTS backend backed by HashMaps keyed by
36/// `(database_id, tid, collection, …)` tuples.
37///
38/// Uses `RefCell` for interior mutability so the `FtsBackend` trait
39/// can use `&self` uniformly (redb has its own transactional isolation).
40#[derive(Debug, Default)]
41pub struct MemoryBackend {
42    /// `(database_id, tid, collection, term) → posting list`.
43    postings: RefCell<HashMap<QuadKey, Vec<Posting>>>,
44    /// `(database_id, tid, collection, doc_id) → token count`.
45    doc_lengths: RefCell<HashMap<DocLenKey, u32>>,
46    /// `(database_id, tid, collection) → (doc_count, total_token_sum)`.
47    stats: RefCell<HashMap<TripleKey, (u32, u64)>>,
48    /// `(database_id, tid, collection, subkey) → blob` for docmap, fieldnorms, analyzer, language.
49    meta: RefCell<HashMap<QuadKey, Vec<u8>>>,
50    /// `(database_id, tid, collection, segment_id) → compressed segment bytes`.
51    segments: RefCell<HashMap<QuadKey, Vec<u8>>>,
52}
53
54impl MemoryBackend {
55    pub fn new() -> Self {
56        Self::default()
57    }
58}
59
60fn quad(database_id: u64, tid: u64, collection: &str, sub: &str) -> QuadKey {
61    (database_id, tid, collection.to_string(), sub.to_string())
62}
63
64fn doc_len_key(database_id: u64, tid: u64, collection: &str, doc_id: Surrogate) -> DocLenKey {
65    (database_id, tid, collection.to_string(), doc_id)
66}
67
68fn triple(database_id: u64, tid: u64, collection: &str) -> TripleKey {
69    (database_id, tid, collection.to_string())
70}
71
72impl FtsBackend for MemoryBackend {
73    type Error = MemoryError;
74
75    fn read_postings(
76        &self,
77        database_id: u64,
78        tid: u64,
79        collection: &str,
80        term: &str,
81    ) -> Result<Vec<Posting>, Self::Error> {
82        Ok(self
83            .postings
84            .borrow()
85            .get(&quad(database_id, tid, collection, term))
86            .cloned()
87            .unwrap_or_default())
88    }
89
90    fn write_postings(
91        &self,
92        database_id: u64,
93        tid: u64,
94        collection: &str,
95        term: &str,
96        postings: &[Posting],
97    ) -> Result<(), Self::Error> {
98        let key = quad(database_id, tid, collection, term);
99        let mut map = self.postings.borrow_mut();
100        if postings.is_empty() {
101            map.remove(&key);
102        } else {
103            map.insert(key, postings.to_vec());
104        }
105        Ok(())
106    }
107
108    fn remove_postings(
109        &self,
110        database_id: u64,
111        tid: u64,
112        collection: &str,
113        term: &str,
114    ) -> Result<(), Self::Error> {
115        self.postings
116            .borrow_mut()
117            .remove(&quad(database_id, tid, collection, term));
118        Ok(())
119    }
120
121    fn read_doc_length(
122        &self,
123        database_id: u64,
124        tid: u64,
125        collection: &str,
126        doc_id: Surrogate,
127    ) -> Result<Option<u32>, Self::Error> {
128        Ok(self
129            .doc_lengths
130            .borrow()
131            .get(&doc_len_key(database_id, tid, collection, doc_id))
132            .copied())
133    }
134
135    fn write_doc_length(
136        &self,
137        database_id: u64,
138        tid: u64,
139        collection: &str,
140        doc_id: Surrogate,
141        length: u32,
142    ) -> Result<(), Self::Error> {
143        self.doc_lengths
144            .borrow_mut()
145            .insert(doc_len_key(database_id, tid, collection, doc_id), length);
146        Ok(())
147    }
148
149    fn remove_doc_length(
150        &self,
151        database_id: u64,
152        tid: u64,
153        collection: &str,
154        doc_id: Surrogate,
155    ) -> Result<(), Self::Error> {
156        self.doc_lengths
157            .borrow_mut()
158            .remove(&doc_len_key(database_id, tid, collection, doc_id));
159        Ok(())
160    }
161
162    fn collection_terms(
163        &self,
164        database_id: u64,
165        tid: u64,
166        collection: &str,
167    ) -> Result<Vec<String>, Self::Error> {
168        Ok(self
169            .postings
170            .borrow()
171            .keys()
172            .filter(|(d, t, c, _)| *d == database_id && *t == tid && c == collection)
173            .map(|(_, _, _, term)| term.clone())
174            .collect())
175    }
176
177    fn collection_stats(
178        &self,
179        database_id: u64,
180        tid: u64,
181        collection: &str,
182    ) -> Result<(u32, u64), Self::Error> {
183        Ok(self
184            .stats
185            .borrow()
186            .get(&triple(database_id, tid, collection))
187            .copied()
188            .unwrap_or((0, 0)))
189    }
190
191    fn increment_stats(
192        &self,
193        database_id: u64,
194        tid: u64,
195        collection: &str,
196        doc_len: u32,
197    ) -> Result<(), Self::Error> {
198        let mut stats = self.stats.borrow_mut();
199        let entry = stats
200            .entry(triple(database_id, tid, collection))
201            .or_insert((0, 0));
202        entry.0 += 1;
203        entry.1 += doc_len as u64;
204        Ok(())
205    }
206
207    fn decrement_stats(
208        &self,
209        database_id: u64,
210        tid: u64,
211        collection: &str,
212        doc_len: u32,
213    ) -> Result<(), Self::Error> {
214        let mut stats = self.stats.borrow_mut();
215        let entry = stats
216            .entry(triple(database_id, tid, collection))
217            .or_insert((0, 0));
218        entry.0 = entry.0.saturating_sub(1);
219        entry.1 = entry.1.saturating_sub(doc_len as u64);
220        Ok(())
221    }
222
223    fn read_meta(
224        &self,
225        database_id: u64,
226        tid: u64,
227        collection: &str,
228        subkey: &str,
229    ) -> Result<Option<Vec<u8>>, Self::Error> {
230        Ok(self
231            .meta
232            .borrow()
233            .get(&quad(database_id, tid, collection, subkey))
234            .cloned())
235    }
236
237    fn write_meta(
238        &self,
239        database_id: u64,
240        tid: u64,
241        collection: &str,
242        subkey: &str,
243        value: &[u8],
244    ) -> Result<(), Self::Error> {
245        self.meta
246            .borrow_mut()
247            .insert(quad(database_id, tid, collection, subkey), value.to_vec());
248        Ok(())
249    }
250
251    fn write_segment(
252        &self,
253        database_id: u64,
254        tid: u64,
255        collection: &str,
256        segment_id: &str,
257        data: &[u8],
258    ) -> Result<(), Self::Error> {
259        self.segments.borrow_mut().insert(
260            quad(database_id, tid, collection, segment_id),
261            data.to_vec(),
262        );
263        Ok(())
264    }
265
266    fn read_segment(
267        &self,
268        database_id: u64,
269        tid: u64,
270        collection: &str,
271        segment_id: &str,
272    ) -> Result<Option<Vec<u8>>, Self::Error> {
273        Ok(self
274            .segments
275            .borrow()
276            .get(&quad(database_id, tid, collection, segment_id))
277            .cloned())
278    }
279
280    fn list_segments(
281        &self,
282        database_id: u64,
283        tid: u64,
284        collection: &str,
285    ) -> Result<Vec<String>, Self::Error> {
286        Ok(self
287            .segments
288            .borrow()
289            .keys()
290            .filter(|(d, t, c, _)| *d == database_id && *t == tid && c == collection)
291            .map(|(_, _, _, seg)| seg.clone())
292            .collect())
293    }
294
295    fn remove_segment(
296        &self,
297        database_id: u64,
298        tid: u64,
299        collection: &str,
300        segment_id: &str,
301    ) -> Result<(), Self::Error> {
302        self.segments
303            .borrow_mut()
304            .remove(&quad(database_id, tid, collection, segment_id));
305        Ok(())
306    }
307
308    fn purge_collection(
309        &self,
310        database_id: u64,
311        tid: u64,
312        collection: &str,
313    ) -> Result<usize, Self::Error> {
314        let matches_dtc = |d: u64, t: u64, c: &str| d == database_id && t == tid && c == collection;
315
316        let mut postings = self.postings.borrow_mut();
317        let mut doc_lengths = self.doc_lengths.borrow_mut();
318        let before = postings.len() + doc_lengths.len();
319        postings.retain(|k, _| !matches_dtc(k.0, k.1, &k.2));
320        doc_lengths.retain(|k, _| !matches_dtc(k.0, k.1, &k.2));
321        self.stats
322            .borrow_mut()
323            .remove(&triple(database_id, tid, collection));
324        self.meta
325            .borrow_mut()
326            .retain(|k, _| !matches_dtc(k.0, k.1, &k.2));
327        self.segments
328            .borrow_mut()
329            .retain(|k, _| !matches_dtc(k.0, k.1, &k.2));
330        let after = postings.len() + doc_lengths.len();
331        Ok(before - after)
332    }
333
334    fn purge_tenant(&self, database_id: u64, tid: u64) -> Result<usize, Self::Error> {
335        let matches_dt = |d: u64, t: u64| d == database_id && t == tid;
336
337        let mut postings = self.postings.borrow_mut();
338        let mut doc_lengths = self.doc_lengths.borrow_mut();
339        let before = postings.len() + doc_lengths.len();
340        postings.retain(|k, _| !matches_dt(k.0, k.1));
341        doc_lengths.retain(|k, _| !matches_dt(k.0, k.1));
342        self.stats.borrow_mut().retain(|k, _| !matches_dt(k.0, k.1));
343        self.meta.borrow_mut().retain(|k, _| !matches_dt(k.0, k.1));
344        self.segments
345            .borrow_mut()
346            .retain(|k, _| !matches_dt(k.0, k.1));
347        let after = postings.len() + doc_lengths.len();
348        Ok(before - after)
349    }
350}
351
352#[cfg(test)]
353mod tests {
354    use super::*;
355
356    const DB: u64 = 0;
357    const T: u64 = 1;
358
359    #[test]
360    fn roundtrip_postings() {
361        let backend = MemoryBackend::new();
362        let postings = vec![Posting {
363            doc_id: Surrogate(1),
364            term_freq: 2,
365            positions: vec![0, 5],
366        }];
367        backend
368            .write_postings(DB, T, "col", "hello", &postings)
369            .unwrap();
370
371        let read = backend.read_postings(DB, T, "col", "hello").unwrap();
372        assert_eq!(read.len(), 1);
373        assert_eq!(read[0].doc_id, Surrogate(1));
374    }
375
376    #[test]
377    fn roundtrip_doc_lengths() {
378        let backend = MemoryBackend::new();
379        backend
380            .write_doc_length(DB, T, "col", Surrogate(1), 42)
381            .unwrap();
382        assert_eq!(
383            backend.read_doc_length(DB, T, "col", Surrogate(1)).unwrap(),
384            Some(42)
385        );
386
387        backend
388            .remove_doc_length(DB, T, "col", Surrogate(1))
389            .unwrap();
390        assert_eq!(
391            backend.read_doc_length(DB, T, "col", Surrogate(1)).unwrap(),
392            None
393        );
394    }
395
396    #[test]
397    fn incremental_stats() {
398        let backend = MemoryBackend::new();
399        backend.increment_stats(DB, T, "col", 10).unwrap();
400        backend.increment_stats(DB, T, "col", 20).unwrap();
401        assert_eq!(backend.collection_stats(DB, T, "col").unwrap(), (2, 30));
402
403        backend.decrement_stats(DB, T, "col", 10).unwrap();
404        assert_eq!(backend.collection_stats(DB, T, "col").unwrap(), (1, 20));
405    }
406
407    #[test]
408    fn stats_saturating_sub() {
409        let backend = MemoryBackend::new();
410        backend.decrement_stats(DB, T, "col", 100).unwrap();
411        assert_eq!(backend.collection_stats(DB, T, "col").unwrap(), (0, 0));
412    }
413
414    #[test]
415    fn purge_clears_stats_and_isolates_collections() {
416        let backend = MemoryBackend::new();
417        backend.increment_stats(DB, T, "col", 10).unwrap();
418        backend
419            .write_doc_length(DB, T, "col", Surrogate(1), 10)
420            .unwrap();
421        backend
422            .write_postings(
423                DB,
424                T,
425                "col",
426                "hello",
427                &[Posting {
428                    doc_id: Surrogate(1),
429                    term_freq: 1,
430                    positions: vec![0],
431                }],
432            )
433            .unwrap();
434
435        backend.increment_stats(DB, T, "other", 7).unwrap();
436        backend
437            .write_doc_length(DB, T, "other", Surrogate(1), 7)
438            .unwrap();
439        backend
440            .write_postings(
441                DB,
442                T,
443                "other",
444                "world",
445                &[Posting {
446                    doc_id: Surrogate(1),
447                    term_freq: 1,
448                    positions: vec![0],
449                }],
450            )
451            .unwrap();
452
453        backend.purge_collection(DB, T, "col").unwrap();
454        assert_eq!(backend.collection_stats(DB, T, "col").unwrap(), (0, 0));
455        assert!(
456            backend
457                .read_postings(DB, T, "col", "hello")
458                .unwrap()
459                .is_empty()
460        );
461        assert_eq!(
462            backend.read_doc_length(DB, T, "col", Surrogate(1)).unwrap(),
463            None
464        );
465
466        assert_eq!(backend.collection_stats(DB, T, "other").unwrap(), (1, 7));
467        assert_eq!(
468            backend
469                .read_postings(DB, T, "other", "world")
470                .unwrap()
471                .len(),
472            1
473        );
474        assert_eq!(
475            backend
476                .read_doc_length(DB, T, "other", Surrogate(1))
477                .unwrap(),
478            Some(7)
479        );
480    }
481
482    #[test]
483    fn collection_terms() {
484        let backend = MemoryBackend::new();
485        backend
486            .write_postings(
487                DB,
488                T,
489                "col",
490                "hello",
491                &[Posting {
492                    doc_id: Surrogate(1),
493                    term_freq: 1,
494                    positions: vec![0],
495                }],
496            )
497            .unwrap();
498        backend
499            .write_postings(
500                DB,
501                T,
502                "col",
503                "world",
504                &[Posting {
505                    doc_id: Surrogate(1),
506                    term_freq: 1,
507                    positions: vec![1],
508                }],
509            )
510            .unwrap();
511
512        let mut terms = backend.collection_terms(DB, T, "col").unwrap();
513        terms.sort();
514        assert_eq!(terms, vec!["hello", "world"]);
515    }
516
517    #[test]
518    fn segment_roundtrip() {
519        let backend = MemoryBackend::new();
520        let data = b"compressed segment bytes";
521        backend.write_segment(DB, T, "col", "id1", data).unwrap();
522        assert_eq!(
523            backend.read_segment(DB, T, "col", "id1").unwrap(),
524            Some(data.to_vec())
525        );
526        assert_eq!(backend.read_segment(DB, T, "col", "missing").unwrap(), None);
527    }
528
529    #[test]
530    fn segment_list_filters_by_collection() {
531        let backend = MemoryBackend::new();
532        backend.write_segment(DB, T, "col", "a", b"a").unwrap();
533        backend.write_segment(DB, T, "col", "b", b"b").unwrap();
534        backend.write_segment(DB, T, "other", "c", b"c").unwrap();
535
536        let mut segs = backend.list_segments(DB, T, "col").unwrap();
537        segs.sort();
538        assert_eq!(segs, vec!["a", "b"]);
539
540        let other = backend.list_segments(DB, T, "other").unwrap();
541        assert_eq!(other, vec!["c"]);
542    }
543
544    #[test]
545    fn segment_remove() {
546        let backend = MemoryBackend::new();
547        backend.write_segment(DB, T, "col", "id1", b"data").unwrap();
548        backend.remove_segment(DB, T, "col", "id1").unwrap();
549        assert_eq!(backend.read_segment(DB, T, "col", "id1").unwrap(), None);
550    }
551
552    #[test]
553    fn purge_clears_segments() {
554        let backend = MemoryBackend::new();
555        backend.write_segment(DB, T, "col", "a", b"a").unwrap();
556        backend.write_segment(DB, T, "other", "b", b"b").unwrap();
557
558        backend.purge_collection(DB, T, "col").unwrap();
559        assert!(backend.list_segments(DB, T, "col").unwrap().is_empty());
560        assert_eq!(backend.list_segments(DB, T, "other").unwrap().len(), 1);
561    }
562
563    #[test]
564    fn purge_tenant_isolates_tenants() {
565        let backend = MemoryBackend::new();
566        backend.increment_stats(DB, 1, "col", 5).unwrap();
567        backend.increment_stats(DB, 2, "col", 7).unwrap();
568        backend
569            .write_postings(
570                DB,
571                1,
572                "col",
573                "t",
574                &[Posting {
575                    doc_id: Surrogate(1),
576                    term_freq: 1,
577                    positions: vec![0],
578                }],
579            )
580            .unwrap();
581        backend
582            .write_postings(
583                DB,
584                2,
585                "col",
586                "t",
587                &[Posting {
588                    doc_id: Surrogate(1),
589                    term_freq: 1,
590                    positions: vec![0],
591                }],
592            )
593            .unwrap();
594
595        backend.purge_tenant(DB, 1).unwrap();
596        assert_eq!(backend.collection_stats(DB, 1, "col").unwrap(), (0, 0));
597        assert!(backend.read_postings(DB, 1, "col", "t").unwrap().is_empty());
598        assert_eq!(backend.collection_stats(DB, 2, "col").unwrap(), (1, 7));
599        assert_eq!(backend.read_postings(DB, 2, "col", "t").unwrap().len(), 1);
600    }
601
602    #[test]
603    fn databases_isolated() {
604        let backend = MemoryBackend::new();
605        backend.increment_stats(0, T, "col", 5).unwrap();
606        backend.increment_stats(9, T, "col", 7).unwrap();
607        backend
608            .write_postings(
609                0,
610                T,
611                "col",
612                "t",
613                &[Posting {
614                    doc_id: Surrogate(1),
615                    term_freq: 1,
616                    positions: vec![0],
617                }],
618            )
619            .unwrap();
620        backend
621            .write_postings(
622                9,
623                T,
624                "col",
625                "t",
626                &[Posting {
627                    doc_id: Surrogate(1),
628                    term_freq: 1,
629                    positions: vec![0],
630                }],
631            )
632            .unwrap();
633
634        backend.purge_tenant(0, T).unwrap();
635        assert_eq!(backend.collection_stats(0, T, "col").unwrap(), (0, 0));
636        assert!(backend.read_postings(0, T, "col", "t").unwrap().is_empty());
637        // Same tenant in a different database must be unaffected.
638        assert_eq!(backend.collection_stats(9, T, "col").unwrap(), (1, 7));
639        assert_eq!(backend.read_postings(9, T, "col", "t").unwrap().len(), 1);
640    }
641}
nodedb_fts/backend/memory.rs

nodedb_fts/backend/
memory.rs
