Do not assign a single worker for mpn works

ziesha-network · May 17, 2023 · 602a864 · 602a864
1 parent cd2ced4
commit 602a864
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 11 deletions.
diff --git a/src/mpn/mod.rs b/src/mpn/mod.rs
@@ -62,6 +62,7 @@ pub struct MpnWorkPool {
     final_delta: ZkDeltaPairs,
     works: HashMap<usize, MpnWork>,
     solutions: HashMap<usize, MpnSolution>,
+    assignees: HashMap<Address, Vec<usize>>,
 }
 
 impl MpnWorkPool {
@@ -73,11 +74,20 @@ impl MpnWorkPool {
         remaining
     }
     pub fn get_works(&self, address: Address) -> HashMap<usize, MpnWork> {
-        let mut result: HashMap<usize, MpnWork> = self
-            .remaining_works()
-            .into_iter()
-            .filter(|(_, v)| v.worker.address == address)
-            .collect();
+        if let Some(works) = self.assignees.get(&address) {
+            return works
+                .iter()
+                .filter_map(|i| {
+                    if let Some(w) = self.works.get(i) {
+                        Some((*i, w.clone()))
+                    } else {
+                        None
+                    }
+                })
+                .collect();
+        }
+
+        let mut result: HashMap<usize, MpnWork> = self.remaining_works().into_iter().collect();
 
         if result.is_empty() {
             result = self.remaining_works();
@@ -244,7 +254,6 @@ pub struct MpnWork {
     pub public_inputs: ZkPublicInputs,
     pub data: MpnWorkData,
     pub new_root: ZkCompressedState,
-    pub worker: MpnWorker,
     pub reward: Amount,
 }
 
@@ -341,7 +350,6 @@ pub fn prepare_works<K: KvStore, B: Blockchain<K>>(
             public_inputs,
             new_root,
             data: MpnWorkData::Deposit(transitions),
-            worker: workers[worker_id].clone(),
             reward: deposit_reward,
         });
         worker_id = (worker_id + 1) % workers.len();
@@ -365,7 +373,6 @@ pub fn prepare_works<K: KvStore, B: Blockchain<K>>(
             public_inputs,
             new_root,
             data: MpnWorkData::Withdraw(transitions),
-            worker: workers[worker_id].clone(),
             reward: withdraw_reward,
         });
         worker_id = (worker_id + 1) % workers.len();
@@ -390,7 +397,6 @@ pub fn prepare_works<K: KvStore, B: Blockchain<K>>(
             public_inputs,
             new_root,
             data: MpnWorkData::Update(transitions),
-            worker: workers[worker_id].clone(),
             reward: update_reward,
         });
         worker_id = (worker_id + 1) % workers.len();
@@ -402,6 +408,7 @@ pub fn prepare_works<K: KvStore, B: Blockchain<K>>(
         works: works.into_iter().enumerate().collect(),
         final_delta,
         solutions: HashMap::new(),
+        assignees: HashMap::new(),
     })
 }
 

diff --git a/src/node/heartbeat/generate_block.rs b/src/node/heartbeat/generate_block.rs
@@ -86,8 +86,8 @@ pub async fn generate_block<K: KvStore, B: Blockchain<K>>(
         if let Some(claim) = ctx.validator_claim.clone() {
             if claim.address == ctx.validator_wallet.get_address() {
                 if let Some(work_pool) = &ctx.mpn_work_pool {
-                    for work in work_pool.remaining_works().values() {
-                        log::error!("Prover {} is late!", work.worker.address);
+                    for work in work_pool.remaining_works().keys() {
+                        log::error!("Solution for work {} is late!", work);
                     }
                 }
             }