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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

use std::sync::Arc;

use alloy_primitives::Address;
use brontes_metrics::db_cache::CacheData;
use brontes_types::db::{
    address_metadata::AddressMetadata, address_to_protocol_info::ProtocolInfo,
    searcher::SearcherInfo, token_info::TokenInfo,
};
use moka::{policy::EvictionPolicy, sync::SegmentedCache};

const MEGABYTE: usize = 1024 * 1024;

#[derive(Clone)]
pub struct ReadWriteCache {
    address_meta:      Arc<SegmentedCache<Address, Option<AddressMetadata>, ahash::RandomState>>,
    searcher_eoa:      Arc<SegmentedCache<Address, Option<SearcherInfo>, ahash::RandomState>>,
    searcher_contract: Arc<SegmentedCache<Address, Option<SearcherInfo>, ahash::RandomState>>,
    protocol_info:     Arc<SegmentedCache<Address, Option<ProtocolInfo>, ahash::RandomState>>,
    token_info:        Arc<SegmentedCache<Address, Option<TokenInfo>, ahash::RandomState>>,

    pub metrics: Option<CacheData>,
}

impl ReadWriteCache {
    pub fn new(memory_per_table_mb: usize, metrics: bool) -> Self {
        let metrics = metrics.then(CacheData::default);
        Self {
            metrics,
            address_meta: SegmentedCache::builder(200)
                .eviction_policy(EvictionPolicy::lru())
                .max_capacity(
                    ((memory_per_table_mb * MEGABYTE) / std::mem::size_of::<AddressMetadata>())
                        as u64,
                )
                .build_with_hasher(ahash::RandomState::new())
                .into(),

            searcher_eoa: SegmentedCache::builder(200)
                .eviction_policy(EvictionPolicy::lru())
                .max_capacity(
                    ((memory_per_table_mb * MEGABYTE) / std::mem::size_of::<SearcherInfo>()) as u64,
                )
                .build_with_hasher(ahash::RandomState::new())
                .into(),

            searcher_contract: SegmentedCache::builder(200)
                .eviction_policy(EvictionPolicy::lru())
                .max_capacity(
                    ((memory_per_table_mb * MEGABYTE) / std::mem::size_of::<SearcherInfo>()) as u64,
                )
                .build_with_hasher(ahash::RandomState::new())
                .into(),
            protocol_info: SegmentedCache::builder(200)
                .eviction_policy(EvictionPolicy::lru())
                .max_capacity(
                    ((memory_per_table_mb * MEGABYTE) / std::mem::size_of::<ProtocolInfo>()) as u64,
                )
                .build_with_hasher(ahash::RandomState::new())
                .into(),

            token_info: SegmentedCache::builder(200)
                .eviction_policy(EvictionPolicy::lru())
                .max_capacity(
                    ((memory_per_table_mb * MEGABYTE) / std::mem::size_of::<TokenInfo>()) as u64,
                )
                .build_with_hasher(ahash::RandomState::new())
                .into(),
        }
    }

    fn record_metrics<R, T, TY>(
        &self,
        read: bool,
        name: &str,
        cache: &T,
        f: impl FnOnce(&T) -> R,
    ) -> R {
        if let Some(metrics) = self.metrics.clone() {
            if read {
                metrics.cache_read::<R, TY>(name, || f(cache))
            } else {
                metrics.cache_write::<R, TY>(name, || f(cache))
            }
        } else {
            f(cache)
        }
    }

    pub fn address_meta<R>(
        &self,
        read: bool,
        f: impl FnOnce(&SegmentedCache<Address, Option<AddressMetadata>, ahash::RandomState>) -> R,
    ) -> R {
        self.record_metrics::<R, _, AddressMetadata>(read, "address_meta", &*self.address_meta, f)
    }

    pub fn searcher_contract<R>(
        &self,
        read: bool,
        f: impl FnOnce(&SegmentedCache<Address, Option<SearcherInfo>, ahash::RandomState>) -> R,
    ) -> R {
        self.record_metrics::<R, _, SearcherInfo>(
            read,
            "searcher_contract",
            &*self.searcher_contract,
            f,
        )
    }

    pub fn searcher_eoa<R>(
        &self,
        read: bool,
        f: impl FnOnce(&SegmentedCache<Address, Option<SearcherInfo>, ahash::RandomState>) -> R,
    ) -> R {
        self.record_metrics::<R, _, SearcherInfo>(read, "searcher_eoa", &*self.searcher_eoa, f)
    }

    pub fn protocol_info<R>(
        &self,
        read: bool,
        f: impl FnOnce(&SegmentedCache<Address, Option<ProtocolInfo>, ahash::RandomState>) -> R,
    ) -> R {
        self.record_metrics::<R, _, ProtocolInfo>(read, "protocol_info", &*self.protocol_info, f)
    }

    pub fn token_info<R>(
        &self,
        read: bool,
        f: impl FnOnce(&SegmentedCache<Address, Option<TokenInfo>, ahash::RandomState>) -> R,
    ) -> R {
        self.record_metrics::<R, _, TokenInfo>(read, "token_info", &*self.token_info, f)
    }
}