[compiler-rt] [memprof] Add flag to control profile dump at exit (PR #119452)

[Ellis Hoag via llvm-commits]

https://github.com/ellishg created https://github.com/llvm/llvm-project/pull/119452

Add the `dump_at_exit` flag to control whether or not profiles should be dumped when the program exits. Since we can call `__memprof_profile_dump()` directly, we don't necessarily need to dump profiles at exit.

>From ec382a624af18cb7f0434c24d1c811ca1bde671d Mon Sep 17 00:00:00 2001
From: Ellis Hoag <ellis.sparky.hoag at gmail.com>
Date: Tue, 10 Dec 2024 12:39:14 -0800
Subject: [PATCH] [memprof] Add flag to control profile dump at exit

---
 compiler-rt/lib/memprof/memprof_allocator.cpp    |  3 ++-
 compiler-rt/lib/memprof/memprof_flags.inc        |  4 +++-
 .../test/memprof/TestCases/dump_at_exit.cpp      | 16 ++++++++++++++++
 3 files changed, 21 insertions(+), 2 deletions(-)
 create mode 100644 compiler-rt/test/memprof/TestCases/dump_at_exit.cpp

diff --git a/compiler-rt/lib/memprof/memprof_allocator.cpp b/compiler-rt/lib/memprof/memprof_allocator.cpp
index 19b2b901068246..c3448c22532bb8 100644
--- a/compiler-rt/lib/memprof/memprof_allocator.cpp
+++ b/compiler-rt/lib/memprof/memprof_allocator.cpp
@@ -301,7 +301,8 @@ struct Allocator {
 
   ~Allocator() {
     atomic_store_relaxed(&destructing, 1);
-    FinishAndWrite();
+    if (flags()->dump_at_exit)
+      FinishAndWrite();
   }
 
   static void PrintCallback(const uptr Key, LockedMemInfoBlock *const &Value,
diff --git a/compiler-rt/lib/memprof/memprof_flags.inc b/compiler-rt/lib/memprof/memprof_flags.inc
index 7c5dc091f79351..3a40ba0ab5498b 100644
--- a/compiler-rt/lib/memprof/memprof_flags.inc
+++ b/compiler-rt/lib/memprof/memprof_flags.inc
@@ -38,4 +38,6 @@ MEMPROF_FLAG(bool, allocator_frees_and_returns_null_on_realloc_zero, true,
 MEMPROF_FLAG(bool, print_text, false,
   "If set, prints the heap profile in text format. Else use the raw binary serialization format.")
 MEMPROF_FLAG(bool, print_terse, false,
-             "If set, prints memory profile in a terse format. Only applicable if print_text = true.")
\ No newline at end of file
+             "If set, prints memory profile in a terse format. Only applicable if print_text = true.")
+MEMPROF_FLAG(bool, dump_at_exit, true,
+             "If set, dump profiles when the program terminates.")
diff --git a/compiler-rt/test/memprof/TestCases/dump_at_exit.cpp b/compiler-rt/test/memprof/TestCases/dump_at_exit.cpp
new file mode 100644
index 00000000000000..426849d1cea01e
--- /dev/null
+++ b/compiler-rt/test/memprof/TestCases/dump_at_exit.cpp
@@ -0,0 +1,16 @@
+// RUN: %clangxx_memprof %s -o %t
+
+// RUN: %env_memprof_opts=print_text=true:log_path=stdout:dump_at_exit=false %run %t | count 0
+// RUN: %env_memprof_opts=print_text=true:log_path=stdout:dump_at_exit=true %run %t | FileCheck %s
+
+#include <stdlib.h>
+#include <string.h>
+
+int main() {
+  char *x = (char *)malloc(10);
+  memset(x, 0, 10);
+  free(x);
+  return 0;
+}
+
+// CHECK: Recorded MIBs