Remove memcpy from EmpNetworkAdapter (#193)

Summary:
Pull Request resolved: #193

This diff remove `memcpy` from EmpNetworkAdapter to improve performance. `memcpy` is expensive, as it is essentially adding linear overhead for every single operation.

Changes:

- Add `send_impl` and `recv_impl` functions to IPartyCommunicationAgent
- Refactor classes that implemented IPartyCommunicationAgent to call `send_impl` and `recv_impl` in their respective `send` and `receive` functions
- Refactor EmpNetworkAdapter to use `send_impl` and `recv_impl`

Reviewed By: RuiyuZhu

Differential Revision: D35818108

fbshipit-source-id: a8c6f842df1a58e388dea7fec1fe80ff2fcf0747

fbpcf/engine/communication/IPartyCommunicationAgent.h

@@ -16,8 +16,11 @@
 #error "Machine must be little endian"
 #endif
 
-namespace fbpcf::engine::communication {
+namespace fbpcf::engine::util {
+class EmpNetworkAdapter;
+}
 
+namespace fbpcf::engine::communication {
 /**
  * This is the network API between two parties.
  * NOTE: sendT/receiveT only work when the two parties have the same endianness
@@ -89,6 +92,8 @@ class IPartyCommunicationAgent {
   virtual std::pair<uint64_t, uint64_t> getTrafficStatistics() const = 0;
 
  private:
+  friend class util::EmpNetworkAdapter;
+
   // convert a vector of bits into a vector of bytes
   static std::vector<unsigned char> compressToBytes(
       const std::vector<bool>& bits) {
@@ -125,6 +130,10 @@ class IPartyCommunicationAgent {
     }
     return bits;
   }
+
+  virtual void recvImpl(void* data, int nBytes) = 0;
+
+  virtual void sendImpl(const void* data, int nBytes) = 0;
 };
 
 template <>

fbpcf/engine/communication/InMemoryPartyCommunicationAgentHost.cpp

@@ -9,20 +9,33 @@
 
 namespace fbpcf::engine::communication {
 
+void InMemoryPartyCommunicationAgent::sendImpl(const void* data, int nBytes) {
+  std::vector<unsigned char> buffer(nBytes);
+  memcpy(buffer.data(), data, nBytes);
+  host_.send(myId_, buffer);
+  sentData_ += nBytes;
+}
+
 void InMemoryPartyCommunicationAgent::send(
     const std::vector<unsigned char>& data) {
-  host_.send(myId_, data);
-  sentData_ += data.size();
+  sendImpl(static_cast<const void*>(data.data()), data.size());
 }
 
-std::vector<unsigned char> InMemoryPartyCommunicationAgent::receive(
-    size_t size) {
-  auto result = host_.receive(myId_, size);
-  if (result.size() != size) {
+void InMemoryPartyCommunicationAgent::recvImpl(void* data, int nBytes) {
+  auto result = host_.receive(myId_, nBytes);
+
+  if (result.size() != nBytes) {
     throw std::runtime_error("unexpected message size!");
   }
-  receivedData_ += size;
-  return result;
+  memcpy(data, result.data(), nBytes);
+  receivedData_ += nBytes;
+}
+
+std::vector<unsigned char> InMemoryPartyCommunicationAgent::receive(
+    size_t size) {
+  std::vector<unsigned char> v(size);
+  recvImpl(static_cast<void*>(v.data()), size);
+  return v;
 }
 
 InMemoryPartyCommunicationAgentHost::InMemoryPartyCommunicationAgentHost() {

fbpcf/engine/communication/InMemoryPartyCommunicationAgentHost.h

@@ -45,6 +45,10 @@ class InMemoryPartyCommunicationAgent final : public IPartyCommunicationAgent {
     return {sentData_, receivedData_};
   }
 
+  void recvImpl(void* data, int nBytes) override;
+
+  void sendImpl(const void* data, int nBytes) override;
+
  private:
   InMemoryPartyCommunicationAgentHost& host_;
   int myId_;

fbpcf/engine/communication/SocketPartyCommunicationAgent.cpp

@@ -87,42 +87,49 @@ SocketPartyCommunicationAgent::~SocketPartyCommunicationAgent() {
   }
 }
 
-void SocketPartyCommunicationAgent::send(
-    const std::vector<unsigned char>& data) {
+void SocketPartyCommunicationAgent::sendImpl(const void* data, int nBytes) {
   size_t bytesWritten;
   if (!ssl_) {
-    bytesWritten =
-        fwrite(data.data(), sizeof(unsigned char), data.size(), outgoingPort_);
+    bytesWritten = fwrite(data, sizeof(unsigned char), nBytes, outgoingPort_);
   } else {
-    bytesWritten = SSL_write(ssl_, (void*)data.data(), data.size());
+    bytesWritten = SSL_write(ssl_, data, nBytes);
   }
-  assert(bytesWritten == data.size());
+  assert(bytesWritten == nBytes);
   sentData_ += bytesWritten;
   if (!ssl_) {
     fflush(outgoingPort_);
   }
 }
 
-std::vector<unsigned char> SocketPartyCommunicationAgent::receive(size_t size) {
+void SocketPartyCommunicationAgent::send(
+    const std::vector<unsigned char>& data) {
+  sendImpl(static_cast<const void*>(data.data()), data.size());
+}
+
+void SocketPartyCommunicationAgent::recvImpl(void* data, int nBytes) {
   size_t bytesRead = 0;
-  std::vector<unsigned char> rst(size);
 
   if (!ssl_) {
-    bytesRead = fread(rst.data(), sizeof(unsigned char), size, incomingPort_);
+    bytesRead = fread(data, sizeof(unsigned char), nBytes, incomingPort_);
   } else {
     // fread is blocking, but SSL_read is nonblocking. This discrepancy
     // can cause issues at the application level. We need to make sure that
     // both APIs behave consistently, so here we add a loop to ensure we
     // mimick blocking behavior.
-    while (bytesRead < size) {
+    while (bytesRead < nBytes) {
       bytesRead += SSL_read(
           ssl_,
-          rst.data() + (bytesRead * sizeof(unsigned char)),
-          size - bytesRead);
+          (unsigned char*)data + (bytesRead * sizeof(unsigned char)),
+          nBytes - bytesRead);
     }
   }
-  assert(bytesRead == size);
+  assert(bytesRead == nBytes);
   receivedData_ += bytesRead;
+}
+
+std::vector<unsigned char> SocketPartyCommunicationAgent::receive(size_t size) {
+  std::vector<unsigned char> rst(size);
+  recvImpl(static_cast<void*>(rst.data()), size);
   return rst;
 }
 

fbpcf/engine/communication/SocketPartyCommunicationAgent.h

@@ -57,6 +57,10 @@ class SocketPartyCommunicationAgent final : public IPartyCommunicationAgent {
     return {sentData_, receivedData_};
   }
 
+  void recvImpl(void* data, int nBytes) override;
+
+  void sendImpl(const void* data, int nBytes) override;
+
  private:
   void openServerPort(int portNo);
   void openClientPort(const std::string& serverAddress, int portNo);

fbpcf/engine/util/EmpNetworkAdapter.h

@@ -22,15 +22,12 @@ class EmpNetworkAdapter {
   explicit EmpNetworkAdapter(communication::IPartyCommunicationAgent& agent)
       : agent_(agent) {}
 
-  void send_data(const void* data, int nByte) {
-    std::vector<unsigned char> buffer(nByte);
-    memcpy(buffer.data(), data, nByte);
-    agent_.send(buffer);
+  void send_data(const void* data, int nBytes) {
+    agent_.sendImpl(data, nBytes);
   }
 
-  void recv_data(void* data, int nByte) {
-    auto buffer = agent_.receive(nByte);
-    memcpy(data, buffer.data(), nByte);
+  void recv_data(void* data, int nBytes) {
+    agent_.recvImpl(data, nBytes);
   }
 
   void send_block(const __m128i* data, int nBlock) {