1use byteorder::{ByteOrder, LittleEndian, ReadBytesExt};
9use hashbrown::HashMap;
10use radix::RadixNum;
11use rocksdb::backup::{
12 BackupEngine as DBBackupEngine, BackupEngineOptions as DBBackupEngineOptions,
13 RestoreOptions as DBRestoreOptions,
14};
15use rocksdb::{
16 DB, DBCompactionStyle, DBCompressionType, Env as DBEnv, Error as DBError, FlushOptions,
17 Options as DBOptions, WriteBatch, WriteOptions,
18};
19use std::fmt;
20use std::fs;
21use std::io::{self, Cursor};
22use std::path::{Path, PathBuf};
23use std::str;
24use std::sync::{Arc, Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard};
25use std::thread;
26use std::time::{Duration, SystemTime};
27use std::vec::Drain;
28
29use super::generic::{
30 StoreGeneric, StoreGenericActionBuilder, StoreGenericBuilder, StoreGenericPool,
31};
32use super::identifiers::*;
33use super::item::StoreItemPart;
34use super::keyer::{StoreKeyerBuilder, StoreKeyerHasher, StoreKeyerKey, StoreKeyerPrefix};
35
36#[derive(Clone)]
39pub struct StoreKVPool {
40 pool: Arc<RwLock<HashMap<StoreKVKey, StoreKVBox>>>,
41 kv_store_config: Arc<crate::config::ConfigStoreKV>,
42 store_access_lock: Arc<RwLock<()>>,
43 store_acquire_lock: Arc<Mutex<()>>,
44 store_flush_lock: Arc<Mutex<()>>,
45}
46
47pub struct StoreKVBuilder {
48 kv_store_config: Arc<crate::config::ConfigStoreKV>,
49}
50
51pub struct StoreKV {
52 database: DB,
53 last_used: Arc<RwLock<SystemTime>>,
54 last_flushed: Arc<RwLock<SystemTime>>,
55 pub lock: RwLock<bool>,
56 kv_store_config: Arc<crate::config::ConfigStoreKV>,
57}
58
59pub struct StoreKVActionBuilder<'build> {
60 pub kv_pool: &'build StoreKVPool,
61}
62
63pub struct StoreKVAction<'a> {
64 store: Option<StoreKVBox>,
65 bucket: StoreItemPart<'a>,
66}
67
68#[derive(PartialEq, Eq, Hash, Clone, Copy)]
69pub struct StoreKVKey {
70 collection_hash: StoreKVAtom,
71}
72
73#[derive(PartialEq)]
74pub enum StoreKVAcquireMode {
75 Any,
76 OpenOnly,
77}
78
79type StoreKVAtom = u32;
80type StoreKVBox = Arc<StoreKV>;
81
82const ATOM_HASH_RADIX: usize = 16;
83
84impl StoreKVPool {
85 pub fn new(kv_store_config: Arc<crate::config::ConfigStoreKV>) -> Self {
86 Self {
87 pool: Arc::default(),
88 kv_store_config,
89 store_access_lock: Arc::default(),
90 store_acquire_lock: Arc::default(),
91 store_flush_lock: Arc::default(),
92 }
93 }
94
95 pub fn count(&self) -> usize {
96 self.pool.read().unwrap().len()
97 }
98
99 pub fn lock_read_access<'a>(&'a self) -> RwLockReadGuard<'a, ()> {
100 self.store_access_lock.read().unwrap()
101 }
102
103 pub fn lock_write_access<'a>(&'a self) -> RwLockWriteGuard<'a, ()> {
104 self.store_access_lock.write().unwrap()
105 }
106
107 pub fn acquire(
108 &self,
109 mode: StoreKVAcquireMode,
110 collection: impl AsRef<str>,
111 ) -> Result<Option<StoreKVBox>, ()> {
112 let collection = collection.as_ref();
113 let pool_key = StoreKVKey::from_str(collection);
114
115 let _acquire = self.store_acquire_lock.lock().unwrap();
118
119 let store_pool_read = self.pool.read().unwrap();
121
122 if let Some(store_kv) = store_pool_read.get(&pool_key) {
123 Self::proceed_acquire_cache("kv", collection, pool_key, store_kv).map(Some)
124 } else {
125 tracing::info!(
126 "kv store not in pool for collection: {} {}, opening it",
127 collection,
128 pool_key
129 );
130
131 drop(store_pool_read);
134
135 let can_open_db = if mode == StoreKVAcquireMode::OpenOnly {
137 self.kv_store_config.path(pool_key.collection_hash).exists()
138 } else {
139 true
140 };
141
142 let builder = StoreKVBuilder {
143 kv_store_config: Arc::clone(&self.kv_store_config),
144 };
145
146 if can_open_db {
150 Self::proceed_acquire_open("kv", collection, pool_key, &self.pool, &builder)
151 .map(Some)
152 } else {
153 Ok(None)
154 }
155 }
156 }
157
158 fn close(&self, collection_hash: StoreKVAtom) {
159 tracing::debug!(
160 "closing key-value database for collection: <{:x}>",
161 collection_hash
162 );
163
164 let mut store_pool_write = self.pool.write().unwrap();
165
166 let collection_target = StoreKVKey::from_atom(collection_hash);
167
168 store_pool_write.remove(&collection_target);
169 }
170
171 pub fn janitor(&self) {
172 Self::proceed_janitor(
173 "kv",
174 &self.pool,
175 self.kv_store_config.pool.inactive_after,
176 &self.store_access_lock,
177 )
178 }
179
180 pub fn backup(&self, path: &Path) -> Result<(), io::Error> {
181 tracing::debug!("backing up all kv stores to path: {:?}", path);
182
183 fs::create_dir_all(path)?;
185
186 self.dump_action(
188 "backup",
189 &self.kv_store_config.path,
190 path,
191 &Self::backup_item,
192 )
193 }
194
195 pub fn restore(&self, path: &Path) -> Result<(), io::Error> {
196 tracing::debug!("restoring all kv stores from path: {:?}", path);
197
198 self.dump_action(
200 "restore",
201 path,
202 &self.kv_store_config.path,
203 &Self::restore_item,
204 )
205 }
206
207 pub fn flush(&self, force: bool) {
208 tracing::debug!("scanning for kv store pool items to flush to disk");
209
210 let _flush = self.store_flush_lock.lock().unwrap();
213
214 let mut keys_flush: Vec<StoreKVKey> = Vec::new();
216
217 {
218 let _access = self.store_access_lock.write().unwrap();
221
222 let store_pool_read = self.pool.read().unwrap();
223
224 for (key, store) in &*store_pool_read {
225 let not_flushed_for = store
230 .last_flushed
231 .read()
232 .unwrap()
233 .elapsed()
234 .unwrap_or_else(|err| {
235 tracing::error!(
236 "kv key: {} last flush duration clock issue, zeroing: {}",
237 key,
238 err
239 );
240
241 Duration::from_secs(0)
243 })
244 .as_secs();
245
246 if force || not_flushed_for >= self.kv_store_config.database.flush_after {
247 tracing::info!(
248 "kv key: {} not flushed for: {} seconds, may flush",
249 key,
250 not_flushed_for
251 );
252
253 keys_flush.push(*key);
254 } else {
255 tracing::debug!(
256 "kv key: {} not flushed for: {} seconds, no flush",
257 key,
258 not_flushed_for
259 );
260 }
261 }
262 }
263
264 if keys_flush.is_empty() {
266 tracing::info!("no kv store pool items need to be flushed at the moment");
267
268 return;
269 }
270
271 let mut count_flushed = 0;
273
274 {
275 for key in &keys_flush {
276 {
277 let _access = self.store_access_lock.write().unwrap();
280
281 if let Some(store) = self.pool.read().unwrap().get(key) {
282 tracing::debug!("kv key: {} flush started", key);
283
284 if let Err(err) = store.flush() {
285 tracing::error!("kv key: {} flush failed: {}", key, err);
286 } else {
287 count_flushed += 1;
288
289 tracing::debug!("kv key: {} flush complete", key);
290 }
291
292 *store.last_flushed.write().unwrap() = SystemTime::now();
294 }
295 }
296
297 thread::yield_now();
299 }
300 }
301
302 tracing::info!(
303 "done scanning for kv store pool items to flush to disk (flushed: {})",
304 count_flushed
305 );
306 }
307
308 #[allow(clippy::type_complexity)]
309 fn dump_action(
310 &self,
311 action: &str,
312 read_path: &Path,
313 write_path: &Path,
314 fn_item: &dyn Fn(&Self, &Path, &Path, &str) -> Result<(), io::Error>,
315 ) -> Result<(), io::Error> {
316 for collection in fs::read_dir(read_path)? {
318 let collection = collection?;
319
320 if let (Ok(collection_file_type), Some(collection_name)) =
322 (collection.file_type(), collection.file_name().to_str())
323 {
324 if collection_file_type.is_dir() {
325 tracing::debug!("kv collection ongoing {}: {}", action, collection_name);
326
327 fn_item(self, write_path, &collection.path(), collection_name)?;
328 }
329 }
330 }
331
332 Ok(())
333 }
334
335 fn backup_item(
336 &self,
337 backup_path: &Path,
338 _origin_path: &Path,
339 collection_name: &str,
340 ) -> Result<(), io::Error> {
341 let _access = self.store_access_lock.write().unwrap();
344
345 let kv_backup_path = backup_path.join(collection_name);
347
348 tracing::debug!(
349 "kv collection: {} backing up to path: {:?}",
350 collection_name,
351 kv_backup_path
352 );
353
354 if kv_backup_path.exists() {
356 fs::remove_dir_all(&kv_backup_path)?;
357 }
358
359 fs::create_dir_all(backup_path.join(collection_name))?;
361
362 if let Ok(collection_radix) = RadixNum::from_str(collection_name, ATOM_HASH_RADIX) {
365 if let Ok(collection_hash) = collection_radix.as_decimal() {
366 let origin_kv = StoreKVBuilder {
367 kv_store_config: Arc::clone(&self.kv_store_config),
368 }
369 .open(collection_hash as StoreKVAtom)
370 .map_err(|_| io::Error::other("database open failure"))?;
371
372 let kv_backup_options = DBBackupEngineOptions::new(&kv_backup_path)
374 .map_err(|_| io::Error::other("backup engine options acquire failure"))?;
375 let kv_backup_environment = DBEnv::new()
376 .map_err(|_| io::Error::other("backup engine environment acquire failure"))?;
377
378 let mut kv_backup_engine =
379 DBBackupEngine::open(&kv_backup_options, &kv_backup_environment)
380 .map_err(|_| io::Error::other("backup engine failure"))?;
381
382 kv_backup_engine
384 .create_new_backup(&origin_kv)
385 .map_err(|_| io::Error::other("database backup failure"))?;
386
387 tracing::info!(
388 "kv collection: {} backed up to path: {:?}",
389 collection_name,
390 kv_backup_path
391 );
392 }
393 }
394
395 Ok(())
396 }
397
398 fn restore_item(
399 &self,
400 _backup_path: &Path,
401 origin_path: &Path,
402 collection_name: &str,
403 ) -> Result<(), io::Error> {
404 let _access = self.store_access_lock.write().unwrap();
407
408 tracing::debug!(
409 "kv collection: {} restoring from path: {:?}",
410 collection_name,
411 origin_path
412 );
413
414 if let Ok(collection_radix) = RadixNum::from_str(collection_name, ATOM_HASH_RADIX) {
417 if let Ok(collection_hash) = collection_radix.as_decimal() {
418 self.close(collection_hash as StoreKVAtom);
420
421 let kv_path = self.kv_store_config.path(collection_hash as StoreKVAtom);
423
424 if kv_path.exists() {
426 fs::remove_dir_all(&kv_path)?;
427 }
428
429 fs::create_dir_all(&kv_path)?;
431
432 let kv_backup_options = DBBackupEngineOptions::new(&origin_path)
434 .map_err(|_| io::Error::other("backup engine options acquire failure"))?;
435 let kv_backup_environment = DBEnv::new()
436 .map_err(|_| io::Error::other("backup engine environment acquire failure"))?;
437
438 let mut kv_backup_engine =
439 DBBackupEngine::open(&kv_backup_options, &kv_backup_environment)
440 .map_err(|_| io::Error::other("backup engine failure"))?;
441
442 kv_backup_engine
443 .restore_from_latest_backup(&kv_path, &kv_path, &DBRestoreOptions::default())
444 .map_err(|_| io::Error::other("database restore failure"))?;
445
446 tracing::info!(
447 "kv collection: {} restored to path: {:?} from backup: {:?}",
448 collection_name,
449 kv_path,
450 origin_path
451 );
452 }
453 }
454
455 Ok(())
456 }
457}
458
459impl StoreGenericPool<StoreKVKey, StoreKV, StoreKVBuilder> for StoreKVPool {}
460
461impl StoreKVBuilder {
462 fn open(&self, collection_hash: StoreKVAtom) -> Result<DB, DBError> {
463 tracing::debug!(
464 "opening key-value database for collection: <{:x}>",
465 collection_hash
466 );
467
468 let db_options = self.configure();
470
471 DB::open(&db_options, self.kv_store_config.path(collection_hash))
473 }
474
475 fn configure(&self) -> DBOptions {
476 tracing::debug!("configuring key-value database");
477
478 let db_conf = &self.kv_store_config.database;
479
480 let mut db_options = DBOptions::default();
482
483 db_options.create_if_missing(true);
485 db_options.set_use_fsync(false);
486 db_options.set_compaction_style(DBCompactionStyle::Level);
487 db_options.set_min_write_buffer_number(1);
488 db_options.set_max_write_buffer_number(2);
489
490 db_options.set_compression_type(if db_conf.compress {
492 DBCompressionType::Zstd
493 } else {
494 DBCompressionType::None
495 });
496
497 db_options.set_max_open_files(if let Some(value) = db_conf.max_files {
498 value as i32
499 } else {
500 -1
501 });
502
503 db_options.increase_parallelism(db_conf.parallelism as i32);
504 db_options.set_max_subcompactions(db_conf.max_compactions as u32);
505 db_options.set_max_background_jobs((db_conf.max_compactions + db_conf.max_flushes) as i32);
506 db_options.set_write_buffer_size(db_conf.write_buffer * 1024);
507
508 db_options
509 }
510}
511
512impl crate::config::ConfigStoreKV {
513 fn path(&self, collection_hash: StoreKVAtom) -> PathBuf {
514 self.path.join(format!("{:x}", collection_hash))
515 }
516}
517
518impl StoreGenericBuilder<StoreKVKey, StoreKV> for StoreKVBuilder {
519 fn build(&self, pool_key: StoreKVKey) -> Result<StoreKV, ()> {
520 self.open(pool_key.collection_hash)
521 .map(|db| {
522 let now = SystemTime::now();
523
524 StoreKV {
525 database: db,
526 last_used: Arc::new(RwLock::new(now)),
527 last_flushed: Arc::new(RwLock::new(now)),
528 lock: RwLock::new(false),
529 kv_store_config: Arc::clone(&self.kv_store_config),
530 }
531 })
532 .map_err(|err| {
533 tracing::error!("failed opening kv: {}", err);
534 })
535 }
536}
537
538impl StoreKV {
539 pub fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, DBError> {
540 self.database.get(key)
541 }
542
543 pub fn put(&self, key: &[u8], data: &[u8]) -> Result<(), DBError> {
544 let mut batch = WriteBatch::default();
545
546 batch.put(key, data);
547
548 self.do_write(batch)
549 }
550
551 pub fn delete(&self, key: &[u8]) -> Result<(), DBError> {
552 let mut batch = WriteBatch::default();
553
554 batch.delete(key);
555
556 self.do_write(batch)
557 }
558
559 fn flush(&self) -> Result<(), DBError> {
560 let mut flush_options = FlushOptions::default();
562
563 flush_options.set_wait(true);
564
565 self.database.flush_opt(&flush_options)
567 }
568
569 fn do_write(&self, batch: WriteBatch) -> Result<(), DBError> {
570 let mut write_options = WriteOptions::default();
572
573 if !self.kv_store_config.database.write_ahead_log {
575 tracing::debug!("ignoring wal for kv write");
576
577 write_options.disable_wal(true);
578 } else {
579 tracing::debug!("using wal for kv write");
580
581 write_options.disable_wal(false);
582 }
583
584 self.database.write_opt(batch, &write_options)
586 }
587}
588
589impl StoreGeneric for StoreKV {
590 fn ref_last_used(&self) -> &RwLock<SystemTime> {
591 &self.last_used
592 }
593}
594
595impl<'build> StoreKVActionBuilder<'build> {
596 pub fn access(bucket: StoreItemPart, store: Option<StoreKVBox>) -> StoreKVAction {
597 Self::build(bucket, store)
598 }
599
600 pub fn erase<T: AsRef<str>>(&self, collection: T, bucket: Option<T>) -> Result<u32, ()> {
601 self.dispatch_erase("kv", collection, bucket)
602 }
603
604 fn build(bucket: StoreItemPart, store: Option<StoreKVBox>) -> StoreKVAction {
605 StoreKVAction { store, bucket }
606 }
607}
608
609impl<'build> StoreGenericActionBuilder for StoreKVActionBuilder<'build> {
610 fn proceed_erase_collection(&self, collection_str: &str) -> Result<u32, ()> {
611 let collection_atom = StoreKeyerHasher::to_compact(collection_str);
612 let collection_path = self.kv_pool.kv_store_config.path(collection_atom);
613
614 self.kv_pool.close(collection_atom);
616
617 if collection_path.exists() {
618 tracing::debug!(
619 "kv collection store exists, erasing: {}/* at path: {:?}",
620 collection_str,
621 &collection_path
622 );
623
624 let erase_result = fs::remove_dir_all(&collection_path);
626
627 if erase_result.is_ok() {
628 tracing::debug!("done with kv collection erasure");
629
630 Ok(1)
631 } else {
632 Err(())
633 }
634 } else {
635 tracing::debug!(
636 "kv collection store does not exist, consider already erased: {}/* at path: {:?}",
637 collection_str,
638 &collection_path
639 );
640
641 Ok(0)
642 }
643 }
644
645 fn proceed_erase_bucket(&self, _collection: &str, _bucket: &str) -> Result<u32, ()> {
646 Err(())
649 }
650}
651
652impl<'a> StoreKVAction<'a> {
653 pub fn get_meta_to_value(&self, meta: StoreMetaKey) -> Result<Option<StoreMetaValue>, ()> {
657 if let Some(ref store) = self.store {
658 let store_key = StoreKeyerBuilder::meta_to_value(self.bucket.as_str(), &meta);
659
660 tracing::debug!("store get meta-to-value: {}", store_key);
661
662 match store.get(&store_key.as_bytes()) {
663 Ok(Some(value)) => {
664 tracing::debug!("got meta-to-value: {}", store_key);
665
666 Ok(if let Ok(value) = str::from_utf8(&value) {
667 match meta {
668 StoreMetaKey::IIDIncr => value
669 .parse::<StoreObjectIID>()
670 .ok()
671 .map(StoreMetaValue::IIDIncr)
672 .or(None),
673 }
674 } else {
675 None
676 })
677 }
678 Ok(None) => {
679 tracing::debug!("no meta-to-value found: {}", store_key);
680
681 Ok(None)
682 }
683 Err(err) => {
684 tracing::error!(
685 "error getting meta-to-value: {} with trace: {}",
686 store_key,
687 err
688 );
689
690 Err(())
691 }
692 }
693 } else {
694 Ok(None)
695 }
696 }
697
698 pub fn set_meta_to_value(&self, meta: StoreMetaKey, value: StoreMetaValue) -> Result<(), ()> {
699 if let Some(ref store) = self.store {
700 let store_key = StoreKeyerBuilder::meta_to_value(self.bucket.as_str(), &meta);
701
702 tracing::debug!("store set meta-to-value: {}", store_key);
703
704 let value_string = match value {
705 StoreMetaValue::IIDIncr(iid_incr) => iid_incr.to_string(),
706 };
707
708 store
709 .put(&store_key.as_bytes(), value_string.as_bytes())
710 .or(Err(()))
711 } else {
712 Err(())
713 }
714 }
715
716 pub fn get_term_to_iids(
720 &self,
721 term_hashed: StoreTermHashed,
722 ) -> Result<Option<Vec<StoreObjectIID>>, ()> {
723 if let Some(ref store) = self.store {
724 let store_key = StoreKeyerBuilder::term_to_iids(self.bucket.as_str(), term_hashed);
725
726 tracing::debug!("store get term-to-iids: {}", store_key);
727
728 match store.get(&store_key.as_bytes()) {
729 Ok(Some(value)) => {
730 tracing::debug!(
731 "got term-to-iids: {} with encoded value: {:?}",
732 store_key,
733 &*value
734 );
735
736 Self::decode_u32_list(&*value)
737 .or(Err(()))
738 .map(|value_decoded| {
739 tracing::debug!(
740 "got term-to-iids: {} with decoded value: {:?}",
741 store_key,
742 &value_decoded
743 );
744
745 Some(value_decoded)
746 })
747 }
748 Ok(None) => {
749 tracing::debug!("no term-to-iids found: {}", store_key);
750
751 Ok(None)
752 }
753 Err(err) => {
754 tracing::error!(
755 "error getting term-to-iids: {} with trace: {}",
756 store_key,
757 err
758 );
759
760 Err(())
761 }
762 }
763 } else {
764 Ok(None)
765 }
766 }
767
768 pub fn set_term_to_iids(
769 &self,
770 term_hashed: StoreTermHashed,
771 iids: &[StoreObjectIID],
772 ) -> Result<(), ()> {
773 if let Some(ref store) = self.store {
774 let store_key = StoreKeyerBuilder::term_to_iids(self.bucket.as_str(), term_hashed);
775
776 tracing::debug!("store set term-to-iids: {}", store_key);
777
778 let iids_encoded = Self::encode_u32_list(iids);
780
781 tracing::debug!(
782 "store set term-to-iids: {} with encoded value: {:?}",
783 store_key,
784 iids_encoded
785 );
786
787 store.put(&store_key.as_bytes(), &iids_encoded).or(Err(()))
788 } else {
789 Err(())
790 }
791 }
792
793 pub fn delete_term_to_iids(&self, term_hashed: StoreTermHashed) -> Result<(), ()> {
794 if let Some(ref store) = self.store {
795 let store_key = StoreKeyerBuilder::term_to_iids(self.bucket.as_str(), term_hashed);
796
797 tracing::debug!("store delete term-to-iids: {}", store_key);
798
799 store.delete(&store_key.as_bytes()).or(Err(()))
800 } else {
801 Err(())
802 }
803 }
804
805 pub fn get_oid_to_iid(&self, oid: StoreObjectOID<'a>) -> Result<Option<StoreObjectIID>, ()> {
809 if let Some(ref store) = self.store {
810 let store_key = StoreKeyerBuilder::oid_to_iid(self.bucket.as_str(), oid);
811
812 tracing::debug!("store get oid-to-iid: {}", store_key);
813
814 match store.get(&store_key.as_bytes()) {
815 Ok(Some(value)) => {
816 tracing::debug!(
817 "got oid-to-iid: {} with encoded value: {:?}",
818 store_key,
819 &*value
820 );
821
822 Self::decode_u32(&*value).or(Err(())).map(|value_decoded| {
823 tracing::debug!(
824 "got oid-to-iid: {} with decoded value: {:?}",
825 store_key,
826 &value_decoded
827 );
828
829 Some(value_decoded)
830 })
831 }
832 Ok(None) => {
833 tracing::debug!("no oid-to-iid found: {}", store_key);
834
835 Ok(None)
836 }
837 Err(err) => {
838 tracing::error!(
839 "error getting oid-to-iid: {} with trace: {}",
840 store_key,
841 err
842 );
843
844 Err(())
845 }
846 }
847 } else {
848 Ok(None)
849 }
850 }
851
852 pub fn set_oid_to_iid(&self, oid: StoreObjectOID<'a>, iid: StoreObjectIID) -> Result<(), ()> {
853 if let Some(ref store) = self.store {
854 let store_key = StoreKeyerBuilder::oid_to_iid(self.bucket.as_str(), oid);
855
856 tracing::debug!("store set oid-to-iid: {}", store_key);
857
858 let iid_encoded = Self::encode_u32(iid);
860
861 tracing::debug!(
862 "store set oid-to-iid: {} with encoded value: {:?}",
863 store_key,
864 iid_encoded
865 );
866
867 store.put(&store_key.as_bytes(), &iid_encoded).or(Err(()))
868 } else {
869 Err(())
870 }
871 }
872
873 pub fn delete_oid_to_iid(&self, oid: StoreObjectOID<'a>) -> Result<(), ()> {
874 if let Some(ref store) = self.store {
875 let store_key = StoreKeyerBuilder::oid_to_iid(self.bucket.as_str(), oid);
876
877 tracing::debug!("store delete oid-to-iid: {}", store_key);
878
879 store.delete(&store_key.as_bytes()).or(Err(()))
880 } else {
881 Err(())
882 }
883 }
884
885 pub fn get_iid_to_oid(&self, iid: StoreObjectIID) -> Result<Option<String>, ()> {
889 if let Some(ref store) = self.store {
890 let store_key = StoreKeyerBuilder::iid_to_oid(self.bucket.as_str(), iid);
891
892 tracing::debug!("store get iid-to-oid: {}", store_key);
893
894 match store.get(&store_key.as_bytes()) {
895 Ok(Some(value)) => Ok(str::from_utf8(&value).ok().map(|value| value.to_string())),
896 Ok(None) => Ok(None),
897 Err(_) => Err(()),
898 }
899 } else {
900 Ok(None)
901 }
902 }
903
904 pub fn set_iid_to_oid(&self, iid: StoreObjectIID, oid: StoreObjectOID<'a>) -> Result<(), ()> {
905 if let Some(ref store) = self.store {
906 let store_key = StoreKeyerBuilder::iid_to_oid(self.bucket.as_str(), iid);
907
908 tracing::debug!("store set iid-to-oid: {}", store_key);
909
910 store.put(&store_key.as_bytes(), oid.as_bytes()).or(Err(()))
911 } else {
912 Err(())
913 }
914 }
915
916 pub fn delete_iid_to_oid(&self, iid: StoreObjectIID) -> Result<(), ()> {
917 if let Some(ref store) = self.store {
918 let store_key = StoreKeyerBuilder::iid_to_oid(self.bucket.as_str(), iid);
919
920 tracing::debug!("store delete iid-to-oid: {}", store_key);
921
922 store.delete(&store_key.as_bytes()).or(Err(()))
923 } else {
924 Err(())
925 }
926 }
927
928 pub fn get_iid_to_terms(
932 &self,
933 iid: StoreObjectIID,
934 ) -> Result<Option<Vec<StoreTermHashed>>, ()> {
935 if let Some(ref store) = self.store {
936 let store_key = StoreKeyerBuilder::iid_to_terms(self.bucket.as_str(), iid);
937
938 tracing::debug!("store get iid-to-terms: {}", store_key);
939
940 match store.get(&store_key.as_bytes()) {
941 Ok(Some(value)) => {
942 tracing::debug!(
943 "got iid-to-terms: {} with encoded value: {:?}",
944 store_key,
945 &*value
946 );
947
948 Self::decode_u32_list(&*value)
949 .or(Err(()))
950 .map(|value_decoded| {
951 tracing::debug!(
952 "got iid-to-terms: {} with decoded value: {:?}",
953 store_key,
954 &value_decoded
955 );
956
957 if !value_decoded.is_empty() {
958 Some(value_decoded)
959 } else {
960 None
961 }
962 })
963 }
964 Ok(None) => Ok(None),
965 Err(_) => Err(()),
966 }
967 } else {
968 Ok(None)
969 }
970 }
971
972 pub fn set_iid_to_terms(
973 &self,
974 iid: StoreObjectIID,
975 terms_hashed: &[StoreTermHashed],
976 ) -> Result<(), ()> {
977 if let Some(ref store) = self.store {
978 let store_key = StoreKeyerBuilder::iid_to_terms(self.bucket.as_str(), iid);
979
980 tracing::debug!("store set iid-to-terms: {}", store_key);
981
982 let terms_hashed_encoded = Self::encode_u32_list(terms_hashed);
984
985 tracing::debug!(
986 "store set iid-to-terms: {} with encoded value: {:?}",
987 store_key,
988 terms_hashed_encoded
989 );
990
991 store
992 .put(&store_key.as_bytes(), &terms_hashed_encoded)
993 .or(Err(()))
994 } else {
995 Err(())
996 }
997 }
998
999 pub fn delete_iid_to_terms(&self, iid: StoreObjectIID) -> Result<(), ()> {
1000 if let Some(ref store) = self.store {
1001 let store_key = StoreKeyerBuilder::iid_to_terms(self.bucket.as_str(), iid);
1002
1003 tracing::debug!("store delete iid-to-terms: {}", store_key);
1004
1005 store.delete(&store_key.as_bytes()).or(Err(()))
1006 } else {
1007 Err(())
1008 }
1009 }
1010
1011 pub fn batch_flush_bucket(
1012 &self,
1013 iid: StoreObjectIID,
1014 oid: StoreObjectOID<'a>,
1015 iid_terms_hashed: &[StoreTermHashed],
1016 ) -> Result<u32, ()> {
1017 let mut count = 0;
1018
1019 tracing::debug!(
1020 "store batch flush bucket: {} with hashed terms: {:?}",
1021 iid,
1022 iid_terms_hashed
1023 );
1024
1025 match (
1027 self.delete_oid_to_iid(oid),
1028 self.delete_iid_to_oid(iid),
1029 self.delete_iid_to_terms(iid),
1030 ) {
1031 (Ok(_), Ok(_), Ok(_)) => {
1032 for iid_term in iid_terms_hashed {
1034 if let Ok(Some(mut iid_term_iids)) = self.get_term_to_iids(*iid_term) {
1035 if iid_term_iids.contains(&iid) {
1036 count += 1;
1037
1038 iid_term_iids.retain(|cur_iid| cur_iid != &iid);
1040 }
1041
1042 let is_ok = if iid_term_iids.is_empty() {
1043 self.delete_term_to_iids(*iid_term).is_ok()
1044 } else {
1045 self.set_term_to_iids(*iid_term, &iid_term_iids).is_ok()
1046 };
1047
1048 if !is_ok {
1049 return Err(());
1050 }
1051 }
1052 }
1053
1054 Ok(count)
1055 }
1056 _ => Err(()),
1057 }
1058 }
1059
1060 pub fn batch_truncate_object(
1061 &self,
1062 term_hashed: StoreTermHashed,
1063 term_iids_drain: Drain<StoreObjectIID>,
1064 ) -> Result<u32, ()> {
1065 let mut count = 0;
1066
1067 for term_iid_drain in term_iids_drain {
1068 tracing::debug!("store batch truncate object iid: {}", term_iid_drain);
1069
1070 if let Ok(Some(mut term_iid_drain_terms)) = self.get_iid_to_terms(term_iid_drain) {
1072 count += 1;
1073
1074 term_iid_drain_terms.retain(|cur_term| cur_term != &term_hashed);
1075
1076 if term_iid_drain_terms.is_empty() {
1078 if let Ok(Some(term_iid_drain_oid)) = self.get_iid_to_oid(term_iid_drain) {
1080 if self
1081 .batch_flush_bucket(term_iid_drain, &term_iid_drain_oid, &Vec::new())
1082 .is_err()
1083 {
1084 tracing::error!(
1085 "failed executing store batch truncate object batch-flush-bucket"
1086 );
1087 }
1088 } else {
1089 tracing::error!("failed getting store batch truncate object iid-to-oid");
1090 }
1091 } else {
1092 if self
1094 .set_iid_to_terms(term_iid_drain, &term_iid_drain_terms)
1095 .is_err()
1096 {
1097 tracing::error!("failed setting store batch truncate object iid-to-terms");
1098 }
1099 }
1100 }
1101 }
1102
1103 Ok(count)
1104 }
1105
1106 pub fn batch_erase_bucket(&self) -> Result<u32, ()> {
1107 if let Some(ref store) = self.store {
1108 let (k_meta_to_value, k_term_to_iids, k_oid_to_iid, k_iid_to_oid, k_iid_to_terms) = (
1110 StoreKeyerBuilder::meta_to_value(self.bucket.as_str(), &StoreMetaKey::IIDIncr),
1111 StoreKeyerBuilder::term_to_iids(self.bucket.as_str(), 0),
1112 StoreKeyerBuilder::oid_to_iid(self.bucket.as_str(), ""),
1113 StoreKeyerBuilder::iid_to_oid(self.bucket.as_str(), 0),
1114 StoreKeyerBuilder::iid_to_terms(self.bucket.as_str(), 0),
1115 );
1116
1117 let key_prefixes: [StoreKeyerPrefix; 5] = [
1118 k_meta_to_value.as_prefix(),
1119 k_term_to_iids.as_prefix(),
1120 k_oid_to_iid.as_prefix(),
1121 k_iid_to_oid.as_prefix(),
1122 k_iid_to_terms.as_prefix(),
1123 ];
1124
1125 for key_prefix in &key_prefixes {
1127 tracing::debug!(
1128 "store batch erase bucket: {} for prefix: {:?}",
1129 self.bucket.as_str(),
1130 key_prefix
1131 );
1132
1133 let key_prefix_start: StoreKeyerKey = [
1136 key_prefix[0],
1137 key_prefix[1],
1138 key_prefix[2],
1139 key_prefix[3],
1140 key_prefix[4],
1141 0,
1142 0,
1143 0,
1144 0,
1145 ];
1146 let key_prefix_end: StoreKeyerKey = [
1147 key_prefix[0],
1148 key_prefix[1],
1149 key_prefix[2],
1150 key_prefix[3],
1151 key_prefix[4],
1152 255,
1153 255,
1154 255,
1155 255,
1156 ];
1157
1158 let mut batch = WriteBatch::default();
1160
1161 batch.delete_range(&key_prefix_start, &key_prefix_end);
1162
1163 if let Err(err) = store.do_write(batch) {
1165 tracing::error!(
1166 "failed in store batch erase bucket: {} with error: {}",
1167 self.bucket.as_str(),
1168 err
1169 );
1170 } else {
1171 store.delete(&key_prefix_end).ok();
1175
1176 tracing::debug!(
1177 "succeeded in store batch erase bucket: {}",
1178 self.bucket.as_str()
1179 );
1180 }
1181 }
1182
1183 tracing::info!(
1184 "done processing store batch erase bucket: {}",
1185 self.bucket.as_str()
1186 );
1187
1188 Ok(1)
1189 } else {
1190 Err(())
1191 }
1192 }
1193
1194 fn encode_u32(decoded: u32) -> [u8; 4] {
1195 let mut encoded = [0; 4];
1196
1197 LittleEndian::write_u32(&mut encoded, decoded);
1198
1199 encoded
1200 }
1201
1202 fn decode_u32(encoded: &[u8]) -> Result<u32, ()> {
1203 Cursor::new(encoded).read_u32::<LittleEndian>().or(Err(()))
1204 }
1205
1206 fn encode_u32_list(decoded: &[u32]) -> Vec<u8> {
1207 let mut encoded = Vec::with_capacity(decoded.len() * 4);
1210
1211 for decoded_item in decoded {
1212 encoded.extend(&Self::encode_u32(*decoded_item))
1213 }
1214
1215 encoded
1216 }
1217
1218 fn decode_u32_list(encoded: &[u8]) -> Result<Vec<u32>, ()> {
1219 let mut decoded = Vec::with_capacity(encoded.len() / 4);
1222
1223 for encoded_chunk in encoded.chunks(4) {
1224 if let Ok(decoded_chunk) = Self::decode_u32(encoded_chunk) {
1225 decoded.push(decoded_chunk);
1226 } else {
1227 return Err(());
1228 }
1229 }
1230
1231 Ok(decoded)
1232 }
1233}
1234
1235impl StoreKVKey {
1236 pub fn from_atom(collection_hash: StoreKVAtom) -> StoreKVKey {
1237 StoreKVKey { collection_hash }
1238 }
1239
1240 #[allow(clippy::should_implement_trait)]
1241 pub fn from_str(collection_str: &str) -> StoreKVKey {
1242 StoreKVKey {
1243 collection_hash: StoreKeyerHasher::to_compact(collection_str),
1244 }
1245 }
1246}
1247
1248impl fmt::Display for StoreKVKey {
1249 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1250 write!(f, "<{:x}>", self.collection_hash)
1251 }
1252}
1253
1254#[cfg(test)]
1255mod tests {
1256 use super::*;
1257
1258 #[test]
1259 fn it_acquires_database() {
1260 let kv_store_config = test_kv_store_config();
1261 let kv_pool = StoreKVPool::new(kv_store_config);
1262
1263 assert!(kv_pool.acquire(StoreKVAcquireMode::Any, "c:test:1").is_ok());
1264 }
1265
1266 #[test]
1267 fn it_janitors_database() {
1268 let kv_store_config = test_kv_store_config();
1269 let kv_pool = StoreKVPool::new(kv_store_config);
1270
1271 kv_pool.janitor();
1272 }
1273
1274 #[test]
1275 fn it_proceeds_primitives() {
1276 let kv_store_config = test_kv_store_config();
1277 let kv_pool = StoreKVPool::new(kv_store_config);
1278
1279 let store = kv_pool
1280 .acquire(StoreKVAcquireMode::Any, "c:test:2")
1281 .unwrap()
1282 .unwrap();
1283
1284 assert!(store.get(&[0]).is_ok());
1285 assert!(store.put(&[0], &[1, 0, 0, 0]).is_ok());
1286 assert!(store.delete(&[0]).is_ok());
1287 }
1288
1289 #[test]
1290 fn it_proceeds_actions() {
1291 let kv_store_config = test_kv_store_config();
1292 let kv_pool = StoreKVPool::new(kv_store_config);
1293
1294 let store = kv_pool
1295 .acquire(StoreKVAcquireMode::Any, "c:test:3")
1296 .unwrap();
1297 let action =
1298 StoreKVActionBuilder::access(StoreItemPart::from_str("b:test:3").unwrap(), store);
1299
1300 assert!(action.get_meta_to_value(StoreMetaKey::IIDIncr).is_ok());
1301 assert!(
1302 action
1303 .set_meta_to_value(StoreMetaKey::IIDIncr, StoreMetaValue::IIDIncr(1))
1304 .is_ok()
1305 );
1306
1307 assert!(action.get_term_to_iids(1).is_ok());
1308 assert!(action.set_term_to_iids(1, &[0, 1, 2]).is_ok());
1309 assert!(action.delete_term_to_iids(1).is_ok());
1310
1311 assert!(action.get_oid_to_iid(&"s".to_string()).is_ok());
1312 assert!(action.set_oid_to_iid(&"s".to_string(), 4).is_ok());
1313 assert!(action.delete_oid_to_iid(&"s".to_string()).is_ok());
1314
1315 assert!(action.get_iid_to_oid(4).is_ok());
1316 assert!(action.set_iid_to_oid(4, &"s".to_string()).is_ok());
1317 assert!(action.delete_iid_to_oid(4).is_ok());
1318
1319 assert!(action.get_iid_to_terms(4).is_ok());
1320 assert!(action.set_iid_to_terms(4, &[45402]).is_ok());
1321 assert!(action.delete_iid_to_terms(4).is_ok());
1322 }
1323
1324 #[test]
1325 fn it_encodes_atom() {
1326 assert_eq!(StoreKVAction::encode_u32(0), [0, 0, 0, 0]);
1327 assert_eq!(StoreKVAction::encode_u32(1), [1, 0, 0, 0]);
1328 assert_eq!(StoreKVAction::encode_u32(45402), [90, 177, 0, 0]);
1329 }
1330
1331 #[test]
1332 fn it_decodes_atom() {
1333 assert_eq!(StoreKVAction::decode_u32(&[0, 0, 0, 0]), Ok(0));
1334 assert_eq!(StoreKVAction::decode_u32(&[1, 0, 0, 0]), Ok(1));
1335 assert_eq!(StoreKVAction::decode_u32(&[90, 177, 0, 0]), Ok(45402));
1336 }
1337
1338 #[test]
1339 fn it_encodes_atom_list() {
1340 assert_eq!(
1341 StoreKVAction::encode_u32_list(&[0, 2, 3]),
1342 [0, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0]
1343 );
1344 assert_eq!(StoreKVAction::encode_u32_list(&[45402]), [90, 177, 0, 0]);
1345 }
1346
1347 #[test]
1348 fn it_decodes_atom_list() {
1349 assert_eq!(
1350 StoreKVAction::decode_u32_list(&[0, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0]),
1351 Ok(vec![0, 2, 3])
1352 );
1353 assert_eq!(
1354 StoreKVAction::decode_u32_list(&[90, 177, 0, 0]),
1355 Ok(vec![45402])
1356 );
1357 }
1358
1359 fn test_kv_store_config() -> Arc<crate::config::ConfigStoreKV> {
1360 Arc::new(
1361 config::Config::builder()
1362 .add_source(config::File::from_str(
1363 crate::config::tests::defaults_toml(),
1364 config::FileFormat::Toml,
1365 ))
1366 .build()
1367 .unwrap()
1368 .get::<crate::config::ConfigStoreKV>("store.kv")
1369 .unwrap(),
1370 )
1371 }
1372}
1373
1374#[cfg(all(feature = "benchmark", test))]
1375mod benches {
1376 extern crate test;
1377
1378 use super::*;
1379 use test::Bencher;
1380
1381 #[bench]
1382 fn bench_encode_atom(b: &mut Bencher) {
1383 b.iter(|| StoreKVAction::encode_u32(0));
1384 }
1385
1386 #[bench]
1387 fn bench_decode_atom(b: &mut Bencher) {
1388 let encoded_atom = [0, 0, 0, 0];
1389
1390 b.iter(|| StoreKVAction::decode_u32(&encoded_atom));
1391 }
1392
1393 #[bench]
1394 fn bench_encode_atom_list(b: &mut Bencher) {
1395 let atom_list = [0, 2, 3];
1396
1397 b.iter(|| StoreKVAction::encode_u32_list(&atom_list));
1398 }
1399
1400 #[bench]
1401 fn bench_decode_atom_list(b: &mut Bencher) {
1402 let encoded_atom_list = [0, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0];
1403
1404 b.iter(|| StoreKVAction::decode_u32_list(&encoded_atom_list));
1405 }
1406}
1407
1408impl fmt::Debug for StoreKVPool {
1411 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1412 use crate::util::fmt::{AsPrettyMutex, AsPrettyRwLock};
1413
1414 let Self {
1416 pool,
1417 store_access_lock,
1418 store_acquire_lock,
1419 store_flush_lock,
1420 kv_store_config: _kv_store_config,
1423 } = self;
1424
1425 f.debug_struct("StoreKVPool")
1426 .field("pool", &AsPrettyRwLock(pool))
1427 .field("store_access_lock", &AsPrettyRwLock(store_access_lock))
1428 .field("store_acquire_lock", &AsPrettyMutex(store_acquire_lock))
1429 .field("store_flush_lock", &AsPrettyMutex(store_flush_lock))
1430 .finish_non_exhaustive()
1431 }
1432}
1433
1434impl fmt::Debug for StoreKVKey {
1435 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1436 fmt::Display::fmt(&self, f)
1437 }
1438}
1439
1440impl fmt::Debug for StoreKV {
1441 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1442 use crate::util::fmt::AsPrettyRwLock;
1443
1444 let Self {
1446 database,
1447 last_used,
1448 last_flushed,
1449 lock,
1450 kv_store_config: _kv_store_config,
1453 } = self;
1454
1455 f.debug_struct("StoreKV")
1456 .field("database", database)
1457 .field("last_used", &AsPrettyRwLock(last_used))
1458 .field("last_flushed", &AsPrettyRwLock(last_flushed))
1459 .field("lock", &AsPrettyRwLock(lock))
1460 .finish_non_exhaustive()
1461 }
1462}