3 files changed, 56 insertions, 5 deletions

diff --git a/compiler-rt/include/sanitizer/lsan_interface.h b/compiler-rt/include/sanitizer/lsan_interface.h
index 97a88a5021e..96ddc84efb9 100644
--- a/compiler-rt/include/sanitizer/lsan_interface.h
+++ b/compiler-rt/include/sanitizer/lsan_interface.h
@@ -20,15 +20,23 @@
 extern "C" {
 #endif
   // Allocations made between calls to __lsan_disable() and __lsan_enable() will
-  // be treated as non-leaks. Disable/enable pairs can be nested.
+  // be treated as non-leaks. Disable/enable pairs may be nested.
   void __lsan_disable();
   void __lsan_enable();
   // The heap object into which p points will be treated as a non-leak.
   void __lsan_ignore_object(const void *p);
   // The user may optionally provide this function to disallow leak checking
-  // for the program it is linked into. Note: this function may be called late,
-  // after all the global destructors.
+  // for the program it is linked into (if the return value is non-zero). This
+  // function must be defined as returning a constant value; any behavior beyond
+  // that is unsupported.
   int __lsan_is_turned_off();
+  // Calling this function makes LSan enter the leak checking phase immediately.
+  // Use this if normal end-of-process leak checking happens too late (e.g. if
+  // you have intentional memory leaks in your shutdown code). Calling this
+  // function overrides end-of-process leak checking; it must be called at
+  // most once per process. This function will terminate the process if there
+  // are memory leaks and the exit_code flag is non-zero. 
+  void __lsan_do_leak_check();
 #ifdef __cplusplus
 }  // extern "C"
 
diff --git a/compiler-rt/lib/lsan/lit_tests/TestCases/do_leak_check_override.cc b/compiler-rt/lib/lsan/lit_tests/TestCases/do_leak_check_override.cc
new file mode 100644
index 00000000000..440ba42a54c
--- /dev/null
+++ b/compiler-rt/lib/lsan/lit_tests/TestCases/do_leak_check_override.cc
@@ -0,0 +1,36 @@
+// Test for __lsan_do_leak_check(). We test it by making the leak check run
+// before global destructors, which also tests compatibility with HeapChecker's
+// "normal" mode (LSan runs in "strict" mode by default).
+// RUN: LSAN_BASE="use_stacks=0:use_registers=0"
+// RUN: %clangxx_lsan %s -o %t
+// RUN: LSAN_OPTIONS=$LSAN_BASE %t 2>&1 | FileCheck --check-prefix=CHECK-strict %s
+// RUN: LSAN_OPTIONS=$LSAN_BASE %t foo 2>&1 | FileCheck --check-prefix=CHECK-normal %s
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <sanitizer/lsan_interface.h>
+
+struct LeakyGlobal {
+  LeakyGlobal() {
+    p = malloc(1337);
+  }
+  ~LeakyGlobal() {
+    p = 0;
+  }
+  void *p;
+};
+
+LeakyGlobal leaky_global;
+
+int main(int argc, char *argv[]) {
+  // Register leak check to run before global destructors.
+  if (argc > 1)
+    atexit(&__lsan_do_leak_check);
+  void *p = malloc(666);
+  printf("Test alloc: %p\n", p);
+  printf("Test alloc in leaky global: %p\n", leaky_global.p);
+  return 0;
+}
+
+// CHECK-strict: SUMMARY: LeakSanitizer: 2003 byte(s) leaked in 2 allocation(s)
+// CHECK-normal: SUMMARY: LeakSanitizer: 666 byte(s) leaked in 1 allocation(s)
diff --git a/compiler-rt/lib/lsan/lsan_common.cc b/compiler-rt/lib/lsan/lsan_common.cc
index fbb7a7457fd..8bb89e748a5 100644
--- a/compiler-rt/lib/lsan/lsan_common.cc
+++ b/compiler-rt/lib/lsan/lsan_common.cc
@@ -125,7 +125,7 @@ void ScanRangeForPointers(uptr begin, uptr end,
   if (pp % alignment)
     pp = pp + alignment - pp % alignment;
   for (; pp + sizeof(void *) <= end; pp += alignment) {  // NOLINT
-    void *p = *reinterpret_cast<void**>(pp);
+    void *p = *reinterpret_cast<void **>(pp);
     if (!CanBeAHeapPointer(reinterpret_cast<uptr>(p))) continue;
     uptr chunk = PointsIntoChunk(p);
     if (!chunk) continue;
@@ -353,7 +353,7 @@ void DoLeakCheck() {
   EnsureMainThreadIDIsCorrect();
   BlockingMutexLock l(&global_mutex);
   static bool already_done;
-  CHECK(!already_done);
+  if (already_done) return;
   already_done = true;
   if (&__lsan_is_turned_off && __lsan_is_turned_off())
     return;
@@ -544,6 +544,13 @@ void __lsan_enable() {
 #endif
 }
 
+SANITIZER_INTERFACE_ATTRIBUTE
+void __lsan_do_leak_check() {
+#if CAN_SANITIZE_LEAKS
+  __lsan::DoLeakCheck();
+#endif
+}
+
 #if !SANITIZER_SUPPORTS_WEAK_HOOKS
 SANITIZER_WEAK_ATTRIBUTE SANITIZER_INTERFACE_ATTRIBUTE
 int __lsan_is_turned_off() {