Custom transactions from JDC removed from mempool during update

roles/jd-server/src/lib/job_declarator/message_handler.rs

@@ -7,14 +7,17 @@ use roles_logic_sv2::{
         ProvideMissingTransactions, ProvideMissingTransactionsSuccess, SubmitSolutionJd,
     },
     parsers::JobDeclaration,
+    utils::Mutex,
 };
-use std::{convert::TryInto, io::Cursor};
+use std::{convert::TryInto, io::Cursor, sync::Arc};
 use stratum_common::bitcoin::{Transaction, Txid};
 pub type SendTo = SendTo_<JobDeclaration<'static>, ()>;
+use crate::mempool::JDsMempool;
+
 use super::{signed_token, TransactionState};
 use roles_logic_sv2::{errors::Error, parsers::PoolMessages as AllMessages};
 use stratum_common::bitcoin::consensus::Decodable;
-use tracing::info;
+use tracing::{debug, info};
 
 use super::JobDeclaratorDownstream;
 
@@ -38,6 +41,53 @@ impl JobDeclaratorDownstream {
     }
 }
 
+pub fn clear_declared_mining_job(
+    mining_job: DeclareMiningJob,
+    mempool: Arc<Mutex<JDsMempool>>,
+) -> Result<(), Error> {
+    // If there is an old declared mining job, remove its transactions from the mempool
+    // Retrieve necessary data from the old job
+    let transactions_to_remove = mining_job.tx_short_hash_list.inner_as_ref();
+    if transactions_to_remove.is_empty() {
+        info!("No transactions to remove from mempool");
+        return Ok(());
+    }
+
+    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()));
+        }
+    }
+    Ok(())
+}
+
 impl ParseClientJobDeclarationMessages for JobDeclaratorDownstream {
     fn handle_allocate_mining_job_token(
         &mut self,
@@ -61,6 +111,11 @@ impl ParseClientJobDeclarationMessages for JobDeclaratorDownstream {
     }
 
     fn handle_declare_mining_job(&mut self, message: DeclareMiningJob) -> Result<SendTo, Error> {
+        // Clone the old declared mining job to retain its data
+        if let Some(old_declare_mining_job_) = self.declared_mining_job.0.clone() {
+            clear_declared_mining_job(old_declare_mining_job_, self.mempool.clone())?;
+        }
+
         // the transactions that are present in the mempool are stored here, that is sent to the
         // mempool which use the rpc client to retrieve the whole data for each transaction.
         // The unknown transactions is a vector that contains the transactions that are not in the

roles/jd-server/src/lib/mempool/mod.rs

@@ -7,7 +7,7 @@ use hashbrown::HashMap;
 use roles_logic_sv2::utils::Mutex;
 use rpc_sv2::{mini_rpc_client, mini_rpc_client::RpcError};
 use std::{convert::TryInto, str::FromStr, sync::Arc};
-use stratum_common::{bitcoin, bitcoin::hash_types::Txid};
+use stratum_common::bitcoin::{self, hash_types::Txid};
 
 #[derive(Clone, Debug)]
 pub struct TransactionWithHash {
@@ -95,31 +95,32 @@ impl JDsMempool {
     }
 
     pub async fn update_mempool(self_: Arc<Mutex<Self>>) -> Result<(), JdsMempoolError> {
-        let mut mempool_ordered: HashMap<Txid, Option<Transaction>> = HashMap::new();
-
         let client = self_
             .safe_lock(|x| x.get_client())?
             .ok_or(JdsMempoolError::NoClient)?;
-
-        let mempool: Vec<String> = client.get_raw_mempool().await?;
-        for id in &mempool {
+        let node_mempool: Vec<String> = client.get_raw_mempool().await?;
+        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())))?;
-
-            let tx = self_.safe_lock(|x| match x.mempool.get(&key_id) {
-                Some(entry) => entry.clone(),
-                None => None,
-            })?;
-
-            mempool_ordered.insert(key_id, tx);
+            node_mempool_deserialized.push(key_id);
         }
 
-        if mempool_ordered.is_empty() {
-            Err(JdsMempoolError::EmptyMempool)
-        } else {
-            let _ = self_.safe_lock(|x| x.mempool = mempool_ordered);
-            Ok(())
-        }
+        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
+            // don't want to remove them, but we can get rid of the others
+            jds_mempool.retain(|_, val| val.is_some());
+            // 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 {
+                Ok(())
+            }
+        })?
     }
 
     pub async fn on_submit(self_: Arc<Mutex<Self>>) -> Result<(), JdsMempoolError> {