From c1105111b39384b959edee9c91ae543d57a5a795 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames@gmail.com>
Date: Fri, 6 Sep 2019 19:21:59 +0000
Subject: [PATCH] [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
---
 llvm/include/llvm/ExecutionEngine/Orc/RPCUtils.h   | 25 ++++++++++++----
 llvm/unittests/ExecutionEngine/Orc/QueueChannel.h  | 34 +++++++++++++++++++++-
 .../unittests/ExecutionEngine/Orc/RPCUtilsTest.cpp | 11 +++++++
 3 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/RPCUtils.h b/llvm/include/llvm/ExecutionEngine/Orc/RPCUtils.h
index be5cea4..4e63a84 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/RPCUtils.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/RPCUtils.h
@@ -338,7 +338,9 @@ public:
       return Err;
 
     // Close the response message.
-    return C.endSendMessage();
+    if (auto Err = C.endSendMessage())
+      return Err;
+    return C.send();
   }
 
   template <typename ChannelT, typename FunctionIdT, typename SequenceNumberT>
@@ -350,7 +352,9 @@ public:
       return Err2;
     if (auto Err2 = serializeSeq(C, std::move(Err)))
       return Err2;
-    return C.endSendMessage();
+    if (auto Err2 = C.endSendMessage())
+      return Err2;
+    return C.send();
   }
 
 };
@@ -378,8 +382,11 @@ public:
                                                                C, *ResultOrErr))
       return Err;
 
-    // Close the response message.
-    return C.endSendMessage();
+    // End the response message.
+    if (auto Err = C.endSendMessage())
+      return Err;
+
+    return C.send();
   }
 
   template <typename ChannelT, typename FunctionIdT, typename SequenceNumberT>
@@ -389,7 +396,9 @@ public:
       return Err;
     if (auto Err2 = C.startSendMessage(ResponseId, SeqNo))
       return Err2;
-    return C.endSendMessage();
+    if (auto Err2 = C.endSendMessage())
+      return Err2;
+    return C.send();
   }
 
 };
@@ -1520,6 +1529,12 @@ public:
       return std::move(Err);
     }
 
+    if (auto Err = this->C.send()) {
+      detail::ResultTraits<typename Func::ReturnType>::consumeAbandoned(
+          std::move(Result));
+      return std::move(Err);
+    }
+
     while (!ReceivedResponse) {
       if (auto Err = this->handleOne()) {
         detail::ResultTraits<typename Func::ReturnType>::consumeAbandoned(
diff --git a/llvm/unittests/ExecutionEngine/Orc/QueueChannel.h b/llvm/unittests/ExecutionEngine/Orc/QueueChannel.h
index 511f038..1909693 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 1f7c88d..8e4c533 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();
 }
 
-- 
2.7.4