Skip to main content

sonic/store/
fst.rs

1// Sonic
2//
3// Fast, lightweight and schema-less search backend
4// Copyright: 2019, Valerian Saliou <valerian@valeriansaliou.name>
5// Copyright: 2026, Rémi Bardon <remi@remibardon.name>
6// License: Mozilla Public License v2.0 (MPL v2.0)
7
8use fst::automaton::AlwaysMatch;
9use fst::set::Stream as FSTStream;
10use fst::{
11    Automaton, Error as FSTError, IntoStreamer, Set as FSTSet, SetBuilder as FSTSetBuilder,
12    Streamer,
13};
14use fst_levenshtein::Levenshtein;
15use fst_regex::Regex;
16use hashbrown::{HashMap, HashSet};
17use indexmap::IndexMap;
18use radix::RadixNum;
19use regex_syntax::escape as regex_escape;
20use std::collections::VecDeque;
21use std::fmt;
22use std::fs::{self, File};
23use std::io::{self, BufRead, BufReader, BufWriter, Write};
24use std::iter::FromIterator;
25use std::path::{Path, PathBuf};
26use std::str;
27use std::sync::{Arc, Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard};
28use std::thread;
29use std::time::{Duration, SystemTime};
30
31use super::generic::{
32    StoreGeneric, StoreGenericActionBuilder, StoreGenericBuilder, StoreGenericPool,
33};
34use super::keyer::StoreKeyerHasher;
35use crate::lexer::ranges::LexerRegexRange;
36use crate::query::QueryMatchScore;
37
38// NOTE: This type cannot be generic over a lifetime as spawning threads would
39//   force it to be `'static`.
40#[derive(Clone)]
41pub struct StoreFSTPool {
42    fst_store_config: Arc<crate::config::ConfigStoreFST>,
43    // NOTE: This shouldn’t be here, but until a big rewrite let’s not care.
44    pub fst_action_config: StoreFSTActionConfig,
45    graph_pool: Arc<RwLock<HashMap<StoreFSTKey, StoreFSTBox>>>,
46    graph_acquire_lock: Arc<Mutex<()>>,
47    graph_rebuild_lock: Arc<Mutex<()>>,
48    graph_access_lock: Arc<RwLock<()>>,
49    graph_consolidate: Arc<RwLock<HashSet<StoreFSTKey>>>,
50}
51
52pub struct StoreFSTBuilder<'build> {
53    fst_store_config: &'build crate::config::ConfigStoreFST,
54    // NOTE: This shouldn’t be here, but until a big rewrite let’s not care.
55    fst_action_config: StoreFSTActionConfig,
56    graph_consolidate: Arc<RwLock<HashSet<StoreFSTKey>>>,
57}
58
59pub struct StoreFST {
60    graph: FSTSet,
61    target: StoreFSTKey,
62    pending: StoreFSTPending,
63    last_used: Arc<RwLock<SystemTime>>,
64    last_consolidated: Arc<RwLock<SystemTime>>,
65    graph_consolidate: Arc<RwLock<HashSet<StoreFSTKey>>>,
66    // NOTE: This shouldn’t be here, but until a big rewrite let’s not care.
67    action_config: StoreFSTActionConfig,
68}
69
70#[derive(Default)]
71pub struct StoreFSTPending {
72    pop: Arc<RwLock<HashSet<Vec<u8>>>>,
73    push: Arc<RwLock<HashSet<Vec<u8>>>>,
74}
75
76pub struct StoreFSTActionBuilder<'build> {
77    pub fst_store_config: &'build crate::config::ConfigStoreFST,
78}
79
80pub struct StoreFSTAction {
81    store: StoreFSTBox,
82}
83
84impl StoreFSTAction {
85    fn config(&self) -> &StoreFSTActionConfig {
86        &self.store.action_config
87    }
88}
89
90#[derive(PartialEq, Eq, Hash, Clone, Copy)]
91pub struct StoreFSTKey {
92    collection_hash: StoreFSTAtom,
93    bucket_hash: StoreFSTAtom,
94}
95
96pub struct StoreFSTMisc;
97
98#[derive(Copy, Clone)]
99enum StoreFSTPathMode {
100    Permanent,
101    Temporary,
102    Backup,
103}
104
105type StoreFSTAtom = u32;
106type StoreFSTBox = Arc<StoreFST>;
107
108#[derive(Debug, Clone, Copy)]
109pub struct StoreFSTActionConfig {
110    pub prefix_matching_enabled: bool,
111    pub fuzzy_matching_enabled: bool,
112}
113
114impl Default for StoreFSTActionConfig {
115    fn default() -> Self {
116        Self {
117            prefix_matching_enabled: true,
118            fuzzy_matching_enabled: true,
119        }
120    }
121}
122
123const WORD_LIMIT_LENGTH: usize = 40;
124const ATOM_HASH_RADIX: usize = 16;
125
126impl StoreFSTPathMode {
127    fn extension(&self) -> &'static str {
128        match self {
129            StoreFSTPathMode::Permanent => ".fst",
130            StoreFSTPathMode::Temporary => ".fst.tmp",
131            StoreFSTPathMode::Backup => ".fst.bck",
132        }
133    }
134}
135
136impl StoreFSTPool {
137    pub fn new(
138        fst_store_config: Arc<crate::config::ConfigStoreFST>,
139        fst_action_config: StoreFSTActionConfig,
140    ) -> Self {
141        Self {
142            fst_store_config,
143            fst_action_config,
144            graph_pool: Arc::default(),
145            graph_acquire_lock: Arc::default(),
146            graph_rebuild_lock: Arc::default(),
147            graph_access_lock: Arc::default(),
148            graph_consolidate: Arc::default(),
149        }
150    }
151
152    pub fn count(&self) -> (usize, usize) {
153        (
154            self.graph_pool.read().unwrap().len(),
155            self.graph_consolidate.read().unwrap().len(),
156        )
157    }
158
159    pub fn lock_read_access<'a>(&'a self) -> RwLockReadGuard<'a, ()> {
160        self.graph_access_lock.read().unwrap()
161    }
162
163    pub fn lock_write_access<'a>(&'a self) -> RwLockWriteGuard<'a, ()> {
164        self.graph_access_lock.write().unwrap()
165    }
166
167    pub fn acquire<T: AsRef<str>>(&self, collection: T, bucket: T) -> Result<StoreFSTBox, ()> {
168        let (collection_str, bucket_str) = (collection.as_ref(), bucket.as_ref());
169
170        let pool_key = StoreFSTKey::from_str(collection_str, bucket_str);
171
172        // Freeze acquire lock, and reference it in context
173        // Notice: this prevents two graphs on the same collection to be opened at the same time.
174        let _acquire = self.graph_acquire_lock.lock().unwrap();
175
176        // Acquire a thread-safe store pool reference in read mode
177        let graph_pool_read = self.graph_pool.read().unwrap();
178
179        if let Some(store_fst) = graph_pool_read.get(&pool_key) {
180            Self::proceed_acquire_cache("fst", collection_str, pool_key, store_fst)
181        } else {
182            tracing::info!(
183                "fst store not in pool for collection: {} <{:x}> / bucket: {} <{:x}>, opening it",
184                collection_str,
185                pool_key.collection_hash,
186                bucket_str,
187                pool_key.bucket_hash
188            );
189
190            // Important: we need to drop the read reference first, to avoid dead-locking \
191            //   when acquiring the RWLock in write mode in this block.
192            drop(graph_pool_read);
193
194            let builder = StoreFSTBuilder {
195                fst_store_config: &self.fst_store_config,
196                graph_consolidate: Arc::clone(&self.graph_consolidate),
197                fst_action_config: self.fst_action_config,
198            };
199
200            Self::proceed_acquire_open("fst", collection_str, pool_key, &self.graph_pool, &builder)
201        }
202    }
203
204    pub fn janitor(&self) {
205        Self::proceed_janitor(
206            "fst",
207            &self.graph_pool,
208            self.fst_store_config.pool.inactive_after,
209            &self.graph_access_lock,
210        )
211    }
212
213    pub fn backup(&self, path: &Path) -> Result<(), io::Error> {
214        tracing::debug!("backing up all fst stores to path: {:?}", path);
215
216        // Create backup directory (full path)
217        fs::create_dir_all(path)?;
218
219        // Proceed dump action (backup)
220        self.dump_action(
221            "backup",
222            StoreFSTPathMode::Permanent,
223            &self.fst_store_config.path,
224            path,
225            &Self::backup_item,
226        )
227    }
228
229    pub fn restore(&self, path: &Path) -> Result<(), io::Error> {
230        tracing::debug!("restoring all fst stores from path: {:?}", path);
231
232        // Proceed dump action (restore)
233        self.dump_action(
234            "restore",
235            StoreFSTPathMode::Backup,
236            path,
237            &self.fst_store_config.path,
238            &Self::restore_item,
239        )
240    }
241
242    pub fn consolidate(&self, force: bool) {
243        tracing::debug!("scanning for fst store pool items to consolidate");
244
245        // Notice: we do not consolidate all items at each tick, we try to even out multiple \
246        //   consolidation tasks over time. This lowers the overall HZ of the tasker system for \
247        //   certain heavy tasks, which is better to spread out consolidation steps over time over \
248        //   a large number of very active buckets.
249
250        // Acquire rebuild lock, and reference it in context
251        // Notice: this prevents two consolidate operations to be executed at the same time.
252        let _rebuild = self.graph_rebuild_lock.lock().unwrap();
253
254        // Exit trap: Register is empty? Abort there.
255        if self.graph_consolidate.read().unwrap().is_empty() {
256            tracing::info!("no fst store pool items to consolidate in register");
257
258            return;
259        }
260
261        // Step 1: List keys to be consolidated
262        let mut keys_consolidate: Vec<StoreFSTKey> = Vec::new();
263
264        {
265            // Acquire access lock (in blocking write mode), and reference it in context
266            // Notice: this prevents store to be acquired from any context
267            let _access = self.graph_access_lock.write().unwrap();
268
269            let (graph_pool_read, graph_consolidate_read) = (
270                self.graph_pool.read().unwrap(),
271                self.graph_consolidate.read().unwrap(),
272            );
273
274            for key in &*graph_consolidate_read {
275                if let Some(store) = graph_pool_read.get(key) {
276                    // Important: be lenient with system clock going back to a past duration, \
277                    //   since we may be running in a virtualized environment where clock is not \
278                    //   guaranteed to be monotonic. This is done to avoid poisoning associated \
279                    //   mutexes by crashing on unwrap().
280                    let not_consolidated_for = store
281                        .last_consolidated
282                        .read()
283                        .unwrap()
284                        .elapsed()
285                        .unwrap_or_else(|err| {
286                            tracing::error!(
287                                "fst key: {} last consolidated duration clock issue, zeroing: {}",
288                                key,
289                                err
290                            );
291
292                            // Assuming a zero seconds fallback duration
293                            Duration::from_secs(0)
294                        })
295                        .as_secs();
296
297                    if force
298                        || not_consolidated_for >= self.fst_store_config.graph.consolidate_after
299                    {
300                        tracing::info!(
301                            "fst key: {} not consolidated for: {} seconds, may consolidate",
302                            key,
303                            not_consolidated_for
304                        );
305
306                        keys_consolidate.push(*key);
307                    } else {
308                        tracing::debug!(
309                            "fst key: {} not consolidated for: {} seconds, no consolidate",
310                            key,
311                            not_consolidated_for
312                        );
313                    }
314                }
315            }
316        }
317
318        // Exit trap: Nothing to consolidate yet? Abort there.
319        if keys_consolidate.is_empty() {
320            tracing::info!("no fst store pool items need to consolidate at the moment");
321
322            return;
323        }
324
325        // Step 2: Clear keys to be consolidated from register
326        {
327            // Acquire access lock (in blocking write mode), and reference it in context
328            // Notice: this prevents store to be acquired from any context
329            let _access = self.graph_access_lock.write().unwrap();
330
331            let mut graph_consolidate_write = self.graph_consolidate.write().unwrap();
332
333            for key in &keys_consolidate {
334                graph_consolidate_write.remove(key);
335
336                tracing::debug!("fst key: {} cleared from consolidate register", key);
337            }
338        }
339
340        // Step 3: Consolidate FSTs, one-by-one (sequential locking; this avoids global locks)
341        let (mut count_moved, mut count_pushed, mut count_popped) = (0, 0, 0);
342
343        {
344            for key in &keys_consolidate {
345                {
346                    // As we may be renaming the FST file, ensure no consumer out of this is \
347                    //   trying to access the FST file as it gets processed. This also waits for \
348                    //   current consumers to finish reading the FST, and prevents any new \
349                    //   consumer from opening it while we are not done there.
350                    let _access = self.graph_access_lock.write().unwrap();
351
352                    let do_close = if let Some(store) = self.graph_pool.read().unwrap().get(key) {
353                        tracing::debug!("fst key: {} consolidate started", key);
354
355                        let consolidate_counts = self.consolidate_item(store);
356
357                        count_moved += consolidate_counts.1;
358                        count_pushed += consolidate_counts.2;
359                        count_popped += consolidate_counts.3;
360
361                        tracing::debug!("fst key: {} consolidate complete", key);
362
363                        // Should close this FST?
364                        consolidate_counts.0
365                    } else {
366                        false
367                    };
368
369                    // Nuke old opened FST?
370                    // Notice: last consolidated date will be bumped to a new date in the future \
371                    //   when a push or pop operation will be done, thus effectively scheduling \
372                    //   a consolidation in the future properly.
373                    // Notice: we remove this one early as to release write lock early
374                    if do_close {
375                        self.graph_pool.write().unwrap().remove(key);
376                    }
377                }
378
379                // Give a bit of time to other threads before continuing (a consolidate operation \
380                //   must not block all other threads until it completes); this method tells the \
381                //   thread scheduler to give a bit of priority to other threads, and get back \
382                //   to this thread's work when other threads are done. On large setups, this \
383                //   loop can starve other threads due to the locks used (unfortunately they \
384                //   are all necessary).
385                thread::yield_now();
386            }
387        }
388
389        tracing::info!(
390            "done scanning for fst store pool items to consolidate (move: {}, push: {}, pop: {})",
391            count_moved,
392            count_pushed,
393            count_popped
394        );
395    }
396
397    #[allow(clippy::type_complexity)]
398    fn dump_action(
399        &self,
400        action: &str,
401        path_mode: StoreFSTPathMode,
402        read_path: &Path,
403        write_path: &Path,
404        fn_item: &dyn Fn(&Self, &Path, &Path, &str, &str) -> Result<(), io::Error>,
405    ) -> Result<(), io::Error> {
406        let fst_extension = path_mode.extension();
407        let fst_extension_len = fst_extension.len();
408
409        // Iterate on FST collections
410        for collection in fs::read_dir(read_path)? {
411            let collection = collection?;
412
413            // Actual collection found?
414            if let (Ok(collection_file_type), Some(collection_name)) =
415                (collection.file_type(), collection.file_name().to_str())
416            {
417                if collection_file_type.is_dir() {
418                    tracing::debug!("fst collection ongoing {}: {}", action, collection_name);
419
420                    // Create write folder for collection
421                    fs::create_dir_all(write_path.join(collection_name))?;
422
423                    // Iterate on FST collection buckets
424                    for bucket in fs::read_dir(read_path.join(collection_name))? {
425                        let bucket = bucket?;
426
427                        // Actual bucket found?
428                        if let (Ok(bucket_file_type), Some(bucket_file_name)) =
429                            (bucket.file_type(), bucket.file_name().to_str())
430                        {
431                            let bucket_file_name_len = bucket_file_name.len();
432
433                            if bucket_file_type.is_file()
434                                && bucket_file_name_len > fst_extension_len
435                                && bucket_file_name.ends_with(fst_extension)
436                            {
437                                // Acquire bucket name (from full file name)
438                                let bucket_name =
439                                    &bucket_file_name[..(bucket_file_name_len - fst_extension_len)];
440
441                                tracing::debug!(
442                                    "fst bucket ongoing {}: {}/{}",
443                                    action,
444                                    collection_name,
445                                    bucket_name
446                                );
447
448                                fn_item(
449                                    self,
450                                    write_path,
451                                    &bucket.path(),
452                                    collection_name,
453                                    bucket_name,
454                                )?;
455                            }
456                        }
457                    }
458                }
459            }
460        }
461
462        Ok(())
463    }
464
465    fn backup_item(
466        &self,
467        backup_path: &Path,
468        _origin_path: &Path,
469        collection_name: &str,
470        bucket_name: &str,
471    ) -> Result<(), io::Error> {
472        // Acquire access lock (in blocking write mode), and reference it in context
473        // Notice: this prevents store to be acquired from any context
474        let _access = self.graph_access_lock.write().unwrap();
475
476        // Generate path to FST backup
477        let fst_backup_path = backup_path.join(collection_name).join(format!(
478            "{}{}",
479            bucket_name,
480            StoreFSTPathMode::Backup.extension()
481        ));
482
483        tracing::debug!(
484            "fst bucket: {}/{} backing up to path: {:?}",
485            collection_name,
486            bucket_name,
487            fst_backup_path
488        );
489
490        // Erase any previously-existing FST backup
491        fs::remove_file(&fst_backup_path).ok();
492
493        // Stream actual FST data to FST backup
494        let backup_fst_file = File::create(&fst_backup_path)?;
495        let mut backup_fst_writer = BufWriter::new(backup_fst_file);
496
497        let mut count_words = 0;
498
499        // Convert names to hashes (as names are hashes encoded as base-16 strings, but we need \
500        //   them as proper integers)
501        if let (Ok(collection_radix), Ok(bucket_radix)) = (
502            RadixNum::from_str(collection_name, ATOM_HASH_RADIX),
503            RadixNum::from_str(bucket_name, ATOM_HASH_RADIX),
504        ) {
505            if let (Ok(collection_hash), Ok(bucket_hash)) =
506                (collection_radix.as_decimal(), bucket_radix.as_decimal())
507            {
508                let origin_fst = StoreFSTBuilder::open(
509                    collection_hash as StoreFSTAtom,
510                    bucket_hash as StoreFSTAtom,
511                    &self.fst_store_config,
512                )
513                .map_err(|_| io::Error::other("graph open failure"))?;
514
515                let mut origin_fst_stream = origin_fst.stream();
516
517                while let Some(word) = origin_fst_stream.next() {
518                    count_words += 1;
519
520                    // Write word, and append a new line
521                    backup_fst_writer.write_all(word)?;
522                    backup_fst_writer.write_all(b"\n")?;
523                }
524
525                tracing::info!(
526                    "fst bucket: {}/{} backed up to path: {:?} ({} words)",
527                    collection_name,
528                    bucket_name,
529                    fst_backup_path,
530                    count_words
531                );
532            }
533        }
534
535        Ok(())
536    }
537
538    fn restore_item(
539        &self,
540        _backup_path: &Path,
541        origin_path: &Path,
542        collection_name: &str,
543        bucket_name: &str,
544    ) -> Result<(), io::Error> {
545        // Acquire access lock (in blocking write mode), and reference it in context
546        // Notice: this prevents store to be acquired from any context
547        let _access = self.graph_access_lock.write().unwrap();
548
549        tracing::debug!(
550            "fst bucket: {}/{} restoring from path: {:?}",
551            collection_name,
552            bucket_name,
553            origin_path
554        );
555
556        // Convert names to hashes (as names are hashes encoded as base-16 strings, but we need \
557        //   them as proper integers)
558        if let (Ok(collection_radix), Ok(bucket_radix)) = (
559            RadixNum::from_str(collection_name, ATOM_HASH_RADIX),
560            RadixNum::from_str(bucket_name, ATOM_HASH_RADIX),
561        ) {
562            if let (Ok(collection_hash), Ok(bucket_hash)) =
563                (collection_radix.as_decimal(), bucket_radix.as_decimal())
564            {
565                // Force a FST store close
566                self.close(collection_hash as StoreFSTAtom, bucket_hash as StoreFSTAtom);
567
568                // Generate path to FST
569                let fst_path = self.fst_store_config.path(
570                    StoreFSTPathMode::Permanent,
571                    collection_hash as StoreFSTAtom,
572                    Some(bucket_hash as StoreFSTAtom),
573                );
574
575                // Remove existing FST data?
576                if fst_path.exists() {
577                    fs::remove_file(&fst_path)?;
578                }
579
580                // Stream backup words to restored FST
581                let fst_writer = BufWriter::new(File::create(&fst_path)?);
582                let fst_backup_reader = BufReader::new(File::open(&origin_path)?);
583
584                let mut fst_builder = FSTSetBuilder::new(fst_writer)
585                    .map_err(|_| io::Error::other("graph restore builder failure"))?;
586
587                for word in fst_backup_reader.lines() {
588                    let word = word?;
589
590                    fst_builder
591                        .insert(word)
592                        .map_err(|_| io::Error::other("graph restore word insert failure"))?;
593                }
594
595                fst_builder
596                    .finish()
597                    .map_err(|_| io::Error::other("graph restore finish failure"))?;
598
599                tracing::info!(
600                    "fst bucket: {}/{} restored to path: {:?} from backup: {:?}",
601                    collection_name,
602                    bucket_name,
603                    fst_path,
604                    origin_path
605                );
606            }
607        }
608
609        Ok(())
610    }
611
612    fn consolidate_item(&self, store: &StoreFSTBox) -> (bool, usize, usize, usize) {
613        let (mut should_close, mut count_moved, mut count_pushed, mut count_popped) =
614            (false, 0, 0, 0);
615
616        // Acquire write references to pending sets
617        let (mut pending_push_write, mut pending_pop_write) = (
618            store.pending.push.write().unwrap(),
619            store.pending.pop.write().unwrap(),
620        );
621
622        // Do consolidate? (any change committed)
623        // Notice: if both pending sets are empty do not consolidate as there may have been a \
624        //   push then a pop of this push, nulling out any committed change.
625        if !(pending_push_write.is_empty() && pending_pop_write.is_empty()) {
626            // Read old FST (or default to empty FST)
627            if let Ok(old_fst) = StoreFSTBuilder::open(
628                store.target.collection_hash,
629                store.target.bucket_hash,
630                &self.fst_store_config,
631            ) {
632                // Initialize the new FST (temporary)
633                let bucket_tmp_path = self.fst_store_config.path(
634                    StoreFSTPathMode::Temporary,
635                    store.target.collection_hash,
636                    Some(store.target.bucket_hash),
637                );
638
639                let bucket_tmp_path_parent = bucket_tmp_path.parent().unwrap();
640
641                if fs::create_dir_all(&bucket_tmp_path_parent).is_ok() {
642                    // Erase any previously-existing temporary FST (eg. process stopped while \
643                    //   writing the temporary FST); there is no guarantee this succeeds.
644                    fs::remove_file(&bucket_tmp_path).ok();
645
646                    if let Ok(tmp_fst_file) = File::create(&bucket_tmp_path) {
647                        let tmp_fst_writer = BufWriter::new(tmp_fst_file);
648
649                        // Create a builder that can be used to insert new key-value pairs.
650                        if let Ok(mut tmp_fst_builder) = FSTSetBuilder::new(tmp_fst_writer) {
651                            // Convert push keys to an ordered vector
652                            // Notice: we must go from a Vec to a VecDeque as to sort values, \
653                            //   which is a requirement for FST insertions.
654                            let mut ordered_push_vec: Vec<&[u8]> =
655                                Vec::from_iter(pending_push_write.iter().map(|item| item.as_ref()));
656
657                            ordered_push_vec.sort();
658
659                            let mut ordered_push: VecDeque<&[u8]> =
660                                VecDeque::from_iter(ordered_push_vec);
661
662                            // Append words not in pop list to new FST (ie. old words minus pop \
663                            //   words)
664                            let mut old_fst_stream = old_fst.stream();
665
666                            'old: while let Some(old_fst_word) = old_fst_stream.next() {
667                                // Append new words from front? (ie. push words)
668                                // Notice: as an FST is ordered, inserts would fail if they are \
669                                //   committed out-of-order. Thus, the only way to check for \
670                                //   order is there.
671                                // Notice: a quick check is done before engaging in the loop, to \
672                                //   prevent any de-optimized jump instruction, as we may call \
673                                //   this code block a lot on large FSTs, and the loop should not \
674                                //   be engaged that often on stabilized FSTs (ie. mature FSTs).
675                                if let Some(push_first_ref) = ordered_push.front() {
676                                    // Engage the loop?
677                                    if *push_first_ref <= old_fst_word {
678                                        while let Some(push_front_ref) = ordered_push.front() {
679                                            if *push_front_ref <= old_fst_word {
680                                                // Pop front item and consume it
681                                                // Notice: as we validated previously that there \
682                                                //   is a front value, this unwrap is safe.
683                                                let push_front = ordered_push.pop_front().unwrap();
684
685                                                if StoreFSTMisc::check_over_limits(
686                                                    tmp_fst_builder.bytes_written() as usize,
687                                                    count_pushed + count_moved,
688                                                    &self.fst_store_config.graph,
689                                                ) {
690                                                    // FST cannot accept more items (limits reached)
691                                                    tracing::warn!(
692                                                        "limit reached on new from old in fst"
693                                                    );
694
695                                                    // Important: stop the main loop (limit reached)
696                                                    break 'old;
697                                                }
698
699                                                if let Err(err) = tmp_fst_builder.insert(push_front)
700                                                {
701                                                    // Could not insert word in FST
702                                                    tracing::error!(
703                                                        "failed inserting new from old in fst: {}",
704                                                        err
705                                                    );
706                                                } else {
707                                                    // Word inserted in FST
708                                                    count_pushed += 1;
709                                                }
710
711                                                // Continue scanning next word (may also come \
712                                                //   before this FST word in order)
713                                                continue;
714                                            }
715
716                                            // Important: stop loop on next front item (always \
717                                            //   the same)
718                                            break;
719                                        }
720                                    }
721                                }
722
723                                // Restore old word (if not popped)
724                                if !pending_pop_write.contains(old_fst_word) {
725                                    if StoreFSTMisc::check_over_limits(
726                                        tmp_fst_builder.bytes_written() as usize,
727                                        count_pushed + count_moved,
728                                        &self.fst_store_config.graph,
729                                    ) {
730                                        // FST cannot accept more items (limits reached)
731                                        tracing::warn!("limit reached on old word in fst");
732
733                                        // Important: stop the main loop (limit reached)
734                                        break 'old;
735                                    }
736
737                                    if let Err(err) = tmp_fst_builder.insert(old_fst_word) {
738                                        // Could not move word to FST
739                                        tracing::error!(
740                                            "failed inserting old word in fst: {}",
741                                            err
742                                        );
743                                    } else {
744                                        // Word moved to FST
745                                        count_moved += 1;
746                                    }
747                                } else {
748                                    count_popped += 1;
749                                }
750                            }
751
752                            // Complete FST with last pushed items
753                            // Notice: this is necessary if the FST was empty, or if we have push \
754                            //   items that come after the last ordered word of the FST.
755                            while let Some(push_front) = ordered_push.pop_front() {
756                                if StoreFSTMisc::check_over_limits(
757                                    tmp_fst_builder.bytes_written() as usize,
758                                    count_pushed + count_moved,
759                                    &self.fst_store_config.graph,
760                                ) {
761                                    // FST cannot accept more items (limits reached)
762                                    tracing::warn!(
763                                        "limit reached on new word from complete in fst"
764                                    );
765
766                                    // Important: stop the main loop (limit reached)
767                                    break;
768                                }
769
770                                if let Err(err) = tmp_fst_builder.insert(push_front) {
771                                    // Could not insert word in FST
772                                    tracing::error!(
773                                        "failed inserting new word from complete in fst: {}",
774                                        err
775                                    );
776                                } else {
777                                    // Word inserted in FST
778                                    count_pushed += 1;
779                                }
780                            }
781
782                            // Finish building new FST
783                            if tmp_fst_builder.finish().is_ok() {
784                                // Should close open store reference to old FST
785                                should_close = true;
786
787                                // Replace old FST with new FST (this nukes the old FST)
788                                // Notice: there is no need to re-open the new FST, as it will be \
789                                //   automatically opened on its next access.
790                                let bucket_final_path = self.fst_store_config.path(
791                                    StoreFSTPathMode::Permanent,
792                                    store.target.collection_hash,
793                                    Some(store.target.bucket_hash),
794                                );
795
796                                // Proceed temporary FST to final FST path rename
797                                if fs::rename(&bucket_tmp_path, &bucket_final_path).is_ok() {
798                                    tracing::info!(
799                                        "done consolidate fst at path: {:?}",
800                                        bucket_final_path
801                                    );
802                                } else {
803                                    tracing::error!(
804                                        "error consolidating fst at path: {:?}",
805                                        bucket_final_path
806                                    );
807                                }
808                            } else {
809                                tracing::error!(
810                                    "error finishing building temporary fst at path: {:?}",
811                                    bucket_tmp_path
812                                );
813                            }
814                        } else {
815                            tracing::error!(
816                                "error starting building temporary fst at path: {:?}",
817                                bucket_tmp_path
818                            );
819                        }
820                    } else {
821                        tracing::error!(
822                            "error initializing temporary fst at path: {:?}",
823                            bucket_tmp_path
824                        );
825                    }
826                } else {
827                    tracing::error!(
828                        "error initializing temporary fst directory at path: {:?}",
829                        bucket_tmp_path_parent
830                    );
831                }
832            } else {
833                tracing::error!("error opening old fst");
834            }
835
836            // Reset all pending sets
837            *pending_push_write = HashSet::new();
838            *pending_pop_write = HashSet::new();
839        }
840
841        (should_close, count_moved, count_pushed, count_popped)
842    }
843
844    fn close(&self, collection_hash: StoreFSTAtom, bucket_hash: StoreFSTAtom) {
845        tracing::debug!(
846            "closing finite-state transducer graph for collection: <{:x}> and bucket: <{:x}>",
847            collection_hash,
848            bucket_hash
849        );
850
851        let bucket_target = StoreFSTKey::from_atom(collection_hash, bucket_hash);
852
853        self.graph_pool.write().unwrap().remove(&bucket_target);
854        self.graph_consolidate
855            .write()
856            .unwrap()
857            .remove(&bucket_target);
858    }
859}
860
861impl<'build> StoreGenericPool<StoreFSTKey, StoreFST, StoreFSTBuilder<'build>> for StoreFSTPool {}
862
863impl<'build> StoreFSTBuilder<'build> {
864    fn open(
865        collection_hash: StoreFSTAtom,
866        bucket_hash: StoreFSTAtom,
867        fst_store_config: &crate::config::ConfigStoreFST,
868    ) -> Result<FSTSet, FSTError> {
869        tracing::debug!(
870            "opening finite-state transducer graph for collection: <{:x}> and bucket: <{:x}>",
871            collection_hash,
872            bucket_hash
873        );
874
875        let collection_bucket_path = fst_store_config.path(
876            StoreFSTPathMode::Permanent,
877            collection_hash,
878            Some(bucket_hash),
879        );
880
881        if collection_bucket_path.exists() {
882            // Open graph at path for collection
883            // Notice: this is unsafe, as loaded memory is a memory-mapped file, that cannot be \
884            //   guaranteed not to be muted while we own a read handle to it. Though, we use \
885            //   higher-level locking mechanisms on all callers of this method, so we are safe.
886            unsafe { FSTSet::from_path(collection_bucket_path) }
887        } else {
888            // FST does not exist on disk, generate an empty FST for now; until a consolidation \
889            //   task occurs and populates the on-disk-FST.
890            let empty_iter: Vec<&str> = Vec::new();
891
892            FSTSet::from_iter(empty_iter)
893        }
894    }
895}
896
897impl crate::config::ConfigStoreFST {
898    fn path(
899        &self,
900        mode: StoreFSTPathMode,
901        collection_hash: StoreFSTAtom,
902        bucket_hash: Option<StoreFSTAtom>,
903    ) -> PathBuf {
904        let mut final_path = self.path.join(format!("{:x}", collection_hash));
905
906        if let Some(bucket_hash) = bucket_hash {
907            final_path = final_path.join(format!("{:x}{}", bucket_hash, mode.extension()));
908        }
909
910        final_path
911    }
912}
913
914impl<'build> StoreGenericBuilder<StoreFSTKey, StoreFST> for StoreFSTBuilder<'build> {
915    fn build(&self, pool_key: StoreFSTKey) -> Result<StoreFST, ()> {
916        Self::open(
917            pool_key.collection_hash,
918            pool_key.bucket_hash,
919            self.fst_store_config,
920        )
921        .map(|graph| {
922            let now = SystemTime::now();
923
924            StoreFST {
925                graph,
926                target: pool_key,
927                pending: StoreFSTPending::default(),
928                last_used: Arc::new(RwLock::new(now)),
929                last_consolidated: Arc::new(RwLock::new(now)),
930                graph_consolidate: Arc::clone(&self.graph_consolidate),
931                action_config: self.fst_action_config,
932            }
933        })
934        .map_err(|err| {
935            tracing::error!("failed opening fst: {}", err);
936        })
937    }
938}
939
940impl StoreFST {
941    pub fn cardinality(&self) -> usize {
942        self.graph.len()
943    }
944
945    pub fn as_stream(&self) -> FSTStream<'_, AlwaysMatch> {
946        self.graph.into_stream()
947    }
948
949    pub fn lookup_begins(&self, word: &str) -> Result<FSTStream<'_, Regex>, ()> {
950        // Notice: this regex maps over an unicode range, for speed reasons at scale. \
951        //   We found out that the 'match any' syntax ('.*') was super-slow. Using the restrictive \
952        //   syntax below divided the cost of eg. a search query by 2. The regex below has been \
953        //   found out to be nearly zero-cost to compile and execute, for whatever reason.
954        // Regex format: '{escaped_word}([{unicode_range}]*)'
955        let mut regex_str = regex_escape(word);
956
957        regex_str.push('(');
958
959        let write_result = LexerRegexRange::from(word)
960            .unwrap_or_default()
961            .write_to(&mut regex_str);
962
963        regex_str.push_str("*)");
964
965        // Regex write failed? (this should not happen)
966        if let Err(err) = write_result {
967            tracing::error!(
968                "could not lookup word in fst via 'begins': {} because regex write failed: {}",
969                word,
970                err
971            );
972
973            return Err(());
974        }
975
976        // Proceed word lookup
977        tracing::debug!(
978            "looking-up word in fst via 'begins': {} with regex: {}",
979            word,
980            regex_str
981        );
982
983        if let Ok(regex) = Regex::new(&regex_str) {
984            Ok(self.graph.search(regex).into_stream())
985        } else {
986            Err(())
987        }
988    }
989
990    pub fn lookup_typos(
991        &self,
992        word: &str,
993        typo_factor: u32,
994    ) -> Result<FSTStream<'_, Levenshtein>, ()> {
995        tracing::debug!(
996            "looking-up word in fst via 'typos': {} with typo factor: {}",
997            word,
998            typo_factor
999        );
1000
1001        if let Ok(fuzzy) = Levenshtein::new(word, typo_factor) {
1002            Ok(self.graph.search(fuzzy).into_stream())
1003        } else {
1004            Err(())
1005        }
1006    }
1007
1008    pub fn should_consolidate(&self) {
1009        // Check if not already scheduled
1010        if !self
1011            .graph_consolidate
1012            .read()
1013            .unwrap()
1014            .contains(&self.target)
1015        {
1016            // Schedule target for next consolidation tick (ie. collection + bucket tuple)
1017            self.graph_consolidate.write().unwrap().insert(self.target);
1018
1019            // Bump 'last consolidated' time, effectively de-bouncing consolidation to a fixed \
1020            //   and predictable tick time in the future.
1021            let mut last_consolidated_value = self.last_consolidated.write().unwrap();
1022
1023            *last_consolidated_value = SystemTime::now();
1024
1025            // Perform an early drop of the lock (frees up write lock early)
1026            drop(last_consolidated_value);
1027
1028            tracing::info!("graph consolidation scheduled on pool key: {}", self.target);
1029        } else {
1030            tracing::debug!(
1031                "graph consolidation already scheduled on pool key: {}",
1032                self.target
1033            );
1034        }
1035    }
1036}
1037
1038impl StoreGeneric for StoreFST {
1039    fn ref_last_used(&self) -> &RwLock<SystemTime> {
1040        &self.last_used
1041    }
1042}
1043
1044impl<'build> StoreFSTActionBuilder<'build> {
1045    pub fn access(store: StoreFSTBox) -> StoreFSTAction {
1046        Self::build(store)
1047    }
1048
1049    fn build(store: StoreFSTBox) -> StoreFSTAction {
1050        StoreFSTAction { store }
1051    }
1052}
1053
1054impl StoreFSTPool {
1055    pub fn erase<T: AsRef<str>>(&self, collection: T, bucket: Option<T>) -> Result<u32, ()> {
1056        self.dispatch_erase("fst", collection, bucket)
1057    }
1058}
1059
1060impl StoreGenericActionBuilder for StoreFSTPool {
1061    fn proceed_erase_collection(&self, collection_str: &str) -> Result<u32, ()> {
1062        let path_mode = StoreFSTPathMode::Permanent;
1063
1064        let collection_atom = StoreKeyerHasher::to_compact(collection_str);
1065        let collection_path = self.fst_store_config.path(path_mode, collection_atom, None);
1066
1067        // Force a FST graph close (on all contained buckets)
1068        // Notice: we first need to scan for opened buckets in-memory, as not all FSTs may be \
1069        //   committed to disk; thus some FST stores that exist in-memory may not exist on-disk.
1070        let mut bucket_atoms: Vec<StoreFSTAtom> = Vec::new();
1071
1072        {
1073            let graph_pool_read = self.graph_pool.read().unwrap();
1074
1075            for target_key in graph_pool_read.keys() {
1076                if target_key.collection_hash == collection_atom {
1077                    bucket_atoms.push(target_key.bucket_hash);
1078                }
1079            }
1080        }
1081
1082        if !bucket_atoms.is_empty() {
1083            tracing::debug!(
1084                "will force-close {} fst buckets for collection: {}",
1085                bucket_atoms.len(),
1086                collection_str
1087            );
1088
1089            let (mut graph_pool_write, mut graph_consolidate_write) = (
1090                self.graph_pool.write().unwrap(),
1091                self.graph_consolidate.write().unwrap(),
1092            );
1093
1094            for bucket_atom in bucket_atoms {
1095                tracing::debug!(
1096                    "fst bucket graph force close for bucket: {}/<{:x}>",
1097                    collection_str,
1098                    bucket_atom
1099                );
1100
1101                let bucket_target = StoreFSTKey::from_atom(collection_atom, bucket_atom);
1102
1103                graph_pool_write.remove(&bucket_target);
1104                graph_consolidate_write.remove(&bucket_target);
1105            }
1106        }
1107
1108        // Remove all FSTs on-disk
1109        if collection_path.exists() {
1110            tracing::debug!(
1111                "fst collection store exists, erasing: {}/* at path: {:?}",
1112                collection_str,
1113                &collection_path
1114            );
1115
1116            // Remove FST graph storage from filesystem
1117            let erase_result = fs::remove_dir_all(&collection_path);
1118
1119            if erase_result.is_ok() {
1120                tracing::debug!("done with fst collection erasure");
1121
1122                Ok(1)
1123            } else {
1124                Err(())
1125            }
1126        } else {
1127            tracing::debug!(
1128                "fst collection store does not exist, consider already erased: {}/* at path: {:?}",
1129                collection_str,
1130                &collection_path
1131            );
1132
1133            Ok(0)
1134        }
1135    }
1136
1137    fn proceed_erase_bucket(&self, collection_str: &str, bucket_str: &str) -> Result<u32, ()> {
1138        tracing::debug!(
1139            "sub-erase on fst bucket: {} for collection: {}",
1140            bucket_str,
1141            collection_str
1142        );
1143
1144        let (collection_atom, bucket_atom) = (
1145            StoreKeyerHasher::to_compact(collection_str),
1146            StoreKeyerHasher::to_compact(bucket_str),
1147        );
1148
1149        let bucket_path = self.fst_store_config.path(
1150            StoreFSTPathMode::Permanent,
1151            collection_atom,
1152            Some(bucket_atom),
1153        );
1154
1155        // Force a FST graph close
1156        self.close(collection_atom, bucket_atom);
1157
1158        // Remove FST on-disk
1159        if bucket_path.exists() {
1160            tracing::debug!(
1161                "fst bucket graph exists, erasing: {}/{} at path: {:?}",
1162                collection_str,
1163                bucket_str,
1164                &bucket_path
1165            );
1166
1167            // Remove FST graph storage from filesystem
1168            let erase_result = fs::remove_file(&bucket_path);
1169
1170            if erase_result.is_ok() {
1171                tracing::debug!("done with fst bucket erasure");
1172
1173                Ok(1)
1174            } else {
1175                Err(())
1176            }
1177        } else {
1178            tracing::debug!(
1179                "fst bucket graph does not exist, consider already erased: {}/{} at path: {:?}",
1180                collection_str,
1181                bucket_str,
1182                &bucket_path
1183            );
1184
1185            Ok(0)
1186        }
1187    }
1188}
1189
1190impl StoreFSTAction {
1191    pub fn push_word(&self, word: &str, fst_store_config: &crate::config::ConfigStoreFST) -> bool {
1192        // Word over limit? (abort, the FST does not perform well over large words)
1193        if Self::word_over_limit(word) {
1194            return false;
1195        }
1196
1197        let word_bytes = word.as_bytes();
1198
1199        // Nuke word from 'pop' set? (void a previous un-consolidated commit)
1200        if self.store.pending.pop.read().unwrap().contains(word_bytes) {
1201            self.store.pending.pop.write().unwrap().remove(word_bytes);
1202        }
1203
1204        // Add word in 'push' set? (only if word is not in FST)
1205        // Notice: also check whether FST is over limits or not from there, to avoid stacking \
1206        //   words that could never be consolidated to final FST anyway.
1207        let graph_fst = self.store.graph.as_fst();
1208
1209        if !self.store.graph.contains(&word)
1210            && !self.store.pending.push.read().unwrap().contains(word_bytes)
1211            && self.store.pending.push.read().unwrap().len() < fst_store_config.graph.max_words
1212            && !StoreFSTMisc::check_over_limits(
1213                graph_fst.size(),
1214                graph_fst.len(),
1215                &fst_store_config.graph,
1216            )
1217        {
1218            self.store
1219                .pending
1220                .push
1221                .write()
1222                .unwrap()
1223                .insert(word_bytes.to_vec());
1224
1225            self.store.should_consolidate();
1226
1227            // Pushed
1228            true
1229        } else {
1230            // Not pushed
1231            false
1232        }
1233    }
1234
1235    pub fn pop_word(&self, word: &str) -> bool {
1236        // Word over limit? (abort, the FST does not perform well over large words)
1237        if Self::word_over_limit(word) {
1238            return false;
1239        }
1240
1241        let word_bytes = word.as_bytes();
1242
1243        // Nuke word from 'push' set? (void a previous un-consolidated commit)
1244        if self.store.pending.push.read().unwrap().contains(word_bytes) {
1245            self.store.pending.push.write().unwrap().remove(word_bytes);
1246        }
1247
1248        // Add word in 'pop' set? (only if word is in FST)
1249        if self.store.graph.contains(word_bytes)
1250            && !self.store.pending.pop.read().unwrap().contains(word_bytes)
1251        {
1252            self.store
1253                .pending
1254                .pop
1255                .write()
1256                .unwrap()
1257                .insert(word_bytes.to_vec());
1258
1259            self.store.should_consolidate();
1260
1261            // Popped
1262            true
1263        } else {
1264            // Not popped
1265            false
1266        }
1267    }
1268
1269    pub fn suggest_words(
1270        &self,
1271        from_word: &str,
1272        // Length before stemming. Useful to apply fuzzy matching rules based
1273        // on user input.
1274        original_word_len: usize,
1275        limit: usize,
1276        max_typo_factor: Option<u32>,
1277    ) -> Option<
1278        impl ExactSizeIterator<Item = (String, QueryMatchScore)> + DoubleEndedIterator + use<>,
1279    > {
1280        // Word over limit? (abort, the FST does not perform well over large words)
1281        if Self::word_over_limit(from_word) {
1282            return None;
1283        }
1284
1285        let mut found_words: IndexMap<String, QueryMatchScore> = IndexMap::with_capacity(limit);
1286
1287        if self.config().prefix_matching_enabled {
1288            // Try to complete provided word
1289            if let Some(stream) = self.lookup_begins(from_word, original_word_len) {
1290                for (word, score) in stream {
1291                    if found_words.contains_key(&word) {
1292                        continue;
1293                    }
1294
1295                    found_words.insert(word, score);
1296
1297                    // Requested limit reached? Stop there.
1298                    if found_words.len() >= limit {
1299                        break;
1300                    }
1301                }
1302            }
1303        }
1304
1305        // Try to fuzzy-suggest other words? (eg. correct typos)
1306        if self.config().fuzzy_matching_enabled && found_words.len() < limit {
1307            // Allow more typos in word as the word gets longer, up to a maximum limit
1308            let max_typo_factor = max_typo_factor.unwrap_or(typo_factor(original_word_len));
1309            let mut typo_factor = 1u32;
1310
1311            // TODO: Rework the Levenshtein query feature to avoid repeating
1312            //   the same query over and over again. Maybe try to see if
1313            //   `fst_levenshtein` can return distances in its response.
1314            while found_words.len() < limit && typo_factor <= max_typo_factor {
1315                let Some(stream) = self.lookup_typos(from_word, typo_factor) else {
1316                    break;
1317                };
1318
1319                for (word, score) in stream {
1320                    if found_words.contains_key(&word) {
1321                        continue;
1322                    }
1323
1324                    found_words.insert(word, score);
1325
1326                    // Requested limit reached? Stop there.
1327                    if found_words.len() >= limit {
1328                        break;
1329                    }
1330                }
1331
1332                typo_factor += 1;
1333            }
1334        }
1335
1336        if !found_words.is_empty() {
1337            Some(found_words.into_iter())
1338        } else {
1339            None
1340        }
1341    }
1342
1343    pub fn lookup_begins(
1344        &self,
1345        from_word: &str,
1346        // Length before stemming. Useful to calculate correct score.
1347        original_word_len: usize,
1348    ) -> Option<impl Iterator<Item = (String, QueryMatchScore)>> {
1349        // Word over limit? (abort, the FST does not perform well over large words)
1350        if Self::word_over_limit(from_word) {
1351            return None;
1352        }
1353
1354        if !self.config().prefix_matching_enabled {
1355            return None;
1356        }
1357
1358        let Ok(stream) = self.store.lookup_begins(from_word) else {
1359            return None;
1360        };
1361
1362        tracing::debug!(
1363            word = ?from_word,
1364            "looking up for word in 'begins' fst stream"
1365        );
1366
1367        Some(FSTStreamIterator(stream).map(move |word| {
1368            // WARN: Calculating distance to original word length might
1369            //   yield weird results when combines with stemming.
1370            let distance: usize = original_word_len.abs_diff(word.len());
1371            let score = u16::try_from(distance).unwrap_or(u16::MAX);
1372            (word, score)
1373        }))
1374    }
1375
1376    pub fn lookup_typos(
1377        &self,
1378        from_word: &str,
1379        typo_factor: u32,
1380    ) -> Option<impl Iterator<Item = (String, QueryMatchScore)>> {
1381        if !self.config().fuzzy_matching_enabled {
1382            return None;
1383        }
1384
1385        let Ok(stream) = self.store.lookup_typos(from_word, typo_factor) else {
1386            return None;
1387        };
1388
1389        tracing::debug!(
1390            word = ?from_word, typo_factor,
1391            "looking up for word in 'typos' fst stream"
1392        );
1393
1394        // NOTE: Returning the same score for every word works only
1395        //   because we re-run `lookup_typos` for increasingly
1396        //   larger typo factors and do not re-insert existing
1397        //   values. As explained in previous TODO, we should try
1398        //   to get the real distance back from `fst_levenshtein`.
1399        let score = u16::try_from(typo_factor).unwrap_or(u16::MAX);
1400
1401        Some(FSTStreamIterator(stream).map(move |word| (word, score)))
1402    }
1403
1404    pub fn list_words(&self, limit: usize, offset: usize) -> Result<Vec<String>, ()> {
1405        let stream = self.store.as_stream();
1406
1407        // Enumerate words from FST stream
1408        match stream
1409            .into_strs()
1410            .map(|words| words.into_iter().skip(offset).take(limit).collect())
1411        {
1412            Err(err) => {
1413                tracing::debug!("conversion of stream failed: {}", err);
1414                Err(())
1415            }
1416            Ok(words) => Ok(words),
1417        }
1418    }
1419
1420    pub fn count_words(&self) -> usize {
1421        self.store.cardinality()
1422    }
1423
1424    fn word_over_limit(word: &str) -> bool {
1425        if word.len() > WORD_LIMIT_LENGTH {
1426            tracing::debug!("got over-limit fst word: {}", word);
1427
1428            true
1429        } else {
1430            false
1431        }
1432    }
1433}
1434
1435/// Allow more typos in word as the word gets longer, up to a maximum limit.
1436pub(crate) fn typo_factor(word_len: usize) -> u32 {
1437    match word_len {
1438        1..=3 => 0,
1439        4..=6 => 1,
1440        7..=9 => 2,
1441        _ => 3,
1442    }
1443}
1444
1445impl StoreFSTMisc {
1446    pub fn count_collection_buckets(
1447        collection: impl AsRef<str>,
1448        fst_store_config: &crate::config::ConfigStoreFST,
1449    ) -> Result<usize, ()> {
1450        let mut count = 0;
1451
1452        let path_mode = StoreFSTPathMode::Permanent;
1453
1454        let collection_atom = StoreKeyerHasher::to_compact(collection.as_ref());
1455        let collection_path = fst_store_config.path(path_mode, collection_atom, None);
1456
1457        if collection_path.exists() {
1458            // Scan collection directory for contained buckets (count them)
1459            if let Ok(entries) = fs::read_dir(&collection_path) {
1460                let fst_extension = path_mode.extension();
1461                let fst_extension_len = fst_extension.len();
1462
1463                for entry in entries.flatten() {
1464                    if let Some(entry_name) = entry.file_name().to_str() {
1465                        let entry_name_len = entry_name.len();
1466
1467                        // FST file found? This is a bucket.
1468                        if entry_name_len > fst_extension_len && entry_name.ends_with(fst_extension)
1469                        {
1470                            count += 1;
1471                        }
1472                    }
1473                }
1474            } else {
1475                tracing::error!("failed reading directory for count: {:?}", collection_path);
1476
1477                return Err(());
1478            }
1479        }
1480
1481        Ok(count)
1482    }
1483
1484    fn check_over_limits(
1485        bytes_count: usize,
1486        words_count: usize,
1487        fst_graph_config: &crate::config::ConfigStoreFSTGraph,
1488    ) -> bool {
1489        // Over bytes limit?
1490        let max_size = fst_graph_config.max_size * 1024;
1491
1492        if bytes_count >= max_size {
1493            tracing::info!(
1494                "fst has exceeded maximum allowed bytes: {} over limit: {}",
1495                bytes_count,
1496                max_size
1497            );
1498
1499            return true;
1500        }
1501
1502        // Over words limit?
1503        if words_count >= fst_graph_config.max_words {
1504            tracing::info!(
1505                "fst has exceeded maximum allowed words: {} over limit: {}",
1506                words_count,
1507                fst_graph_config.max_words
1508            );
1509
1510            return true;
1511        }
1512
1513        // Not over limit
1514        false
1515    }
1516}
1517
1518impl StoreFSTKey {
1519    pub fn from_atom(collection_hash: StoreFSTAtom, bucket_hash: StoreFSTAtom) -> StoreFSTKey {
1520        StoreFSTKey {
1521            collection_hash,
1522            bucket_hash,
1523        }
1524    }
1525
1526    pub fn from_str(collection_str: &str, bucket_str: &str) -> StoreFSTKey {
1527        StoreFSTKey {
1528            collection_hash: StoreKeyerHasher::to_compact(collection_str),
1529            bucket_hash: StoreKeyerHasher::to_compact(bucket_str),
1530        }
1531    }
1532}
1533
1534impl fmt::Display for StoreFSTKey {
1535    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1536        write!(f, "<{:x}>/<{:x}>", self.collection_hash, self.bucket_hash)
1537    }
1538}
1539
1540// MARK: - Helpers
1541
1542#[repr(transparent)]
1543struct FSTStreamIterator<'a, A: Automaton>(fst::set::Stream<'a, A>);
1544
1545impl<'a, A: Automaton> Iterator for FSTStreamIterator<'a, A> {
1546    type Item = String;
1547
1548    fn next(&mut self) -> Option<Self::Item> {
1549        match self.0.next() {
1550            Some(bytes) => match str::from_utf8(bytes) {
1551                Ok(str) => Some(str.to_owned()),
1552                Err(_) => None,
1553            },
1554            None => None,
1555        }
1556    }
1557}
1558
1559// MARK: - Tests
1560
1561#[cfg(test)]
1562mod tests {
1563    use super::*;
1564
1565    #[test]
1566    fn it_acquires_graph() {
1567        let fst_pool = test_fst_pool();
1568
1569        assert!(fst_pool.acquire("c:test:1", "b:test:1").is_ok());
1570    }
1571
1572    #[test]
1573    fn it_janitors_graph() {
1574        let fst_pool = test_fst_pool();
1575
1576        fst_pool.janitor();
1577    }
1578
1579    #[test]
1580    fn it_proceeds_primitives() {
1581        let fst_pool = test_fst_pool();
1582
1583        let store = fst_pool.acquire("c:test:2", "b:test:2").unwrap();
1584
1585        assert!(store.lookup_typos("valerien", 1).is_ok());
1586    }
1587
1588    fn test_fst_pool() -> StoreFSTPool {
1589        let fst_store_config = test_fst_store_config();
1590
1591        StoreFSTPool::new(fst_store_config, Default::default())
1592    }
1593
1594    fn test_fst_store_config() -> Arc<crate::config::ConfigStoreFST> {
1595        Arc::new(
1596            config::Config::builder()
1597                .add_source(config::File::from_str(
1598                    crate::config::tests::defaults_toml(),
1599                    config::FileFormat::Toml,
1600                ))
1601                .build()
1602                .unwrap()
1603                .get::<crate::config::ConfigStoreFST>("store.fst")
1604                .unwrap(),
1605        )
1606    }
1607}
1608
1609// MARK: - Boilerplate
1610
1611impl fmt::Debug for StoreFSTPool {
1612    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1613        use crate::util::fmt::{AsPrettyMutex, AsPrettyRwLock};
1614
1615        // NOTE: Deconstructing to future-proof this function.
1616        let Self {
1617            fst_action_config,
1618            graph_pool,
1619            graph_acquire_lock,
1620            graph_rebuild_lock,
1621            graph_access_lock,
1622            graph_consolidate,
1623            // NOTE: We don’t care about the configuration,
1624            //   we can see it elsewhere if needed.
1625            fst_store_config: _fst_store_config,
1626        } = self;
1627
1628        f.debug_struct("StoreFSTPool")
1629            .field("fst_action_config", fst_action_config)
1630            .field("graph_pool", &AsPrettyRwLock(graph_pool))
1631            .field("graph_acquire_lock", &AsPrettyMutex(graph_acquire_lock))
1632            .field("graph_rebuild_lock", &AsPrettyMutex(graph_rebuild_lock))
1633            .field("graph_access_lock", &AsPrettyRwLock(graph_access_lock))
1634            .field("graph_consolidate", &AsPrettyRwLock(graph_consolidate))
1635            .finish_non_exhaustive()
1636    }
1637}
1638
1639impl fmt::Debug for StoreFSTKey {
1640    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1641        fmt::Display::fmt(&self, f)
1642    }
1643}
1644
1645impl fmt::Debug for StoreFST {
1646    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1647        use crate::util::fmt::AsPrettyRwLock;
1648
1649        // NOTE: Deconstructing to future-proof this function.
1650        let Self {
1651            graph,
1652            target,
1653            pending,
1654            last_used,
1655            last_consolidated,
1656            graph_consolidate,
1657            action_config,
1658        } = self;
1659
1660        f.debug_struct("StoreFST")
1661            .field("graph", graph)
1662            .field("target", target)
1663            .field("pending", pending)
1664            .field("last_used", &AsPrettyRwLock(last_used))
1665            .field("last_consolidated", &AsPrettyRwLock(last_consolidated))
1666            .field("graph_consolidate", &AsPrettyRwLock(graph_consolidate))
1667            .field("action_config", action_config)
1668            .finish()
1669    }
1670}
1671
1672impl fmt::Debug for StoreFSTPending {
1673    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1674        use crate::util::fmt::AsPrettyRwLock;
1675
1676        // NOTE: Deconstructing to future-proof this function.
1677        let Self { pop, push } = self;
1678
1679        f.debug_struct("StoreFSTPending")
1680            .field("pop", &AsPrettyRwLock(pop))
1681            .field("push", &AsPrettyRwLock(push))
1682            .finish()
1683    }
1684}