make sure pir server end task after all result data has been fetched …

…by client (primihub#765)
phoenix20162016 · Mar 21, 2024 · 8278d03 · 8278d03
1 parent b2e5d72
commit 8278d03
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 1 deletion.
diff --git a/src/primihub/kernel/pir/operator/base_pir.h b/src/primihub/kernel/pir/operator/base_pir.h
@@ -66,6 +66,7 @@ class BasePirOperator {
   Options options_;
   std::string key_{"pir_key"};
   std::string response_key_{"response_pir_key"};
+  std::string key_task_end_{"pir_task_end"};
 };
 }  // namespace primihub::pir
 #endif  // SRC_PRIMIHUB_KERNEL_PIR_OPERATOR_BASE_PIR_H_
diff --git a/src/primihub/kernel/pir/operator/keyword_pir_impl/keyword_pir_client.cc b/src/primihub/kernel/pir/operator/keyword_pir_impl/keyword_pir_client.cc
@@ -93,6 +93,12 @@ retcode KeywordPirOperatorClient::OnExecute(const PirDataType& input,
   VLOG(5) << "query_resultquery_resultquery_resultquery_result: "
           << query_result.size();
   ExtractResult(orig_item, query_result, result);
+  {
+    std::string task_end{"SUCCESS"};
+    auto link_ctx = this->GetLinkContext();
+    ret = link_ctx->Send(this->key_task_end_, PeerNode(), task_end);
+    CHECK_RETCODE_WITH_RETVALUE(ret, retcode::FAIL);
+  }
   return retcode::SUCCESS;
 }
 

diff --git a/src/primihub/kernel/pir/operator/keyword_pir_impl/keyword_pir_server.cc b/src/primihub/kernel/pir/operator/keyword_pir_impl/keyword_pir_server.cc
@@ -86,7 +86,13 @@ retcode KeywordPirOperatorServer::OnExecute(const PirDataType& input,
 
   ret = ProcessQuery(sender_db);
   CHECK_RETCODE_WITH_RETVALUE(ret, retcode::FAIL);
-
+  {
+    std::string task_end;
+    auto link_ctx = this->GetLinkContext();
+    ret = link_ctx->Recv(this->key_task_end_, ProxyNode(), &task_end);
+    CHECK_RETCODE_WITH_RETVALUE(ret, retcode::FAIL);
+    LOG(INFO) << "task status: " << task_end;
+  }
   VLOG(5) << "end of execute task";
   return retcode::SUCCESS;
 }