Bug 1722261 - Store {,Racy}RegisteredThread pointer in ThreadRegistrationData - r=canaltinova

This will help with replacing the old {,Racy}RegisteredThread with the new ThreadRegistration classes, while still being able to access the old classes (until they are not needed anymore).

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

tools/profiler/core/RegisteredThread.cpp

@@ -10,6 +10,7 @@
 #include "js/AllocationRecording.h"
 #include "js/ProfilingStack.h"
 #include "js/TraceLoggerAPI.h"
+#include "mozilla/ProfilerThreadRegistrationData.h"
 
 RacyRegisteredThread::RacyRegisteredThread(
     mozilla::profiler::ThreadRegistration& aThreadRegistration,
@@ -60,3 +61,27 @@ void RegisteredThread::SetJSContext(JSContext* aContext) {
 void RegisteredThread::PollJSSampling() {
   mRacyRegisteredThread.mThreadRegistration.mData.PollJSSampling();
 }
+
+const RacyRegisteredThread& mozilla::profiler::
+    ThreadRegistrationUnlockedConstReader::RacyRegisteredThreadCRef() const {
+  MOZ_ASSERT(mRegisteredThread);
+  return mRegisteredThread->RacyRegisteredThread();
+}
+
+RacyRegisteredThread&
+mozilla::profiler::ThreadRegistrationUnlockedConstReaderAndAtomicRW::
+    RacyRegisteredThreadRef() {
+  MOZ_ASSERT(mRegisteredThread);
+  return mRegisteredThread->RacyRegisteredThread();
+}
+
+RegisteredThread& mozilla::profiler::ThreadRegistrationLockedRWFromAnyThread::
+    RegisteredThreadRef() {
+  MOZ_ASSERT(mRegisteredThread);
+  return *mRegisteredThread;
+}
+
+void mozilla::profiler::ThreadRegistrationLockedRWOnThread::SetRegisteredThread(
+    RegisteredThread* aRegisteredThread) {
+  mRegisteredThread = aRegisteredThread;
+}

tools/profiler/core/platform.cpp

@@ -4123,6 +4123,18 @@ static ProfilingStack* locked_register_thread(
 
   TLSRegisteredThread::SetRegisteredThread(aLock, registeredThread.get());
 
+  ThreadRegistry::OffThreadRef::RWFromAnyThreadWithLock lockedRWFromAnyThread =
+      aOffThreadRef.LockedRWFromAnyThread();
+
+  ThreadRegistration::LockedRWOnThread* lockedRWOnThread =
+      lockedRWFromAnyThread.GetLockedRWOnThread();
+  MOZ_RELEASE_ASSERT(
+      lockedRWOnThread,
+      "At the moment, we should only get here when registering the current "
+      "thread (either through profiler_register_thread or from profiler_init "
+      "on the main thread registering that main thread).");
+  lockedRWOnThread->SetRegisteredThread(registeredThread.get());
+
   if (ActivePS::Exists(aLock) && ActivePS::ShouldProfileThread(aLock, info)) {
     registeredThread->RacyRegisteredThread().SetIsBeingProfiled(true);
     nsCOMPtr<nsIEventTarget> eventTarget = registeredThread->GetEventTarget();
@@ -5455,6 +5467,14 @@ static void locked_unregister_thread(PSLockRef lock) {
       ActivePS::UnregisterThread(lock, registeredThread);
     }
 
+    ThreadRegistration::WithOnThreadRef(
+        [](ThreadRegistration::OnThreadRef threadRegistration) {
+          threadRegistration.WithLockedRWOnThread(
+              [](ThreadRegistration::LockedRWOnThread& aRW) {
+                aRW.SetRegisteredThread(nullptr);
+              });
+        });
+
     // Clear the pointer to the RegisteredThread object that we're about to
     // destroy.
     TLSRegisteredThread::ResetRegisteredThread(lock);

tools/profiler/public/ProfilerThreadRegistrationData.h

@@ -203,6 +203,9 @@ class ThreadRegistrationData {
   static const int SLEEPING_OBSERVED = 2;
   // Read&written from thread and suspended thread.
   Atomic<int> mSleep{AWAKE};
+
+  // TODO: Remove when {,Racy}RegisteredThread are removed in a later patch.
+  RegisteredThread* mRegisteredThread;
 };
 
 // Accessing const data from any thread.
@@ -216,6 +219,9 @@ class ThreadRegistrationUnlockedConstReader : public ThreadRegistrationData {
 
   [[nodiscard]] const void* StackTop() const { return mStackTop; }
 
+  // TODO: Remove when {,Racy}RegisteredThread are removed in a later patch.
+  [[nodiscard]] const RacyRegisteredThread& RacyRegisteredThreadCRef() const;
+
  protected:
   ThreadRegistrationUnlockedConstReader(const char* aName,
                                         const void* aStackTop)
@@ -284,6 +290,9 @@ class ThreadRegistrationUnlockedConstReaderAndAtomicRW
 
   [[nodiscard]] bool IsSleeping() const { return mSleep != AWAKE; }
 
+  // TODO: Remove when {,Racy}RegisteredThread are removed in a later patch.
+  [[nodiscard]] RacyRegisteredThread& RacyRegisteredThreadRef();
+
  protected:
   ThreadRegistrationUnlockedConstReaderAndAtomicRW(const char* aName,
                                                    const void* aStackTop)
@@ -389,6 +398,9 @@ class ThreadRegistrationLockedRWFromAnyThread
     mJSSampling = INACTIVE_REQUESTED;
   }
 
+  // TODO: Remove when {,Racy}RegisteredThread are removed in a later patch.
+  [[nodiscard]] RegisteredThread& RegisteredThreadRef();
+
  protected:
   ThreadRegistrationLockedRWFromAnyThread(const char* aName,
                                           const void* aStackTop)
@@ -410,6 +422,9 @@ class ThreadRegistrationLockedRWOnThread
   // Poll to see if JS sampling should be started/stopped.
   void PollJSSampling();
 
+  // TODO: Remove when {,Racy}RegisteredThread are removed in a later patch.
+  void SetRegisteredThread(RegisteredThread* aRegisteredThread);
+
  public:
   ThreadRegistrationLockedRWOnThread(const char* aName, const void* aStackTop)
       : ThreadRegistrationLockedRWFromAnyThread(aName, aStackTop) {}