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