Update sampling profiler thread (#1421)

Summary:

Allow repeated calls to enable the sampling profiler,
updating the thread targeted by the profiler each time.

Previously, the sampling profiler would target only a
single thread. If the runtime was used from another
thread, multiple threads would be trying to use the
runtime in parallel, causing incorrect behavior.

Reviewed By: avp

Differential Revision: D57397716

API/hermes/hermes.cpp

@@ -1326,10 +1326,10 @@ void HermesRuntime::registerForProfiling() {
 #if HERMESVM_SAMPLING_PROFILER_AVAILABLE
   vm::Runtime &runtime = impl(this)->runtime_;
   if (runtime.samplingProfiler) {
-    ::hermes::hermes_fatal(
-        "re-registering HermesVMs for profiling is not allowed");
+    runtime.samplingProfiler->setRuntimeThread();
+  } else {
+    runtime.samplingProfiler = ::hermes::vm::SamplingProfiler::create(runtime);
   }
-  runtime.samplingProfiler = ::hermes::vm::SamplingProfiler::create(runtime);
 #else
   throwHermesNotCompiledWithSamplingProfilerSupport();
 #endif // HERMESVM_SAMPLING_PROFILER_AVAILABLE

API/hermes/hermes.h

@@ -177,7 +177,10 @@ class HERMES_EXPORT HermesRuntime : public jsi::Runtime {
       const DebugFlags &debugFlags);
 #endif
 
-  /// Register this runtime for sampling profiler.
+  /// Register this runtime and thread for sampling profiler. Before using the
+  /// runtime on another thread, invoke this function again from the new thread
+  /// to make the sampling profiler target the new thread (and forget the old
+  /// thread).
   void registerForProfiling();
   /// Unregister this runtime for sampling profiler.
   void unregisterForProfiling();

include/hermes/VM/Profiler/SamplingProfiler.h

@@ -109,9 +109,11 @@ class SamplingProfiler {
   };
 
   /// \return true if this SamplingProfiler belongs to the current running
-  /// thread. Does not acquire any locks, and as such should not be used in
-  /// production.
-  bool belongsToCurrentThread() const;
+  /// thread. The current thread can change (e.g. in the time between
+  /// this function returning and the caller inspecting the value), so the
+  /// usefulness of this method depends upon knowledge of when the runtime
+  /// will switch threads.
+  bool belongsToCurrentThread();
 
   /// \returns the NativeFunctionPtr for \p stackFrame. Caller must hold
   /// runtimeDataLock_.
@@ -149,8 +151,8 @@ class SamplingProfiler {
   /// Max size of sampleStorage_.
   static const int kMaxStackDepth = 500;
 
-  /// Protect data specific to a runtime, such as the sampled stacks and
-  /// domains.
+  /// Protect data specific to a runtime, such as the sampled stacks,
+  /// domains, and thread associated with the runtime.
   std::mutex runtimeDataLock_;
 
  protected:
@@ -274,6 +276,10 @@ class SamplingProfiler {
   /// suspend() that hansn't been resume()d yet.
   void resume();
 
+  /// Inform the sampling profiler that the runtime will now be executing
+  /// bytecode on the current thread.
+  void setRuntimeThread();
+
  protected:
   explicit SamplingProfiler(Runtime &runtime);
 };

lib/VM/Profiler/SamplingProfilerPosix.cpp

@@ -70,10 +70,9 @@ struct SamplingProfilerPosix : SamplingProfiler {
   SamplingProfilerPosix(Runtime &rt);
   ~SamplingProfilerPosix() override;
 
-  /// Thread that this profiler instance represents. This can currently only be
-  /// set from the constructor of SamplingProfiler, so we need to construct a
-  /// new SamplingProfiler every time the runtime is moved to a different
-  /// thread.
+  /// Thread that this profiler instance represents. This can be updated as the
+  /// runtime is invoked on different threads. Must only be accessed while
+  /// holding the runtimeDataLock_.
   pthread_t currentThread_;
 
 #if defined(HERMESVM_ENABLE_LOOM) && defined(__ANDROID__)
@@ -316,7 +315,9 @@ bool Sampler::platformSuspendVMAndWalkStack(SamplingProfiler *profiler) {
   // acquired the updates to domains_.
   self->profilerForSig_.store(profiler, std::memory_order_release);
 
-  // Signal target runtime thread to sample stack.
+  // Signal target runtime thread to sample stack. The runtimeDataLock is
+  // held by the caller, ensuring the runtime won't start to be used on
+  // another thread before sampling begins.
   pthread_kill(posixProfiler->currentThread_, SIGPROF);
 
   // Threading: samplingDoneSem_ will synchronise this thread with the
@@ -470,9 +471,16 @@ std::unique_ptr<SamplingProfiler> SamplingProfiler::create(Runtime &rt) {
   return std::make_unique<sampling_profiler::SamplingProfilerPosix>(rt);
 }
 
-bool SamplingProfiler::belongsToCurrentThread() const {
-  return static_cast<const sampling_profiler::SamplingProfilerPosix *>(this)
-             ->currentThread_ == pthread_self();
+bool SamplingProfiler::belongsToCurrentThread() {
+  auto profiler = static_cast<sampling_profiler::SamplingProfilerPosix *>(this);
+  std::lock_guard<std::mutex> lock(profiler->runtimeDataLock_);
+  return profiler->currentThread_ == pthread_self();
+}
+
+void SamplingProfiler::setRuntimeThread() {
+  auto profiler = static_cast<sampling_profiler::SamplingProfilerPosix *>(this);
+  std::lock_guard<std::mutex> lock(profiler->runtimeDataLock_);
+  profiler->currentThread_ = pthread_self();
 }
 
 } // namespace vm

lib/VM/Profiler/SamplingProfilerWindows.cpp

@@ -26,14 +26,21 @@
 
 namespace hermes {
 namespace vm {
+
+/// Open the current thread and \return a handle to it. The handle must later
+/// be closed with a call to the Win32 CloseHandle API.
+static HANDLE openCurrentThread() {
+  return OpenThread(
+      THREAD_GET_CONTEXT | THREAD_SUSPEND_RESUME | THREAD_QUERY_INFORMATION,
+      false,
+      GetCurrentThreadId());
+}
+
 namespace sampling_profiler {
 namespace {
 struct SamplingProfilerWindows : SamplingProfiler {
   SamplingProfilerWindows(Runtime &rt) : SamplingProfiler(rt) {
-    currentThread_ = OpenThread(
-        THREAD_GET_CONTEXT | THREAD_SUSPEND_RESUME | THREAD_QUERY_INFORMATION,
-        false,
-        GetCurrentThreadId());
+    currentThread_ = openCurrentThread();
 
 #if defined(HERMESVM_ENABLE_LOOM)
     fbloom_profilo_api()->fbloom_register_enable_for_loom_callback(
@@ -48,6 +55,8 @@ struct SamplingProfilerWindows : SamplingProfiler {
     // TODO(T125910634): re-introduce the requirement for destroying the
     // sampling profiler on the same thread in which it was created.
     Sampler::get()->unregisterRuntime(this);
+
+    CloseHandle(currentThread_);
   }
 
 #if defined(HERMESVM_ENABLE_LOOM)
@@ -121,10 +130,8 @@ struct SamplingProfilerWindows : SamplingProfiler {
   bool loomDataPushEnabled_{false};
 #endif // defined(HERMESVM_ENABLE_LOOM)
 
-  /// Thread that this profiler instance represents. This can currently only
-  /// be set from the constructor of SamplingProfiler, so we need to construct
-  /// a new SamplingProfiler every time the runtime is moved to a different
-  /// thread.
+  /// Thread that this profiler instance represents. This can be updated
+  /// by later calls to SetRuntimeThread.
   HANDLE currentThread_;
 };
 } // namespace
@@ -171,7 +178,9 @@ void Sampler::platformPostSampleStack(SamplingProfiler *localProfiler) {
 bool Sampler::platformSuspendVMAndWalkStack(SamplingProfiler *profiler) {
   auto *winProfiler = static_cast<SamplingProfilerWindows *>(profiler);
 
-  // Suspend the JS thread.
+  // Suspend the JS thread. The runtimeDataLock is held by the caller, ensuring
+  // the runtime won't start to be used on another thread before sampling
+  // begins.
   DWORD prevSuspendCount = SuspendThread(winProfiler->currentThread_);
   if (prevSuspendCount == static_cast<DWORD>(-1)) {
     return true;
@@ -200,11 +209,19 @@ std::unique_ptr<SamplingProfiler> SamplingProfiler::create(Runtime &rt) {
   return std::make_unique<sampling_profiler::SamplingProfilerWindows>(rt);
 }
 
-bool SamplingProfiler::belongsToCurrentThread() const {
-  return GetThreadId(
-             static_cast<const sampling_profiler::SamplingProfilerWindows *>(
-                 this)
-                 ->currentThread_) == GetCurrentThreadId();
+bool SamplingProfiler::belongsToCurrentThread() {
+  auto profiler =
+      static_cast<sampling_profiler::SamplingProfilerWindows *>(this);
+  std::lock_guard<std::mutex> lock(profiler->runtimeDataLock_);
+  return GetThreadId(profiler->currentThread_) == GetCurrentThreadId();
+}
+
+void SamplingProfiler::setRuntimeThread() {
+  auto profiler =
+      static_cast<sampling_profiler::SamplingProfilerWindows *>(this);
+  std::lock_guard<std::mutex> lock(profiler->runtimeDataLock_);
+  CloseHandle(profiler->currentThread_);
+  profiler->currentThread_ = openCurrentThread();
 }
 
 } // namespace vm

unittests/VMRuntime/SamplingProfilerTest.cpp

@@ -9,6 +9,7 @@
 
 #if HERMESVM_SAMPLING_PROFILER_AVAILABLE
 
+#include "TestHelpers1.h"
 #include "hermes/VM/Runtime.h"
 
 #include <gtest/gtest.h>
@@ -63,6 +64,34 @@ TEST(SamplingProfilerTest, MultipleProfilers) {
 }
 #endif
 
+TEST(SamplingProfilerTest, RegisterDifferentThread) {
+  constexpr uint32_t kThreadCount = 3;
+
+  auto rt = makeRuntime(withSamplingProfilerEnabled);
+  rt->samplingProfiler->enable();
+
+  for (uint32_t threadNumber = 0; threadNumber < kThreadCount; ++threadNumber) {
+    std::thread([&]() {
+      rt->samplingProfiler->setRuntimeThread();
+      EXPECT_TRUE(rt->samplingProfiler->belongsToCurrentThread());
+    }).join();
+  }
+
+  rt->samplingProfiler->disable();
+}
+
+TEST(SamplingProfilerTest, RegisterIdenticalThread) {
+  auto rt = makeRuntime(withSamplingProfilerEnabled);
+  rt->samplingProfiler->enable();
+
+  rt->samplingProfiler->setRuntimeThread();
+  EXPECT_TRUE(rt->samplingProfiler->belongsToCurrentThread());
+  rt->samplingProfiler->setRuntimeThread();
+  EXPECT_TRUE(rt->samplingProfiler->belongsToCurrentThread());
+
+  rt->samplingProfiler->disable();
+}
+
 } // namespace
 
 #endif // HERMESVM_SAMPLING_PROFILER_AVAILABLE