Skip to main content

kaspa_txscript/
caches.rs

1use indexmap::IndexMap;
2use parking_lot::RwLock;
3use rand::Rng;
4use std::{
5    collections::hash_map::RandomState,
6    hash::BuildHasher,
7    sync::{
8        Arc,
9        atomic::{AtomicU64, Ordering},
10    },
11};
12
13#[derive(Clone)]
14pub struct Cache<TKey: Clone + std::hash::Hash + Eq + Send + Sync, TData: Clone + Send + Sync, S = RandomState> {
15    // We use IndexMap and not HashMap, because it makes it cheaper to remove a random element when the cache is full.
16    map: Arc<RwLock<IndexMap<TKey, TData, S>>>,
17    size: usize,
18    counters: Arc<TxScriptCacheCounters>,
19}
20
21impl<TKey: Clone + std::hash::Hash + Eq + Send + Sync, TData: Clone + Send + Sync, S: BuildHasher + Default> Cache<TKey, TData, S> {
22    pub fn new(size: u64) -> Self {
23        Self::with_counters(size, Default::default())
24    }
25
26    pub fn with_counters(size: u64, counters: Arc<TxScriptCacheCounters>) -> Self {
27        Self {
28            map: Arc::new(RwLock::new(IndexMap::with_capacity_and_hasher(size as usize, S::default()))),
29            size: size as usize,
30            counters,
31        }
32    }
33
34    pub fn clear(&self) {
35        self.map.write().clear();
36    }
37
38    pub(crate) fn get(&self, key: &TKey) -> Option<TData> {
39        self.map.read().get(key).cloned().inspect(|_data| {
40            self.counters.get_counts.fetch_add(1, Ordering::Relaxed);
41        })
42    }
43
44    pub(crate) fn insert(&self, key: TKey, data: TData) {
45        if self.size == 0 {
46            return;
47        }
48        let mut write_guard = self.map.write();
49        if write_guard.len() == self.size {
50            write_guard.swap_remove_index(rand::thread_rng().gen_range(0..self.size));
51        }
52        write_guard.insert(key, data);
53        self.counters.insert_counts.fetch_add(1, Ordering::Relaxed);
54    }
55}
56
57#[derive(Default)]
58pub struct TxScriptCacheCounters {
59    pub insert_counts: AtomicU64,
60    pub get_counts: AtomicU64,
61}
62
63impl TxScriptCacheCounters {
64    pub fn snapshot(&self) -> TxScriptCacheCountersSnapshot {
65        TxScriptCacheCountersSnapshot {
66            insert_counts: self.insert_counts.load(Ordering::Relaxed),
67            get_counts: self.get_counts.load(Ordering::Relaxed),
68        }
69    }
70}
71
72#[derive(Debug, PartialEq, Eq)]
73pub struct TxScriptCacheCountersSnapshot {
74    pub insert_counts: u64,
75    pub get_counts: u64,
76}
77
78impl TxScriptCacheCountersSnapshot {
79    pub fn hit_ratio(&self) -> f64 {
80        if self.insert_counts > 0 { self.get_counts as f64 / self.insert_counts as f64 } else { 0f64 }
81    }
82}
83
84impl core::ops::Sub for &TxScriptCacheCountersSnapshot {
85    type Output = TxScriptCacheCountersSnapshot;
86
87    fn sub(self, rhs: Self) -> Self::Output {
88        Self::Output {
89            insert_counts: self.insert_counts.saturating_sub(rhs.insert_counts),
90            get_counts: self.get_counts.saturating_sub(rhs.get_counts),
91        }
92    }
93}