1use 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#[derive(Clone)]
41pub struct StoreFSTPool {
42 fst_store_config: Arc<crate::config::ConfigStoreFST>,
43 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 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 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 let _acquire = self.graph_acquire_lock.lock().unwrap();
175
176 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 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 fs::create_dir_all(path)?;
218
219 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 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 let _rebuild = self.graph_rebuild_lock.lock().unwrap();
253
254 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 let mut keys_consolidate: Vec<StoreFSTKey> = Vec::new();
263
264 {
265 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 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 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 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 {
327 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 let (mut count_moved, mut count_pushed, mut count_popped) = (0, 0, 0);
342
343 {
344 for key in &keys_consolidate {
345 {
346 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 consolidate_counts.0
365 } else {
366 false
367 };
368
369 if do_close {
375 self.graph_pool.write().unwrap().remove(key);
376 }
377 }
378
379 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 for collection in fs::read_dir(read_path)? {
411 let collection = collection?;
412
413 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 fs::create_dir_all(write_path.join(collection_name))?;
422
423 for bucket in fs::read_dir(read_path.join(collection_name))? {
425 let bucket = bucket?;
426
427 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 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 let _access = self.graph_access_lock.write().unwrap();
475
476 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 fs::remove_file(&fst_backup_path).ok();
492
493 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 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 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 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 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 self.close(collection_hash as StoreFSTAtom, bucket_hash as StoreFSTAtom);
567
568 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 if fst_path.exists() {
577 fs::remove_file(&fst_path)?;
578 }
579
580 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 let (mut pending_push_write, mut pending_pop_write) = (
618 store.pending.push.write().unwrap(),
619 store.pending.pop.write().unwrap(),
620 );
621
622 if !(pending_push_write.is_empty() && pending_pop_write.is_empty()) {
626 if let Ok(old_fst) = StoreFSTBuilder::open(
628 store.target.collection_hash,
629 store.target.bucket_hash,
630 &self.fst_store_config,
631 ) {
632 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 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 if let Ok(mut tmp_fst_builder) = FSTSetBuilder::new(tmp_fst_writer) {
651 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 let mut old_fst_stream = old_fst.stream();
665
666 'old: while let Some(old_fst_word) = old_fst_stream.next() {
667 if let Some(push_first_ref) = ordered_push.front() {
676 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 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 tracing::warn!(
692 "limit reached on new from old in fst"
693 );
694
695 break 'old;
697 }
698
699 if let Err(err) = tmp_fst_builder.insert(push_front)
700 {
701 tracing::error!(
703 "failed inserting new from old in fst: {}",
704 err
705 );
706 } else {
707 count_pushed += 1;
709 }
710
711 continue;
714 }
715
716 break;
719 }
720 }
721 }
722
723 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 tracing::warn!("limit reached on old word in fst");
732
733 break 'old;
735 }
736
737 if let Err(err) = tmp_fst_builder.insert(old_fst_word) {
738 tracing::error!(
740 "failed inserting old word in fst: {}",
741 err
742 );
743 } else {
744 count_moved += 1;
746 }
747 } else {
748 count_popped += 1;
749 }
750 }
751
752 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 tracing::warn!(
763 "limit reached on new word from complete in fst"
764 );
765
766 break;
768 }
769
770 if let Err(err) = tmp_fst_builder.insert(push_front) {
771 tracing::error!(
773 "failed inserting new word from complete in fst: {}",
774 err
775 );
776 } else {
777 count_pushed += 1;
779 }
780 }
781
782 if tmp_fst_builder.finish().is_ok() {
784 should_close = true;
786
787 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 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 *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 unsafe { FSTSet::from_path(collection_bucket_path) }
887 } else {
888 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 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 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 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(®ex_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 if !self
1011 .graph_consolidate
1012 .read()
1013 .unwrap()
1014 .contains(&self.target)
1015 {
1016 self.graph_consolidate.write().unwrap().insert(self.target);
1018
1019 let mut last_consolidated_value = self.last_consolidated.write().unwrap();
1022
1023 *last_consolidated_value = SystemTime::now();
1024
1025 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 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 if collection_path.exists() {
1110 tracing::debug!(
1111 "fst collection store exists, erasing: {}/* at path: {:?}",
1112 collection_str,
1113 &collection_path
1114 );
1115
1116 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 self.close(collection_atom, bucket_atom);
1157
1158 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 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 if Self::word_over_limit(word) {
1194 return false;
1195 }
1196
1197 let word_bytes = word.as_bytes();
1198
1199 if self.store.pending.pop.read().unwrap().contains(word_bytes) {
1201 self.store.pending.pop.write().unwrap().remove(word_bytes);
1202 }
1203
1204 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 true
1229 } else {
1230 false
1232 }
1233 }
1234
1235 pub fn pop_word(&self, word: &str) -> bool {
1236 if Self::word_over_limit(word) {
1238 return false;
1239 }
1240
1241 let word_bytes = word.as_bytes();
1242
1243 if self.store.pending.push.read().unwrap().contains(word_bytes) {
1245 self.store.pending.push.write().unwrap().remove(word_bytes);
1246 }
1247
1248 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 true
1263 } else {
1264 false
1266 }
1267 }
1268
1269 pub fn suggest_words(
1270 &self,
1271 from_word: &str,
1272 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 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 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 if found_words.len() >= limit {
1299 break;
1300 }
1301 }
1302 }
1303 }
1304
1305 if self.config().fuzzy_matching_enabled && found_words.len() < limit {
1307 let max_typo_factor = max_typo_factor.unwrap_or(typo_factor(original_word_len));
1309 let mut typo_factor = 1u32;
1310
1311 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 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 original_word_len: usize,
1348 ) -> Option<impl Iterator<Item = (String, QueryMatchScore)>> {
1349 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 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 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 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
1435pub(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 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 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 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 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 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#[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#[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
1609impl fmt::Debug for StoreFSTPool {
1612 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1613 use crate::util::fmt::{AsPrettyMutex, AsPrettyRwLock};
1614
1615 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 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 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 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}