Bug 1626837 - ProfileChunkedBuffer ChunkManager handling - r=canaltinova

`ProfileChunkedBuffer` can handle zero or one `ProfileBufferChunkManager` at a
time, and can optionally take ownership of the manager.

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

--HG--
extra : moz-landing-system : lando

mozglue/baseprofiler/public/ProfileChunkedBuffer.h

@@ -71,6 +71,25 @@ class ProfileChunkedBuffer {
   explicit ProfileChunkedBuffer(ThreadSafety aThreadSafety)
       : mMutex(aThreadSafety != ThreadSafety::WithoutMutex) {}
 
+  // Start in-session with external chunk manager.
+  ProfileChunkedBuffer(ThreadSafety aThreadSafety,
+                       ProfileBufferChunkManager& aChunkManager)
+      : mMutex(aThreadSafety != ThreadSafety::WithoutMutex) {
+    SetChunkManager(aChunkManager);
+  }
+
+  // Start in-session with owned chunk manager.
+  ProfileChunkedBuffer(ThreadSafety aThreadSafety,
+                       UniquePtr<ProfileBufferChunkManager>&& aChunkManager)
+      : mMutex(aThreadSafety != ThreadSafety::WithoutMutex) {
+    SetChunkManager(std::move(aChunkManager));
+  }
+
+  ~ProfileChunkedBuffer() {
+    // Do proper clean-up by resetting the chunk manager.
+    ResetChunkManager();
+  }
+
   // This cannot change during the lifetime of this buffer, so there's no need
   // to lock.
   [[nodiscard]] bool IsThreadSafe() const { return mMutex.IsActivated(); }
@@ -80,6 +99,40 @@ class ProfileChunkedBuffer {
     return !!mChunkManager;
   }
 
+  // Stop using the current chunk manager.
+  // If we own the current chunk manager, it will be destroyed.
+  // This will always clear currently-held chunks, if any.
+  void ResetChunkManager() {
+    baseprofiler::detail::BaseProfilerMaybeAutoLock lock(mMutex);
+    Unused << ResetChunkManager(lock);
+  }
+
+  // Set the current chunk manager.
+  // The caller is responsible for keeping the chunk manager alive as along as
+  // it's used here (until the next (Re)SetChunkManager, or
+  // ~ProfileChunkedBuffer).
+  void SetChunkManager(ProfileBufferChunkManager& aChunkManager) {
+    baseprofiler::detail::BaseProfilerMaybeAutoLock lock(mMutex);
+    Unused << ResetChunkManager(lock);
+    SetChunkManager(aChunkManager, lock);
+  }
+
+  // Set the current chunk manager, and keep ownership of it.
+  void SetChunkManager(UniquePtr<ProfileBufferChunkManager>&& aChunkManager) {
+    baseprofiler::detail::BaseProfilerMaybeAutoLock lock(mMutex);
+    Unused << ResetChunkManager(lock);
+    mOwnedChunkManager = std::move(aChunkManager);
+    if (mOwnedChunkManager) {
+      SetChunkManager(*mOwnedChunkManager, lock);
+    }
+  }
+
+  // Stop using the current chunk manager, and return it if owned here.
+  [[nodiscard]] UniquePtr<ProfileBufferChunkManager> ExtractChunkManager() {
+    baseprofiler::detail::BaseProfilerMaybeAutoLock lock(mMutex);
+    return ResetChunkManager(lock);
+  }
+
   void Clear() {
     baseprofiler::detail::BaseProfilerMaybeAutoLock lock(mMutex);
     if (MOZ_UNLIKELY(!mChunkManager)) {
@@ -219,6 +272,57 @@ class ProfileChunkedBuffer {
 #endif  // DEBUG
 
  private:
+  [[nodiscard]] UniquePtr<ProfileBufferChunkManager> ResetChunkManager(
+      const baseprofiler::detail::BaseProfilerMaybeAutoLock&) {
+    UniquePtr<ProfileBufferChunkManager> chunkManager;
+    if (mChunkManager) {
+      mChunkManager->ForgetUnreleasedChunks();
+#ifdef DEBUG
+      mChunkManager->DeregisteredFrom(this);
+#endif
+      mChunkManager = nullptr;
+      chunkManager = std::move(mOwnedChunkManager);
+      if (mCurrentChunk) {
+        mCurrentChunk->MarkDone();
+        mCurrentChunk = nullptr;
+      }
+      mNextChunks = nullptr;
+      mNextChunkRangeStart = mRangeEnd;
+      mRangeStart = mRangeEnd;
+      mPushedBlockCount = 0;
+      mClearedBlockCount = 0;
+    }
+    return chunkManager;
+  }
+
+  void SetChunkManager(
+      ProfileBufferChunkManager& aChunkManager,
+      const baseprofiler::detail::BaseProfilerMaybeAutoLock& aLock) {
+    MOZ_ASSERT(!mChunkManager);
+    mChunkManager = &aChunkManager;
+#ifdef DEBUG
+    mChunkManager->RegisteredWith(this);
+#endif
+
+    mChunkManager->SetChunkDestroyedCallback(
+        [this](const ProfileBufferChunk& aChunk) {
+          for (;;) {
+            ProfileBufferIndex rangeStart = mRangeStart;
+            if (MOZ_LIKELY(rangeStart <= aChunk.RangeStart())) {
+              if (MOZ_LIKELY(mRangeStart.compareExchange(
+                      rangeStart,
+                      aChunk.RangeStart() + aChunk.BufferBytes()))) {
+                break;
+              }
+            }
+          }
+          mClearedBlockCount += aChunk.BlockCount();
+        });
+
+    // We start with one chunk right away.
+    SetAndInitializeCurrentChunk(mChunkManager->GetChunk(), aLock);
+  }
+
   [[nodiscard]] size_t SizeOfExcludingThis(
       MallocSizeOf aMallocSizeOf,
       const baseprofiler::detail::BaseProfilerMaybeAutoLock&) const {
@@ -260,6 +364,9 @@ class ProfileChunkedBuffer {
   // It may be owned locally (see below) or externally.
   ProfileBufferChunkManager* mChunkManager = nullptr;
 
+  // Only non-null when we own the current Chunk Manager.
+  UniquePtr<ProfileBufferChunkManager> mOwnedChunkManager;
+
   UniquePtr<ProfileBufferChunk> mCurrentChunk;
 
   UniquePtr<ProfileBufferChunk> mNextChunks;

mozglue/tests/TestBaseProfiler.cpp

@@ -850,6 +850,31 @@ static void TestChunkedBuffer() {
                 "UniquePtr<ProfileBufferChunk>");
   MOZ_RELEASE_ASSERT(!chunks, "Expected no chunks when out-of-session");
 
+  // Use ProfileBufferChunkManagerWithLocalLimit, which will give away
+  // ProfileBufferChunks that can contain 128 bytes, using up to 1KB of memory
+  // (including usable 128 bytes and headers).
+  constexpr size_t bufferMaxSize = 1024;
+  constexpr ProfileChunkedBuffer::Length chunkMinSize = 128;
+  ProfileBufferChunkManagerWithLocalLimit cm(bufferMaxSize, chunkMinSize);
+  cb.SetChunkManager(cm);
+
+  // Let the chunk manager fulfill the initial request for an extra chunk.
+  cm.FulfillChunkRequests();
+
+  MOZ_RELEASE_ASSERT(cm.MaxTotalSize() == bufferMaxSize);
+  MOZ_RELEASE_ASSERT(cb.BufferLength().isSome());
+  MOZ_RELEASE_ASSERT(*cb.BufferLength() == bufferMaxSize);
+
+  // Steal the underlying ProfileBufferChunks from the ProfileChunkedBuffer.
+  chunks = cb.GetAllChunks();
+  MOZ_RELEASE_ASSERT(!!chunks, "Expected at least one chunk");
+  MOZ_RELEASE_ASSERT(!chunks->GetNext(), "Expected only one chunk");
+  const ProfileChunkedBuffer::Length chunkActualSize = chunks->BufferBytes();
+  MOZ_RELEASE_ASSERT(chunkActualSize >= chunkMinSize);
+  MOZ_RELEASE_ASSERT(chunks->RangeStart() == 1);
+  MOZ_RELEASE_ASSERT(chunks->OffsetFirstBlock() == 0);
+  MOZ_RELEASE_ASSERT(chunks->OffsetPastLastBlock() == 0);
+
 #ifdef DEBUG
   // cb.Dump();
 #endif
@@ -860,9 +885,37 @@ static void TestChunkedBuffer() {
   // cb.Dump();
 #endif
 
+  // Reset to out-of-session.
+  cb.ResetChunkManager();
+
+  chunks = cb.GetAllChunks();
+  MOZ_RELEASE_ASSERT(!chunks, "Expected no chunks when out-of-session");
+
   printf("TestChunkedBuffer done\n");
 }
 
+static void TestChunkedBufferSingle() {
+  printf("TestChunkedBufferSingle...\n");
+
+  constexpr ProfileChunkedBuffer::Length chunkMinSize = 128;
+
+  // Create a ProfileChunkedBuffer that will own&use a
+  // ProfileBufferChunkManagerSingle, which will give away one
+  // ProfileBufferChunk that can contain 128 bytes.
+  ProfileChunkedBuffer cbSingle(
+      ProfileChunkedBuffer::ThreadSafety::WithoutMutex,
+      MakeUnique<ProfileBufferChunkManagerSingle>(chunkMinSize));
+
+  MOZ_RELEASE_ASSERT(cbSingle.BufferLength().isSome());
+  MOZ_RELEASE_ASSERT(*cbSingle.BufferLength() >= chunkMinSize);
+
+#ifdef DEBUG
+  // cbSingle.Dump();
+#endif
+
+  printf("TestChunkedBufferSingle done\n");
+}
+
 static void TestModuloBuffer(ModuloBuffer<>& mb, uint32_t MBSize) {
   using MB = ModuloBuffer<>;
 
@@ -2059,6 +2112,7 @@ void TestProfilerDependencies() {
   TestChunkManagerSingle();
   TestChunkManagerWithLocalLimit();
   TestChunkedBuffer();
+  TestChunkedBufferSingle();
   TestModuloBuffer();
   TestBlocksRingBufferAPI();
   TestBlocksRingBufferUnderlyingBufferChanges();