Bug 1774773 - Added atomic locking around marker schema functions - r=florian

To avoid real or apparent races to the list of marker schema functions, some
bits of the atomic count are used to effectively implement a RWLock, so that
any number of threads can add their own new marker schema concurrently, while
the profile-capture reading of the whole list can only be done on its own.

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

mozglue/baseprofiler/core/ProfilerMarkers.cpp

@@ -16,14 +16,28 @@ namespace base_profiler_markers_detail {
 // work with too-small types.)
 using DeserializerTagAtomic = unsigned;
 
+// The atomic sDeserializerCount still also include bits that act as a "RWLock":
+// Whoever can set this bit gets exclusive access to the count and the whole
+// sMarkerTypeFunctions1Based array, guaranteeing that it cannot be modified.
+static constexpr DeserializerTagAtomic scExclusiveLock = 0x80'00'00'00u;
+// Code that wants shared access can add this value, then ensure there is no
+// exclusive lock, after which it's guaranteed that no exclusive lock can be
+// taken until the shared lock count goes back to zero.
+static constexpr DeserializerTagAtomic scSharedLockUnit = 0x00'01'00'00u;
+// This mask isolates the actual count value from the lock bits.
+static constexpr DeserializerTagAtomic scTagMask = 0x00'00'FF'FFu;
+
 // Number of currently-registered deserializers and other marker type functions.
-static Atomic<DeserializerTagAtomic, MemoryOrdering::Relaxed>
+// The high bits contain lock bits, see above.
+static Atomic<DeserializerTagAtomic, MemoryOrdering::ReleaseAcquire>
     sDeserializerCount{0};
 
 // This needs to be big enough to handle all possible marker types. If one day
 // this needs to be higher, the underlying DeserializerTag type will have to be
 // changed.
 static constexpr DeserializerTagAtomic DeserializerMax = 250;
+static_assert(DeserializerMax <= scTagMask,
+              "DeserializerMax doesn't fit in scTagMask");
 
 static_assert(
     DeserializerMax <= std::numeric_limits<Streaming::DeserializerTag>::max(),
@@ -49,7 +63,28 @@ static Streaming::MarkerTypeFunctions
   MOZ_RELEASE_ASSERT(!!aMarkerTypeNameFunction);
   MOZ_RELEASE_ASSERT(!!aMarkerSchemaFunction);
 
-  DeserializerTagAtomic tag = ++sDeserializerCount;
+  // Add a shared lock request, which will prevent future exclusive locking.
+  DeserializerTagAtomic tagWithLock = (sDeserializerCount += scSharedLockUnit);
+
+  // An exclusive locker may have arrived before us, just wait for it to finish.
+  while ((tagWithLock & scExclusiveLock) != 0u) {
+    tagWithLock = sDeserializerCount;
+  }
+
+  MOZ_ASSERT(
+      // This is equivalent to shifting right to only keep the lock counts.
+      tagWithLock / scSharedLockUnit <
+          // This is effectively half of the permissible shared lock range,
+          // that would mean way too many threads doing this work here!
+          scExclusiveLock / scSharedLockUnit / 2,
+      "The shared lock count is getting unexpectedly high, verify the "
+      "algorithm, and tweak constants if needed");
+
+  // Reserve a tag. Even if there are multiple shared-lock holders here, each
+  // one will get a different value, and therefore will access a different part
+  // of the sMarkerTypeFunctions1Based array.
+  const DeserializerTagAtomic tag = ++sDeserializerCount & scTagMask;
+
   MOZ_RELEASE_ASSERT(
       tag <= DeserializerMax,
       "Too many deserializers, consider increasing DeserializerMax. "
@@ -57,6 +92,9 @@ static Streaming::MarkerTypeFunctions
   sMarkerTypeFunctions1Based[tag - 1] = {aDeserializer, aMarkerTypeNameFunction,
                                          aMarkerSchemaFunction};
 
+  // And release our shared lock, to allow exclusive readers.
+  sDeserializerCount -= scSharedLockUnit;
+
   return static_cast<DeserializerTag>(tag);
 }
 
@@ -69,9 +107,38 @@ static Streaming::MarkerTypeFunctions
   return sMarkerTypeFunctions1Based[aTag - 1].mMarkerDataDeserializer;
 }
 
-/* static */ Span<const Streaming::MarkerTypeFunctions>
-Streaming::MarkerTypeFunctionsArray() {
-  return {sMarkerTypeFunctions1Based, sDeserializerCount};
+Streaming::LockedMarkerTypeFunctionsList::LockedMarkerTypeFunctionsList() {
+  for (;;) {
+    const DeserializerTagAtomic count = sDeserializerCount;
+    if ((count & scTagMask) != count) {
+      // Someone already has a lock, loop around.
+      continue;
+    }
+
+    // There are currently no locks, try to add our exclusive lock.
+    if (!sDeserializerCount.compareExchange(count, count | scExclusiveLock)) {
+      // Someone else modified sDeserializerCount since our read, loop around.
+      continue;
+    }
+
+    // We applied our exclusive lock, we can now read the list of functions,
+    // without interference until ~LockedMarkerTypeFunctionsList().
+    // (Note that sDeserializerCount may receive shared lock requests, but the
+    // count won't change.)
+    mMarkerTypeFunctionsSpan = {sMarkerTypeFunctions1Based, count};
+    break;
+  }
+}
+
+Streaming::LockedMarkerTypeFunctionsList::~LockedMarkerTypeFunctionsList() {
+  MOZ_ASSERT(
+      (sDeserializerCount & scExclusiveLock) == scExclusiveLock,
+      "sDeserializerCount should still have the the exclusive lock bit set");
+  MOZ_ASSERT(
+      (sDeserializerCount & scTagMask) ==
+          DeserializerTagAtomic(mMarkerTypeFunctionsSpan.size()),
+      "sDeserializerCount should have the same count since construction");
+  sDeserializerCount &= ~scExclusiveLock;
 }
 
 // Only accessed on the main thread.

mozglue/baseprofiler/core/platform.cpp

@@ -1816,9 +1816,8 @@ static void StreamCategories(SpliceableJSONWriter& aWriter) {
 
 static void StreamMarkerSchema(SpliceableJSONWriter& aWriter) {
   // Get an array view with all registered marker-type-specific functions.
-  Span<const base_profiler_markers_detail::Streaming::MarkerTypeFunctions>
-      markerTypeFunctionsArray =
-          base_profiler_markers_detail::Streaming::MarkerTypeFunctionsArray();
+  base_profiler_markers_detail::Streaming::LockedMarkerTypeFunctionsList
+      markerTypeFunctionsArray;
   // List of streamed marker names, this is used to spot duplicates.
   std::set<std::string> names;
   // Stream the display schema for each different one. (Duplications may come

mozglue/baseprofiler/public/BaseProfilerMarkersDetail.h

@@ -72,7 +72,23 @@ struct Streaming {
       DeserializerTag aTag);
 
   // Retrieve all MarkerTypeFunctions's.
-  MFBT_API static Span<const MarkerTypeFunctions> MarkerTypeFunctionsArray();
+  // While this object lives, no other operations can happen on this list.
+  class LockedMarkerTypeFunctionsList {
+   public:
+    MFBT_API LockedMarkerTypeFunctionsList();
+    MFBT_API ~LockedMarkerTypeFunctionsList();
+
+    LockedMarkerTypeFunctionsList(const LockedMarkerTypeFunctionsList&) =
+        delete;
+    LockedMarkerTypeFunctionsList& operator=(
+        const LockedMarkerTypeFunctionsList&) = delete;
+
+    auto begin() const { return mMarkerTypeFunctionsSpan.begin(); }
+    auto end() const { return mMarkerTypeFunctionsSpan.end(); }
+
+   private:
+    Span<const MarkerTypeFunctions> mMarkerTypeFunctionsSpan;
+  };
 };
 
 // This helper will examine a marker type's `StreamJSONMarkerData` function, see

tools/profiler/core/platform.cpp

@@ -2652,9 +2652,8 @@ static void StreamCategories(SpliceableJSONWriter& aWriter) {
 
 static void StreamMarkerSchema(SpliceableJSONWriter& aWriter) {
   // Get an array view with all registered marker-type-specific functions.
-  Span<const base_profiler_markers_detail::Streaming::MarkerTypeFunctions>
-      markerTypeFunctionsArray =
-          base_profiler_markers_detail::Streaming::MarkerTypeFunctionsArray();
+  base_profiler_markers_detail::Streaming::LockedMarkerTypeFunctionsList
+      markerTypeFunctionsArray;
   // List of streamed marker names, this is used to spot duplicates.
   std::set<std::string> names;
   // Stream the display schema for each different one. (Duplications may come