Bug 1445822 - Make ProfilingStack ref-counted, owned by both profiler and thread - r=mstange

When the profiler shuts down, it destroys all remaining `RegisteredThread`'s,
including the embedded ProfilingStack that stores labels for each thread.
Now, it is possible that some threads may outlive the profiler (e.g., Rayon
threads, which cannot be synchronously shut down), so there could be labels left
when the `ProfilingStack` gets destroyed, which triggers the `stackPointer == 0`
assertion in the destructor.

Removing the assertion would not be enough, because after that, the thread may
still try to pop labels (using a pointer to the now-destroyed `ProfilingStack`
stored in `AutoProfilerLabel`), causing a UAF.
`AutoProfilerLabel` could theoretically check that the profiler is still alive
before trying to pop labels, but that would add a significant cost every time
the `ProfilingStack` is accessed, because the profiler mutex would need to be
locked.

The solution here is to make `ProfilingStack` a separate heap object, owned
both by the profiler through the thread's `RegisteredThread`, and by the thread
itself through the thread-local `sProfilingStackOwnerTLS`.
The first owning link is severed when the profiler shuts down.
The other owning link is severed when the thread unregisters itself (before or
after shutdown).
This way the `ProfilingStack` stays alive as long as needed by the thread, even
if it outlives the profiler.

Implementation detail: `ProfilingStack` is used in other places as a straight
object, so it cannot be made ref-counted itself. Instead we use
`ProfilingStackOwner`, a small ref-counted shell around it.

Differential Revision: https://phabricator.services.mozilla.com/D49141

--HG--
extra : moz-landing-system : lando

tools/profiler/core/RegisteredThread.h

@@ -12,7 +12,8 @@
 
 #include "js/TraceLoggerAPI.h"
 #include "jsapi.h"
-#include "mozilla/UniquePtr.h"
+#include "mozilla/NotNull.h"
+#include "mozilla/RefPtr.h"
 #include "nsIEventTarget.h"
 
 // This class contains the state for a single thread that is accessible without
@@ -23,7 +24,11 @@
 class RacyRegisteredThread final {
  public:
   explicit RacyRegisteredThread(int aThreadId)
-      : mThreadId(aThreadId), mSleep(AWAKE), mIsBeingProfiled(false) {
+      : mProfilingStackOwner(
+            mozilla::MakeNotNull<RefPtr<class ProfilingStackOwner>>()),
+        mThreadId(aThreadId),
+        mSleep(AWAKE),
+        mIsBeingProfiled(false) {
     MOZ_COUNT_CTOR(RacyRegisteredThread);
   }
 
@@ -77,12 +82,18 @@ class RacyRegisteredThread final {
   int ThreadId() const { return mThreadId; }
 
   class ProfilingStack& ProfilingStack() {
-    return mProfilingStack;
+    return mProfilingStackOwner->ProfilingStack();
+  }
+  const class ProfilingStack& ProfilingStack() const {
+    return mProfilingStackOwner->ProfilingStack();
+  }
+
+  class ProfilingStackOwner& ProfilingStackOwner() {
+    return *mProfilingStackOwner;
   }
-  const class ProfilingStack& ProfilingStack() const { return mProfilingStack; }
 
  private:
-  class ProfilingStack mProfilingStack;
+  mozilla::NotNull<RefPtr<class ProfilingStackOwner>> mProfilingStackOwner;
 
   // mThreadId contains the thread ID of the current thread. It is safe to read
   // this from multiple threads concurrently, as it will never be mutated.

tools/profiler/core/platform.cpp

@@ -1192,7 +1192,7 @@ class TLSRegisteredThread {
  public:
   static bool Init(PSLockRef) {
     bool ok1 = sRegisteredThread.init();
-    bool ok2 = AutoProfilerLabel::sProfilingStack.init();
+    bool ok2 = AutoProfilerLabel::sProfilingStackOwnerTLS.init();
     return ok1 && ok2;
   }
 
@@ -1212,16 +1212,41 @@ class TLSRegisteredThread {
   // RacyRegisteredThread() can also be used to get the ProfilingStack, but that
   // is marginally slower because it requires an extra pointer indirection.
   static ProfilingStack* Stack() {
-    return AutoProfilerLabel::sProfilingStack.get();
+    ProfilingStackOwner* profilingStackOwner =
+        AutoProfilerLabel::sProfilingStackOwnerTLS.get();
+    if (!profilingStackOwner) {
+      return nullptr;
+    }
+    return &profilingStackOwner->ProfilingStack();
   }
 
-  static void SetRegisteredThread(PSLockRef,
-                                  class RegisteredThread* aRegisteredThread) {
+  static void SetRegisteredThreadAndAutoProfilerLabelProfilingStack(
+      PSLockRef, class RegisteredThread* aRegisteredThread) {
+    MOZ_RELEASE_ASSERT(
+        aRegisteredThread,
+        "Use ResetRegisteredThread() instead of SetRegisteredThread(nullptr)");
     sRegisteredThread.set(aRegisteredThread);
-    AutoProfilerLabel::sProfilingStack.set(
-        aRegisteredThread
-            ? &aRegisteredThread->RacyRegisteredThread().ProfilingStack()
-            : nullptr);
+    ProfilingStackOwner& profilingStackOwner =
+        aRegisteredThread->RacyRegisteredThread().ProfilingStackOwner();
+    profilingStackOwner.AddRef();
+    AutoProfilerLabel::sProfilingStackOwnerTLS.set(&profilingStackOwner);
+  }
+
+  // Only reset the registered thread. The AutoProfilerLabel's ProfilingStack
+  // is kept, because the thread may not have unregistered itself yet, so it may
+  // still push/pop labels even after the profiler has shut down.
+  static void ResetRegisteredThread(PSLockRef) {
+    sRegisteredThread.set(nullptr);
+  }
+
+  // Reset the AutoProfilerLabels' ProfilingStack, because the thread is
+  // unregistering itself.
+  static void ResetAutoProfilerLabelProfilingStack(PSLockRef) {
+    MOZ_RELEASE_ASSERT(
+        AutoProfilerLabel::sProfilingStackOwnerTLS.get(),
+        "ResetAutoProfilerLabelProfilingStack should only be called once");
+    AutoProfilerLabel::sProfilingStackOwnerTLS.get()->Release();
+    AutoProfilerLabel::sProfilingStackOwnerTLS.set(nullptr);
   }
 
  private:
@@ -1249,7 +1274,14 @@ MOZ_THREAD_LOCAL(RegisteredThread*) TLSRegisteredThread::sRegisteredThread;
 //
 // This second pointer isn't ideal, but does provide a way to satisfy those
 // constraints. TLSRegisteredThread is responsible for updating it.
-MOZ_THREAD_LOCAL(ProfilingStack*) AutoProfilerLabel::sProfilingStack;
+//
+// The (Racy)RegisteredThread and AutoProfilerLabel::sProfilingStackOwnerTLS
+// co-own the thread's ProfilingStack, so whichever is reset second, is
+// responsible for destroying the ProfilingStack; Because MOZ_THREAD_LOCAL
+// doesn't support RefPtr, AddRef&Release are done explicitly in
+// TLSRegisteredThread.
+MOZ_THREAD_LOCAL(ProfilingStackOwner*)
+AutoProfilerLabel::sProfilingStackOwnerTLS;
 
 // The name of the main thread.
 static const char* const kMainThreadName = "GeckoMain";
@@ -3302,7 +3334,8 @@ static ProfilingStack* locked_register_thread(PSLockRef aLock,
   UniquePtr<RegisteredThread> registeredThread = MakeUnique<RegisteredThread>(
       info, NS_GetCurrentThreadNoCreate(), aStackTop);
 
-  TLSRegisteredThread::SetRegisteredThread(aLock, registeredThread.get());
+  TLSRegisteredThread::SetRegisteredThreadAndAutoProfilerLabelProfilingStack(
+      aLock, registeredThread.get());
 
   if (ActivePS::Exists(aLock) && ActivePS::ShouldProfileThread(aLock, info)) {
     registeredThread->RacyRegisteredThread().SetIsBeingProfiled(true);
@@ -3378,18 +3411,21 @@ static void locked_profiler_start(PSLockRef aLock, PowerOfTwo32 aCapacity,
 // This basically duplicates AutoProfilerLabel's constructor.
 static void* MozGlueLabelEnter(const char* aLabel, const char* aDynamicString,
                                void* aSp) {
-  ProfilingStack* profilingStack = AutoProfilerLabel::sProfilingStack.get();
-  if (profilingStack) {
-    profilingStack->pushLabelFrame(aLabel, aDynamicString, aSp,
-                                   JS::ProfilingCategoryPair::OTHER);
+  ProfilingStackOwner* profilingStackOwner =
+      AutoProfilerLabel::sProfilingStackOwnerTLS.get();
+  if (profilingStackOwner) {
+    profilingStackOwner->ProfilingStack().pushLabelFrame(
+        aLabel, aDynamicString, aSp, JS::ProfilingCategoryPair::OTHER);
   }
-  return profilingStack;
+  return profilingStackOwner;
 }
 
 // This basically duplicates AutoProfilerLabel's destructor.
-static void MozGlueLabelExit(void* sProfilingStack) {
-  if (sProfilingStack) {
-    reinterpret_cast<ProfilingStack*>(sProfilingStack)->pop();
+static void MozGlueLabelExit(void* aProfilingStackOwner) {
+  if (aProfilingStackOwner) {
+    reinterpret_cast<ProfilingStackOwner*>(aProfilingStackOwner)
+        ->ProfilingStack()
+        .pop();
   }
 }
 
@@ -3459,6 +3495,7 @@ void profiler_init(void* aStackTop) {
     // indicates that the profiler has initialized successfully.
     CorePS::Create(lock);
 
+    // profiler_init implicitly registers this thread as main thread.
     locked_register_thread(lock, kMainThreadName, aStackTop);
 
     // Platform-specific initialization.
@@ -3627,7 +3664,10 @@ void profiler_shutdown() {
 
     // We just destroyed CorePS and the ThreadInfos it contains, so we can
     // clear this thread's TLSRegisteredThread.
-    TLSRegisteredThread::SetRegisteredThread(lock, nullptr);
+    TLSRegisteredThread::ResetRegisteredThread(lock);
+    // We can also clear the AutoProfilerLabel's ProfilingStack because the
+    // main thread should not use labels after profiler_shutdown.
+    TLSRegisteredThread::ResetAutoProfilerLabelProfilingStack(lock);
 
 #ifdef MOZ_TASK_TRACER
     tasktracer::ShutdownTaskTracer();
@@ -4375,6 +4415,10 @@ void profiler_unregister_thread() {
 
   if (!CorePS::Exists()) {
     // This function can be called after the main thread has already shut down.
+    // We want to reset the AutoProfilerLabel's ProfilingStack pointer (if
+    // needed), because a thread could stay registered after the profiler has
+    // shut down.
+    TLSRegisteredThread::ResetAutoProfilerLabelProfilingStack(lock);
     return;
   }
 
@@ -4394,8 +4438,10 @@ void profiler_unregister_thread() {
     }
 
     // Clear the pointer to the RegisteredThread object that we're about to
-    // destroy.
-    TLSRegisteredThread::SetRegisteredThread(lock, nullptr);
+    // destroy, as well as the AutoProfilerLabel's ProfilingStack because the
+    // thread is unregistering itself and won't need the ProfilingStack anymore.
+    TLSRegisteredThread::ResetRegisteredThread(lock);
+    TLSRegisteredThread::ResetAutoProfilerLabelProfilingStack(lock);
 
     // Remove the thread from the list of registered threads. This deletes the
     // registeredThread object.

tools/profiler/public/GeckoProfiler.h

@@ -1058,6 +1058,35 @@ class MOZ_RAII AutoProfilerThreadWake {
   bool mIssuedWake;
 };
 
+// Ref-counted shell around a `ProfilingStack`, to be used by the owning
+// (Racy)RegisteredThread and AutoProfilerLabel.
+class ProfilingStackOwner {
+ public:
+  class ProfilingStack& ProfilingStack() {
+    return mProfilingStack;
+  }
+
+  // Using hand-rolled ref-counting, to evade leak checking (emergency patch
+  // for bug 1445822).
+  // TODO: Eliminate all/most leaks if possible.
+  void AddRef() const { ++mRefCnt; }
+  void Release() const {
+    MOZ_ASSERT(int32_t(mRefCnt) > 0);
+    if (--mRefCnt == 0) {
+      delete this;
+    }
+  }
+
+ private:
+  ~ProfilingStackOwner() = default;
+
+  class ProfilingStack mProfilingStack;
+
+  mutable Atomic<int32_t, MemoryOrdering::ReleaseAcquire,
+                 recordreplay::Behavior::DontPreserve>
+      mRefCnt;
+};
+
 // This class creates a non-owning ProfilingStack reference. Objects of this
 // class are stack-allocated, and so exist within a thread, and are thus bounded
 // by the lifetime of the thread, which ensures that the references held can't
@@ -1071,7 +1100,9 @@ class MOZ_RAII AutoProfilerLabel {
     MOZ_GUARD_OBJECT_NOTIFIER_INIT;
 
     // Get the ProfilingStack from TLS.
-    Push(sProfilingStack.get(), aLabel, aDynamicString, aCategoryPair, aFlags);
+    ProfilingStackOwner* profilingStackOwner = sProfilingStackOwnerTLS.get();
+    Push(profilingStackOwner ? &profilingStackOwner->ProfilingStack() : nullptr,
+         aLabel, aDynamicString, aCategoryPair, aFlags);
   }
 
   // This is the AUTO_PROFILER_LABEL_FAST variant. It retrieves the
@@ -1115,7 +1146,7 @@ class MOZ_RAII AutoProfilerLabel {
 
  public:
   // See the comment on the definition in platform.cpp for details about this.
-  static MOZ_THREAD_LOCAL(ProfilingStack*) sProfilingStack;
+  static MOZ_THREAD_LOCAL(ProfilingStackOwner*) sProfilingStackOwnerTLS;
 };
 
 class MOZ_RAII AutoProfilerTracing {