check backups before registering

marlinprotocol · Dec 12, 2024 · 4060b7d · 4060b7d
1 parent 7f063da
commit 4060b7d
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 3 deletions.
diff --git 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<Address, HashMap<U256, U256>>, // Generator -> Markets -> Kalypso Points Per Market
 
     withdrawl_requests: HashMap<Address, HashSet<WithdrawlRequest>>,
+
+    stake_backup: HashMap<Address, StakeBackup>,
+}
+
+#[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<StakeBackup> {
+        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
@@ -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(),