wen_restart: Fix the epoch_stakes used in calculation.

wen-restart/proto/wen_restart.proto

@@ -21,9 +21,16 @@ message LastVotedForkSlotsAggregateRecord {
     optional LastVotedForkSlotsAggregateFinal final_result = 2;
 }
 
+message LastVotedForkSlotsEpochInfoRecord {
+    uint64 epoch = 1;
+    uint64 total_active_stake = 2;
+    uint64 total_stake = 3;
+    bool considered_during_exit = 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))