From b801f9bf824f35f04faf62ff33186e4908cdbd31 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20Patron?= <francois.patron@kandy.eu>
Date: Sat, 28 Dec 2024 17:38:21 +0100
Subject: [PATCH] Add last TIG changes

---
 tig-benchmarker/slave.py | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/tig-benchmarker/slave.py b/tig-benchmarker/slave.py
index 82c706e..d2387b6 100644
--- a/tig-benchmarker/slave.py
+++ b/tig-benchmarker/slave.py
@@ -10,7 +10,7 @@ import time
 import zlib
 from threading import Thread
 from common.structs import OutputData, MerkleProof
-from common.merkle_tree import MerkleTree
+from common.merkle_tree import MerkleTree, MerkleHash
 from datetime import datetime, timedelta
 
 logger = logging.getLogger(os.path.splitext(os.path.basename(__file__))[0])
@@ -97,7 +97,11 @@ def send_results(session, master_ip, master_port, tig_worker_path, download_wasm
     try:
         batch_id = READY_BATCH_IDS.pop()
     except KeyError:
-        logger.debug("No pending batches")
+        logger.debug("No batches to send")
+        return
+    
+    if now() - FINISHED_BATCH_IDS.get(batch_id, 0) < 10000:
+        logger.debug(f"Batch {batch_id} submitted recently")
         return
     
     output_folder = f"{output_path}/{batch_id}"
@@ -107,6 +111,7 @@ def send_results(session, master_ip, master_port, tig_worker_path, download_wasm
     if (
         not os.path.exists(f"{output_folder}/result.json")
         or not os.path.exists(f"{output_folder}/data.zlib")
+        or not os.path.exists(f"{output_folder}/hashes.zlib")
     ):
         if os.path.exists(f"{output_folder}/result.json"):
             os.remove(f"{output_folder}/result.json")
@@ -132,22 +137,21 @@ def send_results(session, master_ip, master_port, tig_worker_path, download_wasm
             READY_BATCH_IDS.add(batch_id) # requeue
 
     else:
+        with open(f"{output_folder}/hashes.zlib", "rb") as f:
+            hashes = json.loads(zlib.decompress(f.read()).decode())
         with open(f"{output_folder}/data.zlib", "rb") as f:
-            leafs = [
-                OutputData.from_dict(x) 
-                for x in json.loads(zlib.decompress(f.read()).decode())
-            ]
-            
+            leafs = json.loads(zlib.decompress(f.read()).decode())
+        
         merkle_tree = MerkleTree(
-            [x.to_merkle_hash() for x in leafs],
+            [MerkleHash.from_str(x) for x in hashes],
             batch["batch_size"]
         )
 
         proofs_to_submit = [
-            MerkleProof(
+            dict(
                 leaf=leafs[n - batch["start_nonce"]],
-                branch=merkle_tree.calc_merkle_branch(branch_idx=n - batch["start_nonce"])
-            ).to_dict()
+                branch=merkle_tree.calc_merkle_branch(branch_idx=n - batch["start_nonce"]).to_str()
+            )
             for n in batch["sampled_nonces"]
         ]