Bug 1577658 - DeserializeAfterKindAndStream takes a lambda: (tid) -> writer - r=canaltinova

Previously, DeserializeAfterKindAndStream would take a JSON writer and a thread id, using the thread id (if specified) to only output markers from that thread to the given writer.

Now, DeserializeAfterKindAndStream takes a lambda, which will be called with the thread id found in the marker, that lambda can either return null (nothing to output) or a pointer to a writer, in which case the marker is read and output to the given writer.
This makes DeserializeAfterKindAndStream more flexible, and will allow handling markers from different threads, each possibly being output to different writers.

Also, for simplicity the entry is now always fully read, so there is no need for the caller to do anything. The return bool is therefore unnecessary, and has been removed.

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

mozglue/baseprofiler/core/ProfileBufferEntry.cpp

@@ -807,27 +807,24 @@ void ProfileBuffer::StreamMarkersToJSON(SpliceableJSONWriter& aWriter,
     MOZ_ASSERT(static_cast<ProfileBufferEntry::KindUnderlyingType>(type) <
                static_cast<ProfileBufferEntry::KindUnderlyingType>(
                    ProfileBufferEntry::Kind::MODERN_LIMIT));
-    bool entryWasFullyRead = false;
-
     if (type == ProfileBufferEntry::Kind::Marker) {
-      entryWasFullyRead = ::mozilla::base_profiler_markers_detail::
-          DeserializeAfterKindAndStream(
-              aER, aWriter, aThreadId,
-              [&](ProfileChunkedBuffer& aChunkedBuffer) {
-                ProfilerBacktrace backtrace("", &aChunkedBuffer);
-                backtrace.StreamJSON(aWriter, TimeStamp::ProcessCreation(),
-                                     aUniqueStacks);
-              },
-              // We don't have Rust markers in the mozglue.
-              [&](mozilla::base_profiler_markers_detail::Streaming::
-                      DeserializerTag) {
-                MOZ_ASSERT_UNREACHABLE("No Rust markers in mozglue.");
-              });
-    }
-
-    if (!entryWasFullyRead) {
-      // Not a marker, or marker for another thread.
-      // We probably didn't read the whole entry, so we need to skip to the end.
+      ::mozilla::base_profiler_markers_detail::DeserializeAfterKindAndStream(
+          aER,
+          [&](const BaseProfilerThreadId& aMarkerThreadId) {
+            return (aMarkerThreadId == aThreadId) ? &aWriter : nullptr;
+          },
+          [&](ProfileChunkedBuffer& aChunkedBuffer) {
+            ProfilerBacktrace backtrace("", &aChunkedBuffer);
+            backtrace.StreamJSON(aWriter, TimeStamp::ProcessCreation(),
+                                 aUniqueStacks);
+          },
+          // We don't have Rust markers in the mozglue.
+          [&](mozilla::base_profiler_markers_detail::Streaming::
+                  DeserializerTag) {
+            MOZ_ASSERT_UNREACHABLE("No Rust markers in mozglue.");
+          });
+    } else {
+      // The entry was not a marker, we need to skip to the end.
       aER.SetRemainingBytes(0);
     }
   });

mozglue/baseprofiler/public/BaseProfilerMarkersDetail.h

@@ -295,48 +295,61 @@ ProfileBufferBlockIndex AddMarkerToBuffer(
       aBuffer, aName, aCategory, std::move(aOptions), aTs...);
 }
 
-template <typename StackCallback, typename RustMarkerCallback>
-[[nodiscard]] bool DeserializeAfterKindAndStream(
+// Assuming aEntryReader points right after the entry type (being Marker), this
+// reads the remainder of the marker and outputs it.
+// - GetWriterForThreadCallback, called first, after the thread id is read:
+//     (ThreadId) -> SpliceableJSONWriter* or null
+//   If null, nothing will be output, but aEntryReader will still be read fully.
+// - StackCallback, only called if GetWriterForThreadCallback didn't return
+//   null, and if the marker contains a stack:
+//     (ProfileChunkedBuffer&) -> void
+// - RustMarkerCallback, only called if GetWriterForThreadCallback didn't return
+//   null, and if the marker contains a Rust payload:
+//     (DeserializerTag) -> void
+template <typename GetWriterForThreadCallback, typename StackCallback,
+          typename RustMarkerCallback>
+void DeserializeAfterKindAndStream(
     ProfileBufferEntryReader& aEntryReader,
-    baseprofiler::SpliceableJSONWriter& aWriter,
-    baseprofiler::BaseProfilerThreadId aThreadIdOrUnspecified,
+    GetWriterForThreadCallback&& aGetWriterForThreadCallback,
     StackCallback&& aStackCallback, RustMarkerCallback&& aRustMarkerCallback) {
   // Each entry is made up of the following:
   //   ProfileBufferEntry::Kind::Marker, <- already read by caller
   //   options,                          <- next location in entries
   //   name,
-  //   category,
-  //   deserializer tag,
-  //   payload type (cpp or rust)
   //   payload
   const MarkerOptions options = aEntryReader.ReadObject<MarkerOptions>();
-  if (aThreadIdOrUnspecified.IsSpecified() &&
-      options.ThreadId().ThreadId() != aThreadIdOrUnspecified) {
-    // A specific thread is being read, we're not in it.
-    return false;
+
+  baseprofiler::SpliceableJSONWriter* writer =
+      std::forward<GetWriterForThreadCallback>(aGetWriterForThreadCallback)(
+          options.ThreadId().ThreadId());
+  if (!writer) {
+    // No writer associated with this thread id, drop it.
+    aEntryReader.SetRemainingBytes(0);
+    return;
   }
+
   // Write the information to JSON with the following schema:
   // [name, startTime, endTime, phase, category, data]
-  aWriter.StartArrayElement();
+  writer->StartArrayElement();
   {
-    aWriter.UniqueStringElement(aEntryReader.ReadObject<ProfilerString8View>());
+    writer->UniqueStringElement(aEntryReader.ReadObject<ProfilerString8View>());
 
     const double startTime = options.Timing().GetStartTime();
-    aWriter.TimeDoubleMsElement(startTime);
+    writer->TimeDoubleMsElement(startTime);
 
     const double endTime = options.Timing().GetEndTime();
-    aWriter.TimeDoubleMsElement(endTime);
+    writer->TimeDoubleMsElement(endTime);
 
-    aWriter.IntElement(static_cast<int64_t>(options.Timing().MarkerPhase()));
+    writer->IntElement(static_cast<int64_t>(options.Timing().MarkerPhase()));
 
     MarkerCategory category = aEntryReader.ReadObject<MarkerCategory>();
-    aWriter.IntElement(static_cast<int64_t>(category.GetCategory()));
+    writer->IntElement(static_cast<int64_t>(category.GetCategory()));
 
     if (const auto tag =
             aEntryReader.ReadObject<mozilla::base_profiler_markers_detail::
                                         Streaming::DeserializerTag>();
         tag != 0) {
-      aWriter.StartObjectElement(JSONWriter::SingleLineStyle);
+      writer->StartObjectElement(JSONWriter::SingleLineStyle);
       {
         // Stream "common props".
 
@@ -347,7 +360,7 @@ template <typename StackCallback, typename RustMarkerCallback>
           // `nsContentUtils::GenerateProcessSpecificId`, which is specifically
           // designed to only use 53 of the 64 bits to be lossless when passed
           // into and out of JS as a double.
-          aWriter.DoubleProperty(
+          writer->DoubleProperty(
               "innerWindowID",
               static_cast<double>(options.InnerWindowId().Id()));
         }
@@ -356,9 +369,9 @@ template <typename StackCallback, typename RustMarkerCallback>
         if (ProfileChunkedBuffer* chunkedBuffer =
                 options.Stack().GetChunkedBuffer();
             chunkedBuffer) {
-          aWriter.StartObjectProperty("stack");
+          writer->StartObjectProperty("stack");
           { std::forward<StackCallback>(aStackCallback)(*chunkedBuffer); }
-          aWriter.EndObject();
+          writer->EndObject();
         }
 
         auto payloadType = static_cast<mozilla::MarkerPayloadType>(
@@ -372,24 +385,25 @@ template <typename StackCallback, typename RustMarkerCallback>
                 MarkerDataDeserializer deserializer =
                     mozilla::base_profiler_markers_detail::Streaming::
                         DeserializerForTag(tag);
-
             MOZ_RELEASE_ASSERT(deserializer);
-            deserializer(aEntryReader, aWriter);
+            deserializer(aEntryReader, *writer);
+            MOZ_ASSERT(aEntryReader.RemainingBytes() == 0u);
             break;
           }
           case mozilla::MarkerPayloadType::Rust:
             std::forward<RustMarkerCallback>(aRustMarkerCallback)(tag);
+            MOZ_ASSERT(aEntryReader.RemainingBytes() == 0u);
             break;
           default:
             MOZ_ASSERT_UNREACHABLE("Unknown payload type.");
             break;
         }
       }
-      aWriter.EndObject();
+      writer->EndObject();
     }
   }
-  aWriter.EndArray();
-  return true;
+  writer->EndArray();
+  MOZ_ASSERT(aEntryReader.RemainingBytes() == 0u);
 }
 
 }  // namespace mozilla::base_profiler_markers_detail

mozglue/tests/TestBaseProfiler.cpp

@@ -4368,16 +4368,16 @@ void StreamMarkers(const mozilla::ProfileChunkedBuffer& aBuffer,
           aEntryReader.ReadObject<mozilla::ProfileBufferEntryKind>();
       MOZ_RELEASE_ASSERT(entryKind == mozilla::ProfileBufferEntryKind::Marker);
 
-      const bool success =
-          mozilla::base_profiler_markers_detail::DeserializeAfterKindAndStream(
-              aEntryReader, aWriter,
-              mozilla::baseprofiler::BaseProfilerThreadId{},
-              [&](mozilla::ProfileChunkedBuffer&) {
-                aWriter.StringElement("Real backtrace would be here");
-              },
-              [&](mozilla::base_profiler_markers_detail::Streaming::
-                      DeserializerTag) {});
-      MOZ_RELEASE_ASSERT(success);
+      mozilla::base_profiler_markers_detail::DeserializeAfterKindAndStream(
+          aEntryReader,
+          [&](const mozilla::baseprofiler::BaseProfilerThreadId&) {
+            return &aWriter;
+          },
+          [&](mozilla::ProfileChunkedBuffer&) {
+            aWriter.StringElement("Real backtrace would be here");
+          },
+          [&](mozilla::base_profiler_markers_detail::Streaming::
+                  DeserializerTag) {});
     });
   }
   aWriter.EndArray();

tools/profiler/core/ProfileBufferEntry.cpp

@@ -1279,42 +1279,38 @@ void ProfileBuffer::StreamMarkersToJSON(SpliceableJSONWriter& aWriter,
     MOZ_ASSERT(static_cast<ProfileBufferEntry::KindUnderlyingType>(type) <
                static_cast<ProfileBufferEntry::KindUnderlyingType>(
                    ProfileBufferEntry::Kind::MODERN_LIMIT));
-    bool entryWasFullyRead = false;
-
     if (type == ProfileBufferEntry::Kind::Marker) {
-      entryWasFullyRead =
-          mozilla::base_profiler_markers_detail::DeserializeAfterKindAndStream(
-              aER, aWriter, aThreadId,
-              [&](ProfileChunkedBuffer& aChunkedBuffer) {
-                ProfilerBacktrace backtrace("", &aChunkedBuffer);
-                backtrace.StreamJSON(aWriter, aProcessStartTime, aUniqueStacks);
-              },
-              [&](mozilla::base_profiler_markers_detail::Streaming::
-                      DeserializerTag aTag) {
-                size_t payloadSize = aER.RemainingBytes();
+      mozilla::base_profiler_markers_detail::DeserializeAfterKindAndStream(
+          aER,
+          [&](const ProfilerThreadId& aMarkerThreadId) {
+            return (aMarkerThreadId == aThreadId) ? &aWriter : nullptr;
+          },
+          [&](ProfileChunkedBuffer& aChunkedBuffer) {
+            ProfilerBacktrace backtrace("", &aChunkedBuffer);
+            backtrace.StreamJSON(aWriter, aProcessStartTime, aUniqueStacks);
+          },
+          [&](mozilla::base_profiler_markers_detail::Streaming::DeserializerTag
+                  aTag) {
+            size_t payloadSize = aER.RemainingBytes();
 
-                ProfileBufferEntryReader::DoubleSpanOfConstBytes spans =
-                    aER.ReadSpans(payloadSize);
-                if (MOZ_LIKELY(spans.IsSingleSpan())) {
-                  // Only a single span, we can just refer to it directly
-                  // instead of copying it.
-                  profiler::ffi::gecko_profiler_serialize_marker_for_tag(
-                      aTag, spans.mFirstOrOnly.Elements(), payloadSize,
-                      &aWriter);
-                } else {
-                  // Two spans, we need to concatenate them by copying.
-                  uint8_t* payloadBuffer = new uint8_t[payloadSize];
-                  spans.CopyBytesTo(payloadBuffer);
-                  profiler::ffi::gecko_profiler_serialize_marker_for_tag(
-                      aTag, payloadBuffer, payloadSize, &aWriter);
-                  delete[] payloadBuffer;
-                }
-              });
-    }
-
-    if (!entryWasFullyRead) {
-      // The entry was not a marker, or it was a marker for another thread.
-      // We probably didn't read the whole entry, so we need to skip to the end.
+            ProfileBufferEntryReader::DoubleSpanOfConstBytes spans =
+                aER.ReadSpans(payloadSize);
+            if (MOZ_LIKELY(spans.IsSingleSpan())) {
+              // Only a single span, we can just refer to it directly
+              // instead of copying it.
+              profiler::ffi::gecko_profiler_serialize_marker_for_tag(
+                  aTag, spans.mFirstOrOnly.Elements(), payloadSize, &aWriter);
+            } else {
+              // Two spans, we need to concatenate them by copying.
+              uint8_t* payloadBuffer = new uint8_t[payloadSize];
+              spans.CopyBytesTo(payloadBuffer);
+              profiler::ffi::gecko_profiler_serialize_marker_for_tag(
+                  aTag, payloadBuffer, payloadSize, &aWriter);
+              delete[] payloadBuffer;
+            }
+          });
+    } else {
+      // The entry was not a marker, we need to skip to the end.
       aER.SetRemainingBytes(0);
     }
   });