Bug 1659901 - Rework profiler's TLS accesses - r=mstange

To ensure correct usage of TLSs in the profiler, they are now better encapsulated so that:
- init() is called once and its result is cached. (TLSREGISTEREDThread::Init() doesn't need proof of the PSLock, because it's using thread-safe function-static initializers.)
- get() and set() always init() as needed, or in some particular cases strongly assert that init() was successful.

Also, a null-check was missing in profiler_init_threadmanager().

Depends on D87588

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

tools/profiler/core/platform.cpp

@@ -1444,22 +1444,31 @@ static PSMutex gPSMutex;
 Atomic<uint32_t, MemoryOrdering::Relaxed> RacyFeatures::sActiveAndFeatures(0);
 
 // Each live thread has a RegisteredThread, and we store a reference to it in
-// TLS. This class encapsulates that TLS.
+// TLS. This class encapsulates that TLS, and also handles the associated
+// profiling stack used by AutoProfilerLabel.
 class TLSRegisteredThread {
  public:
-  static bool Init(PSLockRef) {
-    bool ok1 = sRegisteredThread.init();
-    bool ok2 = AutoProfilerLabel::sProfilingStackOwnerTLS.init();
-    return ok1 && ok2;
+  static bool Init() {
+    // Only one call to MOZ_THREAD_LOCAL::init() is needed, we cache the result
+    // for later calls to Init(), in particular before using get() and set().
+    static const bool ok = AutoProfilerLabel::ProfilingStackOwnerTLS::Init() &&
+                           sRegisteredThread.init();
+    return ok;
   }
 
   // Get the entire RegisteredThread. Accesses are guarded by gPSMutex.
   static class RegisteredThread* RegisteredThread(PSLockRef) {
+    if (!Init()) {
+      return nullptr;
+    }
     return sRegisteredThread.get();
   }
 
   // Get only the RacyRegisteredThread. Accesses are not guarded by gPSMutex.
   static class RacyRegisteredThread* RacyRegisteredThread() {
+    if (!Init()) {
+      return nullptr;
+    }
     class RegisteredThread* registeredThread = sRegisteredThread.get();
     return registeredThread ? &registeredThread->RacyRegisteredThread()
                             : nullptr;
@@ -1469,8 +1478,11 @@ 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() {
+    if (!Init()) {
+      return nullptr;
+    }
     ProfilingStackOwner* profilingStackOwner =
-        AutoProfilerLabel::sProfilingStackOwnerTLS.get();
+        AutoProfilerLabel::ProfilingStackOwnerTLS::Get();
     if (!profilingStackOwner) {
       return nullptr;
     }
@@ -1479,6 +1491,9 @@ class TLSRegisteredThread {
 
   static void SetRegisteredThreadAndAutoProfilerLabelProfilingStack(
       PSLockRef, class RegisteredThread* aRegisteredThread) {
+    if (!Init()) {
+      return;
+    }
     MOZ_RELEASE_ASSERT(
         aRegisteredThread,
         "Use ResetRegisteredThread() instead of SetRegisteredThread(nullptr)");
@@ -1486,24 +1501,30 @@ class TLSRegisteredThread {
     ProfilingStackOwner& profilingStackOwner =
         aRegisteredThread->RacyRegisteredThread().ProfilingStackOwner();
     profilingStackOwner.AddRef();
-    AutoProfilerLabel::sProfilingStackOwnerTLS.set(&profilingStackOwner);
+    AutoProfilerLabel::ProfilingStackOwnerTLS::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) {
+    if (!Init()) {
+      return;
+    }
     sRegisteredThread.set(nullptr);
   }
 
   // Reset the AutoProfilerLabels' ProfilingStack, because the thread is
   // unregistering itself.
   static void ResetAutoProfilerLabelProfilingStack(PSLockRef) {
+    if (!Init()) {
+      return;
+    }
     MOZ_RELEASE_ASSERT(
-        AutoProfilerLabel::sProfilingStackOwnerTLS.get(),
+        AutoProfilerLabel::ProfilingStackOwnerTLS::Get(),
         "ResetAutoProfilerLabelProfilingStack should only be called once");
-    AutoProfilerLabel::sProfilingStackOwnerTLS.get()->Release();
-    AutoProfilerLabel::sProfilingStackOwnerTLS.set(nullptr);
+    AutoProfilerLabel::ProfilingStackOwnerTLS::Get()->Release();
+    AutoProfilerLabel::ProfilingStackOwnerTLS::Set(nullptr);
   }
 
  private:
@@ -1532,13 +1553,13 @@ 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.
 //
-// The (Racy)RegisteredThread and AutoProfilerLabel::sProfilingStackOwnerTLS
+// The (Racy)RegisteredThread and AutoProfilerLabel::ProfilingStackOwnerTLS
 // 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;
+AutoProfilerLabel::ProfilingStackOwnerTLS::sProfilingStackOwnerTLS;
 
 void ProfilingStackOwner::DumpStackAndCrash() const {
   fprintf(stderr,
@@ -3817,7 +3838,7 @@ static ProfilingStack* locked_register_thread(PSLockRef aLock,
 
   VTUNE_REGISTER_THREAD(aName);
 
-  if (!TLSRegisteredThread::Init(aLock)) {
+  if (!TLSRegisteredThread::Init()) {
     return nullptr;
   }
 
@@ -3911,7 +3932,7 @@ static void locked_profiler_start(PSLockRef aLock, PowerOfTwo32 aCapacity,
 static void* MozGlueLabelEnter(const char* aLabel, const char* aDynamicString,
                                void* aSp) {
   ProfilingStackOwner* profilingStackOwner =
-      AutoProfilerLabel::sProfilingStackOwnerTLS.get();
+      AutoProfilerLabel::ProfilingStackOwnerTLS::Get();
   if (profilingStackOwner) {
     profilingStackOwner->ProfilingStack().pushLabelFrame(
         aLabel, aDynamicString, aSp, JS::ProfilingCategoryPair::OTHER);
@@ -3956,7 +3977,7 @@ void profiler_init_threadmanager() {
   PSAutoLock lock(gPSMutex);
   RegisteredThread* registeredThread =
       TLSRegisteredThread::RegisteredThread(lock);
-  if (!registeredThread->GetEventTarget()) {
+  if (registeredThread && !registeredThread->GetEventTarget()) {
     registeredThread->ResetMainThread(NS_GetCurrentThreadNoCreate());
   }
 }
@@ -5087,6 +5108,9 @@ ProfilingStack* profiler_register_thread(const char* aName,
         "profiler_register_thread again",
         TextMarkerPayload(text, TimeStamp::NowUnfuzzed()), &lock);
 
+    MOZ_RELEASE_ASSERT(
+        TLSRegisteredThread::Init(),
+        "Thread should not have already been registered without TLS::Init()");
     MOZ_RELEASE_ASSERT(TLSRegisteredThread::RegisteredThread(lock),
                        "TLS should be set when re-registering thread");
     MOZ_RELEASE_ASSERT(
@@ -5118,6 +5142,9 @@ void profiler_unregister_thread() {
   if (RegisteredThread* registeredThread =
           FindCurrentThreadRegisteredThread(lock);
       registeredThread) {
+    MOZ_RELEASE_ASSERT(
+        TLSRegisteredThread::Init(),
+        "Thread should not have been registered without TLS::Init()");
     MOZ_RELEASE_ASSERT(TLSRegisteredThread::RegisteredThread(lock),
                        "TLS should be set when un-registering thread");
     MOZ_RELEASE_ASSERT(

tools/profiler/public/GeckoProfiler.h

@@ -1187,7 +1187,7 @@ class MOZ_RAII AutoProfilerLabel {
                     JS::ProfilingCategoryPair aCategoryPair,
                     uint32_t aFlags = 0) {
     // Get the ProfilingStack from TLS.
-    ProfilingStackOwner* profilingStackOwner = sProfilingStackOwnerTLS.get();
+    ProfilingStackOwner* profilingStackOwner = ProfilingStackOwnerTLS::Get();
     Push(profilingStackOwner ? &profilingStackOwner->ProfilingStack() : nullptr,
          aLabel, aDynamicString, aCategoryPair, aFlags);
   }
@@ -1229,7 +1229,33 @@ class MOZ_RAII AutoProfilerLabel {
 
  public:
   // See the comment on the definition in platform.cpp for details about this.
-  static MOZ_THREAD_LOCAL(ProfilingStackOwner*) sProfilingStackOwnerTLS;
+  class ProfilingStackOwnerTLS {
+   public:
+    static bool Init() {
+      // Only one call to MOZ_THREAD_LOCAL::init() is needed, we cache the
+      // result for later calls to Init(), in particular before using get() and
+      // set().
+      static const bool ok = sProfilingStackOwnerTLS.init();
+      return ok;
+    }
+
+    static ProfilingStackOwner* Get() {
+      if (!Init()) {
+        return nullptr;
+      }
+      return sProfilingStackOwnerTLS.get();
+    }
+
+    static void Set(ProfilingStackOwner* aProfilingStackOwner) {
+      MOZ_RELEASE_ASSERT(Init(),
+                         "ProfilingStackOwnerTLS::Set() should only be called "
+                         "after a successful Init()");
+      sProfilingStackOwnerTLS.set(aProfilingStackOwner);
+    }
+
+   private:
+    static MOZ_THREAD_LOCAL(ProfilingStackOwner*) sProfilingStackOwnerTLS;
+  };
 };
 
 class MOZ_RAII AutoProfilerTracing {