fix cm2020 check hostResource

WeBankBlockchain · Nov 25, 2024 · d8697a4 · d8697a4
1 parent 89b8a6b
commit d8697a4
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 6 deletions.
diff --git a/cpp/wedpr-computing/ppc-psi/src/cm2020-psi/CM2020PSIImpl.cpp b/cpp/wedpr-computing/ppc-psi/src/cm2020-psi/CM2020PSIImpl.cpp
@@ -155,12 +155,11 @@ void CM2020PSIImpl::asyncRunTask()
 
         psi->noticePeerToFinish(task);
     });
-    // check the memory
-    checkHostResource(m_config->minNeededMemoryGB());
-    addPendingTask(taskState);
-
     try
     {
+        addPendingTask(taskState);
+        // check the memory
+        checkHostResource(m_config->minNeededMemoryGB());
         // prepare reader and writer
         auto dataResource = task->selfParty()->dataResource();
         auto reader = loadReader(task->id(), dataResource, DataSchema::Bytes);

diff --git a/cpp/wedpr-computing/ppc-psi/src/ecdh-multi-psi/EcdhMultiPSIImpl.cpp b/cpp/wedpr-computing/ppc-psi/src/ecdh-multi-psi/EcdhMultiPSIImpl.cpp
@@ -131,9 +131,9 @@ void EcdhMultiPSIImpl::asyncRunTask(
             psi->removePartner(taskID);
             psi->removePendingTask(taskID);
         });
+        addPendingTask(taskState);
         // check the memory
         checkHostResource(m_config->minNeededMemoryGB());
-        addPendingTask(taskState);
         // over the peer limit
         if (_task->getAllPeerParties().size() > c_max_peer_size)
         {

diff --git a/cpp/wedpr-computing/ppc-psi/src/labeled-psi/LabeledPSIImpl.cpp b/cpp/wedpr-computing/ppc-psi/src/labeled-psi/LabeledPSIImpl.cpp
@@ -115,9 +115,9 @@ void LabeledPSIImpl::asyncRunTask(
 
                 psi->noticePeerToFinish(_task);
             });
+            addPendingTask(taskState);
             // check the memory
             checkHostResource(m_config->minNeededMemoryGB());
-            addPendingTask(taskState);
 
             auto oprfClient = std::make_shared<EcdhOprfClient>(
                 sizeof(apsi::Item::value_type) + sizeof(apsi::LabelKey), m_config->hash(),