[ORC] Make sure RPC channel-send is called in blocking calls and responses.

ORC-RPC batches calls by default, and the channel's send method must be called
to transfer any buffered calls to the remote. The call to send was missing on
responses and blocking calls in the SingleThreadedRPCEndpoint. This patch adds
the necessary calls and modifies the RPC unit test to check for them.

llvm-svn: 371245

2 files changed, 44 insertions, 1 deletions

diff --git a/llvm/unittests/ExecutionEngine/Orc/QueueChannel.h b/llvm/unittests/ExecutionEngine/Orc/QueueChannel.h
index 511f038dec1..1909693ecb1 100644
--- a/llvm/unittests/ExecutionEngine/Orc/QueueChannel.h
+++ b/llvm/unittests/ExecutionEngine/Orc/QueueChannel.h
@@ -80,6 +80,30 @@ public:
   QueueChannel(QueueChannel&&) = delete;
   QueueChannel& operator=(QueueChannel&&) = delete;
 
+  template <typename FunctionIdT, typename SequenceIdT>
+  Error startSendMessage(const FunctionIdT &FnId, const SequenceIdT &SeqNo) {
+    ++InFlightOutgoingMessages;
+    return orc::rpc::RawByteChannel::startSendMessage(FnId, SeqNo);
+  }
+
+  Error endSendMessage() {
+    --InFlightOutgoingMessages;
+    ++CompletedOutgoingMessages;
+    return orc::rpc::RawByteChannel::endSendMessage();
+  }
+
+  template <typename FunctionIdT, typename SequenceNumberT>
+  Error startReceiveMessage(FunctionIdT &FnId, SequenceNumberT &SeqNo) {
+    ++InFlightIncomingMessages;
+    return orc::rpc::RawByteChannel::startReceiveMessage(FnId, SeqNo);
+  }
+
+  Error endReceiveMessage() {
+    --InFlightIncomingMessages;
+    ++CompletedIncomingMessages;
+    return orc::rpc::RawByteChannel::endReceiveMessage();
+  }
+
   Error readBytes(char *Dst, unsigned Size) override {
     std::unique_lock<std::mutex> Lock(InQueue->getMutex());
     while (Size) {
@@ -112,7 +136,10 @@ public:
     return Error::success();
   }
 
-  Error send() override { return Error::success(); }
+  Error send() override {
+    ++SendCalls;
+    return Error::success();
+  }
 
   void close() {
     auto ChannelClosed = []() { return make_error<QueueChannelClosedError>(); };
@@ -124,6 +151,11 @@ public:
 
   uint64_t NumWritten = 0;
   uint64_t NumRead = 0;
+  std::atomic<size_t> InFlightIncomingMessages{0};
+  std::atomic<size_t> CompletedIncomingMessages{0};
+  std::atomic<size_t> InFlightOutgoingMessages{0};
+  std::atomic<size_t> CompletedOutgoingMessages{0};
+  std::atomic<size_t> SendCalls{0};
 
 private:
 
diff --git a/llvm/unittests/ExecutionEngine/Orc/RPCUtilsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/RPCUtilsTest.cpp
index 1f7c88d93d0..8e4c5330d90 100644
--- a/llvm/unittests/ExecutionEngine/Orc/RPCUtilsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/RPCUtilsTest.cpp
@@ -214,6 +214,17 @@ TEST(DummyRPC, TestCallAsyncVoidBool) {
     EXPECT_FALSE(!!Err) << "Client failed to handle response from void(bool)";
   }
 
+  // The client should have made two calls to send: One implicit call to
+  // negotiate the VoidBool function key, and a second to make the VoidBool
+  // call.
+  EXPECT_EQ(Channels.first->SendCalls, 2U)
+      << "Expected one send call to have been made by client";
+
+  // The server should have made two calls to send: One to send the response to
+  // the negotiate call, and another to send the response to the VoidBool call.
+  EXPECT_EQ(Channels.second->SendCalls, 2U)
+      << "Expected two send calls to have been made by server";
+
   ServerThread.join();
 }