1 files changed, 444 insertions, 0 deletions
diff --git a/mozglue/baseprofiler/public/ProfileBufferChunkManagerWithLocalLimit.h b/mozglue/baseprofiler/public/ProfileBufferChunkManagerWithLocalLimit.h
new file mode 100644
index 0000000000..034279809d
--- /dev/null
+++ b/mozglue/baseprofiler/public/ProfileBufferChunkManagerWithLocalLimit.h
@@ -0,0 +1,444 @@
+/* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* 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 ProfileBufferChunkManagerWithLocalLimit_h
+#define ProfileBufferChunkManagerWithLocalLimit_h
+
+#include "BaseProfiler.h"
+#include "mozilla/BaseProfilerDetail.h"
+#include "mozilla/ProfileBufferChunkManager.h"
+#include "mozilla/ProfileBufferControlledChunkManager.h"
+#include "mozilla/mozalloc.h"
+
+#include <utility>
+
+namespace mozilla {
+
+// Manages the Chunks for this process in a thread-safe manner, with a maximum
+// size per process.
+//
+// "Unreleased" chunks are not owned here, only "released" chunks can be
+// destroyed or recycled when reaching the memory limit, so it is theoretically
+// possible to break that limit, if:
+// - The user of this class doesn't release their chunks, AND/OR
+// - The limit is too small (e.g., smaller than 2 or 3 chunks, which should be
+//   the usual number of unreleased chunks in flight).
+// In this case, it just means that we will use more memory than allowed,
+// potentially risking OOMs. Hopefully this shouldn't happen in real code,
+// assuming that the user is doing the right thing and releasing chunks ASAP,
+// and that the memory limit is reasonably large.
+class ProfileBufferChunkManagerWithLocalLimit final
+    : public ProfileBufferChunkManager,
+      public ProfileBufferControlledChunkManager {
+ public:
+  using Length = ProfileBufferChunk::Length;
+
+  // MaxTotalBytes: Maximum number of bytes allocated in all local Chunks.
+  // ChunkMinBufferBytes: Minimum number of user-available bytes in each Chunk.
+  // Note that Chunks use a bit more memory for their header.
+  explicit ProfileBufferChunkManagerWithLocalLimit(size_t aMaxTotalBytes,
+                                                   Length aChunkMinBufferBytes)
+      : mMaxTotalBytes(aMaxTotalBytes),
+        mChunkMinBufferBytes(aChunkMinBufferBytes) {}
+
+  ~ProfileBufferChunkManagerWithLocalLimit() {
+    if (mUpdateCallback) {
+      // Signal the end of this callback.
+      std::move(mUpdateCallback)(Update(nullptr));
+    }
+  }
+
+  [[nodiscard]] size_t MaxTotalSize() const final {
+    // `mMaxTotalBytes` is `const` so there is no need to lock the mutex.
+    return mMaxTotalBytes;
+  }
+
+  [[nodiscard]] size_t TotalSize() const { return mTotalBytes; }
+
+  [[nodiscard]] UniquePtr<ProfileBufferChunk> GetChunk() final {
+    AUTO_PROFILER_STATS(Local_GetChunk);
+
+    ChunkAndUpdate chunkAndUpdate = [&]() {
+      baseprofiler::detail::BaseProfilerAutoLock lock(mMutex);
+      return GetChunk(lock);
+    }();
+
+    baseprofiler::detail::BaseProfilerAutoLock lock(mUpdateCallbackMutex);
+    if (mUpdateCallback && !chunkAndUpdate.second.IsNotUpdate()) {
+      mUpdateCallback(std::move(chunkAndUpdate.second));
+    }
+
+    return std::move(chunkAndUpdate.first);
+  }
+
+  void RequestChunk(std::function<void(UniquePtr<ProfileBufferChunk>)>&&
+                        aChunkReceiver) final {
+    AUTO_PROFILER_STATS(Local_RequestChunk);
+    baseprofiler::detail::BaseProfilerAutoLock lock(mMutex);
+    if (mChunkReceiver) {
+      // We already have a chunk receiver, meaning a request is pending.
+      return;
+    }
+    // Store the chunk receiver. This indicates that a request is pending, and
+    // it will be handled in the next `FulfillChunkRequests()` call.
+    mChunkReceiver = std::move(aChunkReceiver);
+  }
+
+  void FulfillChunkRequests() final {
+    AUTO_PROFILER_STATS(Local_FulfillChunkRequests);
+    std::function<void(UniquePtr<ProfileBufferChunk>)> chunkReceiver;
+    ChunkAndUpdate chunkAndUpdate = [&]() -> ChunkAndUpdate {
+      baseprofiler::detail::BaseProfilerAutoLock lock(mMutex);
+      if (!mChunkReceiver) {
+        // No receiver means no pending request, we're done.
+        return {};
+      }
+      // Otherwise there is a request, extract the receiver to call below.
+      std::swap(chunkReceiver, mChunkReceiver);
+      MOZ_ASSERT(!mChunkReceiver, "mChunkReceiver should have been emptied");
+      // And allocate the requested chunk. This may fail, it's fine, we're
+      // letting the receiver know about it.
+      AUTO_PROFILER_STATS(Local_FulfillChunkRequests_GetChunk);
+      return GetChunk(lock);
+    }();
+
+    if (chunkReceiver) {
+      {
+        baseprofiler::detail::BaseProfilerAutoLock lock(mUpdateCallbackMutex);
+        if (mUpdateCallback && !chunkAndUpdate.second.IsNotUpdate()) {
+          mUpdateCallback(std::move(chunkAndUpdate.second));
+        }
+      }
+
+      // Invoke callback outside of lock, so that it can use other chunk manager
+      // functions if needed.
+      // Note that this means there could be a race, where another request
+      // happens now and even gets fulfilled before this one is! It should be
+      // rare, and shouldn't be a problem anyway, the user will still get their
+      // requested chunks, new/recycled chunks look the same so their order
+      // doesn't matter.
+      std::move(chunkReceiver)(std::move(chunkAndUpdate.first));
+    }
+  }
+
+  void ReleaseChunk(UniquePtr<ProfileBufferChunk> aChunk) final {
+    if (!aChunk) {
+      return;
+    }
+
+    MOZ_RELEASE_ASSERT(!aChunk->GetNext(), "ReleaseChunk only accepts 1 chunk");
+    MOZ_RELEASE_ASSERT(!aChunk->ChunkHeader().mDoneTimeStamp.IsNull(),
+                       "Released chunk should have a 'Done' timestamp");
+
+    Update update = [&]() {
+      baseprofiler::detail::BaseProfilerAutoLock lock(mMutex);
+      MOZ_ASSERT(mUser, "Not registered yet");
+      // Keep a pointer to the first newly-released chunk, so we can use it to
+      // prepare an update (after `aChunk` is moved-from).
+      const ProfileBufferChunk* const newlyReleasedChunk = aChunk.get();
+      // Transfer the chunk size from the unreleased bucket to the released one.
+      mUnreleasedBufferBytes -= aChunk->BufferBytes();
+      mReleasedBufferBytes += aChunk->BufferBytes();
+      if (!mReleasedChunks) {
+        // No other released chunks at the moment, we're starting the list.
+        MOZ_ASSERT(mReleasedBufferBytes == aChunk->BufferBytes());
+        mReleasedChunks = std::move(aChunk);
+      } else {
+        // Insert aChunk in mReleasedChunks to keep done-timestamp order.
+        const TimeStamp& releasedChunkDoneTimeStamp =
+            aChunk->ChunkHeader().mDoneTimeStamp;
+        if (releasedChunkDoneTimeStamp <
+            mReleasedChunks->ChunkHeader().mDoneTimeStamp) {
+          // aChunk is the oldest -> Insert at the beginning.
+          aChunk->SetLast(std::move(mReleasedChunks));
+          mReleasedChunks = std::move(aChunk);
+        } else {
+          // Go through the already-released chunk list, and insert aChunk
+          // before the first younger released chunk, or at the end.
+          ProfileBufferChunk* chunk = mReleasedChunks.get();
+          for (;;) {
+            ProfileBufferChunk* const nextChunk = chunk->GetNext();
+            if (!nextChunk || releasedChunkDoneTimeStamp <
+                                  nextChunk->ChunkHeader().mDoneTimeStamp) {
+              // Either we're at the last released chunk, or the next released
+              // chunk is younger -> Insert right after this released chunk.
+              chunk->InsertNext(std::move(aChunk));
+              break;
+            }
+            chunk = nextChunk;
+          }
+        }
+      }
+
+      return Update(mUnreleasedBufferBytes, mReleasedBufferBytes,
+                    mReleasedChunks.get(), newlyReleasedChunk);
+    }();
+
+    baseprofiler::detail::BaseProfilerAutoLock lock(mUpdateCallbackMutex);
+    if (mUpdateCallback && !update.IsNotUpdate()) {
+      mUpdateCallback(std::move(update));
+    }
+  }
+
+  void SetChunkDestroyedCallback(
+      std::function<void(const ProfileBufferChunk&)>&& aChunkDestroyedCallback)
+      final {
+    baseprofiler::detail::BaseProfilerAutoLock lock(mMutex);
+    MOZ_ASSERT(mUser, "Not registered yet");
+    mChunkDestroyedCallback = std::move(aChunkDestroyedCallback);
+  }
+
+  [[nodiscard]] UniquePtr<ProfileBufferChunk> GetExtantReleasedChunks() final {
+    UniquePtr<ProfileBufferChunk> chunks;
+    size_t unreleasedBufferBytes = [&]() {
+      baseprofiler::detail::BaseProfilerAutoLock lock(mMutex);
+      MOZ_ASSERT(mUser, "Not registered yet");
+      mReleasedBufferBytes = 0;
+      chunks = std::move(mReleasedChunks);
+      return mUnreleasedBufferBytes;
+    }();
+
+    baseprofiler::detail::BaseProfilerAutoLock lock(mUpdateCallbackMutex);
+    if (mUpdateCallback) {
+      mUpdateCallback(Update(unreleasedBufferBytes, 0, nullptr, nullptr));
+    }
+
+    return chunks;
+  }
+
+  void ForgetUnreleasedChunks() final {
+    Update update = [&]() {
+      baseprofiler::detail::BaseProfilerAutoLock lock(mMutex);
+      MOZ_ASSERT(mUser, "Not registered yet");
+      mUnreleasedBufferBytes = 0;
+      return Update(0, mReleasedBufferBytes, mReleasedChunks.get(), nullptr);
+    }();
+    baseprofiler::detail::BaseProfilerAutoLock lock(mUpdateCallbackMutex);
+    if (mUpdateCallback) {
+      mUpdateCallback(std::move(update));
+    }
+  }
+
+  [[nodiscard]] size_t SizeOfExcludingThis(
+      MallocSizeOf aMallocSizeOf) const final {
+    baseprofiler::detail::BaseProfilerAutoLock lock(mMutex);
+    return SizeOfExcludingThis(aMallocSizeOf, lock);
+  }
+
+  [[nodiscard]] size_t SizeOfIncludingThis(
+      MallocSizeOf aMallocSizeOf) const final {
+    baseprofiler::detail::BaseProfilerAutoLock lock(mMutex);
+    MOZ_ASSERT(mUser, "Not registered yet");
+    return aMallocSizeOf(this) + SizeOfExcludingThis(aMallocSizeOf, lock);
+  }
+
+  void SetUpdateCallback(UpdateCallback&& aUpdateCallback) final {
+    {
+      baseprofiler::detail::BaseProfilerAutoLock lock(mUpdateCallbackMutex);
+      if (mUpdateCallback) {
+        // Signal the end of the previous callback.
+        std::move(mUpdateCallback)(Update(nullptr));
+        mUpdateCallback = nullptr;
+      }
+    }
+
+    if (aUpdateCallback) {
+      Update initialUpdate = [&]() {
+        baseprofiler::detail::BaseProfilerAutoLock lock(mMutex);
+        return Update(mUnreleasedBufferBytes, mReleasedBufferBytes,
+                      mReleasedChunks.get(), nullptr);
+      }();
+
+      baseprofiler::detail::BaseProfilerAutoLock lock(mUpdateCallbackMutex);
+      MOZ_ASSERT(!mUpdateCallback, "Only one update callback allowed");
+      mUpdateCallback = std::move(aUpdateCallback);
+      mUpdateCallback(std::move(initialUpdate));
+    }
+  }
+
+  void DestroyChunksAtOrBefore(TimeStamp aDoneTimeStamp) final {
+    MOZ_ASSERT(!aDoneTimeStamp.IsNull());
+    baseprofiler::detail::BaseProfilerAutoLock lock(mMutex);
+    for (;;) {
+      if (!mReleasedChunks) {
+        // We don't own any released chunks (anymore), we're done.
+        break;
+      }
+      if (mReleasedChunks->ChunkHeader().mDoneTimeStamp > aDoneTimeStamp) {
+        // The current chunk is strictly after the given timestamp, we're done.
+        break;
+      }
+      // We've found a chunk at or before the timestamp, discard it.
+      DiscardOldestReleasedChunk(lock);
+    }
+  }
+
+ protected:
+  const ProfileBufferChunk* PeekExtantReleasedChunksAndLock() final
+      MOZ_CAPABILITY_ACQUIRE(mMutex) {
+    mMutex.Lock();
+    MOZ_ASSERT(mUser, "Not registered yet");
+    return mReleasedChunks.get();
+  }
+  void UnlockAfterPeekExtantReleasedChunks() final
+      MOZ_CAPABILITY_RELEASE(mMutex) {
+    mMutex.Unlock();
+  }
+
+ private:
+  size_t MaybeRecycleChunkAndGetDeallocatedSize(
+      UniquePtr<ProfileBufferChunk>&& chunk,
+      const baseprofiler::detail::BaseProfilerAutoLock& aLock) {
+    // Try to recycle big-enough chunks. (All chunks should have the same size,
+    // but it's a cheap test and may allow future adjustments based on actual
+    // data rate.)
+    if (chunk->BufferBytes() >= mChunkMinBufferBytes) {
+      // We keep up to two recycled chunks at any time.
+      if (!mRecycledChunks) {
+        mRecycledChunks = std::move(chunk);
+        return 0;
+      } else if (!mRecycledChunks->GetNext()) {
+        mRecycledChunks->InsertNext(std::move(chunk));
+        return 0;
+      }
+    }
+    return moz_malloc_usable_size(chunk.get());
+  }
+
+  UniquePtr<ProfileBufferChunk> TakeRecycledChunk(
+      const baseprofiler::detail::BaseProfilerAutoLock& aLock) {
+    UniquePtr<ProfileBufferChunk> recycled;
+    if (mRecycledChunks) {
+      recycled = std::exchange(mRecycledChunks, mRecycledChunks->ReleaseNext());
+      recycled->MarkRecycled();
+    }
+    return recycled;
+  }
+
+  void DiscardOldestReleasedChunk(
+      const baseprofiler::detail::BaseProfilerAutoLock& aLock) {
+    MOZ_ASSERT(!!mReleasedChunks);
+    UniquePtr<ProfileBufferChunk> oldest =
+        std::exchange(mReleasedChunks, mReleasedChunks->ReleaseNext());
+    mReleasedBufferBytes -= oldest->BufferBytes();
+    if (mChunkDestroyedCallback) {
+      // Inform the user that we're going to destroy this chunk.
+      mChunkDestroyedCallback(*oldest);
+    }
+
+    mTotalBytes -=
+        MaybeRecycleChunkAndGetDeallocatedSize(std::move(oldest), aLock);
+  }
+
+  using ChunkAndUpdate = std::pair<UniquePtr<ProfileBufferChunk>, Update>;
+  [[nodiscard]] ChunkAndUpdate GetChunk(
+      const baseprofiler::detail::BaseProfilerAutoLock& aLock) {
+    MOZ_ASSERT(mUser, "Not registered yet");
+    // After this function, the total memory consumption will be the sum of:
+    // - Bytes from released (i.e., full) chunks,
+    // - Bytes from unreleased (still in use) chunks,
+    // - Bytes from the chunk we want to create/recycle. (Note that we don't
+    //   count the extra bytes of chunk header, and of extra allocation ability,
+    //   for the new chunk, as it's assumed to be negligible compared to the
+    //   total memory limit.)
+    // If this total is higher than the local limit, we'll want to destroy
+    // the oldest released chunks until we're under the limit; if any, we may
+    // recycle one of them to avoid a deallocation followed by an allocation.
+    while (mReleasedBufferBytes + mUnreleasedBufferBytes +
+                   mChunkMinBufferBytes >=
+               mMaxTotalBytes &&
+           !!mReleasedChunks) {
+      // We have reached the local limit, discard the oldest released chunk.
+      DiscardOldestReleasedChunk(aLock);
+    }
+
+    // Extract the recycled chunk, if any.
+    ChunkAndUpdate chunkAndUpdate{TakeRecycledChunk(aLock), Update()};
+    UniquePtr<ProfileBufferChunk>& chunk = chunkAndUpdate.first;
+
+    if (!chunk) {
+      // No recycled chunk -> Create a chunk now. (This could still fail.)
+      chunk = ProfileBufferChunk::Create(mChunkMinBufferBytes);
+      mTotalBytes += moz_malloc_usable_size(chunk.get());
+    }
+
+    if (chunk) {
+      // We do have a chunk (recycled or new), record its size as "unreleased".
+      mUnreleasedBufferBytes += chunk->BufferBytes();
+
+      chunkAndUpdate.second =
+          Update(mUnreleasedBufferBytes, mReleasedBufferBytes,
+                 mReleasedChunks.get(), nullptr);
+    }
+
+    return chunkAndUpdate;
+  }
+
+  [[nodiscard]] size_t SizeOfExcludingThis(
+      MallocSizeOf aMallocSizeOf,
+      const baseprofiler::detail::BaseProfilerAutoLock&) const {
+    MOZ_ASSERT(mUser, "Not registered yet");
+    size_t size = 0;
+    if (mReleasedChunks) {
+      size += mReleasedChunks->SizeOfIncludingThis(aMallocSizeOf);
+    }
+    if (mRecycledChunks) {
+      size += mRecycledChunks->SizeOfIncludingThis(aMallocSizeOf);
+    }
+    // Note: Missing size of std::function external resources (if any).
+    return size;
+  }
+
+  // Maxumum number of bytes that should be used by all unreleased and released
+  // chunks. Note that only released chunks can be destroyed here, so it is the
+  // responsibility of the user to properly release their chunks when possible.
+  const size_t mMaxTotalBytes;
+
+  // Minimum number of bytes that new chunks should be able to store.
+  // Used when calling `ProfileBufferChunk::Create()`.
+  const Length mChunkMinBufferBytes;
+
+  // Mutex guarding the following members.
+  mutable baseprofiler::detail::BaseProfilerMutex mMutex;
+
+  // Number of bytes currently held in chunks that have been given away (through
+  // `GetChunk` or `RequestChunk`) and not released yet.
+  size_t mUnreleasedBufferBytes = 0;
+
+  // Number of bytes currently held in chunks that have been released and stored
+  // in `mReleasedChunks` below.
+  size_t mReleasedBufferBytes = 0;
+
+  // Total allocated size (used to substract it from memory counters).
+  size_t mTotalBytes = 0;
+
+  // List of all released chunks. The oldest one should be at the start of the
+  // list, and may be destroyed or recycled when the memory limit is reached.
+  UniquePtr<ProfileBufferChunk> mReleasedChunks;
+
+  // This may hold chunks that were released then slated for destruction, they
+  // will be reused next time an allocation would have been needed.
+  UniquePtr<ProfileBufferChunk> mRecycledChunks;
+
+  // Optional callback used to notify the user when a chunk is about to be
+  // destroyed or recycled. (The data content is always destroyed, but the chunk
+  // container may be reused.)
+  std::function<void(const ProfileBufferChunk&)> mChunkDestroyedCallback;
+
+  // Callback set from `RequestChunk()`, until it is serviced in
+  // `FulfillChunkRequests()`. There can only be one request in flight.
+  std::function<void(UniquePtr<ProfileBufferChunk>)> mChunkReceiver;
+
+  // Separate mutex guarding mUpdateCallback, so that it may be invoked outside
+  // of the main buffer `mMutex`.
+  mutable baseprofiler::detail::BaseProfilerMutex mUpdateCallbackMutex;
+
+  UpdateCallback mUpdateCallback;
+};
+
+}  // namespace mozilla
+
+#endif  // ProfileBufferChunkManagerWithLocalLimit_h