From 8c8220faaf6e528f37f166f533368cc746fbf016 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Wed, 10 Apr 2024 17:01:43 -0600
Subject: [PATCH] v1.18: default staked client in LocalCluster (backport of
 #716) (#724)

* default staked client in LocalCluster (#716)

* default staked client in LocalCluster

* fix underflow

(cherry picked from commit e91a5e27444af47c63946a971244c1075bbc99bf)

# Conflicts:
#	local-cluster/tests/local_cluster.rs

* resolve conflicts

---------

Co-authored-by: Andrew Fitzgerald <apfitzge@gmail.com>
---
 local-cluster/src/local_cluster.rs | 56 +++++++++++++++++++++++-------
 1 file changed, 44 insertions(+), 12 deletions(-)

diff --git a/local-cluster/src/local_cluster.rs b/local-cluster/src/local_cluster.rs
index 9d1b483d85fdd3..1ffd935a82892a 100644
--- a/local-cluster/src/local_cluster.rs
+++ b/local-cluster/src/local_cluster.rs
@@ -2,6 +2,7 @@ use {
     crate::{
         cluster::{Cluster, ClusterValidatorInfo, ValidatorInfo},
         cluster_tests,
+        integration_tests::DEFAULT_NODE_STAKE,
         validator_configs::*,
     },
     itertools::izip,
@@ -45,7 +46,7 @@ use {
         transaction::Transaction,
     },
     solana_stake_program::stake_state,
-    solana_streamer::socket::SocketAddrSpace,
+    solana_streamer::{socket::SocketAddrSpace, streamer::StakedNodes},
     solana_tpu_client::tpu_client::{
         DEFAULT_TPU_CONNECTION_POOL_SIZE, DEFAULT_TPU_ENABLE_UDP, DEFAULT_TPU_USE_QUIC,
     },
@@ -57,7 +58,7 @@ use {
         collections::HashMap,
         io::{Error, ErrorKind, Result},
         iter,
-        net::UdpSocket,
+        net::{IpAddr, Ipv4Addr, UdpSocket},
         path::{Path, PathBuf},
         sync::{Arc, RwLock},
     },
@@ -186,6 +187,46 @@ impl LocalCluster {
     pub fn new(config: &mut ClusterConfig, socket_addr_space: SocketAddrSpace) -> Self {
         assert_eq!(config.validator_configs.len(), config.node_stakes.len());
 
+        let connection_cache = match config.tpu_use_quic {
+            true => {
+                let client_keypair = Keypair::new();
+                let stake = DEFAULT_NODE_STAKE;
+
+                for validator_config in config.validator_configs.iter_mut() {
+                    let mut overrides = HashMap::new();
+                    overrides.insert(client_keypair.pubkey(), stake);
+                    validator_config.staked_nodes_overrides = Arc::new(RwLock::new(overrides));
+                }
+
+                assert!(
+                    config.tpu_use_quic,
+                    "no support for staked override forwarding without quic"
+                );
+
+                let total_stake = config.node_stakes.iter().sum::<u64>();
+                let stakes = HashMap::from([
+                    (client_keypair.pubkey(), stake),
+                    (Pubkey::new_unique(), total_stake.saturating_sub(stake)),
+                ]);
+                let staked_nodes = Arc::new(RwLock::new(StakedNodes::new(
+                    Arc::new(stakes),
+                    HashMap::<Pubkey, u64>::default(), // overrides
+                )));
+
+                Arc::new(ConnectionCache::new_with_client_options(
+                    "connection_cache_local_cluster_quic_staked",
+                    config.tpu_connection_pool_size,
+                    None,
+                    Some((&client_keypair, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)))),
+                    Some((&staked_nodes, &client_keypair.pubkey())),
+                ))
+            }
+            false => Arc::new(ConnectionCache::with_udp(
+                "connection_cache_local_cluster_udp",
+                config.tpu_connection_pool_size,
+            )),
+        };
+
         let mut validator_keys = {
             if let Some(ref keys) = config.validator_keys {
                 assert_eq!(config.validator_configs.len(), keys.len());
@@ -318,16 +359,7 @@ impl LocalCluster {
             entry_point_info: leader_contact_info,
             validators,
             genesis_config,
-            connection_cache: match config.tpu_use_quic {
-                true => Arc::new(ConnectionCache::new_quic(
-                    "connection_cache_local_cluster_quic",
-                    config.tpu_connection_pool_size,
-                )),
-                false => Arc::new(ConnectionCache::with_udp(
-                    "connection_cache_local_cluster_udp",
-                    config.tpu_connection_pool_size,
-                )),
-            },
+            connection_cache,
         };
 
         let node_pubkey_to_vote_key: HashMap<Pubkey, Arc<Keypair>> = keys_in_genesis