Allow to reset malloc hooks.

This is used to prevent the additional indirection even after heap
profiling has finished, preventing any performance impact on processes
that are not currently being profiled.

Test: m
Test: flash sailfish
Test: try tearing down & re-enabling hooks

Bug: 120186127

Change-Id: Idc5988111a47870d2c093fd6a017b47e65f5616b

3 files changed, 58 insertions, 26 deletions

diff --git a/libc/bionic/malloc_common.cpp b/libc/bionic/malloc_common.cpp
index d1aa1eaf8..fc65e15a9 100644
--- a/libc/bionic/malloc_common.cpp
+++ b/libc/bionic/malloc_common.cpp
@@ -562,13 +562,7 @@ static void ClearGlobalFunctions() {
   }
 }
 
-static void* LoadSharedLibrary(const char* shared_lib, const char* prefix, MallocDispatch* dispatch_table) {
-  void* impl_handle = dlopen(shared_lib, RTLD_NOW | RTLD_LOCAL);
-  if (impl_handle == nullptr) {
-    error_log("%s: Unable to open shared library %s: %s", getprogname(), shared_lib, dlerror());
-    return nullptr;
-  }
-
+static bool InitSharedLibrary(void* impl_handle, const char* shared_lib, const char* prefix, MallocDispatch* dispatch_table) {
   static constexpr const char* names[] = {
     "initialize",
     "finalize",
@@ -583,48 +577,61 @@ static void* LoadSharedLibrary(const char* shared_lib, const char* prefix, Mallo
     g_functions[i] = dlsym(impl_handle, symbol);
     if (g_functions[i] == nullptr) {
       error_log("%s: %s routine not found in %s", getprogname(), symbol, shared_lib);
-      dlclose(impl_handle);
       ClearGlobalFunctions();
-      return nullptr;
+      return false;
     }
   }
 
   if (!InitMallocFunctions(impl_handle, dispatch_table, prefix)) {
-    dlclose(impl_handle);
     ClearGlobalFunctions();
+    return false;
+  }
+  return true;
+}
+
+static void* LoadSharedLibrary(const char* shared_lib, const char* prefix, MallocDispatch* dispatch_table) {
+  void* impl_handle = dlopen(shared_lib, RTLD_NOW | RTLD_LOCAL);
+  if (impl_handle == nullptr) {
+    error_log("%s: Unable to open shared library %s: %s", getprogname(), shared_lib, dlerror());
     return nullptr;
   }
 
+  if (!InitSharedLibrary(impl_handle, shared_lib, prefix, dispatch_table)) {
+    dlclose(impl_handle);
+    impl_handle = nullptr;
+  }
+
   return impl_handle;
 }
 
 // A function pointer to heapprofds init function. Used to re-initialize
 // heapprofd. This will start a new profiling session and tear down the old
 // one in case it is still active.
-static _Atomic init_func_t g_heapprofd_init_func = nullptr;
+static _Atomic (void*) g_heapprofd_handle = nullptr;
 
 static void install_hooks(libc_globals* globals, const char* options,
                           const char* prefix, const char* shared_lib) {
-  init_func_t init_func = atomic_load(&g_heapprofd_init_func);
-  if (init_func != nullptr) {
-    init_func(&__libc_malloc_default_dispatch, &gMallocLeakZygoteChild, options);
-    info_log("%s: malloc %s re-enabled", getprogname(), prefix);
-    return;
-  }
-
   MallocDispatch dispatch_table;
-  void* impl_handle = LoadSharedLibrary(shared_lib, prefix, &dispatch_table);
-  if (impl_handle == nullptr) {
-    return;
+
+  void* impl_handle = atomic_load(&g_heapprofd_handle);
+  if (impl_handle != nullptr) {
+    if (!InitSharedLibrary(impl_handle, shared_lib, prefix, &dispatch_table)) {
+      return;
+    }
+  } else {
+    impl_handle = LoadSharedLibrary(shared_lib, prefix, &dispatch_table);
+    if (impl_handle == nullptr) {
+      return;
+    }
   }
-  init_func = reinterpret_cast<init_func_t>(g_functions[FUNC_INITIALIZE]);
+  init_func_t init_func = reinterpret_cast<init_func_t>(g_functions[FUNC_INITIALIZE]);
   if (!init_func(&__libc_malloc_default_dispatch, &gMallocLeakZygoteChild, options)) {
+    error_log("%s: failed to enable malloc %s", getprogname(), prefix);
     dlclose(impl_handle);
     ClearGlobalFunctions();
     return;
   }
 
-  atomic_store(&g_heapprofd_init_func, init_func);
   // We assign free  first explicitly to prevent the case where we observe a
   // alloc, but miss the corresponding free because of initialization order.
   //
@@ -636,6 +643,7 @@ static void install_hooks(libc_globals* globals, const char* options,
   // _Atomic. Assigning to an _Atomic is an atomic_store operation.
   // The assignment is done in declaration order.
   globals->malloc_dispatch = dispatch_table;
+  atomic_store(&g_heapprofd_handle, impl_handle);
 
   info_log("%s: malloc %s enabled", getprogname(), prefix);
 
@@ -761,6 +769,18 @@ extern "C" void MaybeInstallInitHeapprofdHook(int) {
 // =============================================================================
 
 #if !defined(LIBC_STATIC)
+bool MallocDispatchReset() {
+  if (!atomic_exchange(&g_heapprofd_init_in_progress, true)) {
+    __libc_globals.mutate([](libc_globals* globals) {
+      globals->malloc_dispatch = __libc_malloc_default_dispatch;
+    });
+    atomic_store(&g_heapprofd_init_in_progress, false);
+    return true;
+  }
+  errno = EAGAIN;
+  return false;
+}
+
 // Marks this process as a profileable zygote child.
 bool HandleInitZygoteChildProfiling() {
   atomic_store_explicit(&gMallocZygoteChildProfileable, true,
@@ -777,6 +797,10 @@ bool HandleInitZygoteChildProfiling() {
 
 #else
 
+bool MallocDispatchReset() {
+  return true;
+}
+
 bool HandleInitZygoteChildProfiling() {
   return true;
 }
@@ -791,6 +815,13 @@ bool android_mallopt(int opcode, void* arg, size_t arg_size) {
     }
     return HandleInitZygoteChildProfiling();
   }
+  if (opcode == M_RESET_HOOKS) {
+    if (arg != nullptr || arg_size != 0) {
+      errno = EINVAL;
+      return false;
+    }
+    return MallocDispatchReset();
+  }
 
   errno = ENOTSUP;
   return false;
diff --git a/libc/private/bionic_malloc.h b/libc/private/bionic_malloc.h
index a9fa22d83..5f4a75d7f 100644
--- a/libc/private/bionic_malloc.h
+++ b/libc/private/bionic_malloc.h
@@ -32,11 +32,13 @@
 
 // Opcodes for android_mallopt.
 
-// Marks the calling process as a profileable zygote child, possibly
-// initializing profiling infrastructure.
 enum {
+  // Marks the calling process as a profileable zygote child, possibly
+  // initializing profiling infrastructure.
   M_INIT_ZYGOTE_CHILD_PROFILING = 1,
 #define M_INIT_ZYGOTE_CHILD_PROFILING M_INIT_ZYGOTE_CHILD_PROFILING
+  M_RESET_HOOKS = 2,
+#define M_RESET_HOOKS M_RESET_HOOKS
 };
 
 // Manipulates bionic-specific handling of memory allocation APIs such as
diff --git a/libc/symbol_ordering b/libc/symbol_ordering
index 5b365f060..b672b3557 100644
--- a/libc/symbol_ordering
+++ b/libc/symbol_ordering
@@ -86,7 +86,6 @@ __malloc_hook
 __realloc_hook
 __free_hook
 __memalign_hook
-_ZL21g_heapprofd_init_func
 je_malloc_conf
 malloc_initializer
 a0