1 files changed, 388 insertions, 0 deletions
diff --git a/mozglue/baseprofiler/public/BaseProfileJSONWriter.h b/mozglue/baseprofiler/public/BaseProfileJSONWriter.h
new file mode 100644
index 0000000000..5dcf06f3f3
--- /dev/null
+++ b/mozglue/baseprofiler/public/BaseProfileJSONWriter.h
@@ -0,0 +1,388 @@
+/* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#ifndef BASEPROFILEJSONWRITER_H
+#define BASEPROFILEJSONWRITER_H
+
+#include "mozilla/HashFunctions.h"
+#include "mozilla/HashTable.h"
+#include "mozilla/JSONWriter.h"
+#include "mozilla/TimeStamp.h"
+#include "mozilla/UniquePtr.h"
+
+#include <functional>
+#include <ostream>
+#include <string_view>
+
+namespace mozilla {
+namespace baseprofiler {
+
+class SpliceableJSONWriter;
+
+// On average, profile JSONs are large enough such that we want to avoid
+// reallocating its buffer when expanding. Additionally, the contents of the
+// profile are not accessed until the profile is entirely written. For these
+// reasons we use a chunked writer that keeps an array of chunks, which is
+// concatenated together after writing is finished.
+class ChunkedJSONWriteFunc final : public JSONWriteFunc {
+ public:
+  friend class SpliceableJSONWriter;
+
+  ChunkedJSONWriteFunc() : mChunkPtr{nullptr}, mChunkEnd{nullptr} {
+    AllocChunk(kChunkSize);
+  }
+
+  bool IsEmpty() const {
+    MOZ_ASSERT_IF(!mChunkPtr, !mChunkEnd && mChunkList.length() == 0 &&
+                                  mChunkLengths.length() == 0);
+    return !mChunkPtr;
+  }
+
+  void Write(const Span<const char>& aStr) override {
+    MOZ_ASSERT(mChunkPtr >= mChunkList.back().get() && mChunkPtr <= mChunkEnd);
+    MOZ_ASSERT(mChunkEnd >= mChunkList.back().get() + mChunkLengths.back());
+    MOZ_ASSERT(*mChunkPtr == '\0');
+
+    // Most strings to be written are small, but subprocess profiles (e.g.,
+    // from the content process in e10s) may be huge. If the string is larger
+    // than a chunk, allocate its own chunk.
+    char* newPtr;
+    if (aStr.size() >= kChunkSize) {
+      AllocChunk(aStr.size() + 1);
+      newPtr = mChunkPtr + aStr.size();
+    } else {
+      newPtr = mChunkPtr + aStr.size();
+      if (newPtr >= mChunkEnd) {
+        AllocChunk(kChunkSize);
+        newPtr = mChunkPtr + aStr.size();
+      }
+    }
+
+    memcpy(mChunkPtr, aStr.data(), aStr.size());
+    *newPtr = '\0';
+    mChunkPtr = newPtr;
+    mChunkLengths.back() += aStr.size();
+  }
+  void CopyDataIntoLazilyAllocatedBuffer(
+      const std::function<char*(size_t)>& aAllocator) const {
+    // Request a buffer for the full content plus a null terminator.
+    MOZ_ASSERT(mChunkLengths.length() == mChunkList.length());
+    size_t totalLen = 1;
+    for (size_t i = 0; i < mChunkLengths.length(); i++) {
+      MOZ_ASSERT(strlen(mChunkList[i].get()) == mChunkLengths[i]);
+      totalLen += mChunkLengths[i];
+    }
+    char* ptr = aAllocator(totalLen);
+
+    if (!ptr) {
+      // Failed to allocate memory.
+      return;
+    }
+
+    for (size_t i = 0; i < mChunkList.length(); i++) {
+      size_t len = mChunkLengths[i];
+      memcpy(ptr, mChunkList[i].get(), len);
+      ptr += len;
+    }
+    *ptr = '\0';
+  }
+  UniquePtr<char[]> CopyData() const {
+    UniquePtr<char[]> c;
+    CopyDataIntoLazilyAllocatedBuffer([&](size_t allocationSize) {
+      c = MakeUnique<char[]>(allocationSize);
+      return c.get();
+    });
+    return c;
+  }
+  void Take(ChunkedJSONWriteFunc&& aOther) {
+    for (size_t i = 0; i < aOther.mChunkList.length(); i++) {
+      MOZ_ALWAYS_TRUE(mChunkLengths.append(aOther.mChunkLengths[i]));
+      MOZ_ALWAYS_TRUE(mChunkList.append(std::move(aOther.mChunkList[i])));
+    }
+    mChunkPtr = mChunkList.back().get() + mChunkLengths.back();
+    mChunkEnd = mChunkPtr;
+    aOther.mChunkPtr = nullptr;
+    aOther.mChunkEnd = nullptr;
+    aOther.mChunkList.clear();
+    aOther.mChunkLengths.clear();
+  }
+
+ private:
+  void AllocChunk(size_t aChunkSize) {
+    MOZ_ASSERT(mChunkLengths.length() == mChunkList.length());
+    UniquePtr<char[]> newChunk = MakeUnique<char[]>(aChunkSize);
+    mChunkPtr = newChunk.get();
+    mChunkEnd = mChunkPtr + aChunkSize;
+    *mChunkPtr = '\0';
+    MOZ_ALWAYS_TRUE(mChunkLengths.append(0));
+    MOZ_ALWAYS_TRUE(mChunkList.append(std::move(newChunk)));
+  }
+
+  static const size_t kChunkSize = 4096 * 512;
+
+  // Pointer for writing inside the current chunk.
+  //
+  // The current chunk is always at the back of mChunkList, i.e.,
+  // mChunkList.back() <= mChunkPtr <= mChunkEnd.
+  char* mChunkPtr;
+
+  // Pointer to the end of the current chunk.
+  //
+  // The current chunk is always at the back of mChunkList, i.e.,
+  // mChunkEnd >= mChunkList.back() + mChunkLengths.back().
+  char* mChunkEnd;
+
+  // List of chunks and their lengths.
+  //
+  // For all i, the length of the string in mChunkList[i] is
+  // mChunkLengths[i].
+  Vector<UniquePtr<char[]>> mChunkList;
+  Vector<size_t> mChunkLengths;
+};
+
+struct OStreamJSONWriteFunc final : public JSONWriteFunc {
+  explicit OStreamJSONWriteFunc(std::ostream& aStream) : mStream(aStream) {}
+
+  void Write(const Span<const char>& aStr) override {
+    std::string_view sv(aStr.data(), aStr.size());
+    mStream << sv;
+  }
+
+  std::ostream& mStream;
+};
+
+class UniqueJSONStrings;
+
+class SpliceableJSONWriter : public JSONWriter {
+ public:
+  explicit SpliceableJSONWriter(UniquePtr<JSONWriteFunc> aWriter)
+      : JSONWriter(std::move(aWriter)) {}
+
+  void StartBareList(CollectionStyle aStyle = MultiLineStyle) {
+    StartCollection(scEmptyString, scEmptyString, aStyle);
+  }
+
+  void EndBareList() { EndCollection(scEmptyString); }
+
+  // This function must be used to correctly stream timestamps in profiles.
+  // Null timestamps don't output anything.
+  void TimeProperty(const Span<const char>& aName, const TimeStamp& aTime) {
+    if (!aTime.IsNull()) {
+      DoubleProperty(aName,
+                     (aTime - TimeStamp::ProcessCreation()).ToMilliseconds());
+    }
+  }
+
+  void NullElements(uint32_t aCount) {
+    for (uint32_t i = 0; i < aCount; i++) {
+      NullElement();
+    }
+  }
+
+  void Splice(const Span<const char>& aStr) {
+    Separator();
+    WriteFunc()->Write(aStr);
+    mNeedComma[mDepth] = true;
+  }
+
+  void Splice(const char* aStr, size_t aLen) {
+    Separator();
+    WriteFunc()->Write(Span<const char>(aStr, aLen));
+    mNeedComma[mDepth] = true;
+  }
+
+  // Splice the given JSON directly in, without quoting.
+  void SplicedJSONProperty(const Span<const char>& aMaybePropertyName,
+                           const Span<const char>& aJsonValue) {
+    Scalar(aMaybePropertyName, aJsonValue);
+  }
+
+  void CopyAndSplice(const ChunkedJSONWriteFunc& aFunc) {
+    Separator();
+    for (size_t i = 0; i < aFunc.mChunkList.length(); i++) {
+      WriteFunc()->Write(
+          Span<const char>(aFunc.mChunkList[i].get(), aFunc.mChunkLengths[i]));
+    }
+    mNeedComma[mDepth] = true;
+  }
+
+  // Takes the chunks from aFunc and write them. If move is not possible
+  // (e.g., using OStreamJSONWriteFunc), aFunc's chunks are copied and its
+  // storage cleared.
+  virtual void TakeAndSplice(ChunkedJSONWriteFunc&& aFunc) {
+    Separator();
+    for (size_t i = 0; i < aFunc.mChunkList.length(); i++) {
+      WriteFunc()->Write(
+          Span<const char>(aFunc.mChunkList[i].get(), aFunc.mChunkLengths[i]));
+    }
+    aFunc.mChunkPtr = nullptr;
+    aFunc.mChunkEnd = nullptr;
+    aFunc.mChunkList.clear();
+    aFunc.mChunkLengths.clear();
+    mNeedComma[mDepth] = true;
+  }
+
+  // Set (or reset) the pointer to a UniqueJSONStrings.
+  void SetUniqueStrings(UniqueJSONStrings& aUniqueStrings) {
+    MOZ_RELEASE_ASSERT(!mUniqueStrings);
+    mUniqueStrings = &aUniqueStrings;
+  }
+
+  // Set (or reset) the pointer to a UniqueJSONStrings.
+  void ResetUniqueStrings() {
+    MOZ_RELEASE_ASSERT(mUniqueStrings);
+    mUniqueStrings = nullptr;
+  }
+
+  // Add `aStr` to the unique-strings list (if not already there), and write its
+  // index as a named object property.
+  inline void UniqueStringProperty(const Span<const char>& aName,
+                                   const Span<const char>& aStr);
+
+  // Add `aStr` to the unique-strings list (if not already there), and write its
+  // index as an array element.
+  inline void UniqueStringElement(const Span<const char>& aStr);
+
+ private:
+  UniqueJSONStrings* mUniqueStrings = nullptr;
+};
+
+class SpliceableChunkedJSONWriter final : public SpliceableJSONWriter {
+ public:
+  explicit SpliceableChunkedJSONWriter()
+      : SpliceableJSONWriter(MakeUnique<ChunkedJSONWriteFunc>()) {}
+
+  // Access the ChunkedJSONWriteFunc as reference-to-const, usually to copy data
+  // out.
+  const ChunkedJSONWriteFunc& ChunkedWriteFunc() const {
+    MOZ_ASSERT(!mTaken);
+    // The WriteFunc was non-fallibly allocated as a ChunkedJSONWriteFunc in the
+    // only constructor above, so it's safe to cast to ChunkedJSONWriteFunc*.
+    return *static_cast<const ChunkedJSONWriteFunc*>(WriteFunc());
+  }
+
+  // Access the ChunkedJSONWriteFunc as rvalue-reference, usually to take its
+  // data out. This writer shouldn't be used anymore after this.
+  ChunkedJSONWriteFunc&& TakeChunkedWriteFunc() {
+#ifdef DEBUG
+    MOZ_ASSERT(!mTaken);
+    mTaken = true;
+#endif  //
+    // The WriteFunc was non-fallibly allocated as a ChunkedJSONWriteFunc in the
+    // only constructor above, so it's safe to cast to ChunkedJSONWriteFunc*.
+    return std::move(*static_cast<ChunkedJSONWriteFunc*>(WriteFunc()));
+  }
+
+  // Adopts the chunks from aFunc without copying.
+  void TakeAndSplice(ChunkedJSONWriteFunc&& aFunc) override {
+    MOZ_ASSERT(!mTaken);
+    Separator();
+    // The WriteFunc was non-fallibly allocated as a ChunkedJSONWriteFunc in the
+    // only constructor above, so it's safe to cast to ChunkedJSONWriteFunc*.
+    static_cast<ChunkedJSONWriteFunc*>(WriteFunc())->Take(std::move(aFunc));
+    mNeedComma[mDepth] = true;
+  }
+
+#ifdef DEBUG
+ private:
+  bool mTaken = false;
+#endif  //
+};
+
+class JSONSchemaWriter {
+  JSONWriter& mWriter;
+  uint32_t mIndex;
+
+ public:
+  explicit JSONSchemaWriter(JSONWriter& aWriter) : mWriter(aWriter), mIndex(0) {
+    aWriter.StartObjectProperty("schema",
+                                SpliceableJSONWriter::SingleLineStyle);
+  }
+
+  void WriteField(const Span<const char>& aName) {
+    mWriter.IntProperty(aName, mIndex++);
+  }
+
+  template <size_t Np1>
+  void WriteField(const char (&aName)[Np1]) {
+    WriteField(Span<const char>(aName, Np1 - 1));
+  }
+
+  ~JSONSchemaWriter() { mWriter.EndObject(); }
+};
+
+// This class helps create an indexed list of unique strings, and inserts the
+// index as a JSON value. The collected list of unique strings can later be
+// inserted as a JSON array.
+// This can be useful for elements/properties with many repeated strings.
+//
+// With only JSONWriter w,
+// `w.WriteElement("a"); w.WriteElement("b"); w.WriteElement("a");`
+// when done inside a JSON array, will generate:
+// `["a", "b", "c"]`
+//
+// With UniqueStrings u,
+// `u.WriteElement(w, "a"); u.WriteElement(w, "b"); u.WriteElement(w, "a");`
+// when done inside a JSON array, will generate:
+// `[0, 1, 0]`
+// and later, `u.SpliceStringTableElements(w)` (inside a JSON array), will
+// output the corresponding indexed list of unique strings:
+// `["a", "b"]`
+class UniqueJSONStrings {
+ public:
+  // Start an empty list of unique strings.
+  MFBT_API explicit UniqueJSONStrings(
+      JSONWriter::CollectionStyle aStyle = JSONWriter::MultiLineStyle);
+
+  // Start with a copy of the strings from another list.
+  MFBT_API explicit UniqueJSONStrings(
+      const UniqueJSONStrings& aOther,
+      JSONWriter::CollectionStyle aStyle = JSONWriter::MultiLineStyle);
+
+  MFBT_API ~UniqueJSONStrings();
+
+  // Add `aStr` to the list (if not already there), and write its index as a
+  // named object property.
+  void WriteProperty(JSONWriter& aWriter, const Span<const char>& aName,
+                     const Span<const char>& aStr) {
+    aWriter.IntProperty(aName, GetOrAddIndex(aStr));
+  }
+
+  // Add `aStr` to the list (if not already there), and write its index as an
+  // array element.
+  void WriteElement(JSONWriter& aWriter, const Span<const char>& aStr) {
+    aWriter.IntElement(GetOrAddIndex(aStr));
+  }
+
+  // Splice all collected unique strings into an array. This should only be done
+  // once, and then this UniqueStrings shouldn't be used anymore.
+  MFBT_API void SpliceStringTableElements(SpliceableJSONWriter& aWriter);
+
+ private:
+  // If `aStr` is already listed, return its index.
+  // Otherwise add it to the list and return the new index.
+  MFBT_API uint32_t GetOrAddIndex(const Span<const char>& aStr);
+
+  SpliceableChunkedJSONWriter mStringTableWriter;
+  HashMap<HashNumber, uint32_t> mStringHashToIndexMap;
+};
+
+void SpliceableJSONWriter::UniqueStringProperty(const Span<const char>& aName,
+                                                const Span<const char>& aStr) {
+  MOZ_RELEASE_ASSERT(mUniqueStrings);
+  mUniqueStrings->WriteProperty(*this, aName, aStr);
+}
+
+// Add `aStr` to the list (if not already there), and write its index as an
+// array element.
+void SpliceableJSONWriter::UniqueStringElement(const Span<const char>& aStr) {
+  MOZ_RELEASE_ASSERT(mUniqueStrings);
+  mUniqueStrings->WriteElement(*this, aStr);
+}
+
+}  // namespace baseprofiler
+}  // namespace mozilla
+
+#endif  // BASEPROFILEJSONWRITER_H