6 files changed, 58 insertions, 39 deletions
diff --git a/lib/asan/asan_rtl.cc b/lib/asan/asan_rtl.cc
index ad5a8eb98..e246abd5e 100644
--- a/lib/asan/asan_rtl.cc
+++ b/lib/asan/asan_rtl.cc
@@ -393,7 +393,7 @@ static void AsanInitInternal() {
   AsanDoesNotSupportStaticLinkage();
 
   // Install tool-specific callbacks in sanitizer_common.
-  SetDieCallback(AsanDie);
+  AddDieCallback(AsanDie);
   SetCheckFailedCallback(AsanCheckFailed);
   SetPrintfAndReportCallback(AppendToErrorMessageBuffer);
 
diff --git a/lib/dfsan/dfsan.cc b/lib/dfsan/dfsan.cc
index 3b1e6c431..8cd65f799 100644
--- a/lib/dfsan/dfsan.cc
+++ b/lib/dfsan/dfsan.cc
@@ -416,7 +416,7 @@ static void dfsan_init(int argc, char **argv, char **envp) {
   // Register the fini callback to run when the program terminates successfully
   // or it is killed by the runtime.
   Atexit(dfsan_fini);
-  SetDieCallback(dfsan_fini);
+  AddDieCallback(dfsan_fini);
 
   __dfsan_label_info[kInitializingLabel].desc = "<init label>";
 }
diff --git a/lib/msan/msan.cc b/lib/msan/msan.cc
index 0ba7f32bb..1b116815f 100644
--- a/lib/msan/msan.cc
+++ b/lib/msan/msan.cc
@@ -375,7 +375,7 @@ void __msan_init() {
   msan_init_is_running = 1;
   SanitizerToolName = "MemorySanitizer";
 
-  SetDieCallback(MsanDie);
+  AddDieCallback(MsanDie);
   InitTlsSize();
 
   CacheBinaryName();
diff --git a/lib/sanitizer_common/sanitizer_common.cc b/lib/sanitizer_common/sanitizer_common.cc
index fb0cee0f4..4ffa5b8be 100644
--- a/lib/sanitizer_common/sanitizer_common.cc
+++ b/lib/sanitizer_common/sanitizer_common.cc
@@ -105,23 +105,43 @@ uptr stoptheworld_tracer_pid = 0;
 // writing to the same log file.
 uptr stoptheworld_tracer_ppid = 0;
 
-static DieCallbackType InternalDieCallback, UserDieCallback;
-void SetDieCallback(DieCallbackType callback) {
-  InternalDieCallback = callback;
+static const int kMaxNumOfInternalDieCallbacks = 5;
+static DieCallbackType InternalDieCallbacks[kMaxNumOfInternalDieCallbacks];
+
+bool AddDieCallback(DieCallbackType callback) {
+  for (int i = 0; i < kMaxNumOfInternalDieCallbacks; i++) {
+    if (InternalDieCallbacks[i] == nullptr) {
+      InternalDieCallbacks[i] = callback;
+      return true;
+    }
+  }
+  return false;
 }
-void SetUserDieCallback(DieCallbackType callback) {
-  UserDieCallback = callback;
+
+bool RemoveDieCallback(DieCallbackType callback) {
+  for (int i = 0; i < kMaxNumOfInternalDieCallbacks; i++) {
+    if (InternalDieCallbacks[i] == callback) {
+      for (int j = i + 1; j < kMaxNumOfInternalDieCallbacks; j++)
+        InternalDieCallbacks[j - 1] = InternalDieCallbacks[j];
+      InternalDieCallbacks[kMaxNumOfInternalDieCallbacks - 1] = nullptr;
+      return true;
+    }
+  }
+  return false;
 }
 
-DieCallbackType GetDieCallback() {
-  return InternalDieCallback;
+static DieCallbackType UserDieCallback;
+void SetUserDieCallback(DieCallbackType callback) {
+  UserDieCallback = callback;
 }
 
 void NORETURN Die() {
   if (UserDieCallback)
     UserDieCallback();
-  if (InternalDieCallback)
-    InternalDieCallback();
+  for (int i = kMaxNumOfInternalDieCallbacks - 1; i >= 0; i--) {
+    if (InternalDieCallbacks[i])
+      InternalDieCallbacks[i]();
+  }
   internal__exit(common_flags()->exitcode);
 }
 
diff --git a/lib/sanitizer_common/sanitizer_common.h b/lib/sanitizer_common/sanitizer_common.h
index 9aa69c5a4..b63228866 100644
--- a/lib/sanitizer_common/sanitizer_common.h
+++ b/lib/sanitizer_common/sanitizer_common.h
@@ -319,9 +319,16 @@ bool SanitizerGetThreadName(char *name, int max_len);
 // Specific tools may override behavior of "Die" and "CheckFailed" functions
 // to do tool-specific job.
 typedef void (*DieCallbackType)(void);
-void SetDieCallback(DieCallbackType);
-void SetUserDieCallback(DieCallbackType);
-DieCallbackType GetDieCallback();
+
+// It's possible to add several callbacks that would be run when "Die" is
+// called. The callbacks will be run in the opposite order. The tools are
+// strongly recommended to setup all callbacks during initialization, when there
+// is only a single thread.
+bool AddDieCallback(DieCallbackType callback);
+bool RemoveDieCallback(DieCallbackType callback);
+
+void SetUserDieCallback(DieCallbackType callback);
+
 typedef void (*CheckFailedCallbackType)(const char *, int, const char *,
                                        u64, u64);
 void SetCheckFailedCallback(CheckFailedCallbackType callback);
diff --git a/lib/sanitizer_common/sanitizer_stoptheworld_linux_libcdep.cc b/lib/sanitizer_common/sanitizer_stoptheworld_linux_libcdep.cc
index d09baacda..5ea227667 100644
--- a/lib/sanitizer_common/sanitizer_stoptheworld_linux_libcdep.cc
+++ b/lib/sanitizer_common/sanitizer_stoptheworld_linux_libcdep.cc
@@ -211,7 +211,19 @@ static ThreadSuspender *thread_suspender_instance = NULL;
 static const int kSyncSignals[] = { SIGABRT, SIGILL, SIGFPE, SIGSEGV, SIGBUS,
                                     SIGXCPU, SIGXFSZ };
 
-static DieCallbackType old_die_callback;
+static void TracerThreadDieCallback() {
+  // Generally a call to Die() in the tracer thread should be fatal to the
+  // parent process as well, because they share the address space.
+  // This really only works correctly if all the threads are suspended at this
+  // point. So we correctly handle calls to Die() from within the callback, but
+  // not those that happen before or after the callback. Hopefully there aren't
+  // a lot of opportunities for that to happen...
+  ThreadSuspender *inst = thread_suspender_instance;
+  if (inst != NULL && stoptheworld_tracer_pid == internal_getpid()) {
+    inst->KillAllThreads();
+    thread_suspender_instance = NULL;
+  }
+}
 
 // Signal handler to wake up suspended threads when the tracer thread dies.
 static void TracerThreadSignalHandler(int signum, void *siginfo, void *uctx) {
@@ -224,32 +236,13 @@ static void TracerThreadSignalHandler(int signum, void *siginfo, void *uctx) {
       inst->KillAllThreads();
     else
       inst->ResumeAllThreads();
-    SetDieCallback(old_die_callback);
-    old_die_callback = NULL;
+    RAW_CHECK(RemoveDieCallback(TracerThreadDieCallback));
     thread_suspender_instance = NULL;
     atomic_store(&inst->arg->done, 1, memory_order_relaxed);
   }
   internal__exit((signum == SIGABRT) ? 1 : 2);
 }
 
-static void TracerThreadDieCallback() {
-  // Generally a call to Die() in the tracer thread should be fatal to the
-  // parent process as well, because they share the address space.
-  // This really only works correctly if all the threads are suspended at this
-  // point. So we correctly handle calls to Die() from within the callback, but
-  // not those that happen before or after the callback. Hopefully there aren't
-  // a lot of opportunities for that to happen...
-  ThreadSuspender *inst = thread_suspender_instance;
-  if (inst != NULL && stoptheworld_tracer_pid == internal_getpid()) {
-    inst->KillAllThreads();
-    thread_suspender_instance = NULL;
-  }
-  if (old_die_callback)
-    old_die_callback();
-  SetDieCallback(old_die_callback);
-  old_die_callback = NULL;
-}
-
 // Size of alternative stack for signal handlers in the tracer thread.
 static const int kHandlerStackSize = 4096;
 
@@ -267,8 +260,7 @@ static int TracerThread(void* argument) {
   tracer_thread_argument->mutex.Lock();
   tracer_thread_argument->mutex.Unlock();
 
-  old_die_callback = GetDieCallback();
-  SetDieCallback(TracerThreadDieCallback);
+  RAW_CHECK(AddDieCallback(TracerThreadDieCallback));
 
   ThreadSuspender thread_suspender(internal_getppid(), tracer_thread_argument);
   // Global pointer for the signal handler.
@@ -302,7 +294,7 @@ static int TracerThread(void* argument) {
     thread_suspender.ResumeAllThreads();
     exit_code = 0;
   }
-  SetDieCallback(old_die_callback);
+  RAW_CHECK(RemoveDieCallback(TracerThreadDieCallback));
   thread_suspender_instance = NULL;
   atomic_store(&tracer_thread_argument->done, 1, memory_order_relaxed);
   return exit_code;