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;
36
37#[derive(Clone)]
40pub struct StoreFSTPool {
41 fst_store_config: Arc<crate::config::ConfigStoreFST>,
42 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 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 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 let _acquire = self.graph_acquire_lock.lock().unwrap();
174
175 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 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 fs::create_dir_all(path)?;
217
218 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 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 let _rebuild = self.graph_rebuild_lock.lock().unwrap();
252
253 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 let mut keys_consolidate: Vec<StoreFSTKey> = Vec::new();
262
263 {
264 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 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 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 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 {
326 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 let (mut count_moved, mut count_pushed, mut count_popped) = (0, 0, 0);
341
342 {
343 for key in &keys_consolidate {
344 {
345 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 consolidate_counts.0
364 } else {
365 false
366 };
367
368 if do_close {
374 self.graph_pool.write().unwrap().remove(key);
375 }
376 }
377
378 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 for collection in fs::read_dir(read_path)? {
410 let collection = collection?;
411
412 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 fs::create_dir_all(write_path.join(collection_name))?;
421
422 for bucket in fs::read_dir(read_path.join(collection_name))? {
424 let bucket = bucket?;
425
426 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 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 let _access = self.graph_access_lock.write().unwrap();
474
475 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 fs::remove_file(&fst_backup_path).ok();
491
492 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 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 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 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 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 self.close(collection_hash as StoreFSTAtom, bucket_hash as StoreFSTAtom);
566
567 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 if fst_path.exists() {
576 fs::remove_file(&fst_path)?;
577 }
578
579 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 let (mut pending_push_write, mut pending_pop_write) = (
617 store.pending.push.write().unwrap(),
618 store.pending.pop.write().unwrap(),
619 );
620
621 if !(pending_push_write.is_empty() && pending_pop_write.is_empty()) {
625 if let Ok(old_fst) = StoreFSTBuilder::open(
627 store.target.collection_hash,
628 store.target.bucket_hash,
629 &self.fst_store_config,
630 ) {
631 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 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 if let Ok(mut tmp_fst_builder) = FSTSetBuilder::new(tmp_fst_writer) {
650 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 let mut old_fst_stream = old_fst.stream();
664
665 'old: while let Some(old_fst_word) = old_fst_stream.next() {
666 if let Some(push_first_ref) = ordered_push.front() {
675 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 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 tracing::warn!(
691 "limit reached on new from old in fst"
692 );
693
694 break 'old;
696 }
697
698 if let Err(err) = tmp_fst_builder.insert(push_front)
699 {
700 tracing::error!(
702 "failed inserting new from old in fst: {}",
703 err
704 );
705 } else {
706 count_pushed += 1;
708 }
709
710 continue;
713 }
714
715 break;
718 }
719 }
720 }
721
722 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 tracing::warn!("limit reached on old word in fst");
731
732 break 'old;
734 }
735
736 if let Err(err) = tmp_fst_builder.insert(old_fst_word) {
737 tracing::error!(
739 "failed inserting old word in fst: {}",
740 err
741 );
742 } else {
743 count_moved += 1;
745 }
746 } else {
747 count_popped += 1;
748 }
749 }
750
751 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 tracing::warn!(
762 "limit reached on new word from complete in fst"
763 );
764
765 break;
767 }
768
769 if let Err(err) = tmp_fst_builder.insert(push_front) {
770 tracing::error!(
772 "failed inserting new word from complete in fst: {}",
773 err
774 );
775 } else {
776 count_pushed += 1;
778 }
779 }
780
781 if tmp_fst_builder.finish().is_ok() {
783 should_close = true;
785
786 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 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 *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 unsafe { FSTSet::from_path(collection_bucket_path) }
886 } else {
887 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 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 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 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(®ex_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 if !self
1010 .graph_consolidate
1011 .read()
1012 .unwrap()
1013 .contains(&self.target)
1014 {
1015 self.graph_consolidate.write().unwrap().insert(self.target);
1017
1018 let mut last_consolidated_value = self.last_consolidated.write().unwrap();
1021
1022 *last_consolidated_value = SystemTime::now();
1023
1024 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 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 if collection_path.exists() {
1109 tracing::debug!(
1110 "fst collection store exists, erasing: {}/* at path: {:?}",
1111 collection_str,
1112 &collection_path
1113 );
1114
1115 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 self.close(collection_atom, bucket_atom);
1156
1157 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 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 if Self::word_over_limit(word) {
1193 return false;
1194 }
1195
1196 let word_bytes = word.as_bytes();
1197
1198 if self.store.pending.pop.read().unwrap().contains(word_bytes) {
1200 self.store.pending.pop.write().unwrap().remove(word_bytes);
1201 }
1202
1203 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 true
1228 } else {
1229 false
1231 }
1232 }
1233
1234 pub fn pop_word(&self, word: &str) -> bool {
1235 if Self::word_over_limit(word) {
1237 return false;
1238 }
1239
1240 let word_bytes = word.as_bytes();
1241
1242 if self.store.pending.push.read().unwrap().contains(word_bytes) {
1244 self.store.pending.push.write().unwrap().remove(word_bytes);
1245 }
1246
1247 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 true
1262 } else {
1263 false
1265 }
1266 }
1267
1268 pub fn suggest_words(
1269 &self,
1270 from_word: &str,
1271 original_word_len: usize,
1274 limit: usize,
1275 max_typo_factor: Option<u32>,
1276 ) -> Option<impl ExactSizeIterator<Item = (String, u16)> + DoubleEndedIterator + use<>> {
1277 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 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 if found_words.len() >= limit {
1296 break;
1297 }
1298 }
1299 }
1300 }
1301
1302 if self.config().fuzzy_matching_enabled && found_words.len() < limit {
1304 let max_typo_factor = max_typo_factor.unwrap_or(typo_factor(original_word_len));
1306 let mut typo_factor = 1u32;
1307
1308 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 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 original_word_len: usize,
1345 ) -> Option<impl Iterator<Item = (String, u16)>> {
1346 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 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 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 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
1432pub(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 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 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 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 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 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#[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#[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
1606impl fmt::Debug for StoreFSTPool {
1609 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1610 use crate::util::fmt::{AsPrettyMutex, AsPrettyRwLock};
1611
1612 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 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 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 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}