1 files changed, 35 insertions, 15 deletions

diff --git a/llvm/lib/Support/Unix/Signals.inc b/llvm/lib/Support/Unix/Signals.inc
index 5e0cde4a81ed..f68374d29f02 100644
--- a/llvm/lib/Support/Unix/Signals.inc
+++ b/llvm/lib/Support/Unix/Signals.inc
@@ -82,18 +82,15 @@ using namespace llvm;
 static RETSIGTYPE SignalHandler(int Sig);  // defined below.
 static RETSIGTYPE InfoSignalHandler(int Sig);  // defined below.
 
-static void DefaultPipeSignalFunction() {
-  exit(EX_IOERR);
-}
-
 using SignalHandlerFunctionType = void (*)();
 /// The function to call if ctrl-c is pressed.
 static std::atomic<SignalHandlerFunctionType> InterruptFunction =
     ATOMIC_VAR_INIT(nullptr);
 static std::atomic<SignalHandlerFunctionType> InfoSignalFunction =
     ATOMIC_VAR_INIT(nullptr);
-static std::atomic<SignalHandlerFunctionType> PipeSignalFunction =
-    ATOMIC_VAR_INIT(DefaultPipeSignalFunction);
+/// The function to call on SIGPIPE (one-time use only).
+static std::atomic<SignalHandlerFunctionType> OneShotPipeSignalFunction =
+    ATOMIC_VAR_INIT(nullptr);
 
 namespace {
 /// Signal-safe removal of files.
@@ -212,7 +209,7 @@ static StringRef Argv0;
 /// if there is, it's not our direct responsibility. For whatever reason, our
 /// continued execution is no longer desirable.
 static const int IntSigs[] = {
-  SIGHUP, SIGINT, SIGPIPE, SIGTERM, SIGUSR2
+  SIGHUP, SIGINT, SIGTERM, SIGUSR2
 };
 
 /// Signals that represent that we have a bug, and our prompt termination has
@@ -243,7 +240,7 @@ static const int InfoSigs[] = {
 
 static const size_t NumSigs =
     array_lengthof(IntSigs) + array_lengthof(KillSigs) +
-    array_lengthof(InfoSigs);
+    array_lengthof(InfoSigs) + 1 /* SIGPIPE */;
 
 
 static std::atomic<unsigned> NumRegisteredSignals = ATOMIC_VAR_INIT(0);
@@ -328,6 +325,8 @@ static void RegisterHandlers() { // Not signal-safe.
     registerHandler(S, SignalKind::IsKill);
   for (auto S : KillSigs)
     registerHandler(S, SignalKind::IsKill);
+  if (OneShotPipeSignalFunction)
+    registerHandler(SIGPIPE, SignalKind::IsKill);
   for (auto S : InfoSigs)
     registerHandler(S, SignalKind::IsInfo);
 }
@@ -346,6 +345,22 @@ static void RemoveFilesToRemove() {
   FileToRemoveList::removeAllFiles(FilesToRemove);
 }
 
+void sys::CleanupOnSignal(uintptr_t Context) {
+  int Sig = (int)Context;
+
+  if (llvm::is_contained(InfoSigs, Sig)) {
+    InfoSignalHandler(Sig);
+    return;
+  }
+
+  RemoveFilesToRemove();
+
+  if (llvm::is_contained(IntSigs, Sig) || Sig == SIGPIPE)
+    return;
+
+  llvm::sys::RunSignalHandlers();
+}
+
 // The signal handler that runs.
 static RETSIGTYPE SignalHandler(int Sig) {
   // Restore the signal behavior to default, so that the program actually
@@ -362,16 +377,16 @@ static RETSIGTYPE SignalHandler(int Sig) {
   {
     RemoveFilesToRemove();
 
+    if (Sig == SIGPIPE)
+      if (auto OldOneShotPipeFunction =
+              OneShotPipeSignalFunction.exchange(nullptr))
+        return OldOneShotPipeFunction();
+
     if (std::find(std::begin(IntSigs), std::end(IntSigs), Sig)
         != std::end(IntSigs)) {
       if (auto OldInterruptFunction = InterruptFunction.exchange(nullptr))
         return OldInterruptFunction();
 
-      // Send a special return code that drivers can check for, from sysexits.h.
-      if (Sig == SIGPIPE)
-        if (SignalHandlerFunctionType CurrentPipeFunction = PipeSignalFunction)
-          CurrentPipeFunction();
-
       raise(Sig);   // Execute the default handler.
       return;
    }
@@ -410,11 +425,16 @@ void llvm::sys::SetInfoSignalFunction(void (*Handler)()) {
   RegisterHandlers();
 }
 
-void llvm::sys::SetPipeSignalFunction(void (*Handler)()) {
-  PipeSignalFunction.exchange(Handler);
+void llvm::sys::SetOneShotPipeSignalFunction(void (*Handler)()) {
+  OneShotPipeSignalFunction.exchange(Handler);
   RegisterHandlers();
 }
 
+void llvm::sys::DefaultOneShotPipeSignalHandler() {
+  // Send a special return code that drivers can check for, from sysexits.h.
+  exit(EX_IOERR);
+}
+
 // The public API
 bool llvm::sys::RemoveFileOnSignal(StringRef Filename,
                                    std::string* ErrMsg) {