Skip to main content

nodedb_fts/index/writer/
maintenance.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Segment publication and collection purging.
4
5use super::FtsIndex;
6use crate::{
7    backend::FtsBackend,
8    index::error::FtsIndexError,
9    lsm::{compaction, segment::writer as seg_writer},
10    scope::IndexScope,
11};
12use std::sync::atomic::Ordering;
13use tracing::debug;
14
15impl<B: FtsBackend> FtsIndex<B> {
16    /// Flush one index's memtable postings to an immutable segment of that
17    /// index in the backend. Every other index keeps its memtable state.
18    ///
19    /// The postings leave the memtable only after the segment is written, so
20    /// a failed flush loses nothing.
21    ///
22    /// Calling this before serializing the index guarantees that all posting
23    /// data written since the last spill threshold is captured in the backend's
24    /// segment storage rather than the in-memory memtable.  Callers that
25    /// checkpoint the index (e.g., NodeDB-Lite flush) must call this once per
26    /// active index before persisting.
27    pub fn flush_memtable<'a>(
28        &self,
29        database_id: u64,
30        tid: u64,
31        index: impl Into<IndexScope<'a>>,
32    ) -> Result<(), FtsIndexError<B::Error>> {
33        let index = index.into();
34        let segment = self
35            .memtable
36            .with_scope_postings(database_id, tid, index, |postings| {
37                (!postings.is_empty()).then(|| seg_writer::flush_postings_to_segment(postings))
38            })
39            .flatten();
40        let Some(segment_bytes) = segment else {
41            self.memtable.drain_scope(database_id, tid, index);
42            return Ok(());
43        };
44        let segment_bytes = segment_bytes?;
45
46        // A new segment id is above every id the index holds. The counter
47        // restarts with the process, so it is raised past the stored ids:
48        // reusing an id would overwrite a segment and inherit its delete set.
49        let floor = self
50            .backend
51            .list_segments(database_id, tid, index)
52            .map_err(FtsIndexError::backend)?
53            .iter()
54            .map(|id| compaction::parse_segment_number(id))
55            .max()
56            .map_or(0, |n| n.saturating_add(1));
57        self.next_segment_id.fetch_max(floor, Ordering::Relaxed);
58        let seg_id = self.next_segment_id.fetch_add(1, Ordering::Relaxed);
59        let id = compaction::segment_id(seg_id, 0);
60        self.backend
61            .write_segment(database_id, tid, index, &id, &segment_bytes)
62            .map_err(FtsIndexError::backend)?;
63        self.memtable.drain_scope(database_id, tid, index);
64
65        debug!(
66            database_id,
67            tid,
68            collection = index.collection(),
69            field = index.field_key(),
70            seg_id,
71            bytes = segment_bytes.len(),
72            "flushed memtable to segment"
73        );
74        Ok(())
75    }
76
77    /// Flush every index the memtable holds, each to its own segment.
78    pub fn flush_all_memtables(&self) -> Result<(), FtsIndexError<B::Error>> {
79        for scope in self.memtable.scopes() {
80            self.flush_memtable(scope.database_id, scope.tid, scope.index())?;
81        }
82        Ok(())
83    }
84
85    /// Purge every index of a collection. Returns count of removed entries.
86    pub fn purge_collection(
87        &self,
88        database_id: u64,
89        tid: u64,
90        collection: &str,
91    ) -> Result<usize, B::Error> {
92        self.memtable.drain_collection(database_id, tid, collection);
93        self.backend.purge_collection(database_id, tid, collection)
94    }
95
96    /// Purge all entries for a `(database_id, tenant)` across every collection.
97    pub fn purge_tenant(&self, database_id: u64, tid: u64) -> Result<usize, B::Error> {
98        self.memtable.drain_tenant(database_id, tid);
99        self.backend.purge_tenant(database_id, tid)
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use crate::backend::memory::MemoryBackend;
107    use crate::test_support::test_governor;
108    use crate::{
109        FtsSearchParams,
110        block::CompactPosting,
111        lsm::memtable::{Memtable, MemtableConfig},
112        posting::{Bm25Params, QueryMode},
113    };
114    use nodedb_types::Surrogate;
115    use std::sync::atomic::AtomicU64;
116    const DB: u64 = 0;
117    const T: u64 = 1;
118    const DOCS: IndexScope<'static> = IndexScope::document("docs");
119    fn make_index() -> FtsIndex<MemoryBackend> {
120        FtsIndex::new(MemoryBackend::new(), test_governor())
121    }
122
123    fn hits(idx: &FtsIndex<MemoryBackend>, index: IndexScope<'_>, query: &str) -> Vec<u32> {
124        let mut ids: Vec<u32> = idx
125            .search(
126                DB,
127                T,
128                index,
129                FtsSearchParams {
130                    query,
131                    top_k: usize::MAX,
132                    fuzzy_enabled: false,
133                    mode: QueryMode::Or,
134                    prefilter: None,
135                },
136            )
137            .unwrap()
138            .into_iter()
139            .map(|r| r.doc_id.0)
140            .collect();
141        ids.sort_unstable();
142        ids
143    }
144
145    #[test]
146    fn flush_propagates_term_too_long_as_typed_error() {
147        let backend = MemoryBackend::new();
148        let idx = FtsIndex {
149            backend,
150            bm25_params: Bm25Params::default(),
151            memtable: Memtable::new(MemtableConfig {
152                max_postings: 1,
153                max_terms: 1,
154            }),
155            next_segment_id: AtomicU64::new(1),
156            governor: test_governor(),
157        };
158
159        // Insert a single posting under a term whose byte length exceeds the u16 segment-format cap. Bypasses the analyzer (which would tokenize away most pathological inputs); we want to exercise the flush-path boundary check directly.
160        let oversize_term = "x".repeat(crate::lsm::segment::format::MAX_TERM_LEN + 1);
161        idx.memtable.insert(
162            DB,
163            T,
164            DOCS,
165            &oversize_term,
166            CompactPosting {
167                doc_id: Surrogate(1),
168                term_freq: 1,
169                fieldnorm: 1,
170                positions: vec![0],
171            },
172        );
173        idx.memtable.record_doc(DB, T, DOCS, Surrogate(1), 1);
174
175        let err = idx
176            .flush_memtable(DB, T, "docs")
177            .expect_err("flush must reject oversize term");
178        match err {
179            FtsIndexError::TermTooLong { len, max } => {
180                assert_eq!(len, oversize_term.len());
181                assert_eq!(max, crate::lsm::segment::format::MAX_TERM_LEN);
182            }
183            other => panic!("expected TermTooLong, got {other:?}"),
184        }
185        assert_eq!(
186            idx.memtable.get_postings(DB, T, DOCS, &oversize_term).len(),
187            1,
188            "a failed flush keeps the postings in the memtable"
189        );
190    }
191    #[test]
192    fn memtable_flush_on_threshold() {
193        let backend = MemoryBackend::new();
194        let idx = FtsIndex {
195            backend,
196            bm25_params: Bm25Params::default(),
197            memtable: Memtable::new(MemtableConfig {
198                max_postings: 5,
199                max_terms: 100,
200            }),
201            next_segment_id: AtomicU64::new(1),
202            governor: test_governor(),
203        };
204
205        idx.index_document(
206            DB,
207            T,
208            "docs",
209            Surrogate(1),
210            "alpha bravo charlie delta echo foxtrot",
211        )
212        .unwrap();
213
214        assert!(idx.memtable.is_empty());
215        let segments = idx.backend.list_segments(DB, T, DOCS).unwrap();
216        assert!(!segments.is_empty(), "segment should have been written");
217        assert_eq!(hits(&idx, DOCS, "charlie"), vec![1]);
218    }
219    #[test]
220    fn purge_collection_preserves_others() {
221        let idx = make_index();
222        let a = IndexScope::document("col_a");
223        let a_title = IndexScope::field("col_a", "title").unwrap();
224        let b = IndexScope::document("col_b");
225        idx.index_document(DB, T, a, Surrogate(1), "alpha bravo")
226            .unwrap();
227        idx.index_document(DB, T, a_title, Surrogate(1), "alpha")
228            .unwrap();
229        idx.index_document(DB, T, b, Surrogate(1), "delta echo")
230            .unwrap();
231
232        idx.purge_collection(DB, T, "col_a").unwrap();
233        assert_eq!(idx.backend.collection_stats(DB, T, a).unwrap(), (0, 0));
234        assert_eq!(
235            idx.backend.collection_stats(DB, T, a_title).unwrap(),
236            (0, 0)
237        );
238        assert!(idx.backend.collection_stats(DB, T, b).unwrap().0 > 0);
239
240        assert!(!idx.memtable.get_postings(DB, T, b, "delta").is_empty());
241        assert!(idx.memtable.get_postings(DB, T, a, "alpha").is_empty());
242        assert!(
243            idx.memtable
244                .get_postings(DB, T, a_title, "alpha")
245                .is_empty()
246        );
247    }
248
249    /// Flushing one collection's index writes only that index's postings
250    /// to its segment. The other collection keeps its memtable postings and
251    /// stats, and both stay searchable.
252    #[test]
253    fn flush_of_one_collection_leaves_the_other_intact() {
254        let idx = make_index();
255        let a = IndexScope::document("col_a");
256        let b = IndexScope::document("col_b");
257        idx.index_document(DB, T, a, Surrogate(1), "alpha bravo")
258            .unwrap();
259        idx.index_document(DB, T, b, Surrogate(2), "alpha charlie")
260            .unwrap();
261        let b_stats = idx.memtable.stats(DB, T, b);
262
263        idx.flush_memtable(DB, T, a).unwrap();
264
265        assert!(idx.memtable.terms(DB, T, a).is_empty());
266        assert_eq!(idx.backend.list_segments(DB, T, a).unwrap().len(), 1);
267        assert!(idx.backend.list_segments(DB, T, b).unwrap().is_empty());
268        assert_eq!(idx.memtable.get_postings(DB, T, b, "alpha").len(), 1);
269        assert_eq!(idx.memtable.stats(DB, T, b), b_stats);
270        assert_eq!(b_stats, (1, 2));
271
272        assert_eq!(hits(&idx, a, "alpha"), vec![1]);
273        assert_eq!(hits(&idx, b, "alpha"), vec![2]);
274        assert!(hits(&idx, a, "charlie").is_empty());
275        assert_eq!(idx.index_stats(DB, T, a).unwrap().0, 1);
276        assert_eq!(idx.index_stats(DB, T, b).unwrap().0, 1);
277    }
278
279    /// Stats of a field index count only the documents that hold that field.
280    #[test]
281    fn stats_are_per_index() {
282        let idx = make_index();
283        let title = IndexScope::field("docs", "title").unwrap();
284        idx.index_document(DB, T, DOCS, Surrogate(1), "rust book")
285            .unwrap();
286        idx.index_document(DB, T, DOCS, Surrogate(2), "java book")
287            .unwrap();
288        idx.index_document(DB, T, title, Surrogate(1), "rust")
289            .unwrap();
290
291        assert_eq!(idx.index_stats(DB, T, DOCS).unwrap(), (2, 2.0));
292        assert_eq!(idx.index_stats(DB, T, title).unwrap(), (1, 1.0));
293        assert_eq!(idx.memtable.stats(DB, T, DOCS), (2, 4));
294        assert_eq!(idx.memtable.stats(DB, T, title), (1, 1));
295    }
296
297    /// A low spill threshold flushes every index to its own segment, and
298    /// every index stays searchable.
299    #[test]
300    fn low_threshold_spills_each_index_to_its_own_segment() {
301        let idx = FtsIndex::with_memtable_config(
302            MemoryBackend::new(),
303            MemtableConfig {
304                max_postings: 4,
305                max_terms: 100,
306            },
307            test_governor(),
308        );
309        let title = IndexScope::field("docs", "title").unwrap();
310        let other = IndexScope::document("other");
311        idx.index_document(DB, T, title, Surrogate(1), "rust")
312            .unwrap();
313        idx.index_document(DB, T, other, Surrogate(2), "rust")
314            .unwrap();
315        idx.index_document(DB, T, DOCS, Surrogate(1), "rust handbook")
316            .unwrap();
317
318        assert!(idx.memtable.is_empty());
319        for index in [title, other, DOCS] {
320            assert_eq!(idx.backend.list_segments(DB, T, index).unwrap().len(), 1);
321        }
322        assert_eq!(hits(&idx, title, "rust"), vec![1]);
323        assert_eq!(hits(&idx, other, "rust"), vec![2]);
324        assert_eq!(hits(&idx, DOCS, "handbook"), vec![1]);
325        assert!(hits(&idx, title, "handbook").is_empty());
326    }
327
328    #[test]
329    fn with_config_sets_bm25_params_and_memtable_thresholds() {
330        let params = Bm25Params { k1: 2.0, b: 0.5 };
331        let idx = FtsIndex::with_config(
332            MemoryBackend::new(),
333            params,
334            MemtableConfig {
335                max_postings: 1,
336                max_terms: 100,
337            },
338            test_governor(),
339        );
340        assert_eq!(idx.bm25_params.k1, 2.0);
341        assert_eq!(idx.bm25_params.b, 0.5);
342
343        idx.index_document(DB, T, DOCS, Surrogate(1), "alpha")
344            .unwrap();
345        assert!(idx.memtable.is_empty());
346        assert_eq!(idx.backend.list_segments(DB, T, DOCS).unwrap().len(), 1);
347        assert_eq!(hits(&idx, DOCS, "alpha"), vec![1]);
348    }
349}