[llvm] 17a0858 - [ORC] Propagate errors to handlers when sendMessage fails.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2021-10-11T14:23:50-07:00
New Revision: 17a0858f9d17a5d0d1b70b0e409e59a9f96967ca

URL: https://github.com/llvm/llvm-project/commit/17a0858f9d17a5d0d1b70b0e409e59a9f96967ca
DIFF: https://github.com/llvm/llvm-project/commit/17a0858f9d17a5d0d1b70b0e409e59a9f96967ca.diff

LOG: [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.

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h
    llvm/lib/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.cpp
    llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h b/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h
index 2786f76c26a73..27cd92968fe55 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h
@@ -49,6 +49,7 @@ class ExecutorProcessControl {
     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 62f4ff8a34090..64fc717b7b56c 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 fc7d613629dbb..0f198c75f8657 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));
   }
 }