From 17a0858f9d17a5d0d1b70b0e409e59a9f96967ca Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames@gmail.com>
Date: Mon, 11 Oct 2021 13:42:39 -0700
Subject: [PATCH] [ORC] Propagate errors to handlers when sendMessage fails.

In SimpleRemoteEPC, calls to from callWrapperAsync to sendMessage may fail.
The handlers may or may not be sent failure messages by handleDisconnect,
depending on when that method is run. This patch adds a check for an un-failed
handler, and if it finds one sends it a failure message.
---
 .../llvm/ExecutionEngine/Orc/ExecutorProcessControl.h |  1 +
 .../Orc/Shared/SimpleRemoteEPCUtils.cpp               |  5 +++++
 llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp      | 19 +++++++++++++++++++
 3 files changed, 25 insertions(+)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h b/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h
index 2786f76..27cd929 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h
@@ -49,6 +49,7 @@ public:
     friend class ExecutorProcessControl;
   public:
     IncomingWFRHandler() = default;
+    explicit operator bool() const { return !!H; }
     void operator()(shared::WrapperFunctionResult WFR) { H(std::move(WFR)); }
   private:
     template <typename FnT> IncomingWFRHandler(FnT &&Fn)
diff --git a/llvm/lib/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.cpp b/llvm/lib/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.cpp
index 62f4ff8..64fc717 100644
--- a/llvm/lib/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.cpp
@@ -238,6 +238,11 @@ void FDSimpleRemoteEPCTransport::listenLoop() {
     }
   } while (true);
 
+  // Attempt to close FDs, set Disconnected to true so that subsequent
+  // sendMessage calls fail.
+  disconnect();
+
+  // Call up to the client to handle the disconnection.
   C.handleDisconnect(std::move(Err));
 }
 
diff --git a/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp b/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp
index fc7d613..0f198c7 100644
--- a/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp
@@ -67,6 +67,25 @@ void SimpleRemoteEPC::callWrapperAsync(ExecutorAddr WrapperFnAddr,
 
   if (auto Err = sendMessage(SimpleRemoteEPCOpcode::CallWrapper, SeqNo,
                              WrapperFnAddr, ArgBuffer)) {
+    IncomingWFRHandler H;
+
+    // We just registered OnComplete, but there may be a race between this
+    // thread returning from sendMessage and handleDisconnect being called from
+    // the transport's listener thread. If handleDisconnect gets there first
+    // then it will have failed 'H' for us. If we get there first (or if
+    // handleDisconnect already ran) then we need to take care of it.
+    {
+      std::lock_guard<std::mutex> Lock(SimpleRemoteEPCMutex);
+      auto I = PendingCallWrapperResults.find(SeqNo);
+      if (I != PendingCallWrapperResults.end()) {
+        H = std::move(I->second);
+        PendingCallWrapperResults.erase(I);
+      }
+    }
+
+    if (H)
+      H(shared::WrapperFunctionResult::createOutOfBandError("disconnecting"));
+
     getExecutionSession().reportError(std::move(Err));
   }
 }
-- 
2.7.4