From 9e80add2cfa9bb03bbe77cca9adec18137204538 Mon Sep 17 00:00:00 2001
From: Enna1 <xumingjie.enna1@bytedance.com>
Date: Sat, 8 Oct 2022 09:55:40 +0800
Subject: [PATCH] [memprof] dump memprof profile when receive deadly signals

Currently memprof profile is dumped when program exits (call `FinishAndWrite()` in ~Allocator) or `__memprof_profile_dump` is manually called.
For programs that never exit (e.g. server-side application), it will be useful to dump memprof profile when specific signal is received.
This patch installs a signal handler for deadly signals(SIGSEGV, SIGBUS, SIGABRT, SIGILL, SIGTRAP, SIGFPE) like we do in other sanitizers. In the signal handler `__memprof_profile_dump` is called to dump memprof profile.

Reviewed By: tejohnson

Differential Revision: https://reviews.llvm.org/D134795
---
 compiler-rt/lib/memprof/memprof_rtl.cpp        |  9 +++++++++
 .../memprof_profile_dump_on_abort.cpp          | 18 ++++++++++++++++++
 2 files changed, 27 insertions(+)
 create mode 100644 compiler-rt/test/memprof/TestCases/memprof_profile_dump_on_abort.cpp

diff --git a/compiler-rt/lib/memprof/memprof_rtl.cpp b/compiler-rt/lib/memprof/memprof_rtl.cpp
index d568a075c3e1..d30b80304f6b 100644
--- a/compiler-rt/lib/memprof/memprof_rtl.cpp
+++ b/compiler-rt/lib/memprof/memprof_rtl.cpp
@@ -50,6 +50,14 @@ static void MemprofDie() {
   }
 }
 
+static void MemprofOnDeadlySignal(int signo, void *siginfo, void *context) {
+  // We call StartReportDeadlySignal not HandleDeadlySignal so we get the
+  // deadly signal message to stderr but no writing to the profile output file
+  StartReportDeadlySignal();
+  __memprof_profile_dump();
+  Die();
+}
+
 static void CheckUnwind() {
   GET_STACK_TRACE(kStackTraceMax, common_flags()->fast_unwind_on_check);
   stack.Print();
@@ -183,6 +191,7 @@ static void MemprofInitInternal() {
   InitializeShadowMemory();
 
   TSDInit(PlatformTSDDtor);
+  InstallDeadlySignalHandlers(MemprofOnDeadlySignal);
 
   InitializeAllocator();
 
diff --git a/compiler-rt/test/memprof/TestCases/memprof_profile_dump_on_abort.cpp b/compiler-rt/test/memprof/TestCases/memprof_profile_dump_on_abort.cpp
new file mode 100644
index 000000000000..5df611b7e606
--- /dev/null
+++ b/compiler-rt/test/memprof/TestCases/memprof_profile_dump_on_abort.cpp
@@ -0,0 +1,18 @@
+// RUN: %clangxx_memprof %s -o %t
+
+// RUN: %env_memprof_opts=print_text=true:log_path=stdout:handle_abort=1 not %run %t 2>&1 | FileCheck --check-prefix=CHECK-TEXT %s
+
+#include <stdlib.h>
+#include <string.h>
+int main(int argc, char **argv) {
+  char *x = (char *)malloc(10);
+  memset(x, 0, 10);
+  free(x);
+  abort();
+  return 0;
+}
+
+// CHECK-TEXT: MemProfiler:DEADLYSIGNAL
+// CHECK-TEXT: Recorded MIBs (incl. live on exit):
+// CHECK-TEXT: Memory allocation stack id
+// CHECK-TEXT: Stack for id