diff --git a/matching_engine/src/generator_lib/generator_store.rs b/matching_engine/src/generator_lib/generator_store.rs index a2f4013..6f17d41 100644 --- a/matching_engine/src/generator_lib/generator_store.rs +++ b/matching_engine/src/generator_lib/generator_store.rs @@ -303,6 +303,14 @@ pub struct GeneratorStore { kalypso_points_per_market: HashMap>, // Generator -> Markets -> Kalypso Points Per Market withdrawl_requests: HashMap>, + + stake_backup: HashMap, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StakeBackup { + pub native_stake: TokenTracker, + pub symbiotic_stake: TokenTracker, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd)] @@ -465,6 +473,7 @@ impl GeneratorStore { kalypso_points: HashMap::new(), kalypso_points_per_market: HashMap::new(), withdrawl_requests: HashMap::new(), + stake_backup: HashMap::new(), } } @@ -580,7 +589,22 @@ impl GeneratorStore { } pub fn remove_by_address(&mut self, address: &Address) { - self.generators.remove(address); + if let Some(generator) = self.generators.get(address) { + let native_stake = generator.total_native_stake.clone(); + let symbiotic_stake = generator.total_symbiotic_stake.clone(); + + let backup = StakeBackup { + native_stake, + symbiotic_stake, + }; + + self.stake_backup.insert(address.clone(), backup); + self.generators.remove(address); + } + } + + pub fn get_stake_backup(&self, address: &Address) -> Option { + self.stake_backup.get(address).cloned() } pub fn add_extra_stake( diff --git a/matching_engine/src/log_processor/gr.rs b/matching_engine/src/log_processor/gr.rs index 235fa8e..83c3d72 100644 --- a/matching_engine/src/log_processor/gr.rs +++ b/matching_engine/src/log_processor/gr.rs @@ -63,11 +63,26 @@ pub async fn process_generator_registry_logs( .await .unwrap(); + //check if any native or symbiotic stake already existed + let stake_backup = generator_store.get_stake_backup(&address); + let generator = generator_store::Generator { address, reward_address: generator_data.0, - total_native_stake: TokenTracker::new(), - total_symbiotic_stake: TokenTracker::new(), + total_native_stake: { + if stake_backup.is_some() { + stake_backup.clone().unwrap().native_stake.clone() + } else { + TokenTracker::new() + } + }, + total_symbiotic_stake: { + if stake_backup.is_some() { + stake_backup.unwrap().symbiotic_stake.clone() + } else { + TokenTracker::new() + } + }, sum_of_compute_allocations: 0.into(), compute_consumed: 0.into(), native_stake_locked: TokenTracker::new(),