Skip to main content

tycho_rpc/endpoint/proto/
cache.rs

1use std::sync::{Arc, OnceLock};
2
3use arc_swap::ArcSwapOption;
4use bytes::Bytes;
5use tycho_types::boc::{Boc, BocRepr};
6use tycho_types::cell::Cell;
7use tycho_types::models::{Block, BlockId, BlockchainConfig};
8use tycho_types::prelude::HashBytes;
9use tycho_util::FastHasherState;
10
11use super::extractor::{ProtoOkResponse, RawProtoOkResponse};
12use super::make_response_block_id;
13use super::protos::rpc::response;
14
15pub struct ProtoEndpointCache {
16    libraries: moka::sync::Cache<HashBytes, RawProtoOkResponse, FastHasherState>,
17    key_block_proofs: moka::sync::Cache<u32, RawProtoOkResponse, FastHasherState>,
18    latest_key_block: ArcSwapOption<RawProtoOkResponse>,
19    blockchain_config: ArcSwapOption<RawProtoOkResponse>,
20}
21
22impl Default for ProtoEndpointCache {
23    fn default() -> Self {
24        Self {
25            libraries: moka::sync::Cache::builder()
26                .max_capacity(100)
27                .build_with_hasher(Default::default()),
28            key_block_proofs: moka::sync::Cache::builder()
29                .max_capacity(10)
30                .build_with_hasher(Default::default()),
31            latest_key_block: Default::default(),
32            blockchain_config: Default::default(),
33        }
34    }
35}
36
37impl ProtoEndpointCache {
38    pub fn insert_library_cell_response(
39        &self,
40        hash_bytes: HashBytes,
41        cell: Option<Cell>,
42    ) -> RawProtoOkResponse {
43        static EMPTY: OnceLock<RawProtoOkResponse> = OnceLock::new();
44
45        match cell {
46            None => EMPTY
47                .get_or_init(|| {
48                    make_cached(response::Result::GetLibraryCell(response::GetLibraryCell {
49                        cell: None,
50                    }))
51                })
52                .clone(),
53            Some(cell) => {
54                let res = make_cached(response::Result::GetLibraryCell(response::GetLibraryCell {
55                    cell: Some(Bytes::from(Boc::encode(cell))),
56                }));
57                self.libraries.insert(hash_bytes, res.clone());
58                res
59            }
60        }
61    }
62
63    pub fn get_library_cell_response(&self, hash: &HashBytes) -> Option<RawProtoOkResponse> {
64        self.libraries.get(hash)
65    }
66
67    pub fn insert_key_block_proof_response(
68        &self,
69        seqno: u32,
70        proof: Option<(BlockId, Bytes)>,
71    ) -> RawProtoOkResponse {
72        static EMPTY: OnceLock<RawProtoOkResponse> = OnceLock::new();
73
74        match proof {
75            None => EMPTY
76                .get_or_init(|| {
77                    make_cached(response::Result::GetKeyBlockProof(
78                        response::GetKeyBlockProof { key_block: None },
79                    ))
80                })
81                .clone(),
82            Some((block_id, proof)) => {
83                let res = make_cached(response::Result::GetKeyBlockProof(
84                    response::GetKeyBlockProof {
85                        key_block: Some(response::KeyBlockProof {
86                            block_id: Some(make_response_block_id(block_id)),
87                            proof,
88                        }),
89                    },
90                ));
91                self.key_block_proofs.insert(seqno, res.clone());
92                res
93            }
94        }
95    }
96
97    pub fn get_key_block_proof_response(&self, seqno: u32) -> Option<RawProtoOkResponse> {
98        self.key_block_proofs.get(&seqno)
99    }
100
101    pub fn load_latest_key_block(&self) -> arc_swap::Guard<Option<Arc<RawProtoOkResponse>>> {
102        self.latest_key_block.load()
103    }
104
105    pub fn load_blockchain_config(&self) -> arc_swap::Guard<Option<Arc<RawProtoOkResponse>>> {
106        self.blockchain_config.load()
107    }
108
109    pub fn handle_config(&self, global_id: i32, seqno: u32, config: &BlockchainConfig) {
110        self.blockchain_config.store(match BocRepr::encode(config) {
111            Ok(config) => Some(Arc::new(make_cached(
112                response::Result::GetBlockchainConfig(response::GetBlockchainConfig {
113                    global_id,
114                    seqno,
115                    config: config.into(),
116                }),
117            ))),
118            Err(e) => {
119                tracing::error!("failed to serialize blockchain config proto: {e:?}");
120                None
121            }
122        });
123    }
124
125    pub fn handle_key_block(&self, block: &Block) {
126        self.latest_key_block.store(match BocRepr::encode(block) {
127            Ok(block) => Some(Arc::new(make_cached(response::Result::GetLatestKeyBlock(
128                response::GetLatestKeyBlock {
129                    block: block.into(),
130                },
131            )))),
132            Err(e) => {
133                tracing::error!("failed to serialize key block proto: {e:?}");
134                None
135            }
136        });
137    }
138}
139
140fn make_cached(res: response::Result) -> RawProtoOkResponse {
141    ProtoOkResponse::new(res).into_raw()
142}