ecdh psi avoid blocking in MaskPeer (#233)

Add non-block sendbatch for mask_peer send: Attention: add addition api temporary to avoid affecting ecdh_3pc_psi, if ecdh_3pc_psi has the same problem, old api can be modified directly.
secretflow · Jul 6, 2023 · 420dfd7 · 420dfd7
1 parent 38c5d3f
commit 420dfd7
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 5 deletions.
diff --git a/libspu/psi/core/ecdh_psi.cc b/libspu/psi/core/ecdh_psi.cc
@@ -113,7 +113,8 @@ void EcdhPsiContext::MaskPeer(
     // Should send out the dual masked items to peer.
     if (PeerCanTouchResults()) {
       const auto tag = fmt::format("ECDHPSI:Y^B^A:{}", batch_count);
-      SendDualMaskedBatch(dual_masked_peers, batch_count, tag);
+      // call non-block to avoid blocking each other with MaskSelf
+      SendDualMaskedBatchNonBlock(dual_masked_peers, batch_count, tag);
     }
     if (peer_items.empty()) {
       SPDLOG_INFO("MaskPeer:{}--finished, batch_count={}, peer_item_count={}",
@@ -156,10 +157,8 @@ void EcdhPsiContext::RecvDualMaskedSelf(
 namespace {
 
 template <typename T>
-void SendBatchImpl(const std::vector<T>& batch_items,
-                   const std::shared_ptr<yacl::link::Context>& link_ctx,
-                   std::string_view type, int32_t batch_idx,
-                   std::string_view tag) {
+PsiDataBatch BatchData(const std::vector<T>& batch_items, std::string_view type,
+                       int32_t batch_idx) {
   PsiDataBatch batch;
   batch.is_last_batch = batch_items.empty();
   batch.item_num = batch_items.size();
@@ -171,12 +170,30 @@ void SendBatchImpl(const std::vector<T>& batch_items,
       batch.flatten_bytes.append(item);
     }
   }
+  return batch;
+}
 
+template <typename T>
+void SendBatchImpl(const std::vector<T>& batch_items,
+                   const std::shared_ptr<yacl::link::Context>& link_ctx,
+                   std::string_view type, int32_t batch_idx,
+                   std::string_view tag) {
+  auto batch = BatchData<T>(batch_items, type, batch_idx);
   link_ctx->SendAsyncThrottled(
       link_ctx->NextRank(), IcPsiBatchSerializer::Serialize(std::move(batch)),
       tag);
 }
 
+template <typename T>
+void SendBatchNonBlockImpl(const std::vector<T>& batch_items,
+                           const std::shared_ptr<yacl::link::Context>& link_ctx,
+                           std::string_view type, int32_t batch_idx,
+                           std::string_view tag) {
+  auto batch = BatchData<T>(batch_items, type, batch_idx);
+  link_ctx->SendAsync(link_ctx->NextRank(),
+                      IcPsiBatchSerializer::Serialize(std::move(batch)), tag);
+}
+
 void RecvBatchImpl(const std::shared_ptr<yacl::link::Context>& link_ctx,
                    int32_t batch_idx, std::string_view tag,
                    std::vector<std::string>* items) {
@@ -222,6 +239,13 @@ void EcdhPsiContext::SendDualMaskedBatch(
   SendBatchImpl(batch_items, dual_mask_link_ctx_, "dual.enc", batch_idx, tag);
 }
 
+void EcdhPsiContext::SendDualMaskedBatchNonBlock(
+    const std::vector<std::string>& batch_items, int32_t batch_idx,
+    std::string_view tag) {
+  SendBatchNonBlockImpl(batch_items, dual_mask_link_ctx_, "dual.enc", batch_idx,
+                        tag);
+}
+
 void EcdhPsiContext::RecvDualMaskedBatch(std::vector<std::string>* items,
                                          int32_t batch_idx,
                                          std::string_view tag) {

diff --git a/libspu/psi/core/ecdh_psi.h b/libspu/psi/core/ecdh_psi.h
@@ -113,6 +113,10 @@ class EcdhPsiContext {
   void SendDualMaskedBatch(const std::vector<std::string_view>& batch_items,
                            int32_t batch_idx, std::string_view tag = "");
 
+  void SendDualMaskedBatchNonBlock(const std::vector<std::string>& batch_items,
+                                   int32_t batch_idx,
+                                   std::string_view tag = "");
+
   void RecvDualMaskedBatch(std::vector<std::string>* items, int32_t batch_idx,
                            std::string_view tag = "");