Bug 1716274 - profiler_add/remove_state_change_callback - r=canaltinova

Callbacks can be registered to be called at specific profiler state changes.
This may be useful to make sure some markers are recorded at the end of the profile time range, if some information would be useful to always have available (and it doesn't fit in meta-information or elsewhere).

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

tools/profiler/core/platform.cpp

@@ -1463,6 +1463,94 @@ uint32_t ActivePS::sNextGeneration = 0;
 // The mutex that guards accesses to CorePS and ActivePS.
 static PSMutex gPSMutex;
 
+static PSMutex gProfilerStateChangeMutex;
+
+struct IdentifiedProfilingStateChangeCallback {
+  ProfilingStateSet mProfilingStateSet;
+  ProfilingStateChangeCallback mProfilingStateChangeCallback;
+  uintptr_t mUniqueIdentifier;
+
+  explicit IdentifiedProfilingStateChangeCallback(
+      ProfilingStateSet aProfilingStateSet,
+      ProfilingStateChangeCallback&& aProfilingStateChangeCallback,
+      uintptr_t aUniqueIdentifier)
+      : mProfilingStateSet(aProfilingStateSet),
+        mProfilingStateChangeCallback(aProfilingStateChangeCallback),
+        mUniqueIdentifier(aUniqueIdentifier) {}
+};
+using IdentifiedProfilingStateChangeCallbackUPtr =
+    UniquePtr<IdentifiedProfilingStateChangeCallback>;
+
+static Vector<IdentifiedProfilingStateChangeCallbackUPtr>
+    mIdentifiedProfilingStateChangeCallbacks;
+
+void profiler_add_state_change_callback(
+    ProfilingStateSet aProfilingStateSet,
+    ProfilingStateChangeCallback&& aCallback,
+    uintptr_t aUniqueIdentifier /* = 0 */) {
+  gPSMutex.AssertCurrentThreadDoesNotOwn();
+  PSAutoLock lock(gProfilerStateChangeMutex);
+
+#ifdef DEBUG
+  // Check if a non-zero id is not already used. Bug forgive it in non-DEBUG
+  // builds; in the worst case they may get removed too early.
+  if (aUniqueIdentifier != 0) {
+    for (const IdentifiedProfilingStateChangeCallbackUPtr& idedCallback :
+         mIdentifiedProfilingStateChangeCallbacks) {
+      MOZ_ASSERT(idedCallback->mUniqueIdentifier != aUniqueIdentifier);
+    }
+  }
+#endif  // DEBUG
+
+  if (aProfilingStateSet.contains(ProfilingState::AlreadyActive) &&
+      profiler_is_active()) {
+    aCallback(ProfilingState::AlreadyActive);
+  }
+
+  (void)mIdentifiedProfilingStateChangeCallbacks.append(
+      MakeUnique<IdentifiedProfilingStateChangeCallback>(
+          aProfilingStateSet, std::move(aCallback), aUniqueIdentifier));
+}
+
+// Remove the callback with the given identifier.
+void profiler_remove_state_change_callback(uintptr_t aUniqueIdentifier) {
+  MOZ_ASSERT(aUniqueIdentifier != 0);
+  if (aUniqueIdentifier == 0) {
+    // Forgive zero in non-DEBUG builds.
+    return;
+  }
+
+  gPSMutex.AssertCurrentThreadDoesNotOwn();
+  PSAutoLock lock(gProfilerStateChangeMutex);
+
+  mIdentifiedProfilingStateChangeCallbacks.eraseIf(
+      [aUniqueIdentifier](
+          const IdentifiedProfilingStateChangeCallbackUPtr& aIdedCallback) {
+        if (aIdedCallback->mUniqueIdentifier != aUniqueIdentifier) {
+          return false;
+        }
+        if (aIdedCallback->mProfilingStateSet.contains(
+                ProfilingState::RemovingCallback)) {
+          aIdedCallback->mProfilingStateChangeCallback(
+              ProfilingState::RemovingCallback);
+        }
+        return true;
+      });
+}
+
+static void invoke_profiler_state_change_callbacks(
+    ProfilingState aProfilingState) {
+  gPSMutex.AssertCurrentThreadDoesNotOwn();
+  PSAutoLock lock(gProfilerStateChangeMutex);
+
+  for (const IdentifiedProfilingStateChangeCallbackUPtr& idedCallback :
+       mIdentifiedProfilingStateChangeCallbacks) {
+    if (idedCallback->mProfilingStateSet.contains(aProfilingState)) {
+      idedCallback->mProfilingStateChangeCallback(aProfilingState);
+    }
+  }
+}
+
 Atomic<uint32_t, MemoryOrdering::Relaxed> RacyFeatures::sActiveAndFeatures(0);
 
 // Each live thread has a RegisteredThread, and we store a reference to it in
@@ -3122,6 +3210,10 @@ bool profiler_stream_json_for_this_process(
 
   const auto preRecordedMetaInformation = PreRecordMetaInformation();
 
+  if (profiler_is_active()) {
+    invoke_profiler_state_change_callbacks(ProfilingState::GeneratingProfile);
+  }
+
   PSAutoLock lock(gPSMutex);
 
   if (!ActivePS::Exists(lock)) {
@@ -4475,6 +4567,8 @@ void profiler_init(void* aStackTop) {
   ActivePS::SetMemoryCounter(mozilla::profiler::install_memory_hooks());
 #endif
 
+  invoke_profiler_state_change_callbacks(ProfilingState::Started);
+
   // We do this with gPSMutex unlocked. The comment in profiler_stop() explains
   // why.
   NotifyProfilerStarted(capacity, duration, interval, features, filters.begin(),
@@ -4496,6 +4590,11 @@ void profiler_shutdown(IsFastShutdown aIsFastShutdown) {
   MOZ_RELEASE_ASSERT(NS_IsMainThread());
   MOZ_RELEASE_ASSERT(CorePS::Exists());
 
+  if (profiler_is_active()) {
+    invoke_profiler_state_change_callbacks(ProfilingState::Stopping);
+  }
+  invoke_profiler_state_change_callbacks(ProfilingState::ShuttingDown);
+
   const auto preRecordedMetaInformation = PreRecordMetaInformation();
 
   ProfilerParent::ProfilerWillStopIfStarted();
@@ -5028,6 +5127,10 @@ void profiler_start(PowerOfTwo32 aCapacity, double aInterval,
 
     // Reset the current state if the profiler is running.
     if (ActivePS::Exists(lock)) {
+      // Note: Not invoking callbacks with ProfilingState::Stopping, because
+      // we're under lock, and also it would not be useful: Any profiling data
+      // will be discarded, and we're immediately restarting the profiler below
+      // and then notifying ProfilingState::Started.
       samplerThread = locked_profiler_stop(lock);
     }
 
@@ -5041,6 +5144,8 @@ void profiler_start(PowerOfTwo32 aCapacity, double aInterval,
   ActivePS::SetMemoryCounter(mozilla::profiler::install_memory_hooks());
 #endif
 
+  invoke_profiler_state_change_callbacks(ProfilingState::Started);
+
   // We do these operations with gPSMutex unlocked. The comments in
   // profiler_stop() explain why.
   if (samplerThread) {
@@ -5075,6 +5180,10 @@ void profiler_ensure_started(PowerOfTwo32 aCapacity, double aInterval,
       if (!ActivePS::Equals(lock, aCapacity, aDuration, aInterval, aFeatures,
                             aFilters, aFilterCount, aActiveTabID)) {
         // Stop and restart with different settings.
+        // Note: Not invoking callbacks with ProfilingState::Stopping, because
+        // we're under lock, and also it would not be useful: Any profiling data
+        // will be discarded, and we're immediately restarting the profiler
+        // below and then notifying ProfilingState::Started.
         samplerThread = locked_profiler_stop(lock);
         locked_profiler_start(lock, aCapacity, aInterval, aFeatures, aFilters,
                               aFilterCount, aActiveTabID, aDuration);
@@ -5097,6 +5206,8 @@ void profiler_ensure_started(PowerOfTwo32 aCapacity, double aInterval,
   }
 
   if (startedProfiler) {
+    invoke_profiler_state_change_callbacks(ProfilingState::Started);
+
     NotifyProfilerStarted(aCapacity, aDuration, aInterval, aFeatures, aFilters,
                           aFilterCount, aActiveTabID);
   }
@@ -5166,6 +5277,10 @@ void profiler_stop() {
 
   MOZ_RELEASE_ASSERT(CorePS::Exists());
 
+  if (profiler_is_active()) {
+    invoke_profiler_state_change_callbacks(ProfilingState::Stopping);
+  }
+
   ProfilerParent::ProfilerWillStopIfStarted();
 
 #if defined(MOZ_REPLACE_MALLOC) && defined(MOZ_PROFILER_MEMORY)
@@ -5233,6 +5348,8 @@ void profiler_pause() {
 
   MOZ_RELEASE_ASSERT(CorePS::Exists());
 
+  invoke_profiler_state_change_callbacks(ProfilingState::Pausing);
+
   {
     PSAutoLock lock(gPSMutex);
 
@@ -5287,6 +5404,8 @@ void profiler_resume() {
   // gPSMutex must be unlocked when we notify, to avoid potential deadlocks.
   ProfilerParent::ProfilerResumed();
   NotifyObservers("profiler-resumed");
+
+  invoke_profiler_state_change_callbacks(ProfilingState::Resumed);
 }
 
 bool profiler_is_sampling_paused() {

tools/profiler/public/ProfilerState.h

@@ -13,6 +13,11 @@
 #ifndef ProfilerState_h
 #define ProfilerState_h
 
+#include <mozilla/DefineEnum.h>
+#include <mozilla/EnumSet.h>
+
+#include <functional>
+
 //---------------------------------------------------------------------------
 // Profiler features
 //---------------------------------------------------------------------------
@@ -100,6 +105,61 @@ struct ProfilerFeature {
 #undef DECLARE
 };
 
+// clang-format off
+MOZ_DEFINE_ENUM_CLASS(ProfilingState,(
+  // A callback will be invoked ...
+  AlreadyActive,     // if the profiler is active when the callback is added.
+  RemovingCallback,  // when the callback is removed.
+  Started,           // after the profiler has started.
+  Pausing,           // before the profiler is paused.
+  Resumed,           // after the profiler has resumed.
+  GeneratingProfile, // before a profile is created.
+  Stopping,          // before the profiler stops (unless restarting afterward).
+  ShuttingDown       // before the profiler is shut down.
+));
+// clang-format on
+
+inline static const char* ProfilingStateToString(
+    ProfilingState aProfilingState) {
+  switch (aProfilingState) {
+    case ProfilingState::AlreadyActive:
+      return "Profiler already active";
+    case ProfilingState::RemovingCallback:
+      return "Callback being removed";
+    case ProfilingState::Started:
+      return "Profiler started";
+    case ProfilingState::Pausing:
+      return "Profiler pausing";
+    case ProfilingState::Resumed:
+      return "Profiler resumed";
+    case ProfilingState::GeneratingProfile:
+      return "Generating profile";
+    case ProfilingState::Stopping:
+      return "Profiler stopping";
+    case ProfilingState::ShuttingDown:
+      return "Profiler shutting down";
+    default:
+      MOZ_ASSERT_UNREACHABLE("Unexpected ProfilingState enum value");
+      return "?";
+  }
+}
+
+using ProfilingStateSet = mozilla::EnumSet<ProfilingState>;
+
+constexpr ProfilingStateSet AllProfilingStates() {
+  ProfilingStateSet set;
+  using Value = std::underlying_type_t<ProfilingState>;
+  for (Value stateValue = 0;
+       stateValue <= static_cast<Value>(kHighestProfilingState); ++stateValue) {
+    set += static_cast<ProfilingState>(stateValue);
+  }
+  return set;
+}
+
+// Type of callbacks to be invoked at certain state changes.
+// It must NOT call profiler_add/remove_state_change_callback().
+using ProfilingStateChangeCallback = std::function<void(ProfilingState)>;
+
 #ifndef MOZ_GECKO_PROFILER
 
 inline bool profiler_is_active() { return false; }
@@ -108,6 +168,12 @@ inline bool profiler_thread_is_being_profiled() { return false; }
 inline bool profiler_is_active_and_thread_is_registered() { return false; }
 inline bool profiler_feature_active(uint32_t aFeature) { return false; }
 inline bool profiler_is_locked_on_current_thread() { return false; }
+inline void profiler_add_state_change_callback(
+    ProfilingStateSet aProfilingStateSet,
+    ProfilingStateChangeCallback&& aCallback, uintptr_t aUniqueIdentifier = 0) {
+}
+inline void profiler_remove_state_change_callback(uintptr_t aUniqueIdentifier) {
+}
 
 #else  // !MOZ_GECKO_PROFILER
 
@@ -318,6 +384,17 @@ inline bool profiler_is_main_thread() {
 // deadlock.
 bool profiler_is_locked_on_current_thread();
 
+// Install a callback to be invoked at any of the given profiling state changes.
+// An optional non-zero identifier may be given, to allow later removal of the
+// callback, the caller is responsible for making sure it's really unique (e.g.,
+// by using a pointer to an object it owns.)
+void profiler_add_state_change_callback(
+    ProfilingStateSet aProfilingStateSet,
+    ProfilingStateChangeCallback&& aCallback, uintptr_t aUniqueIdentifier = 0);
+
+// Remove the callback with the given non-zero identifier.
+void profiler_remove_state_change_callback(uintptr_t aUniqueIdentifier);
+
 #endif  // MOZ_GECKO_PROFILER
 
 #endif  // ProfilerState_h

tools/profiler/tests/gtest/GeckoProfiler.cpp

@@ -21,8 +21,10 @@
 #include "json/json.h"
 #include "mozilla/Atomics.h"
 #include "mozilla/BlocksRingBuffer.h"
+#include "mozilla/DataMutex.h"
 #include "mozilla/ProfileBufferEntrySerializationGeckoExtensions.h"
 #include "mozilla/ProfileJSONWriter.h"
+#include "mozilla/ScopeExit.h"
 #include "mozilla/UniquePtrExtensions.h"
 #include "mozilla/net/HttpBaseChannel.h"
 #include "nsIChannelEventSink.h"
@@ -2284,6 +2286,106 @@ TEST(GeckoProfiler, PostSamplingCallback)
       [&](SamplingState) { ASSERT_TRUE(false); }));
 }
 
+TEST(GeckoProfiler, ProfilingStateCallback)
+{
+  const char* filters[] = {"GeckoMain"};
+
+  ASSERT_TRUE(!profiler_is_active());
+
+  struct ProfilingStateAndId {
+    ProfilingState mProfilingState;
+    int mId;
+  };
+  DataMutex<Vector<ProfilingStateAndId>> states{"Profiling states"};
+  auto CreateCallback = [&states](int id) {
+    return [id, &states](ProfilingState aProfilingState) {
+      auto lockedStates = states.Lock();
+      ASSERT_TRUE(
+          lockedStates->append(ProfilingStateAndId{aProfilingState, id}));
+    };
+  };
+  auto CheckStatesIsEmpty = [&states]() {
+    auto lockedStates = states.Lock();
+    EXPECT_TRUE(lockedStates->empty());
+  };
+  auto CheckStatesOnlyContains = [&states](ProfilingState aProfilingState,
+                                           int aId) {
+    auto lockedStates = states.Lock();
+    EXPECT_EQ(lockedStates->length(), 1u);
+    if (lockedStates->length() >= 1u) {
+      EXPECT_EQ((*lockedStates)[0].mProfilingState, aProfilingState);
+      EXPECT_EQ((*lockedStates)[0].mId, aId);
+    }
+    lockedStates->clear();
+  };
+
+  profiler_add_state_change_callback(AllProfilingStates(), CreateCallback(1),
+                                     1);
+  // This is in case of error, and it also exercises the (allowed) removal of
+  // unknown callback ids.
+  auto cleanup1 = mozilla::MakeScopeExit(
+      []() { profiler_remove_state_change_callback(1); });
+  CheckStatesIsEmpty();
+
+  profiler_start(PROFILER_DEFAULT_ENTRIES, PROFILER_DEFAULT_INTERVAL,
+                 ProfilerFeature::StackWalk, filters, MOZ_ARRAY_LENGTH(filters),
+                 0);
+
+  CheckStatesOnlyContains(ProfilingState::Started, 1);
+
+  profiler_add_state_change_callback(AllProfilingStates(), CreateCallback(2),
+                                     2);
+  // This is in case of error, and it also exercises the (allowed) removal of
+  // unknown callback ids.
+  auto cleanup2 = mozilla::MakeScopeExit(
+      []() { profiler_remove_state_change_callback(2); });
+  CheckStatesOnlyContains(ProfilingState::AlreadyActive, 2);
+
+  profiler_remove_state_change_callback(2);
+  CheckStatesOnlyContains(ProfilingState::RemovingCallback, 2);
+  // Note: The actual removal is effectively tested below, by not seeing any
+  // more invocations of the 2nd callback.
+
+  ASSERT_EQ(WaitForSamplingState(), SamplingState::SamplingCompleted);
+  UniquePtr<char[]> profileCompleted = profiler_get_profile();
+  CheckStatesOnlyContains(ProfilingState::GeneratingProfile, 1);
+  JSONOutputCheck(profileCompleted.get(), [](const Json::Value& aRoot) {});
+
+  profiler_pause();
+  CheckStatesOnlyContains(ProfilingState::Pausing, 1);
+  UniquePtr<char[]> profilePaused = profiler_get_profile();
+  CheckStatesOnlyContains(ProfilingState::GeneratingProfile, 1);
+  JSONOutputCheck(profilePaused.get(), [](const Json::Value& aRoot) {});
+
+  profiler_resume();
+  CheckStatesOnlyContains(ProfilingState::Resumed, 1);
+  ASSERT_EQ(WaitForSamplingState(), SamplingState::SamplingCompleted);
+  UniquePtr<char[]> profileResumed = profiler_get_profile();
+  CheckStatesOnlyContains(ProfilingState::GeneratingProfile, 1);
+  JSONOutputCheck(profileResumed.get(), [](const Json::Value& aRoot) {});
+
+  // This effectively stops the profiler before restarting it, but
+  // ProfilingState::Stopping is not notified. See `profiler_start` for details.
+  profiler_start(PROFILER_DEFAULT_ENTRIES, PROFILER_DEFAULT_INTERVAL,
+                 ProfilerFeature::StackWalk | ProfilerFeature::NoStackSampling,
+                 filters, MOZ_ARRAY_LENGTH(filters), 0);
+  CheckStatesOnlyContains(ProfilingState::Started, 1);
+  ASSERT_EQ(WaitForSamplingState(), SamplingState::NoStackSamplingCompleted);
+  UniquePtr<char[]> profileNoStacks = profiler_get_profile();
+  CheckStatesOnlyContains(ProfilingState::GeneratingProfile, 1);
+  JSONOutputCheck(profileNoStacks.get(), [](const Json::Value& aRoot) {});
+
+  profiler_stop();
+  CheckStatesOnlyContains(ProfilingState::Stopping, 1);
+  ASSERT_TRUE(!profiler_is_active());
+
+  profiler_remove_state_change_callback(1);
+  CheckStatesOnlyContains(ProfilingState::RemovingCallback, 1);
+
+  // Note: ProfilingState::ShuttingDown cannot be tested here, and the profiler
+  // can only be shut down once per process.
+}
+
 TEST(GeckoProfiler, BaseProfilerHandOff)
 {
   const char* filters[] = {"GeckoMain"};