1 files changed, 54 insertions, 23 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;