snarkos_node_router/helpers/
cache.rs

1// Copyright (c) 2019-2025 Provable Inc.
2// This file is part of the snarkOS library.
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at:
7
8// http://www.apache.org/licenses/LICENSE-2.0
9
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use crate::messages::BlockRequest;
17use snarkvm::prelude::{Network, puzzle::SolutionID};
18
19use core::hash::Hash;
20use linked_hash_map::LinkedHashMap;
21#[cfg(feature = "locktick")]
22use locktick::parking_lot::RwLock;
23#[cfg(not(feature = "locktick"))]
24use parking_lot::RwLock;
25use std::{
26    collections::{HashMap, HashSet, VecDeque},
27    net::{IpAddr, SocketAddr},
28};
29use time::{Duration, OffsetDateTime};
30
31/// The maximum number of items to store in a cache map.
32const MAX_CACHE_SIZE: usize = 1 << 17;
33
34/// A helper containing the peer IP and solution ID.
35type SolutionKey<N> = (SocketAddr, SolutionID<N>);
36/// A helper containing the peer IP and transaction ID.
37type TransactionKey<N> = (SocketAddr, <N as Network>::TransactionID);
38
39#[derive(Debug)]
40pub struct Cache<N: Network> {
41    /// The map of peer connections to their recent timestamps.
42    seen_inbound_connections: RwLock<HashMap<IpAddr, VecDeque<OffsetDateTime>>>,
43    /// The map of peer IPs to their recent timestamps.
44    seen_inbound_messages: RwLock<HashMap<SocketAddr, VecDeque<OffsetDateTime>>>,
45    /// The map of peer IPs to their recent timestamps.
46    seen_inbound_puzzle_requests: RwLock<HashMap<SocketAddr, VecDeque<OffsetDateTime>>>,
47    /// The map of peer IPs to their recent timestamps.
48    seen_inbound_block_requests: RwLock<HashMap<SocketAddr, VecDeque<OffsetDateTime>>>,
49    /// The map of solution IDs to their last seen timestamp.
50    seen_inbound_solutions: RwLock<LinkedHashMap<SolutionKey<N>, OffsetDateTime>>,
51    /// The map of transaction IDs to their last seen timestamp.
52    seen_inbound_transactions: RwLock<LinkedHashMap<TransactionKey<N>, OffsetDateTime>>,
53    /// The map of peer IPs to their block requests.
54    seen_outbound_block_requests: RwLock<HashMap<SocketAddr, HashSet<BlockRequest>>>,
55    /// The map of peer IPs to the number of puzzle requests.
56    seen_outbound_puzzle_requests: RwLock<HashMap<SocketAddr, u32>>,
57    /// The map of solution IDs to their last seen timestamp.
58    seen_outbound_solutions: RwLock<LinkedHashMap<SolutionKey<N>, OffsetDateTime>>,
59    /// The map of transaction IDs to their last seen timestamp.
60    seen_outbound_transactions: RwLock<LinkedHashMap<TransactionKey<N>, OffsetDateTime>>,
61    /// The map of peer IPs to the number of sent peer requests.
62    seen_outbound_peer_requests: RwLock<HashMap<SocketAddr, u32>>,
63}
64
65impl<N: Network> Default for Cache<N> {
66    /// Initializes a new instance of the cache.
67    fn default() -> Self {
68        Self::new()
69    }
70}
71
72impl<N: Network> Cache<N> {
73    const INBOUND_BLOCK_REQUEST_INTERVAL: i64 = 60;
74    const INBOUND_PUZZLE_REQUEST_INTERVAL: i64 = 60;
75
76    /// Initializes a new instance of the cache.
77    pub fn new() -> Self {
78        Self {
79            seen_inbound_connections: Default::default(),
80            seen_inbound_messages: Default::default(),
81            seen_inbound_puzzle_requests: Default::default(),
82            seen_inbound_block_requests: Default::default(),
83            seen_inbound_solutions: RwLock::new(LinkedHashMap::with_capacity(MAX_CACHE_SIZE)),
84            seen_inbound_transactions: RwLock::new(LinkedHashMap::with_capacity(MAX_CACHE_SIZE)),
85            seen_outbound_block_requests: Default::default(),
86            seen_outbound_puzzle_requests: Default::default(),
87            seen_outbound_solutions: RwLock::new(LinkedHashMap::with_capacity(MAX_CACHE_SIZE)),
88            seen_outbound_transactions: RwLock::new(LinkedHashMap::with_capacity(MAX_CACHE_SIZE)),
89            seen_outbound_peer_requests: Default::default(),
90        }
91    }
92}
93
94impl<N: Network> Cache<N> {
95    /// Inserts a new timestamp for the given peer connection, returning the number of recent connection requests.
96    pub fn insert_inbound_connection(&self, peer_ip: IpAddr, interval_in_secs: i64) -> usize {
97        Self::retain_and_insert(&self.seen_inbound_connections, peer_ip, interval_in_secs)
98    }
99
100    /// Inserts a new timestamp for the given peer message, returning the number of recent messages.
101    pub fn insert_inbound_message(&self, peer_ip: SocketAddr, interval_in_secs: i64) -> usize {
102        Self::retain_and_insert(&self.seen_inbound_messages, peer_ip, interval_in_secs)
103    }
104
105    /// Inserts a new timestamp for the given peer IP, returning the number of recent requests.
106    pub fn insert_inbound_puzzle_request(&self, peer_ip: SocketAddr) -> usize {
107        Self::retain_and_insert(&self.seen_inbound_puzzle_requests, peer_ip, Self::INBOUND_PUZZLE_REQUEST_INTERVAL)
108    }
109
110    /// Inserts a new timestamp for the given peer IP, returning the number of recent block requests.
111    pub fn insert_inbound_block_request(&self, peer_ip: SocketAddr) -> usize {
112        Self::retain_and_insert(&self.seen_inbound_block_requests, peer_ip, Self::INBOUND_BLOCK_REQUEST_INTERVAL)
113    }
114
115    /// Inserts a solution ID into the cache, returning the previously seen timestamp if it existed.
116    pub fn insert_inbound_solution(&self, peer_ip: SocketAddr, solution_id: SolutionID<N>) -> Option<OffsetDateTime> {
117        Self::refresh_and_insert(&self.seen_inbound_solutions, (peer_ip, solution_id))
118    }
119
120    /// Inserts a transaction ID into the cache, returning the previously seen timestamp if it existed.
121    pub fn insert_inbound_transaction(
122        &self,
123        peer_ip: SocketAddr,
124        transaction: N::TransactionID,
125    ) -> Option<OffsetDateTime> {
126        Self::refresh_and_insert(&self.seen_inbound_transactions, (peer_ip, transaction))
127    }
128}
129
130impl<N: Network> Cache<N> {
131    /// Returns `true` if the cache contains any inbound block requests for the given peer.
132    pub fn contains_inbound_block_request(&self, peer_ip: &SocketAddr) -> bool {
133        Self::retain(&self.seen_inbound_block_requests, *peer_ip, Self::INBOUND_BLOCK_REQUEST_INTERVAL) > 0
134    }
135
136    /// Returns the number of recent block requests for the given peer.
137    pub fn num_outbound_block_requests(&self, peer_ip: &SocketAddr) -> usize {
138        self.seen_outbound_block_requests.read().get(peer_ip).map(|r| r.len()).unwrap_or(0)
139    }
140
141    /// Returns `true` if the cache contains the given block request for the specified peer.
142    pub fn contains_outbound_block_request(&self, peer_ip: &SocketAddr, request: &BlockRequest) -> bool {
143        self.seen_outbound_block_requests.read().get(peer_ip).map(|r| r.contains(request)).unwrap_or(false)
144    }
145
146    /// Inserts the block request for the given peer IP, returning the number of recent requests.
147    pub fn insert_outbound_block_request(&self, peer_ip: SocketAddr, request: BlockRequest) -> usize {
148        let mut map_write = self.seen_outbound_block_requests.write();
149        let requests = map_write.entry(peer_ip).or_default();
150        requests.insert(request);
151        requests.len()
152    }
153
154    /// Removes the block request for the given peer IP, returning `true` if the request was present.
155    pub fn remove_outbound_block_request(&self, peer_ip: SocketAddr, request: &BlockRequest) -> bool {
156        let mut map_write = self.seen_outbound_block_requests.write();
157        if let Some(requests) = map_write.get_mut(&peer_ip) { requests.remove(request) } else { false }
158    }
159
160    /// Returns `true` if the cache contains a puzzle request from the given peer.
161    pub fn contains_outbound_puzzle_request(&self, peer_ip: &SocketAddr) -> bool {
162        self.seen_outbound_puzzle_requests.read().get(peer_ip).map(|r| *r > 0).unwrap_or(false)
163    }
164
165    /// Increment the peer IP's number of puzzle requests, returning the updated number of puzzle requests.
166    pub fn increment_outbound_puzzle_requests(&self, peer_ip: SocketAddr) -> u32 {
167        Self::increment_counter(&self.seen_outbound_puzzle_requests, peer_ip)
168    }
169
170    /// Decrement the peer IP's number of puzzle requests, returning the updated number of puzzle requests.
171    pub fn decrement_outbound_puzzle_requests(&self, peer_ip: SocketAddr) -> u32 {
172        Self::decrement_counter(&self.seen_outbound_puzzle_requests, peer_ip)
173    }
174
175    /// Inserts a solution ID into the cache, returning the previously seen timestamp if it existed.
176    pub fn insert_outbound_solution(&self, peer_ip: SocketAddr, solution_id: SolutionID<N>) -> Option<OffsetDateTime> {
177        Self::refresh_and_insert(&self.seen_outbound_solutions, (peer_ip, solution_id))
178    }
179
180    /// Inserts a transaction ID into the cache, returning the previously seen timestamp if it existed.
181    pub fn insert_outbound_transaction(
182        &self,
183        peer_ip: SocketAddr,
184        transaction: N::TransactionID,
185    ) -> Option<OffsetDateTime> {
186        Self::refresh_and_insert(&self.seen_outbound_transactions, (peer_ip, transaction))
187    }
188
189    /// Returns `true` if the cache contains a peer request from the given peer.
190    pub fn contains_outbound_peer_request(&self, peer_ip: SocketAddr) -> bool {
191        self.seen_outbound_peer_requests.read().get(&peer_ip).map(|r| *r > 0).unwrap_or(false)
192    }
193
194    /// Increment the peer IP's number of peer requests, returning the updated number of peer requests.
195    pub fn increment_outbound_peer_requests(&self, peer_ip: SocketAddr) -> u32 {
196        Self::increment_counter(&self.seen_outbound_peer_requests, peer_ip)
197    }
198
199    /// Decrement the peer IP's number of peer requests, returning the updated number of peer requests.
200    pub fn decrement_outbound_peer_requests(&self, peer_ip: SocketAddr) -> u32 {
201        Self::decrement_counter(&self.seen_outbound_peer_requests, peer_ip)
202    }
203
204    /// Removes all cache entries applicable to the given key.
205    pub fn clear_peer_entries(&self, peer_ip: SocketAddr) {
206        self.seen_outbound_block_requests.write().remove(&peer_ip);
207    }
208}
209
210impl<N: Network> Cache<N> {
211    /// Insert a new timestamp for the given key, returning the number of recent entries.
212    fn retain_and_insert<K: Eq + Hash + Clone>(
213        map: &RwLock<HashMap<K, VecDeque<OffsetDateTime>>>,
214        key: K,
215        interval_in_secs: i64,
216    ) -> usize {
217        // Fetch the current timestamp.
218        let now = OffsetDateTime::now_utc();
219
220        let mut map_write = map.write();
221        // Load the entry for the key.
222        let timestamps = map_write.entry(key).or_default();
223        // Insert the new timestamp.
224        timestamps.push_back(now);
225        // Retain only the timestamps that are within the recent interval.
226        while timestamps.front().map_or(false, |t| now - *t > Duration::seconds(interval_in_secs)) {
227            timestamps.pop_front();
228        }
229        // Return the frequency of recent requests.
230        timestamps.len()
231    }
232
233    /// Returns the number of recent entries.
234    fn retain<K: Eq + Hash + Clone>(
235        map: &RwLock<HashMap<K, VecDeque<OffsetDateTime>>>,
236        key: K,
237        interval_in_secs: i64,
238    ) -> usize {
239        // Fetch the current timestamp.
240        let now = OffsetDateTime::now_utc();
241
242        let mut map_write = map.write();
243        // Load the entry for the key.
244        let timestamps = map_write.entry(key).or_default();
245        // Retain only the timestamps that are within the recent interval.
246        while timestamps.front().map_or(false, |t| now - *t > Duration::seconds(interval_in_secs)) {
247            timestamps.pop_front();
248        }
249        // Return the frequency of recent requests.
250        timestamps.len()
251    }
252
253    /// Increments the key's counter in the map, returning the updated counter.
254    fn increment_counter<K: Hash + Eq>(map: &RwLock<HashMap<K, u32>>, key: K) -> u32 {
255        let mut map_write = map.write();
256        // Load the entry for the key, and increment the counter.
257        let entry = map_write.entry(key).or_default();
258        *entry = entry.saturating_add(1);
259        // Return the updated counter.
260        *entry
261    }
262
263    /// Decrements the key's counter in the map, returning the updated counter.
264    fn decrement_counter<K: Copy + Hash + Eq>(map: &RwLock<HashMap<K, u32>>, key: K) -> u32 {
265        let mut map_write = map.write();
266        // Load the entry for the key, and decrement the counter.
267        let entry = map_write.entry(key).or_default();
268        let value = entry.saturating_sub(1);
269        // If the entry is 0, remove the entry.
270        if *entry == 0 {
271            map_write.remove(&key);
272        } else {
273            *entry = value;
274        }
275        // Return the updated counter.
276        value
277    }
278
279    /// Updates the map by enforcing the maximum cache size.
280    fn refresh<K: Eq + Hash, V>(map: &RwLock<LinkedHashMap<K, V>>) {
281        let mut map_write = map.write();
282        while map_write.len() >= MAX_CACHE_SIZE {
283            map_write.pop_front();
284        }
285    }
286
287    /// Updates the map by enforcing the maximum cache size, and inserts the given key.
288    /// Returns the previously seen timestamp if it existed.
289    fn refresh_and_insert<K: Eq + Hash>(
290        map: &RwLock<LinkedHashMap<K, OffsetDateTime>>,
291        key: K,
292    ) -> Option<OffsetDateTime> {
293        // Insert the key, and return the previous timestamp if it existed.
294        let previous_timestamp = map.write().insert(key, OffsetDateTime::now_utc());
295        // Refresh the cache.
296        Self::refresh(map);
297        // Return the previous timestamp.
298        previous_timestamp
299    }
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305    use snarkvm::prelude::MainnetV0;
306
307    use std::net::Ipv4Addr;
308
309    type CurrentNetwork = MainnetV0;
310
311    #[test]
312    fn test_inbound_block_request() {
313        let cache = Cache::<CurrentNetwork>::default();
314        let peer_ip = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 1234);
315
316        // Check that the cache is empty.
317        assert_eq!(cache.seen_inbound_block_requests.read().len(), 0);
318
319        // Insert a block request..
320        assert_eq!(cache.insert_inbound_block_request(peer_ip), 1);
321
322        // Check that the cache contains the block request.
323        assert!(cache.contains_inbound_block_request(&peer_ip));
324
325        // Insert another block request for the same peer.
326        assert_eq!(cache.insert_inbound_block_request(peer_ip), 2);
327
328        // Check that the cache contains the block requests.
329        assert!(cache.contains_inbound_block_request(&peer_ip));
330    }
331
332    #[test]
333    fn test_inbound_solution() {
334        let cache = Cache::<CurrentNetwork>::default();
335        let peer_ip = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 1234);
336        let solution_id = SolutionID::<CurrentNetwork>::from(123456789);
337
338        // Check that the cache is empty.
339        assert_eq!(cache.seen_inbound_solutions.read().len(), 0);
340
341        // Insert a solution.
342        assert!(cache.insert_inbound_solution(peer_ip, solution_id).is_none());
343
344        // Check that the cache contains the solution.
345        assert_eq!(cache.seen_inbound_solutions.read().len(), 1);
346
347        // Insert the same solution again.
348        assert!(cache.insert_inbound_solution(peer_ip, solution_id).is_some());
349
350        // Check that the cache still contains the solution.
351        assert_eq!(cache.seen_inbound_solutions.read().len(), 1);
352    }
353
354    #[test]
355    fn test_inbound_transaction() {
356        let cache = Cache::<CurrentNetwork>::default();
357        let peer_ip = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 1234);
358        let transaction = Default::default();
359
360        // Check that the cache is empty.
361        assert_eq!(cache.seen_inbound_transactions.read().len(), 0);
362
363        // Insert a transaction.
364        assert!(cache.insert_inbound_transaction(peer_ip, transaction).is_none());
365
366        // Check that the cache contains the transaction.
367        assert_eq!(cache.seen_inbound_transactions.read().len(), 1);
368
369        // Insert the same transaction again.
370        assert!(cache.insert_inbound_transaction(peer_ip, transaction).is_some());
371
372        // Check that the cache still contains the transaction.
373        assert_eq!(cache.seen_inbound_transactions.read().len(), 1);
374    }
375
376    #[test]
377    fn test_outbound_solution() {
378        let cache = Cache::<CurrentNetwork>::default();
379        let peer_ip = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 1234);
380        let solution_id = SolutionID::<CurrentNetwork>::from(123456789);
381
382        // Check that the cache is empty.
383        assert_eq!(cache.seen_outbound_solutions.read().len(), 0);
384
385        // Insert a solution.
386        assert!(cache.insert_outbound_solution(peer_ip, solution_id).is_none());
387
388        // Check that the cache contains the solution.
389        assert_eq!(cache.seen_outbound_solutions.read().len(), 1);
390
391        // Insert the same solution again.
392        assert!(cache.insert_outbound_solution(peer_ip, solution_id).is_some());
393
394        // Check that the cache still contains the solution.
395        assert_eq!(cache.seen_outbound_solutions.read().len(), 1);
396    }
397
398    #[test]
399    fn test_outbound_transaction() {
400        let cache = Cache::<CurrentNetwork>::default();
401        let peer_ip = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 1234);
402        let transaction = Default::default();
403
404        // Check that the cache is empty.
405        assert_eq!(cache.seen_outbound_transactions.read().len(), 0);
406
407        // Insert a transaction.
408        assert!(cache.insert_outbound_transaction(peer_ip, transaction).is_none());
409
410        // Check that the cache contains the transaction.
411        assert_eq!(cache.seen_outbound_transactions.read().len(), 1);
412
413        // Insert the same transaction again.
414        assert!(cache.insert_outbound_transaction(peer_ip, transaction).is_some());
415
416        // Check that the cache still contains the transaction.
417        assert_eq!(cache.seen_outbound_transactions.read().len(), 1);
418    }
419
420    #[test]
421    fn test_outbound_peer_request() {
422        let cache = Cache::<CurrentNetwork>::default();
423        let peer_ip = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 1234);
424
425        // Check the cache is empty.
426        assert!(cache.seen_outbound_peer_requests.read().is_empty());
427        assert!(!cache.contains_outbound_peer_request(peer_ip));
428
429        // Increment the peer requests.
430        assert_eq!(cache.increment_outbound_peer_requests(peer_ip), 1);
431
432        // Check the cache contains the peer request.
433        assert!(cache.contains_outbound_peer_request(peer_ip));
434
435        // Increment the peer requests again for the same peer IP.
436        assert_eq!(cache.increment_outbound_peer_requests(peer_ip), 2);
437
438        // Check the cache still contains the peer request.
439        assert!(cache.contains_outbound_peer_request(peer_ip));
440
441        // Decrement the peer requests.
442        assert_eq!(cache.decrement_outbound_peer_requests(peer_ip), 1);
443
444        // Decrement the peer requests again.
445        assert_eq!(cache.decrement_outbound_peer_requests(peer_ip), 0);
446
447        // Check the cache is empty.
448        assert!(!cache.contains_outbound_peer_request(peer_ip));
449    }
450}
snarkos_node_router/helpers/cache.rs

snarkos_node_router/helpers/
cache.rs
