wen_restart: Fix the epoch_stakes used in calculation. (anza-xyz#2376)

* wen_restart: Fix the epoch_stakes used in calculation.

* Fix a bad merge.

* Remove EpochStakesCache, it only caches epoch stakes from root_bank, better to just keep root_bank around.

* Split aggregate into smaller functions.

* Switch to node_id_to_stake which is simpler.

* Rename update_slots_stake_map and switch to epoch_total_stake().

* Remove unnecessary utility functions.

* Do not modify epoch_info_vec, just init it with two epochs we will consider.

* Switch to epoch_node_id_to_stake()

* Add test for divergence at Epoch boundary.

* Make linter happy.

* - wait for the new Epoch if > 1/3 of the validators voted for some
  slot in the new Epoch
- switch to voted_percent and voted_for_this_epoch_percent

* Fix a bad merge.

* Fix a bad merge.

* Change constant format.

* Do not loop through the whole table.

* Address reviewer feedback.

* Address reviewer comments.

wen-restart/Cargo.toml

@@ -34,6 +34,7 @@ solana-entry = { workspace = true }
 solana-logger = { workspace = true }
 solana-runtime = { workspace = true, features = ["dev-context-only-utils"] }
 solana-streamer = { workspace = true }
+solana-vote = { workspace = true }
 tempfile = { workspace = true }
 
 [build-dependencies]

wen-restart/proto/wen_restart.proto

@@ -21,9 +21,16 @@ message LastVotedForkSlotsAggregateRecord {
     optional LastVotedForkSlotsAggregateFinal final_result = 2;
 }
 
+message LastVotedForkSlotsEpochInfoRecord {
+    uint64 epoch = 1;
+    uint64 total_stake = 2;
+    uint64 actively_voting_stake = 3;
+    uint64 actively_voting_for_this_epoch_stake = 4;
+}
+
 message LastVotedForkSlotsAggregateFinal {
     map<uint64, uint64> slots_stake_map = 1;
-    uint64 total_active_stake = 2;
+    repeated LastVotedForkSlotsEpochInfoRecord epoch_infos = 2;
 }
 
 message HeaviestForkRecord {

wen-restart/src/heaviest_fork_aggregate.rs

@@ -15,7 +15,8 @@ pub(crate) struct HeaviestForkAggregate {
     supermajority_threshold: f64,
     my_shred_version: u16,
     my_pubkey: Pubkey,
-    // TODO(wen): using local root's EpochStakes, need to fix if crossing Epoch boundary.
+    // We use the epoch_stakes of the Epoch our heaviest bank is in. Proceed and exit only if
+    // enough validator agree with me.
     epoch_stakes: EpochStakes,
     heaviest_forks: HashMap<Pubkey, RestartHeaviestFork>,
     block_stake_map: HashMap<(Slot, Hash), u64>,
@@ -171,7 +172,6 @@ impl HeaviestForkAggregate {
         Some(record)
     }
 
-    // TODO(wen): use better epoch stake and add a test later.
     pub(crate) fn total_active_stake(&self) -> u64 {
         self.active_peers.iter().fold(0, |sum: u64, pubkey| {
             sum.saturating_add(self.epoch_stakes.node_id_to_stake(pubkey).unwrap_or(0))