From 4ddfa774621a72ec2c572ed7e6b663f2b4863674 Mon Sep 17 00:00:00 2001
From: musitdev <philippe.delrieu@free.fr>
Date: Mon, 2 Oct 2023 10:28:54 +0200
Subject: [PATCH] set schedule epoch plus one to schedule one epoch after

---
 stake_aggregate/src/epoch.rs           | 23 ++++++++++-------------
 stake_aggregate/src/leader_schedule.rs | 11 ++++++++++-
 2 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/stake_aggregate/src/epoch.rs b/stake_aggregate/src/epoch.rs
index cab2a7c..e35ee22 100644
--- a/stake_aggregate/src/epoch.rs
+++ b/stake_aggregate/src/epoch.rs
@@ -73,6 +73,7 @@ impl CurrentEpochSlotState {
     pub fn current_epoch_end_slot(&self) -> Slot {
         self.next_epoch_start_slot - 1
     }
+
     pub fn process_new_slot(
         &mut self,
         new_slot: &SubscribeUpdateSlot,
@@ -94,7 +95,11 @@ impl CurrentEpochSlotState {
                 let mut diff = new_slot.slot - self.current_slot.confirmed_slot;
                 //First epoch slot, index is 0 so remove 1 from diff.
                 if self.first_epoch_slot {
-                    change_epoch(&mut self.current_epoch, new_slot.slot);
+                    //calculate next epoch data
+                    self.current_epoch.epoch += 1;
+                    //slot can be non consecutif, use diff.
+                    self.current_epoch.slot_index = 0;
+                    self.current_epoch.absolute_slot = new_slot.slot;
                     log::info!(
                         "change_epoch calculated next epoch:{:?} at slot:{}",
                         self.current_epoch,
@@ -140,9 +145,10 @@ impl CurrentEpochSlotState {
             self.first_epoch_slot = true;
 
             //start leader schedule calculus
-            //switch to next epoch to calculate schedule at next epoch.
-            let mut schedule_epoch = self.current_epoch.clone();
-            change_epoch(&mut schedule_epoch, self.current_slot.confirmed_slot);
+            //switch to 2 next epoch to calculate schedule at next epoch.
+            //at current epoch change the schedule is calculated for the next epoch.
+            let schedule_epoch = crate::leader_schedule::next_schedule_epoch(&self.current_epoch);
+            let schedule_epoch = crate::leader_schedule::next_schedule_epoch(&schedule_epoch);
             Some(LeaderScheduleEvent::InitLeaderschedule(schedule_epoch))
         } else {
             None
@@ -150,15 +156,6 @@ impl CurrentEpochSlotState {
     }
 }
 
-fn change_epoch(current_epoch: &mut EpochInfo, current_slot: Slot) {
-    current_epoch.epoch += 1;
-    //slot can be non consecutif, use diff.
-    current_epoch.slot_index = current_epoch
-        .slot_index
-        .saturating_sub(current_epoch.slots_in_epoch);
-    current_epoch.absolute_slot = current_slot;
-}
-
 #[derive(Default, Debug, Clone)]
 pub struct CurrentSlot {
     pub processed_slot: u64,
diff --git a/stake_aggregate/src/leader_schedule.rs b/stake_aggregate/src/leader_schedule.rs
index bd76844..f4bae42 100644
--- a/stake_aggregate/src/leader_schedule.rs
+++ b/stake_aggregate/src/leader_schedule.rs
@@ -21,7 +21,7 @@ use std::time::Duration;
 use std::time::{SystemTime, UNIX_EPOCH};
 use tokio::task::JoinHandle;
 
-const SCHEDULE_STAKE_BASE_FILE_NAME: &str = "aggregate_export_votestake_";
+const SCHEDULE_STAKE_BASE_FILE_NAME: &str = "aggregate_export_votestake";
 
 pub const MAX_EPOCH_VALUE: u64 = 18446744073709551615;
 
@@ -57,6 +57,15 @@ impl From<SavedStake> for EpochStake {
     }
 }
 
+pub fn next_schedule_epoch(current_epoch: &EpochInfo) -> EpochInfo {
+    let mut next_epoch_info = current_epoch.clone();
+    next_epoch_info.epoch += 1;
+    next_epoch_info.slot_index = 0;
+    next_epoch_info.absolute_slot =
+        current_epoch.absolute_slot + current_epoch.slots_in_epoch - current_epoch.slot_index;
+    next_epoch_info
+}
+
 pub fn bootstrap_leader_schedule(
     current_file_patch: &str,
     next_file_patch: &str,