1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

use indexmap::IndexMap;
use parking_lot::RwLock;
use rand::Rng;
use std::{
    collections::hash_map::RandomState,
    hash::BuildHasher,
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc,
    },
};

#[derive(Clone)]
pub struct Cache<TKey: Clone + std::hash::Hash + Eq + Send + Sync, TData: Clone + Send + Sync, S = RandomState> {
    // We use IndexMap and not HashMap, because it makes it cheaper to remove a random element when the cache is full.
    map: Arc<RwLock<IndexMap<TKey, TData, S>>>,
    size: usize,
    counters: Arc<TxScriptCacheCounters>,
}

impl<TKey: Clone + std::hash::Hash + Eq + Send + Sync, TData: Clone + Send + Sync, S: BuildHasher + Default> Cache<TKey, TData, S> {
    pub fn new(size: u64) -> Self {
        Self::with_counters(size, Default::default())
    }

    pub fn with_counters(size: u64, counters: Arc<TxScriptCacheCounters>) -> Self {
        Self {
            map: Arc::new(RwLock::new(IndexMap::with_capacity_and_hasher(size as usize, S::default()))),
            size: size as usize,
            counters,
        }
    }

    pub(crate) fn get(&self, key: &TKey) -> Option<TData> {
        self.map.read().get(key).cloned().map(|data| {
            self.counters.get_counts.fetch_add(1, Ordering::Relaxed);
            data
        })
    }

    pub(crate) fn insert(&self, key: TKey, data: TData) {
        if self.size == 0 {
            return;
        }
        let mut write_guard = self.map.write();
        if write_guard.len() == self.size {
            write_guard.swap_remove_index(rand::thread_rng().gen_range(0..self.size));
        }
        write_guard.insert(key, data);
        self.counters.insert_counts.fetch_add(1, Ordering::Relaxed);
    }
}

#[derive(Default)]
pub struct TxScriptCacheCounters {
    pub insert_counts: AtomicU64,
    pub get_counts: AtomicU64,
}

impl TxScriptCacheCounters {
    pub fn snapshot(&self) -> TxScriptCacheCountersSnapshot {
        TxScriptCacheCountersSnapshot {
            insert_counts: self.insert_counts.load(Ordering::Relaxed),
            get_counts: self.get_counts.load(Ordering::Relaxed),
        }
    }
}

#[derive(Debug, PartialEq, Eq)]
pub struct TxScriptCacheCountersSnapshot {
    pub insert_counts: u64,
    pub get_counts: u64,
}

impl TxScriptCacheCountersSnapshot {
    pub fn hit_ratio(&self) -> f64 {
        if self.insert_counts > 0 {
            self.get_counts as f64 / self.insert_counts as f64
        } else {
            0f64
        }
    }
}

impl core::ops::Sub for &TxScriptCacheCountersSnapshot {
    type Output = TxScriptCacheCountersSnapshot;

    fn sub(self, rhs: Self) -> Self::Output {
        Self::Output {
            insert_counts: self.insert_counts.checked_sub(rhs.insert_counts).unwrap_or_default(),
            get_counts: self.get_counts.checked_sub(rhs.get_counts).unwrap_or_default(),
        }
    }
}