fix: ensure our ENR cache has enough room for max sessions + remove temp fixes (#1739)

KolbyML · web-flow · commit faddf7879fb4 · 2025-04-02T03:05:57.000-06:00
diff --git a/bin/trin/src/cli.rs b/bin/trin/src/cli.rs
@@ -17,7 +17,7 @@ use ethportal_api::{
 };
 use portalnet::{
     bootnodes::Bootnodes,
-    config::{PortalnetConfig, NODE_ADDR_CACHE_CAPACITY},
+    config::{PortalnetConfig, DISCV5_SESSION_CACHE_CAPACITY},
     constants::{
         DEFAULT_DISCOVERY_PORT, DEFAULT_NETWORK, DEFAULT_UTP_TRANSFER_LIMIT,
         DEFAULT_WEB3_HTTP_ADDRESS, DEFAULT_WEB3_IPC_PATH, DEFAULT_WEB3_WS_PORT,
@@ -433,7 +433,7 @@ impl TrinConfig {
             bootnodes: self.bootnodes.to_enrs(self.network.network()),
             no_stun: self.no_stun,
             no_upnp: self.no_upnp,
-            node_addr_cache_capacity: NODE_ADDR_CACHE_CAPACITY,
+            discv5_session_cache_capacity: DISCV5_SESSION_CACHE_CAPACITY,
             disable_poke: self.disable_poke,
             trusted_block_root: self.trusted_block_root,
             utp_transfer_limit: self.utp_transfer_limit,
diff --git a/bin/trin/src/lib.rs b/bin/trin/src/lib.rs
@@ -77,18 +77,7 @@ pub async fn run_trin(
     // Initialize and spawn uTP socket
     let (utp_talk_reqs_tx, utp_talk_reqs_rx) = mpsc::unbounded_channel();
 
-    // Set the enr_cache_capacity to the maximum uTP limit between all active networks. This is
-    // a trade off between memory usage and increased searches from the networks for each Enr.
-    // utp_transfer_limit is 2x as it would be utp_transfer_limit for incoming and
-    // utp_transfer_limit for outgoing
-    let enr_cache_capacity =
-        portalnet_config.utp_transfer_limit * 2 * trin_config.portal_subnetworks.len();
-    let discv5_utp_socket = Discv5UdpSocket::new(
-        Arc::clone(&discovery),
-        utp_talk_reqs_rx,
-        header_oracle.clone(),
-        enr_cache_capacity,
-    );
+    let discv5_utp_socket = Discv5UdpSocket::new(Arc::clone(&discovery), utp_talk_reqs_rx);
     let utp_socket = UtpSocket::with_socket(discv5_utp_socket);
     let utp_socket = Arc::new(utp_socket);
 
diff --git a/crates/portalnet/src/config.rs b/crates/portalnet/src/config.rs
@@ -5,10 +5,9 @@ use ethportal_api::types::{enr::Enr, network::Network};
 
 use crate::{bootnodes::Bootnodes, constants::DEFAULT_UTP_TRANSFER_LIMIT};
 
-/// Capacity of the cache for observed `NodeAddress` values.
-/// Provides capacity for 32 full k-buckets. This capacity will be shared among all active portal
-/// subnetworks.
-pub const NODE_ADDR_CACHE_CAPACITY: usize = discv5::kbucket::MAX_NODES_PER_BUCKET * 32;
+/// Discv5 session cache capacity.
+/// Provides capacity for 1000 nodes, to match Discv5's default session_cache_capacity value.
+pub const DISCV5_SESSION_CACHE_CAPACITY: usize = 1000;
 
 #[derive(Clone)]
 pub struct PortalnetConfig {
@@ -18,10 +17,10 @@ pub struct PortalnetConfig {
     pub bootnodes: Vec<Enr>,
     pub no_stun: bool,
     pub no_upnp: bool,
-    pub node_addr_cache_capacity: usize,
+    pub discv5_session_cache_capacity: usize,
     pub disable_poke: bool,
     pub trusted_block_root: Option<B256>,
-    // the max number of concurrent utp transfers
+    /// the max number of concurrent utp transfers
     pub utp_transfer_limit: usize,
 }
 
@@ -35,7 +34,7 @@ impl Default for PortalnetConfig {
             bootnodes: Bootnodes::default().to_enrs(Network::Mainnet),
             no_stun: false,
             no_upnp: false,
-            node_addr_cache_capacity: NODE_ADDR_CACHE_CAPACITY,
+            discv5_session_cache_capacity: DISCV5_SESSION_CACHE_CAPACITY,
             disable_poke: false,
             trusted_block_root: None,
             utp_transfer_limit: DEFAULT_UTP_TRANSFER_LIMIT,
diff --git a/crates/portalnet/src/discovery.rs b/crates/portalnet/src/discovery.rs
@@ -23,9 +23,8 @@ use ethportal_api::{
 };
 use lru::LruCache;
 use parking_lot::RwLock;
-use tokio::sync::{mpsc, RwLock as TokioRwLock};
+use tokio::sync::mpsc;
 use tracing::{debug, info, warn};
-use trin_validation::oracle::HeaderOracle;
 use utp_rs::{
     peer::{ConnectionPeer, Peer},
     udp::AsyncUdpSocket,
@@ -159,6 +158,11 @@ impl Discovery {
 
         let discv5_config = ConfigBuilder::new(listen_config)
             .request_timeout(Duration::from_secs(3))
+            // Set the session cache capacity to match our node address cache capacity. If our cache
+            // is smaller then the session cache capacity, it can lead to problems where we can't
+            // send replies to nodes that we have a session with, as we wouldn't have enough room in
+            // to store all the Enr's from all of our current established connections.
+            .session_cache_capacity(portal_config.discv5_session_cache_capacity)
             .build();
         let discv5 = Discv5::new(enr, enr_key, discv5_config)
             .map_err(|e| format!("Failed to create discv5 instance: {e}"))?;
@@ -173,7 +177,9 @@ impl Discovery {
                 .map_err(|e| format!("Failed to add bootnode enr: {e}"))?;
         }
 
-        let node_addr_cache = LruCache::new(portal_config.node_addr_cache_capacity);
+        // We set the cache capacity to double the node address cache capacity pad for
+        // inconsistencies between the two caches.
+        let node_addr_cache = LruCache::new(portal_config.discv5_session_cache_capacity * 2);
         let node_addr_cache = Arc::new(RwLock::new(node_addr_cache));
 
         Ok(Self {
@@ -213,16 +219,6 @@ impl Discovery {
                         let _ = talk_req_tx.send(talk_req).await;
                     }
                     Event::SessionEstablished(enr, socket_addr) => {
-                        // TODO: this is a temporary fix to prevent caching of eth2 nodes
-                        // and will be updated to a more stable solution as soon as it
-                        // validates the theory of what is causing the issue on mainnet.
-                        if enr.get_decodable::<String>(ENR_PORTAL_CLIENT_KEY).is_none() {
-                            debug!(
-                                enr = ?enr,
-                                "discv5 session established with node that does not have a portal client key, not caching"
-                            );
-                            continue;
-                        }
                         if let Some(old) = node_addr_cache.write().put(
                             enr.node_id(),
                             NodeAddress {
@@ -290,8 +286,20 @@ impl Discovery {
     }
 
     /// Looks up the ENR for `node_id`.
+    ///
+    /// First, it checks the cache for the `NodeAddress` and returns the `Enr` if found.
+    /// If not found, it queries the Discv5 routing table for the `node_id` and returns the `Enr` if
+    /// found.
     pub fn find_enr(&self, node_id: &NodeId) -> Option<Enr> {
-        self.discv5.find_enr(node_id)
+        if let Some(enr) = self.cached_node_addr(node_id) {
+            return Some(enr.enr);
+        }
+
+        if let Some(enr) = self.discv5.find_enr(node_id) {
+            return Some(enr);
+        }
+
+        None
     }
 
     /// Adds `enr` to the discv5 routing table.
@@ -363,24 +371,16 @@ impl Discovery {
 pub struct Discv5UdpSocket {
     talk_request_receiver: mpsc::UnboundedReceiver<TalkRequest>,
     discv5: Arc<Discovery>,
-    enr_cache: Arc<TokioRwLock<LruCache<NodeId, Enr>>>,
-    header_oracle: Arc<TokioRwLock<HeaderOracle>>,
 }
 
 impl Discv5UdpSocket {
     pub fn new(
         discv5: Arc<Discovery>,
         talk_request_receiver: mpsc::UnboundedReceiver<TalkRequest>,
-        header_oracle: Arc<TokioRwLock<HeaderOracle>>,
-        enr_cache_capacity: usize,
     ) -> Self {
-        let enr_cache = LruCache::new(enr_cache_capacity);
-        let enr_cache = Arc::new(TokioRwLock::new(enr_cache));
         Self {
             discv5,
             talk_request_receiver,
-            enr_cache,
-            header_oracle,
         }
     }
 }
@@ -437,16 +437,15 @@ impl AsyncUdpSocket<UtpPeer> for Discv5UdpSocket {
         let peer_id = *peer.id();
         let peer_enr = peer.peer().cloned();
         let discv5 = Arc::clone(&self.discv5);
-        let enr_cache = Arc::clone(&self.enr_cache);
-        let header_oracle = Arc::clone(&self.header_oracle);
         let data = buf.to_vec();
         tokio::spawn(async move {
             let enr = match peer_enr {
                 Some(enr) => enr.0,
-                None => match find_enr(&peer_id, &discv5, enr_cache, header_oracle).await {
-                    Ok(enr) => enr,
-                    Err(err) => {
-                        warn!(%err, "unable to send uTP talk request, ENR not found");
+                None => match discv5.find_enr(&peer_id) {
+                    Some(enr) => enr,
+                    None => {
+                        debug!(node_id = %peer_id, "uTP packet to unknown target");
+                        warn!( "unable to send uTP talk request, ENR not found for talk req destination");
                         return;
                     }
                 },
@@ -483,54 +482,3 @@ impl AsyncUdpSocket<UtpPeer> for Discv5UdpSocket {
         }
     }
 }
-
-async fn find_enr(
-    node_id: &NodeId,
-    discv5: &Arc<Discovery>,
-    enr_cache: Arc<TokioRwLock<LruCache<NodeId, Enr>>>,
-    header_oracle: Arc<TokioRwLock<HeaderOracle>>,
-) -> io::Result<Enr> {
-    if let Some(cached_enr) = enr_cache.write().await.get(node_id).cloned() {
-        return Ok(cached_enr);
-    }
-
-    if let Some(enr) = discv5.find_enr(node_id) {
-        enr_cache.write().await.put(*node_id, enr.clone());
-        return Ok(enr);
-    }
-
-    if let Some(enr) = discv5.cached_node_addr(node_id) {
-        enr_cache.write().await.put(*node_id, enr.enr.clone());
-        return Ok(enr.enr);
-    }
-
-    let history_jsonrpc_tx = header_oracle.read().await.history_jsonrpc_tx();
-    if let Ok(history_jsonrpc_tx) = history_jsonrpc_tx {
-        if let Ok(enr) = HeaderOracle::history_get_enr(node_id, history_jsonrpc_tx).await {
-            enr_cache.write().await.put(*node_id, enr.clone());
-            return Ok(enr);
-        }
-    }
-
-    let state_jsonrpc_tx = header_oracle.read().await.state_jsonrpc_tx();
-    if let Ok(state_jsonrpc_tx) = state_jsonrpc_tx {
-        if let Ok(enr) = HeaderOracle::state_get_enr(node_id, state_jsonrpc_tx).await {
-            enr_cache.write().await.put(*node_id, enr.clone());
-            return Ok(enr);
-        }
-    }
-
-    let beacon_jsonrpc_tx = header_oracle.read().await.beacon_jsonrpc_tx();
-    if let Ok(beacon_jsonrpc_tx) = beacon_jsonrpc_tx {
-        if let Ok(enr) = HeaderOracle::beacon_get_enr(node_id, beacon_jsonrpc_tx).await {
-            enr_cache.write().await.put(*node_id, enr.clone());
-            return Ok(enr);
-        }
-    }
-
-    debug!(node_id = %node_id, "uTP packet to unknown target");
-    Err(io::Error::new(
-        io::ErrorKind::Other,
-        "ENR not found for talk req destination",
-    ))
-}
diff --git a/crates/portalnet/src/overlay/service/manager.rs b/crates/portalnet/src/overlay/service/manager.rs
@@ -1037,9 +1037,9 @@ impl<
         // if we're unable to find the ENR for the source node we throw an error
         // since the enr is required for the accept queue, and it is expected to be present
         let enr = self.find_enr(source).ok_or_else(|| {
-            OverlayRequestError::AcceptError(
-                "handle_offer: unable to find ENR for NodeId".to_string(),
-            )
+            OverlayRequestError::AcceptError(format!(
+                "handle_offer: unable to find ENR for NodeId: source={source:?}"
+            ))
         })?;
         for (i, key) in content_keys.iter().enumerate() {
             // Accept content if within radius and not already present in the data store.
@@ -2281,19 +2281,14 @@ impl<
 
     /// Returns an ENR if one is known for the given NodeId.
     pub fn find_enr(&self, node_id: &NodeId) -> Option<Enr> {
-        // Check whether we know this node id in our X's Portal Network's routing table.
-        if let Some(node) = self.kbuckets.entry(*node_id).present_or_pending() {
-            return Some(node.enr);
-        }
-
-        // Check whether this node id is in our discv5 routing table
+        // Check whether this node id is in our enr_session_cache or discv5 routing table
         if let Some(enr) = self.discovery.find_enr(node_id) {
             return Some(enr);
         }
 
-        // Check whether this node id is in our discovery ENR cache
-        if let Some(node_addr) = self.discovery.cached_node_addr(node_id) {
-            return Some(node_addr.enr);
+        // Check whether we know this node id in our X's Portal Network's routing table.
+        if let Some(node) = self.kbuckets.entry(*node_id).present_or_pending() {
+            return Some(node.enr);
         }
 
         // Check the existing find node queries for the ENR.
@@ -2473,14 +2468,11 @@ mod tests {
     use parking_lot::lock_api::Mutex;
     use rstest::*;
     use serial_test::serial;
-    use tokio::{
-        sync::{mpsc::unbounded_channel, RwLock as TokioRwLock},
-        time::timeout,
-    };
+    use tokio::{sync::mpsc::unbounded_channel, time::timeout};
     use tokio_test::{assert_pending, assert_ready, task};
     use trin_metrics::portalnet::PORTALNET_METRICS;
     use trin_storage::{DistanceFunction, MemoryContentStore};
-    use trin_validation::{oracle::HeaderOracle, validator::MockValidator};
+    use trin_validation::validator::MockValidator;
 
     use super::*;
     use crate::{
@@ -2510,15 +2502,9 @@ mod tests {
         };
         let discovery = Arc::new(Discovery::new(portal_config, MAINNET.clone()).unwrap());
 
-        let header_oracle = HeaderOracle::default();
-        let header_oracle = Arc::new(TokioRwLock::new(header_oracle));
         let (_utp_talk_req_tx, utp_talk_req_rx) = unbounded_channel();
-        let discv5_utp = crate::discovery::Discv5UdpSocket::new(
-            Arc::clone(&discovery),
-            utp_talk_req_rx,
-            header_oracle,
-            50,
-        );
+        let discv5_utp =
+            crate::discovery::Discv5UdpSocket::new(Arc::clone(&discovery), utp_talk_req_rx);
         let utp_socket = utp_rs::socket::UtpSocket::with_socket(discv5_utp);
         let metrics = OverlayMetricsReporter {
             overlay_metrics: PORTALNET_METRICS.overlay(),
diff --git a/crates/portalnet/tests/overlay.rs b/crates/portalnet/tests/overlay.rs
@@ -26,11 +26,11 @@ use portalnet::{
     },
 };
 use tokio::{
-    sync::{mpsc, mpsc::unbounded_channel, RwLock as TokioRwLock},
+    sync::{mpsc, mpsc::unbounded_channel},
     time::{self, Duration},
 };
 use trin_storage::{ContentStore, DistanceFunction, MemoryContentStore};
-use trin_validation::{oracle::HeaderOracle, validator::MockValidator};
+use trin_validation::validator::MockValidator;
 use utp_rs::socket::UtpSocket;
 
 async fn init_overlay(
@@ -49,11 +49,8 @@ async fn init_overlay(
     let store = MemoryContentStore::new(node_id, DistanceFunction::Xor);
     let store = Arc::new(Mutex::new(store));
 
-    let header_oracle = HeaderOracle::default();
-    let header_oracle = Arc::new(TokioRwLock::new(header_oracle));
     let (_utp_talk_req_tx, utp_talk_req_rx) = unbounded_channel();
-    let discv5_utp =
-        Discv5UdpSocket::new(Arc::clone(&discovery), utp_talk_req_rx, header_oracle, 50);
+    let discv5_utp = Discv5UdpSocket::new(Arc::clone(&discovery), utp_talk_req_rx);
     let utp_socket = Arc::new(UtpSocket::with_socket(discv5_utp));
 
     let validator = Arc::new(MockValidator {});
diff --git a/testing/utp/src/lib.rs b/testing/utp/src/lib.rs
@@ -25,7 +25,6 @@ use tokio::sync::{
     mpsc::{self, Receiver},
     RwLock,
 };
-use trin_validation::oracle::HeaderOracle;
 use utp_rs::{conn::ConnectionConfig, peer::Peer, socket::UtpSocket};
 
 use crate::rpc::RpcServer;
@@ -179,15 +178,9 @@ pub async fn run_test_app(
     let enr = discovery.local_enr();
     let discovery = Arc::new(discovery);
 
-    let header_oracle = HeaderOracle::default();
-    let header_oracle = Arc::new(RwLock::new(header_oracle));
     let (utp_talk_req_tx, utp_talk_req_rx) = mpsc::unbounded_channel();
-    let discv5_utp_socket = portalnet::discovery::Discv5UdpSocket::new(
-        Arc::clone(&discovery),
-        utp_talk_req_rx,
-        header_oracle,
-        50,
-    );
+    let discv5_utp_socket =
+        portalnet::discovery::Discv5UdpSocket::new(Arc::clone(&discovery), utp_talk_req_rx);
     let utp_socket = utp_rs::socket::UtpSocket::with_socket(discv5_utp_socket);
     let utp_socket = Arc::new(utp_socket);
 
