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();
 }