optimize memory for ecdh-mul-psi

WeBankBlockchain · Nov 14, 2024 · 2fa4fa2 · 2fa4fa2
1 parent 6eab9ed
commit 2fa4fa2
Show file tree

Hide file tree

Showing 6 changed files with 45 additions and 37 deletions.
diff --git a/cpp/ppc-framework/io/DataBatch.h b/cpp/ppc-framework/io/DataBatch.h
@@ -201,14 +201,13 @@ class DataBatch
 
     std::shared_ptr<std::vector<DataType>>& mutableData() { return m_data; }
 
-    template <typename T>
     void release()
     {
         if (!m_data)
         {
             return;
         }
-        std::vector<T>().swap(*m_data);
+        m_data.reset();
     }
 
 private:

diff --git a/cpp/wedpr-computing/ppc-psi/src/ecdh-multi-psi/EcdhMultiCache.cpp b/cpp/wedpr-computing/ppc-psi/src/ecdh-multi-psi/EcdhMultiCache.cpp
@@ -128,7 +128,7 @@ std::vector<std::pair<uint64_t, bcos::bytes>> MasterCache::encryptIntersection(
             }
         });
     // Note: release the m_intersecCipher, make share it not been used after released
-    releaseIntersecCipher();
+    releaseAll();
     return cipherData;
 }
 
@@ -170,6 +170,7 @@ bool CalculatorCache::tryToFinalize()
         }
     }
     m_cacheState = CacheState::Finalized;
+    releaseDataAfterFinalize();
     ECDH_MULTI_LOG(INFO) << LOG_DESC("tryToFinalize:  compute intersection success")
                          << printCacheState()
                          << LOG_KV("intersectionSize", m_intersectionResult.size())

diff --git a/cpp/wedpr-computing/ppc-psi/src/ecdh-multi-psi/EcdhMultiCache.h b/cpp/wedpr-computing/ppc-psi/src/ecdh-multi-psi/EcdhMultiCache.h
@@ -38,18 +38,7 @@ class MasterCache
     {}
     virtual ~MasterCache()
     {
-        m_intersecCipher.clear();
-        m_finishedPartners.clear();
-        m_calculatorCipher.clear();
-        m_partnerToCipher.clear();
-        m_calculatorCipherSeqs.clear();
-        m_partnerCipherSeqs.clear();
-        std::vector<bcos::bytes>().swap(m_intersecCipher);
-        std::set<std::string>().swap(m_finishedPartners);
-        std::map<uint32_t, bcos::bytes>().swap(m_calculatorCipher);
-        std::map<std::string, std::set<bcos::bytes>>().swap(m_partnerToCipher);
-        std::set<uint32_t>().swap(m_calculatorCipherSeqs);
-        std::map<std::string, std::set<uint32_t>>().swap(m_partnerCipherSeqs);
+        releaseAll();
         MallocExtension::instance()->ReleaseFreeMemory();
         ECDH_MULTI_LOG(INFO) << LOG_DESC("the master cipher datacache destroyed ")
                              << LOG_KV("taskID", m_taskState->task()->id());
@@ -96,10 +85,24 @@ class MasterCache
         return false;
     }
 
-    void releaseIntersecCipher()
+    void releaseAll()
     {
-        std::vector<bcos::bytes>.swap(m_intersecCipher);
-        std::vector<uint32_t>.swap(m_intersecCipherIndex);
+        m_intersecCipher.clear();
+        m_intersecCipherIndex.clear();
+        m_calculatorCipher.clear();
+        m_partnerToCipher.clear();
+        // release the intersection information
+        std::vector<bcos::bytes>().swap(m_intersecCipher);
+        std::vector<uint32_t>().swap(m_intersecCipherIndex);
+
+        // release the calculator information
+        std::map<uint32_t, bcos::bytes>().swap(m_calculatorCipher);
+        // release the parterner cipher
+        std::map<std::string, std::set<bcos::bytes>>().swap(m_partnerToCipher);
+
+        MallocExtension::instance()->ReleaseFreeMemory();
+        ECDH_MULTI_LOG(INFO) << LOG_DESC("releaseAll")
+                             << LOG_KV("taskID", m_taskState->task()->id());
     }
 
 private:
@@ -139,16 +142,10 @@ class CalculatorCache
     {}
     virtual ~CalculatorCache()
     {
-        m_CipherDataFromCalculatorSubSeq.clear();
-        m_calculatorIntersectionSubSeq.clear();
-        m_receivedMasterCipher.clear();
-        m_masterCipher.clear();
-        m_intersectionCipher.clear();
-        std::set<uint32_t>().swap(m_CipherDataFromCalculatorSubSeq);
-        std::set<uint32_t>().swap(m_calculatorIntersectionSubSeq);
-        std::set<uint32_t>().swap(m_receivedMasterCipher);
-        std::set<bcos::bytes>().swap(m_masterCipher);
-        std::map<uint32_t, bcos::bytes>().swap(m_intersectionCipher);
+        releaseDataAfterFinalize();
+
+        m_intersectionResult.clear();
+        std::vector<bcos::bytes>().swap(m_intersectionResult);
         MallocExtension::instance()->ReleaseFreeMemory();
         ECDH_MULTI_LOG(INFO) << LOG_DESC("the calculator cipher datacache destroyed")
                              << LOG_KV("taskID", m_taskState->task()->id());
@@ -198,6 +195,22 @@ class CalculatorCache
 
     void syncIntersections();
 
+    void releaseDataAfterFinalize()
+    {
+        for (auto const& it : m_plainData)
+        {
+            it->release();
+        }
+        m_masterCipher.clear();
+        std::set<bcos::bytes>().swap(m_masterCipher);
+
+        m_intersectionCipher.clear();
+        std::map<uint32_t, bcos::bytes>().swap(m_intersectionCipher);
+        MallocExtension::instance()->ReleaseFreeMemory();
+        ECDH_MULTI_LOG(INFO) << LOG_DESC("releaseDataAfterFinalize")
+                             << LOG_KV("taskID", m_taskState->task()->id());
+    }
+
 private:
     TaskState::Ptr m_taskState;
     bool m_syncResult;
@@ -208,11 +221,6 @@ class CalculatorCache
     bcos::SharedMutex x_plainData;
 
 
-    // store the cipher-data of the calculator
-    std::set<uint32_t> m_CipherDataFromCalculatorSubSeq;
-    std::set<uint32_t> m_calculatorIntersectionSubSeq;
-
-
     // the cipher from the master
     std::set<bcos::bytes> m_masterCipher;
     // the seqs of the data received from master
@@ -226,4 +234,4 @@ class CalculatorCache
     bool m_receiveIntersection = false;
     std::vector<bcos::bytes> m_intersectionResult;
 };
-}  // namespace ppc::psi
+}  // namespace ppc::psi
diff --git a/cpp/wedpr-computing/ppc-psi/src/ecdh-multi-psi/core/EcdhMultiPSICalculator.cpp b/cpp/wedpr-computing/ppc-psi/src/ecdh-multi-psi/core/EcdhMultiPSICalculator.cpp
@@ -167,7 +167,7 @@ void EcdhMultiPSICalculator::loadAndEncrypt(std::string _taskID, bcos::bytes _ra
                     seq);
                 dataOffset += dataBatch->size();
             }
-            dataBatch->release<bcos::bytes>();
+            dataBatch->release();
         } while (!m_taskState->sqlReader());
     }
     catch (std::exception& e)

diff --git a/cpp/wedpr-computing/ppc-psi/src/ecdh-multi-psi/core/EcdhMultiPSIMaster.cpp b/cpp/wedpr-computing/ppc-psi/src/ecdh-multi-psi/core/EcdhMultiPSIMaster.cpp
@@ -188,7 +188,7 @@ void EcdhMultiPSIMaster::blindData()
                     }
                 });
             // can release databatch after encrypted
-            dataBatch->release<bcos::bytes>();
+            dataBatch->release();
             ECDH_MASTER_LOG(INFO) << LOG_DESC("encrypt data success")
                                   << LOG_KV("dataSize", cipher.size()) << LOG_KV("task", m_taskID)
                                   << LOG_KV("timecost", (utcSteadyTime() - startT));

diff --git a/cpp/wedpr-computing/ppc-psi/src/ecdh-multi-psi/core/EcdhMultiPSIPartner.cpp b/cpp/wedpr-computing/ppc-psi/src/ecdh-multi-psi/core/EcdhMultiPSIPartner.cpp
@@ -138,7 +138,7 @@ void EcdhMultiPSIPartner::onReceiveRandomA(bcos::bytesPointer _randA)
                 << LOG_DESC("onReceiveRandomA: send cipher data to master success")
                 << LOG_KV("size", dataBatch->size()) << printTaskInfo(m_taskState->task())
                 << LOG_KV("seq", seq);
-            dataBatch->release<bcos::bytes>();
+            dataBatch->release();
         } while (!m_taskState->sqlReader());
     }
     catch (std::exception& e)