5 files changed, 61 insertions, 9 deletions
diff --git a/llvm/include/llvm/Support/InitLLVM.h b/llvm/include/llvm/Support/InitLLVM.h
index 8069859a3e0..3be8d6b6d2e 100644
--- a/llvm/include/llvm/Support/InitLLVM.h
+++ b/llvm/include/llvm/Support/InitLLVM.h
@@ -17,7 +17,8 @@
 // the following one-time initializations:
 //
 //  1. Setting up a signal handler so that pretty stack trace is printed out
-//     if a process crashes.
+//     if a process crashes. A signal handler that exits when a failed write to
+//     a pipe occurs may optionally be installed: this is on-by-default.
 //
 //  2. Set up the global new-handler which is called when a memory allocation
 //     attempt fails.
@@ -32,9 +33,11 @@
 namespace llvm {
 class InitLLVM {
 public:
-  InitLLVM(int &Argc, const char **&Argv);
-  InitLLVM(int &Argc, char **&Argv)
-      : InitLLVM(Argc, const_cast<const char **&>(Argv)) {}
+  InitLLVM(int &Argc, const char **&Argv,
+           bool InstallPipeSignalExitHandler = true);
+  InitLLVM(int &Argc, char **&Argv, bool InstallPipeSignalExitHandler = true)
+      : InitLLVM(Argc, const_cast<const char **&>(Argv),
+                 InstallPipeSignalExitHandler) {}
 
   ~InitLLVM();
 
diff --git a/llvm/include/llvm/Support/Signals.h b/llvm/include/llvm/Support/Signals.h
index a6b215a2431..804bc875441 100644
--- a/llvm/include/llvm/Support/Signals.h
+++ b/llvm/include/llvm/Support/Signals.h
@@ -84,6 +84,28 @@ namespace sys {
   /// function.  Note also that the handler may be executed on a different
   /// thread on some platforms.
   void SetInfoSignalFunction(void (*Handler)());
+
+  /// Registers a function to be called in a "one-shot" manner when a pipe
+  /// signal is delivered to the process (i.e., on a failed write to a pipe).
+  /// After the pipe signal is handled once, the handler is unregistered.
+  ///
+  /// The LLVM signal handling code will not install any handler for the pipe
+  /// signal unless one is provided with this API (see \ref
+  /// DefaultOneShotPipeSignalHandler). This handler must be provided before
+  /// any other LLVM signal handlers are installed: the \ref InitLLVM
+  /// constructor has a flag that can simplify this setup.
+  ///
+  /// Note that the handler is not allowed to call any non-reentrant
+  /// functions.  A null handler pointer disables the current installed
+  /// function.  Note also that the handler may be executed on a
+  /// different thread on some platforms.
+  ///
+  /// This is a no-op on Windows.
+  void SetOneShotPipeSignalFunction(void (*Handler)());
+
+  /// On Unix systems, this function exits with an "IO error" exit code.
+  /// This is a no-op on Windows.
+  void DefaultOneShotPipeSignalHandler();
 } // End sys namespace
 } // End llvm namespace
 
diff --git a/llvm/lib/Support/InitLLVM.cpp b/llvm/lib/Support/InitLLVM.cpp
index 0d7d7fcc8cb..bb9b569d2de 100644
--- a/llvm/lib/Support/InitLLVM.cpp
+++ b/llvm/lib/Support/InitLLVM.cpp
@@ -21,7 +21,11 @@
 using namespace llvm;
 using namespace llvm::sys;
 
-InitLLVM::InitLLVM(int &Argc, const char **&Argv) : StackPrinter(Argc, Argv) {
+InitLLVM::InitLLVM(int &Argc, const char **&Argv,
+                   bool InstallPipeSignalExitHandler)
+    : StackPrinter(Argc, Argv) {
+  if (InstallPipeSignalExitHandler)
+    sys::SetOneShotPipeSignalFunction(sys::DefaultOneShotPipeSignalHandler);
   sys::PrintStackTraceOnErrorSignal(Argv[0]);
   install_out_of_memory_new_handler();
 
diff --git a/llvm/lib/Support/Unix/Signals.inc b/llvm/lib/Support/Unix/Signals.inc
index be05eabfb2e..8c26fa9b8f2 100644
--- a/llvm/lib/Support/Unix/Signals.inc
+++ b/llvm/lib/Support/Unix/Signals.inc
@@ -88,6 +88,9 @@ static std::atomic<SignalHandlerFunctionType> InterruptFunction =
     ATOMIC_VAR_INIT(nullptr);
 static std::atomic<SignalHandlerFunctionType> InfoSignalFunction =
     ATOMIC_VAR_INIT(nullptr);
+/// 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.
@@ -206,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
@@ -237,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);
@@ -322,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);
 }
@@ -361,9 +366,10 @@ static RETSIGTYPE SignalHandler(int Sig) {
       if (auto OldInterruptFunction = InterruptFunction.exchange(nullptr))
         return OldInterruptFunction();
 
-      // Send a special return code that drivers can check for, from sysexits.h.
       if (Sig == SIGPIPE)
-        exit(EX_IOERR);
+        if (auto OldOneShotPipeFunction =
+                OneShotPipeSignalFunction.exchange(nullptr))
+          return OldOneShotPipeFunction();
 
       raise(Sig);   // Execute the default handler.
       return;
@@ -403,6 +409,16 @@ void llvm::sys::SetInfoSignalFunction(void (*Handler)()) {
   RegisterHandlers();
 }
 
+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) {
diff --git a/llvm/lib/Support/Windows/Signals.inc b/llvm/lib/Support/Windows/Signals.inc
index 6a820ef22b1..08a3616427b 100644
--- a/llvm/lib/Support/Windows/Signals.inc
+++ b/llvm/lib/Support/Windows/Signals.inc
@@ -560,6 +560,13 @@ void llvm::sys::SetInfoSignalFunction(void (*Handler)()) {
   // Unimplemented.
 }
 
+void llvm::sys::SetOneShotPipeSignalFunction(void (*Handler)()) {
+  // Unimplemented.
+}
+
+void llvm::sys::DefaultOneShotPipeSignalHandler() {
+  // Unimplemented.
+}
 
 /// Add a function to be called when a signal is delivered to the process. The
 /// handler can have a cookie passed to it to identify what instance of the