snarkos_node_router/helpers/
cache.rs

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

snarkos_node_router/helpers/
cache.rs
