From b30e5eb01805a61d7c2ab8b19504c619626b6750 Mon Sep 17 00:00:00 2001
From: Houdini <nonsoamadi@aol.com>
Date: Fri, 12 Jul 2024 18:26:37 +0100
Subject: [PATCH] Closes Issue #840

Modify the update_mempool function to update the jds mempool to save transactions declared by the downstream and get rid of others.
Create a clear_declared_mining_job to remove old fat transactions before a new mining job is declared.
Manage every unwrap by returning the proper error class and error message
---
 .../src/lib/job_declarator/message_handler.rs | 42 ++++++++++++++-----
 roles/jd-server/src/lib/mempool/mod.rs        | 15 +++----
 2 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/roles/jd-server/src/lib/job_declarator/message_handler.rs b/roles/jd-server/src/lib/job_declarator/message_handler.rs
index bb372cca86..33b71a8323 100644
--- a/roles/jd-server/src/lib/job_declarator/message_handler.rs
+++ b/roles/jd-server/src/lib/job_declarator/message_handler.rs
@@ -52,18 +52,40 @@ pub fn clear_declared_mining_job(
         info!("No transactions to remove from mempool");
         return Ok(());
     }
-    let clear_transactions = |jds_mempool: &mut JDsMempool| {
-        for txid in transactions_to_remove {
-            match jds_mempool.mempool.remove(txid) {
-                Some(_) => info!("Transaction {:?} removed from mempool", txid),
-                None => debug!("Transaction {:?} not found in mempool", txid),
-            };
+
+    let nonce = mining_job.tx_short_hash_nonce;
+
+    for short_id in transactions_to_remove {
+        let result = mempool.safe_lock(|mempool_| -> Result<(), Error> {
+            // Try to manage this unwrap, we use .ok_or() method to return the proper error
+            let short_ids_map = mempool_
+                .to_short_ids(nonce)
+                .ok_or(Error::JDSMissingTransactions)?;
+            let transaction_with_hash = short_ids_map
+                .get(short_id);
+
+            match transaction_with_hash {
+                Some(transaction_with_hash) => {
+                    let txid = transaction_with_hash.id;
+                    match mempool_.mempool.remove(&txid) {
+                        Some(transaction) => {
+                            debug!("Fat transaction {:?} in job with request id {:?} removed from mempool", transaction, mining_job.request_id);
+                            info!("Fat transaction {:?} in job with request id {:?} removed from mempool", txid, mining_job.request_id);
+                        },
+                        None => info!("Thin transaction {:?} in job with request id {:?} removed from mempool", txid, mining_job.request_id),
+                    }
+                },
+                None => debug!("Transaction with short id {:?} not found in mempool while clearing old jobs", short_id),
+            }
+            Ok(())  // Explicitly return Ok(()) inside the closure for proper flow control
+        });
+
+        // Propagate any error from the closure
+        if let Err(err) = result {
+            return Err(Error::PoisonLock(err.to_string()));
         }
-    };
-    match mempool.safe_lock(clear_transactions) {
-        Ok(_) => Ok(()),
-        Err(e) => Err(Error::PoisonLock(e.to_string())),
     }
+    Ok(())
 }
 
 impl ParseClientJobDeclarationMessages for JobDeclaratorDownstream {
diff --git a/roles/jd-server/src/lib/mempool/mod.rs b/roles/jd-server/src/lib/mempool/mod.rs
index 9536c69659..7c88a8638d 100644
--- a/roles/jd-server/src/lib/mempool/mod.rs
+++ b/roles/jd-server/src/lib/mempool/mod.rs
@@ -99,12 +99,13 @@ impl JDsMempool {
             .safe_lock(|x| x.get_client())?
             .ok_or(JdsMempoolError::NoClient)?;
         let node_mempool: Vec<String> = client.get_raw_mempool().await?;
-        let node_mempool_deserialized: Result<Vec<Txid>, JdsMempoolError> = node_mempool
-            .iter()
-            .map(|id| Txid::from_str(id).map_err(|err| JdsMempoolError::Rpc(RpcError::Deserialization(err.to_string()))))
-            .collect();
-        let node_mempool_deserialized = node_mempool_deserialized?;
-        
+        let mut node_mempool_deserialized: Vec<Txid> = vec![];
+        for id in &node_mempool {
+            let key_id = Txid::from_str(id)
+                .map_err(|err| JdsMempoolError::Rpc(RpcError::Deserialization(err.to_string())))?;
+            node_mempool_deserialized.push(key_id);
+        }
+
         self_.safe_lock(|x| {
             let jds_mempool = &mut x.mempool;
             // the fat transactions in the jds-mempool are those declared by some downstream and we
@@ -113,7 +114,7 @@ impl JDsMempool {
             // here we add all the new transactions
             for id in &node_mempool_deserialized {
                 jds_mempool.entry(*id).or_insert(None);
-            };
+            }
             if jds_mempool.is_empty() {
                 Err(JdsMempoolError::EmptyMempool)
             } else {