6 files changed, 70 insertions, 12 deletions
diff --git a/include/sanitizer/asan_interface.h b/include/sanitizer/asan_interface.h
index 435391486..7763389ab 100644
--- a/include/sanitizer/asan_interface.h
+++ b/include/sanitizer/asan_interface.h
@@ -114,8 +114,7 @@ extern "C" {
   // Returns the old value.
   int __asan_set_error_exit_code(int exit_code);
 
-  // Sets the callback to be called right before death on error.
-  // Passing 0 will unset the callback.
+  // Deprecated. Call __sanitizer_set_death_callback instead.
   void __asan_set_death_callback(void (*callback)(void));
 
   void __asan_set_error_report_callback(void (*callback)(const char*));
diff --git a/include/sanitizer/common_interface_defs.h b/include/sanitizer/common_interface_defs.h
index 9cb5ad88d..4c715c66b 100644
--- a/include/sanitizer/common_interface_defs.h
+++ b/include/sanitizer/common_interface_defs.h
@@ -120,6 +120,9 @@ extern "C" {
   // Print the stack trace leading to this call. Useful for debugging user code.
   void __sanitizer_print_stack_trace();
 
+  // Sets the callback to be called right before death on error.
+  // Passing 0 will unset the callback.
+  void __sanitizer_set_death_callback(void (*callback)(void));
 #ifdef __cplusplus
 }  // extern "C"
 #endif
diff --git a/lib/asan/asan_rtl.cc b/lib/asan/asan_rtl.cc
index 3c36bf636..bb27cca74 100644
--- a/lib/asan/asan_rtl.cc
+++ b/lib/asan/asan_rtl.cc
@@ -56,8 +56,6 @@ static void AsanDie() {
   }
   if (common_flags()->coverage)
     __sanitizer_cov_dump();
-  if (death_callback)
-    death_callback();
   if (flags()->abort_on_error)
     Abort();
   internal__exit(flags()->exitcode);
@@ -341,7 +339,6 @@ void ParseExtraActivationFlags() {
 // -------------------------- Globals --------------------- {{{1
 int asan_inited;
 bool asan_init_is_running;
-void (*death_callback)(void);
 
 #if !ASAN_FIXED_MAPPING
 uptr kHighMemEnd, kMidMemBeg, kMidMemEnd;
@@ -767,7 +764,7 @@ void NOINLINE __asan_handle_no_return() {
 }
 
 void NOINLINE __asan_set_death_callback(void (*callback)(void)) {
-  death_callback = callback;
+  SetUserDieCallback(callback);
 }
 
 // Initialize as requested from instrumented application code.
diff --git a/lib/sanitizer_common/sanitizer_common.cc b/lib/sanitizer_common/sanitizer_common.cc
index c3740f24a..7bb4f4581 100644
--- a/lib/sanitizer_common/sanitizer_common.cc
+++ b/lib/sanitizer_common/sanitizer_common.cc
@@ -94,19 +94,23 @@ uptr stoptheworld_tracer_pid = 0;
 // writing to the same log file.
 uptr stoptheworld_tracer_ppid = 0;
 
-static DieCallbackType DieCallback;
+static DieCallbackType InternalDieCallback, UserDieCallback;
 void SetDieCallback(DieCallbackType callback) {
-  DieCallback = callback;
+  InternalDieCallback = callback;
+}
+void SetUserDieCallback(DieCallbackType callback) {
+  UserDieCallback = callback;
 }
 
 DieCallbackType GetDieCallback() {
-  return DieCallback;
+  return InternalDieCallback;
 }
 
 void NORETURN Die() {
-  if (DieCallback) {
-    DieCallback();
-  }
+  if (UserDieCallback)
+    UserDieCallback();
+  if (InternalDieCallback)
+    InternalDieCallback();
   internal__exit(1);
 }
 
@@ -280,4 +284,9 @@ void __sanitizer_set_report_path(const char *path) {
 void __sanitizer_report_error_summary(const char *error_summary) {
   Printf("%s\n", error_summary);
 }
+
+SANITIZER_INTERFACE_ATTRIBUTE
+void __sanitizer_set_death_callback(void (*callback)(void)) {
+  SetUserDieCallback(callback);
+}
 }  // extern "C"
diff --git a/lib/sanitizer_common/sanitizer_common.h b/lib/sanitizer_common/sanitizer_common.h
index c00dce66b..0c0d02913 100644
--- a/lib/sanitizer_common/sanitizer_common.h
+++ b/lib/sanitizer_common/sanitizer_common.h
@@ -245,6 +245,7 @@ bool SanitizerGetThreadName(char *name, int max_len);
 // to do tool-specific job.
 typedef void (*DieCallbackType)(void);
 void SetDieCallback(DieCallbackType);
+void SetUserDieCallback(DieCallbackType);
 DieCallbackType GetDieCallback();
 typedef void (*CheckFailedCallbackType)(const char *, int, const char *,
                                        u64, u64);
diff --git a/test/sanitizer_common/TestCases/Linux/sanitizer_set_death_callback_test.cc b/test/sanitizer_common/TestCases/Linux/sanitizer_set_death_callback_test.cc
new file mode 100644
index 000000000..ee3d59c72
--- /dev/null
+++ b/test/sanitizer_common/TestCases/Linux/sanitizer_set_death_callback_test.cc
@@ -0,0 +1,49 @@
+// RUN: %clangxx -O2 %s -o %t && not %run %t 2>&1 | FileCheck %s
+// Check __sanitizer_set_death_callback. Not all sanitizers implement it yet.
+// XFAIL: lsan
+// XFAIL: tsan
+
+#include <sanitizer/common_interface_defs.h>
+#include <stdio.h>
+#include <pthread.h>
+
+volatile char *zero = 0;
+
+void Death() {
+  fprintf(stderr, "DEATH CALLBACK EXECUTED\n");
+}
+// CHECK: DEATH CALLBACK EXECUTED
+
+int global[10];
+volatile char *sink;
+
+void *Thread(void *x) {
+  global[0]++;
+  return x;
+}
+
+__attribute__((noinline))
+void MaybeInit(int *uninitialized) {
+  if (zero)
+    *uninitialized = 1;
+}
+
+__attribute__((noinline))
+void Leak() {
+  sink = new char[100];  // trigger lsan report.
+}
+
+int main(int argc, char **argv) {
+  int uninitialized;
+  __sanitizer_set_death_callback(Death);
+  MaybeInit(&uninitialized);
+  if (uninitialized)  // trigger msan report.
+    global[0] = 77;
+  pthread_t t;
+  pthread_create(&t, 0, Thread, 0);
+  global[0]++;           // trigger tsan report.
+  pthread_join(t, 0);
+  global[argc + 10]++;   // trigger asan report.
+  Leak();
+  sink = 0;
+}