Reland "Fix GWP hooks not being restored, leading to crashes."

If the DispatchReset fails, the subsequent iteration has the wrong
idea of what the "original" table is, and if a subsequent DispatchReset
succeeds it unhooks them.

Repro in https://r.android.com/1767868.

Bug: 193012939
Bug: 189776979
Change-Id: I30445c053fcb785669f75d9c83056926d850edce
diff --git a/libc/bionic/malloc_heapprofd.cpp b/libc/bionic/malloc_heapprofd.cpp
index 741b45e..9c6ccb4 100644
--- a/libc/bionic/malloc_heapprofd.cpp
+++ b/libc/bionic/malloc_heapprofd.cpp
@@ -237,8 +237,6 @@
     // heapprofd client initialization.
     MallocHeapprofdState expected2 = kHookInstalled;
     if (atomic_compare_exchange_strong(&gHeapprofdState, &expected,
-          kInstallingEphemeralHook) ||
-        atomic_compare_exchange_strong(&gHeapprofdState, &expected2,
           kInstallingEphemeralHook)) {
       const MallocDispatch* default_dispatch = GetDefaultDispatchTable();
 
@@ -248,14 +246,8 @@
       // initialized, allocations may need to be serviced. There are three
       // possible configurations:
 
-      if (default_dispatch == nullptr) {
-        //  1. No malloc hooking has been done (heapprofd, GWP-ASan, etc.). In
-        //  this case, everything but malloc() should come from the system
-        //  allocator.
-        atomic_store(&gPreviousDefaultDispatchTable, nullptr);
-        gEphemeralDispatch = *NativeAllocatorDispatch();
-      } else if (DispatchIsGwpAsan(default_dispatch)) {
-        //  2. GWP-ASan was installed. We should use GWP-ASan for everything but
+      if (DispatchIsGwpAsan(default_dispatch)) {
+        //  1. GWP-ASan was installed. We should use GWP-ASan for everything but
         //  malloc() in the interim period before heapprofd is properly
         //  installed. After heapprofd is finished installing, we will use
         //  GWP-ASan as heapprofd's backing allocator to allow heapprofd and
@@ -263,8 +255,16 @@
         atomic_store(&gPreviousDefaultDispatchTable, default_dispatch);
         gEphemeralDispatch = *default_dispatch;
       } else {
+        // Either,
+        // 2. No malloc hooking has been done (heapprofd, GWP-ASan, etc.). In
+        // this case, everything but malloc() should come from the system
+        // allocator.
+        //
+        // or,
+        //
         // 3. It may be possible at this point in time that heapprofd is
-        // *already* the default dispatch, and as such we don't want to use
+        // *already* the default dispatch, and when it was initialized there
+        // was no default dispatch installed. As such we don't want to use
         // heapprofd as the backing store for itself (otherwise infinite
         // recursion occurs). We will use the system allocator functions. Note:
         // We've checked that no other malloc interceptors are being used by
@@ -273,23 +273,41 @@
         atomic_store(&gPreviousDefaultDispatchTable, nullptr);
         gEphemeralDispatch = *NativeAllocatorDispatch();
       }
-
-      // Now, replace the malloc function so that the next call to malloc() will
-      // initialize heapprofd.
-      gEphemeralDispatch.malloc = MallocInitHeapprofdHook;
-
-      // And finally, install these new malloc-family interceptors.
-      __libc_globals.mutate([](libc_globals* globals) {
-        atomic_store(&globals->default_dispatch_table, &gEphemeralDispatch);
-        if (!MallocLimitInstalled()) {
-          atomic_store(&globals->current_dispatch_table, &gEphemeralDispatch);
-        }
-      });
-      atomic_store(&gHeapprofdState, kEphemeralHookInstalled);
+    } else if (atomic_compare_exchange_strong(&gHeapprofdState, &expected2,
+                                              kInstallingEphemeralHook)) {
+      // if we still have hook installed, we can reuse the previous
+      // decision. THIS IS REQUIRED FOR CORRECTNESS, because otherwise the
+      // following can happen
+      // 1. Assume DispatchIsGwpAsan(default_dispatch)
+      // 2. This function is ran, sets gPreviousDefaultDispatchTable to
+      //    GWP ASan.
+      // 3. The sessions ends, DispatchReset FAILS due to a race. Now
+      //    heapprofd hooks are default dispatch.
+      // 4. We re-enter this function later. If we did NOT look at the
+      //    previously recorded gPreviousDefaultDispatchTable, we would
+      //    incorrectly reach case 3. below.
+      // 5. The session ends, DispatchReset now resets the hooks to the
+      //    system allocator. This is incorrect.
+      const MallocDispatch* prev_dispatch =
+        atomic_load(&gPreviousDefaultDispatchTable);
+      gEphemeralDispatch = prev_dispatch ? *prev_dispatch : *NativeAllocatorDispatch();
     } else {
       error_log("%s: heapprofd: failed to transition kInitialState -> kInstallingEphemeralHook. "
           "current state (possible race): %d", getprogname(), expected2);
+      return;
     }
+    // Now, replace the malloc function so that the next call to malloc() will
+    // initialize heapprofd.
+    gEphemeralDispatch.malloc = MallocInitHeapprofdHook;
+
+    // And finally, install these new malloc-family interceptors.
+    __libc_globals.mutate([](libc_globals* globals) {
+      atomic_store(&globals->default_dispatch_table, &gEphemeralDispatch);
+      if (!MallocLimitInstalled()) {
+        atomic_store(&globals->current_dispatch_table, &gEphemeralDispatch);
+      }
+    });
+    atomic_store(&gHeapprofdState, kEphemeralHookInstalled);
   });
   // Otherwise, we're racing against malloc_limit's enable logic (at most once
   // per process, and a niche feature). This is highly unlikely, so simply give
@@ -335,16 +353,15 @@
     return;
   }
 
+  FinishInstallHooks(globals, nullptr, kHeapprofdPrefix);
+}
+
+void HeapprofdInstallHooksAtInit(libc_globals *globals) {
   // Before we set the new default_dispatch_table in FinishInstallHooks, save
   // the previous dispatch table. If DispatchReset() gets called later, we want
   // to be able to restore the dispatch. We're still under
   // MaybeModifyGlobals locks at this point.
   atomic_store(&gPreviousDefaultDispatchTable, GetDefaultDispatchTable());
-
-  FinishInstallHooks(globals, nullptr, kHeapprofdPrefix);
-}
-
-void HeapprofdInstallHooksAtInit(libc_globals* globals) {
   MaybeModifyGlobals(kWithoutLock, [globals] {
     MallocHeapprofdState expected = kInitialState;
     if (atomic_compare_exchange_strong(&gHeapprofdState, &expected, kInstallingHook)) {