Use LRU in connection-cache (#24109)

Switch to using LRU for connection-cache

Cargo.lock

@@ -2384,9 +2384,9 @@ dependencies = [
 
 [[package]]
 name = "lru"
-version = "0.7.3"
+version = "0.7.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fcb87f3080f6d1d69e8c564c0fcfde1d7aa8cc451ce40cae89479111f03bc0eb"
+checksum = "32613e41de4c47ab04970c348ca7ae7382cf116625755af070b008a15516a889"
 dependencies = [
  "hashbrown",
 ]
@@ -4622,6 +4622,7 @@ dependencies = [
  "jsonrpc-http-server",
  "lazy_static",
  "log",
+ "lru",
  "quinn",
  "rand 0.7.3",
  "rand_chacha 0.2.2",

client/Cargo.toml

@@ -25,6 +25,7 @@ itertools = "0.10.2"
 jsonrpc-core = "18.0.0"
 lazy_static = "1.4.0"
 log = "0.4.14"
+lru = "0.7.5"
 quinn = "0.8.0"
 rand = "0.7.0"
 rand_chacha = "0.2.2"

client/src/connection_cache.rs

@@ -3,10 +3,10 @@ use {
         quic_client::QuicTpuConnection, tpu_connection::TpuConnection, udp_client::UdpTpuConnection,
     },
     lazy_static::lazy_static,
+    lru::LruCache,
     solana_net_utils::VALIDATOR_PORT_RANGE,
     solana_sdk::{transaction::VersionedTransaction, transport::TransportError},
     std::{
-        collections::{hash_map::Entry, BTreeMap, HashMap},
         net::{IpAddr, Ipv4Addr, SocketAddr},
         sync::{Arc, Mutex},
     },
@@ -22,26 +22,14 @@ enum Connection {
 }
 
 struct ConnMap {
-    // Keeps track of the connection associated with an addr and the last time it was used
+    map: LruCache<SocketAddr, Connection>,
-    map: HashMap<SocketAddr, (Connection, u64)>,
-    // Helps to find the least recently used connection. The search and inserts are O(log(n))
-    // but since we're bounding the size of the collections, this should be constant
-    // (and hopefully negligible) time. In theory, we can do this in constant time
-    // with a queue implemented as a doubly-linked list (and all the
-    // HashMap entries holding a "pointer" to the corresponding linked-list node),
-    // so we can push, pop and bump a used connection back to the end of the queue in O(1) time, but
-    // that seems non-"Rust-y" and low bang/buck. This is still pretty terrible though...
-    last_used_times: BTreeMap<u64, SocketAddr>,
-    ticks: u64,
     use_quic: bool,
 }
 
 impl ConnMap {
     pub fn new() -> Self {
         Self {
-            map: HashMap::new(),
+            map: LruCache::new(MAX_CONNECTIONS),
-            last_used_times: BTreeMap::new(),
-            ticks: 0,
             use_quic: false,
         }
     }
@@ -60,54 +48,29 @@ pub fn set_use_quic(use_quic: bool) {
     map.set_use_quic(use_quic);
 }
 
-#[allow(dead_code)]
 // TODO: see https://github.com/solana-labs/solana/issues/23661
 // remove lazy_static and optimize and refactor this
 fn get_connection(addr: &SocketAddr) -> Connection {
     let mut map = (*CONNECTION_MAP).lock().unwrap();
-    let ticks = map.ticks;
+
-    let use_quic = map.use_quic;
+    match map.map.get(addr) {
-    let (conn, target_ticks) = match map.map.entry(*addr) {
+        Some(connection) => connection.clone(),
-        Entry::Occupied(mut entry) => {
+        None => {
-            let mut pair = entry.get_mut();
-            let old_ticks = pair.1;
-            pair.1 = ticks;
-            (pair.0.clone(), old_ticks)
-        }
-        Entry::Vacant(entry) => {
             let (_, send_socket) = solana_net_utils::bind_in_range(
                 IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
                 VALIDATOR_PORT_RANGE,
             )
             .unwrap();
-            let conn = if use_quic {
+            let connection = if map.use_quic {
                 Connection::Quic(Arc::new(QuicTpuConnection::new(send_socket, *addr)))
             } else {
                 Connection::Udp(Arc::new(UdpTpuConnection::new(send_socket, *addr)))
             };
 
-            entry.insert((conn.clone(), ticks));
+            map.map.put(*addr, connection.clone());
-            (conn, ticks)
+            connection
         }
-    };
-
-    let num_connections = map.map.len();
-    if num_connections > MAX_CONNECTIONS {
-        let (old_ticks, target_addr) = {
-            let (old_ticks, target_addr) = map.last_used_times.iter().next().unwrap();
-            (*old_ticks, *target_addr)
-        };
-        map.map.remove(&target_addr);
-        map.last_used_times.remove(&old_ticks);
     }
-
-    if target_ticks != ticks {
-        map.last_used_times.remove(&target_ticks);
-    }
-    map.last_used_times.insert(ticks, *addr);
-
-    map.ticks += 1;
-    conn
 }
 
 // TODO: see https://github.com/solana-labs/solana/issues/23851
@@ -220,11 +183,11 @@ mod tests {
         {
             let map = (*CONNECTION_MAP).lock().unwrap();
             addrs.iter().for_each(|a| {
-                let conn = map.map.get(a).expect("Address not found");
+                let conn = map.map.peek(a).expect("Address not found");
-                assert!(a.ip() == ip(conn.0.clone()));
+                assert!(a.ip() == ip(conn.clone()));
             });
 
-            assert!(map.map.get(&first_addr).is_none());
+            assert!(map.map.peek(&first_addr).is_none());
         }
 
         // Test that get_connection updates which connection is next up for eviction
@@ -237,7 +200,7 @@ mod tests {
         get_connection(&get_addr(&mut rng));
 
         let map = (*CONNECTION_MAP).lock().unwrap();
-        assert!(map.map.get(&addrs[0]).is_some());
+        assert!(map.map.peek(&addrs[0]).is_some());
-        assert!(map.map.get(&addrs[1]).is_none());
+        assert!(map.map.peek(&addrs[1]).is_none());
     }
 }

gossip/src/cluster_info_metrics.rs

@@ -571,13 +571,13 @@ pub(crate) fn submit_gossip_stats(
         .pull
         .votes
         .into_iter()
-        .map(|(slot, num_votes)| (*slot, *num_votes))
+        .map(|(slot, num_votes)| (slot, num_votes))
         .chain(
             crds_stats
                 .push
                 .votes
                 .into_iter()
-                .map(|(slot, num_votes)| (*slot, *num_votes)),
+                .map(|(slot, num_votes)| (slot, num_votes)),
         )
         .into_grouping_map()
         .aggregate(|acc, _slot, num_votes| Some(acc.unwrap_or_default() + num_votes));

programs/bpf/Cargo.lock

@@ -1636,6 +1636,15 @@ dependencies = [
  "cfg-if 1.0.0",
 ]
 
+[[package]]
+name = "lru"
+version = "0.7.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32613e41de4c47ab04970c348ca7ae7382cf116625755af070b008a15516a889"
+dependencies = [
+ "hashbrown",
+]
+
 [[package]]
 name = "matches"
 version = "0.1.9"
@@ -3368,6 +3377,7 @@ dependencies = [
  "jsonrpc-core",
  "lazy_static",
  "log",
+ "lru",
  "quinn",
  "rand 0.7.3",
  "rand_chacha 0.2.2",