1use crate::{cache::CachePolicy, db::DB, errors::StoreError};
2
3use super::prelude::{Cache, DbKey, DbWriter};
4use kaspa_utils::mem_size::MemSizeEstimator;
5use rocksdb::{Direction, IterateBounds, IteratorMode, ReadOptions};
6use serde::{Serialize, de::DeserializeOwned};
7use std::{collections::hash_map::RandomState, error::Error, hash::BuildHasher, marker::PhantomData, sync::Arc};
8
9#[derive(Clone)]
33pub struct CachedDbAccess<TKey, TData, S = RandomState, TLegacy = TData>
34where
35 TKey: Clone + std::hash::Hash + Eq + Send + Sync,
36 TData: Clone + Send + Sync + MemSizeEstimator,
37{
38 db: Arc<DB>,
39
40 cache: Cache<TKey, TData, S>,
42
43 prefix: Vec<u8>,
45
46 version_suffix: Option<u8>,
50
51 _legacy: PhantomData<fn() -> TLegacy>,
54}
55
56pub type KeyDataResult<TData> = Result<(Box<[u8]>, TData), Box<dyn Error>>;
57
58impl<TKey, TData, S, TLegacy> CachedDbAccess<TKey, TData, S, TLegacy>
59where
60 TKey: Clone + std::hash::Hash + Eq + Send + Sync,
61 TData: Clone + Send + Sync + MemSizeEstimator,
62 S: BuildHasher + Default,
63{
64 pub fn new(db: Arc<DB>, cache_policy: CachePolicy, prefix: Vec<u8>) -> Self {
65 Self { db, cache: Cache::new(cache_policy), prefix, version_suffix: None, _legacy: PhantomData }
66 }
67
68 pub fn new_with_version_suffix(db: Arc<DB>, cache_policy: CachePolicy, prefix: Vec<u8>, version_suffix: u8) -> Self {
83 Self { db, cache: Cache::new(cache_policy), prefix, version_suffix: Some(version_suffix), _legacy: PhantomData }
84 }
85
86 pub fn read_from_cache(&self, key: TKey) -> Option<TData>
87 where
88 TKey: Copy + AsRef<[u8]>,
89 {
90 self.cache.get(&key)
91 }
92
93 pub fn has(&self, key: TKey) -> Result<bool, StoreError>
94 where
95 TKey: Clone + AsRef<[u8]>,
96 {
97 if self.cache.contains_key(&key) {
98 return Ok(true);
99 }
100 if self.version_suffix.is_some() {
101 let scan_prefix = DbKey::new(&self.prefix, key);
108 let mut read_opts = ReadOptions::default();
109 read_opts.set_iterate_range(rocksdb::PrefixRange(scan_prefix.as_ref()));
110 let mut iter = self.db.iterator_opt(IteratorMode::From(scan_prefix.as_ref(), Direction::Forward), read_opts);
111 return match iter.next() {
112 Some(Ok(_)) => Ok(true),
113 Some(Err(e)) => Err(e.into()),
114 None => Ok(false),
115 };
116 }
117 Ok(self.db.get_pinned(DbKey::new(&self.prefix, key))?.is_some())
118 }
119
120 pub fn read(&self, key: TKey) -> Result<TData, StoreError>
121 where
122 TKey: Clone + AsRef<[u8]> + ToString,
123 TData: DeserializeOwned,
124 TLegacy: DeserializeOwned + Into<TData>,
125 {
126 if let Some(data) = self.cache.get(&key) {
127 return Ok(data);
128 }
129 if let Some(suffix) = self.version_suffix {
130 return self.read_versioned(key, suffix);
131 }
132 let db_key = DbKey::new(&self.prefix, key.clone());
133 if let Some(slice) = self.db.get_pinned(&db_key)? {
134 let data: TData = bincode::deserialize(&slice)?;
135 self.cache.insert(key, data.clone());
136 Ok(data)
137 } else {
138 Err(StoreError::KeyNotFound(db_key))
139 }
140 }
141
142 fn read_versioned(&self, key: TKey, suffix: u8) -> Result<TData, StoreError>
189 where
190 TKey: Clone + AsRef<[u8]> + ToString,
191 TData: DeserializeOwned,
192 TLegacy: DeserializeOwned + Into<TData>,
193 {
194 let scan_prefix = DbKey::new(&self.prefix, key.clone());
197 let scan_bytes = scan_prefix.as_ref();
198 let scan_len = scan_bytes.len();
199
200 let mut read_opts = ReadOptions::default();
201 read_opts.set_iterate_range(rocksdb::PrefixRange(scan_bytes));
202 let mut iter = self.db.iterator_opt(IteratorMode::From(scan_bytes, Direction::Forward), read_opts);
203
204 let Some(first) = iter.next() else {
207 return Err(StoreError::KeyNotFound(scan_prefix));
208 };
209 let (found_key, found_bytes) = first?;
210
211 let data: TData = if found_key.len() == scan_len {
212 let legacy_value: TLegacy = bincode::deserialize(&found_bytes)?;
215 legacy_value.into()
216 } else if found_key.len() == scan_len + 1 {
217 let tail = found_key[scan_len];
222 if tail != suffix {
223 return Err(StoreError::DataInconsistency(format!(
224 "unexpected version byte {tail:#x} for key {scan_prefix}, expected {suffix:#x}"
225 )));
226 }
227 bincode::deserialize(&found_bytes)?
228 } else {
229 return Err(StoreError::DataInconsistency(format!(
231 "unexpected row key length {} for {scan_prefix}, expected {} (legacy) or {} (versioned)",
232 found_key.len(),
233 scan_len,
234 scan_len + 1,
235 )));
236 };
237
238 self.cache.insert(key, data.clone());
239 Ok(data)
240 }
241
242 pub fn has_with_fallback(&self, fallback_prefix: &[u8], key: TKey) -> Result<bool, StoreError>
253 where
254 TKey: Clone + AsRef<[u8]>,
255 {
256 if self.cache.contains_key(&key) {
257 Ok(true)
258 } else {
259 let db_key = DbKey::new(&self.prefix, key.clone());
260 if self.db.get_pinned(&db_key)?.is_some() {
261 Ok(true)
262 } else {
263 let db_key = DbKey::new(fallback_prefix, key.clone());
264 Ok(self.db.get_pinned(&db_key)?.is_some())
265 }
266 }
267 }
268
269 pub fn read_with_fallback<TFallbackDeser>(&self, fallback_prefix: &[u8], key: TKey) -> Result<TData, StoreError>
272 where
273 TKey: Clone + AsRef<[u8]> + ToString,
274 TData: DeserializeOwned,
275 TFallbackDeser: DeserializeOwned + Into<TData>,
276 {
277 if let Some(data) = self.cache.get(&key) {
278 Ok(data)
279 } else {
280 let db_key = DbKey::new(&self.prefix, key.clone());
281 if let Some(slice) = self.db.get_pinned(&db_key)? {
282 let data: TData = bincode::deserialize(&slice)?;
283 self.cache.insert(key, data.clone());
284 Ok(data)
285 } else {
286 let db_key = DbKey::new(fallback_prefix, key.clone());
287 if let Some(slice) = self.db.get_pinned(&db_key)? {
288 let data: TFallbackDeser = bincode::deserialize(&slice)?;
289 let data: TData = data.into();
290 self.cache.insert(key, data.clone());
291 Ok(data)
292 } else {
293 Err(StoreError::KeyNotFound(db_key))
294 }
295 }
296 }
297 }
298
299 pub fn iterator(&self) -> impl Iterator<Item = KeyDataResult<TData>> + '_
307 where
308 TKey: Clone + AsRef<[u8]>,
309 TData: DeserializeOwned, {
311 let prefix_key = DbKey::prefix_only(&self.prefix);
312 let mut read_opts = ReadOptions::default();
313 read_opts.set_iterate_range(rocksdb::PrefixRange(prefix_key.as_ref()));
314 self.db.iterator_opt(IteratorMode::From(prefix_key.as_ref(), Direction::Forward), read_opts).map(move |iter_result| {
315 match iter_result {
316 Ok((key, data_bytes)) => match bincode::deserialize(&data_bytes) {
317 Ok(data) => Ok((key[prefix_key.prefix_len()..].into(), data)),
318 Err(e) => Err(e.into()),
319 },
320 Err(e) => Err(e.into()),
321 }
322 })
323 }
324
325 pub fn write(&self, mut writer: impl DbWriter, key: TKey, data: TData) -> Result<(), StoreError>
326 where
327 TKey: Clone + AsRef<[u8]>,
328 TData: Serialize,
329 {
330 let bin_data = bincode::serialize(&data)?;
331 self.cache.insert(key.clone(), data);
332 writer.put(self.build_write_key(key), bin_data)?;
333 Ok(())
334 }
335
336 pub fn write_many(
337 &self,
338 mut writer: impl DbWriter,
339 iter: &mut (impl Iterator<Item = (TKey, TData)> + Clone),
340 ) -> Result<(), StoreError>
341 where
342 TKey: Clone + AsRef<[u8]>,
343 TData: Serialize,
344 {
345 let iter_clone = iter.clone();
346 self.cache.insert_many(iter);
347 for (key, data) in iter_clone {
348 let bin_data = bincode::serialize(&data)?;
349 writer.put(self.build_write_key(key), bin_data)?;
350 }
351 Ok(())
352 }
353
354 pub fn write_many_without_cache(
356 &self,
357 mut writer: impl DbWriter,
358 iter: &mut impl Iterator<Item = (TKey, TData)>,
359 ) -> Result<(), StoreError>
360 where
361 TKey: Clone + AsRef<[u8]>,
362 TData: Serialize,
363 {
364 for (key, data) in iter {
365 let bin_data = bincode::serialize(&data)?;
366 writer.put(self.build_write_key(key), bin_data)?;
367 }
368 self.cache.remove_all();
370 Ok(())
371 }
372
373 pub fn delete(&self, mut writer: impl DbWriter, key: TKey) -> Result<(), StoreError>
374 where
375 TKey: Clone + AsRef<[u8]>,
376 {
377 self.cache.remove(&key);
378 self.delete_both_layouts(&mut writer, key)
379 }
380
381 pub fn delete_many(&self, mut writer: impl DbWriter, key_iter: &mut (impl Iterator<Item = TKey> + Clone)) -> Result<(), StoreError>
382 where
383 TKey: Clone + AsRef<[u8]>,
384 {
385 let key_iter_clone = key_iter.clone();
386 self.cache.remove_many(key_iter);
387 for key in key_iter_clone {
388 self.delete_both_layouts(&mut writer, key)?;
389 }
390 Ok(())
391 }
392
393 fn build_write_key(&self, key: TKey) -> DbKey
397 where
398 TKey: Clone + AsRef<[u8]>,
399 {
400 let mut db_key = DbKey::new(&self.prefix, key);
401 if let Some(suffix) = self.version_suffix {
402 db_key.add_suffix([suffix]);
403 }
404 db_key
405 }
406
407 fn delete_both_layouts(&self, writer: &mut impl DbWriter, key: TKey) -> Result<(), StoreError>
412 where
413 TKey: Clone + AsRef<[u8]>,
414 {
415 if self.version_suffix.is_some() {
416 writer.delete(self.build_write_key(key.clone()))?;
417 }
418 writer.delete(DbKey::new(&self.prefix, key))?;
419 Ok(())
420 }
421
422 pub fn delete_all(&self, mut writer: impl DbWriter) -> Result<(), StoreError>
424 where
425 TKey: Clone + AsRef<[u8]>,
426 {
427 self.cache.remove_all();
428 let db_key = DbKey::prefix_only(&self.prefix);
429 let (from, to) = rocksdb::PrefixRange(db_key.as_ref()).into_bounds();
430 writer.delete_range(from.unwrap(), to.unwrap())?;
431 Ok(())
432 }
433
434 pub fn seek_iterator(
439 &self,
440 bucket: Option<&[u8]>, seek_from: Option<TKey>, limit: usize, skip_first: bool, ) -> impl Iterator<Item = KeyDataResult<TData>> + '_
445 where
446 TKey: Clone + AsRef<[u8]>,
447 TData: DeserializeOwned,
448 {
449 let db_key = bucket.map_or_else(
450 move || DbKey::prefix_only(&self.prefix),
451 move |bucket| {
452 let mut key = DbKey::prefix_only(&self.prefix);
453 key.add_bucket(bucket);
454 key
455 },
456 );
457
458 let mut read_opts = ReadOptions::default();
459 read_opts.set_iterate_range(rocksdb::PrefixRange(db_key.as_ref()));
460
461 let mut db_iterator = match seek_from {
462 Some(seek_key) => {
463 self.db.iterator_opt(IteratorMode::From(DbKey::new(&self.prefix, seek_key).as_ref(), Direction::Forward), read_opts)
464 }
465 None => self.db.iterator_opt(IteratorMode::Start, read_opts),
466 };
467
468 if skip_first {
469 db_iterator.next();
470 }
471
472 db_iterator.take(limit).map(move |item| match item {
473 Ok((key_bytes, value_bytes)) => match bincode::deserialize::<TData>(value_bytes.as_ref()) {
474 Ok(value) => Ok((key_bytes[db_key.prefix_len()..].into(), value)),
475 Err(err) => Err(err.into()),
476 },
477 Err(err) => Err(err.into()),
478 })
479 }
480
481 pub fn prefix(&self) -> &[u8] {
482 &self.prefix
483 }
484}
485
486#[cfg(test)]
487mod tests {
488 use super::*;
489 use crate::{
490 create_temp_db,
491 prelude::{BatchDbWriter, ConnBuilder, DirectDbWriter},
492 };
493 use kaspa_hashes::Hash;
494 use rocksdb::WriteBatch;
495
496 #[test]
497 fn test_delete_all() {
498 let (_lifetime, db) = create_temp_db!(ConnBuilder::default().with_files_limit(10));
499 let access = CachedDbAccess::<Hash, u64>::new(db.clone(), CachePolicy::Count(2), vec![1, 2]);
500
501 access.write_many(DirectDbWriter::new(&db), &mut (0..16).map(|i| (i.into(), 2))).unwrap();
502 assert_eq!(16, access.iterator().count());
503 access.delete_all(DirectDbWriter::new(&db)).unwrap();
504 assert_eq!(0, access.iterator().count());
505
506 access.write_many(DirectDbWriter::new(&db), &mut (0..16).map(|i| (i.into(), 2))).unwrap();
507 assert_eq!(16, access.iterator().count());
508 let mut batch = WriteBatch::default();
509 access.delete_all(BatchDbWriter::new(&mut batch)).unwrap();
510 assert_eq!(16, access.iterator().count());
511 db.write(batch).unwrap();
512 assert_eq!(0, access.iterator().count());
513 }
514
515 #[test]
516 fn test_read_with_fallback() {
517 let (_lifetime, db) = create_temp_db!(ConnBuilder::default().with_files_limit(10));
518 let primary_prefix = vec![1];
519 let fallback_prefix = vec![2];
520 let access = CachedDbAccess::<Hash, u64>::new(db.clone(), CachePolicy::Count(10), primary_prefix);
521 let fallback_access = CachedDbAccess::<Hash, u64>::new(db.clone(), CachePolicy::Count(10), fallback_prefix.clone());
522
523 let key: Hash = 1.into();
524 let value = 100;
525
526 fallback_access.write(DirectDbWriter::new(&db), key, value).unwrap();
528
529 let result = access.read_with_fallback::<u64>(&fallback_prefix, key).unwrap();
531 assert_eq!(result, value);
532
533 assert_eq!(access.read_from_cache(key).unwrap(), value);
535 }
536
537 #[test]
538 fn test_has_with_fallback() {
539 let (_lifetime, db) = create_temp_db!(ConnBuilder::default().with_files_limit(10));
540 let primary_prefix = vec![1];
541 let fallback_prefix = vec![2];
542 let access = CachedDbAccess::<Hash, u64>::new(db.clone(), CachePolicy::Count(10), primary_prefix);
543 let fallback_access = CachedDbAccess::<Hash, u64>::new(db.clone(), CachePolicy::Count(10), fallback_prefix.clone());
544
545 let key_in_fallback: Hash = 1.into();
546 let key_not_found: Hash = 2.into();
547
548 fallback_access.write(DirectDbWriter::new(&db), key_in_fallback, 100).unwrap();
550
551 assert!(access.has_with_fallback(&fallback_prefix, key_in_fallback).unwrap());
553
554 assert!(!access.has_with_fallback(&fallback_prefix, key_not_found).unwrap());
556 }
557
558 mod versioned {
559 use super::*;
565 use crate::errors::StoreErrorPredicates;
566 use serde::{Deserialize, Serialize};
567
568 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
570 struct Current {
571 value: u64,
572 extra: Option<u32>,
573 }
574
575 impl MemSizeEstimator for Current {}
576
577 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
581 struct Legacy {
582 value: u64,
583 }
584
585 impl From<Legacy> for Current {
586 fn from(l: Legacy) -> Self {
587 Current { value: l.value, extra: None }
588 }
589 }
590
591 const STORE_PREFIX: u8 = 42;
592 const VERSION_SUFFIX: u8 = 1;
593
594 fn versioned_access(db: Arc<DB>) -> CachedDbAccess<Hash, Current, RandomState, Legacy> {
595 CachedDbAccess::new_with_version_suffix(db, CachePolicy::Count(16), vec![STORE_PREFIX], VERSION_SUFFIX)
596 }
597
598 fn legacy_raw_key(hash: Hash) -> Vec<u8> {
600 let mut bytes = vec![STORE_PREFIX];
601 bytes.extend_from_slice(hash.as_bytes().as_ref());
602 bytes
603 }
604
605 fn versioned_raw_key(hash: Hash, suffix: u8) -> Vec<u8> {
607 let mut bytes = legacy_raw_key(hash);
608 bytes.push(suffix);
609 bytes
610 }
611
612 #[test]
613 fn versioned_round_trip() {
614 let (_lifetime, db) = create_temp_db!(ConnBuilder::default().with_files_limit(10));
615 let access = versioned_access(db.clone());
616
617 let hash: Hash = 7.into();
618 let value = Current { value: 1_000_000, extra: Some(42) };
619
620 access.write(DirectDbWriter::new(&db), hash, value.clone()).unwrap();
621
622 assert!(db.get_pinned(versioned_raw_key(hash, VERSION_SUFFIX)).unwrap().is_some());
624 assert!(db.get_pinned(legacy_raw_key(hash)).unwrap().is_none());
625
626 assert_eq!(access.read(hash).unwrap(), value);
628
629 assert!(access.has(hash).unwrap());
631
632 access.delete(DirectDbWriter::new(&db), hash).unwrap();
634 assert!(!access.has(hash).unwrap());
635 assert!(access.read(hash).is_err());
636 }
637
638 #[test]
639 fn versioned_read_decodes_legacy_row() {
640 let (_lifetime, db) = create_temp_db!(ConnBuilder::default().with_files_limit(10));
641 let access = versioned_access(db.clone());
642
643 let hash: Hash = 13.into();
644 let legacy = Legacy { value: 99 };
645 let legacy_bytes = bincode::serialize(&legacy).unwrap();
646
647 db.put(legacy_raw_key(hash), legacy_bytes).unwrap();
650
651 let current = access.read(hash).unwrap();
655 assert_eq!(current, Current { value: 99, extra: None });
656
657 assert!(access.has(hash).unwrap());
659
660 access.delete(DirectDbWriter::new(&db), hash).unwrap();
663 assert!(db.get_pinned(legacy_raw_key(hash)).unwrap().is_none());
664 assert!(!access.has(hash).unwrap());
665 }
666
667 #[test]
668 fn versioned_read_rejects_unknown_suffix_byte() {
669 let (_lifetime, db) = create_temp_db!(ConnBuilder::default().with_files_limit(10));
670 let access = versioned_access(db.clone());
671
672 let hash: Hash = 21.into();
673 let bytes = bincode::serialize(&Current { value: 1, extra: None }).unwrap();
674
675 db.put(versioned_raw_key(hash, 0xAB), bytes).unwrap();
678
679 let err = access.read(hash).unwrap_err();
680 assert!(matches!(err, StoreError::DataInconsistency(_)), "expected DataInconsistency, got {err:?}");
681 }
682
683 #[test]
684 fn versioned_read_rejects_unexpected_key_length() {
685 let (_lifetime, db) = create_temp_db!(ConnBuilder::default().with_files_limit(10));
686 let access = versioned_access(db.clone());
687
688 let hash: Hash = 34.into();
689 let bytes = bincode::serialize(&Current { value: 1, extra: None }).unwrap();
690
691 let mut bogus_key = legacy_raw_key(hash);
693 bogus_key.push(VERSION_SUFFIX);
694 bogus_key.push(0x00);
695 db.put(bogus_key, bytes).unwrap();
696
697 let err = access.read(hash).unwrap_err();
698 assert!(matches!(err, StoreError::DataInconsistency(_)), "expected DataInconsistency, got {err:?}");
699 }
700
701 #[test]
702 fn versioned_delete_clears_both_layouts() {
703 let (_lifetime, db) = create_temp_db!(ConnBuilder::default().with_files_limit(10));
704 let access = versioned_access(db.clone());
705
706 let hash: Hash = 55.into();
707
708 db.put(legacy_raw_key(hash), bincode::serialize(&Legacy { value: 7 }).unwrap()).unwrap();
713 db.put(versioned_raw_key(hash, VERSION_SUFFIX), bincode::serialize(&Current { value: 7, extra: Some(1) }).unwrap())
714 .unwrap();
715
716 access.delete(DirectDbWriter::new(&db), hash).unwrap();
717
718 assert!(db.get_pinned(legacy_raw_key(hash)).unwrap().is_none());
719 assert!(db.get_pinned(versioned_raw_key(hash, VERSION_SUFFIX)).unwrap().is_none());
720 }
721
722 #[test]
723 fn versioned_has_guards_inserts_against_legacy_rows() {
724 let (_lifetime, db) = create_temp_db!(ConnBuilder::default().with_files_limit(10));
729 let access = versioned_access(db.clone());
730
731 let hash: Hash = 89.into();
732 db.put(legacy_raw_key(hash), bincode::serialize(&Legacy { value: 11 }).unwrap()).unwrap();
733
734 assert!(access.has(hash).unwrap(), "versioned has() must observe legacy rows");
735
736 assert!(!access.has(Hash::from_u64_word(12345)).unwrap());
738 }
739
740 #[test]
741 fn versioned_read_from_missing_key_is_not_found() {
742 let (_lifetime, db) = create_temp_db!(ConnBuilder::default().with_files_limit(10));
743 let access = versioned_access(db);
744
745 let err = access.read(Hash::from_u64_word(999)).unwrap_err();
746 assert!(err.is_key_not_found(), "expected KeyNotFound, got {err:?}");
747 }
748 }
749}