Bug 1667581 - Ensure there's only one nsSegmentedBuffer::FreeOMT pending at any time r=sg

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

xpcom/io/nsSegmentedBuffer.cpp

@@ -134,6 +134,16 @@ void nsSegmentedBuffer::FreeOMT(std::function<void()>&& aTask) {
     return;
   }
 
+  if (mFreeOMT) {
+    // There is a runnable pending which will handle this object
+    if (mFreeOMT->AddTask(std::move(aTask)) > 1) {
+      return;
+    }
+  } else {
+    mFreeOMT = MakeRefPtr<FreeOMTPointers>();
+    mFreeOMT->AddTask(std::move(aTask));
+  }
+
   if (!mIOThread) {
     mIOThread = do_GetService(NS_STREAMTRANSPORTSERVICE_CONTRACTID);
   }
@@ -141,8 +151,24 @@ void nsSegmentedBuffer::FreeOMT(std::function<void()>&& aTask) {
   // During the shutdown we are not able to obtain the IOThread and/or the
   // dispatching of runnable fails.
   if (!mIOThread || NS_FAILED(mIOThread->Dispatch(NS_NewRunnableFunction(
-                        "nsSegmentedBuffer::FreeOMT", aTask)))) {
-    aTask();
+                        "nsSegmentedBuffer::FreeOMT",
+                        [obj = mFreeOMT]() { obj->FreeAll(); })))) {
+    mFreeOMT->FreeAll();
+  }
+}
+
+void nsSegmentedBuffer::FreeOMTPointers::FreeAll() {
+  // Take all the tasks from the object. If AddTask is called after we release
+  // the lock, then another runnable will be dispatched for that task. This is
+  // necessary to avoid blocking the main thread while memory is being freed.
+  nsTArray<std::function<void()>> tasks = [this]() {
+    auto t = mTasks.Lock();
+    return std::move(*t);
+  }();
+
+  // Finally run all the tasks to free memory.
+  for (auto& task : tasks) {
+    task();
   }
 }
 

xpcom/io/nsSegmentedBuffer.h

@@ -12,7 +12,7 @@
 #include "nsCOMPtr.h"
 #include "nsDebug.h"
 #include "nsError.h"
-#include "mozilla/Mutex.h"
+#include "mozilla/DataMutex.h"
 
 class nsIEventTarget;
 
@@ -83,10 +83,36 @@ class nsSegmentedBuffer {
   int32_t mLastSegmentIndex;
 
  private:
+  class FreeOMTPointers {
+    NS_INLINE_DECL_THREADSAFE_REFCOUNTING(FreeOMTPointers)
+
+   public:
+    FreeOMTPointers() : mTasks("nsSegmentedBuffer::FreeOMTPointers") {}
+
+    void FreeAll();
+
+    // Adds a task to the array. Returns the size of the array.
+    size_t AddTask(std::function<void()>&& aTask) {
+      auto tasks = mTasks.Lock();
+      tasks->AppendElement(std::move(aTask));
+      return tasks->Length();
+    }
+
+   private:
+    ~FreeOMTPointers() = default;
+
+    mozilla::DataMutex<nsTArray<std::function<void()>>> mTasks;
+  };
+
   void FreeOMT(void* aPtr);
   void FreeOMT(std::function<void()>&& aTask);
 
   nsCOMPtr<nsIEventTarget> mIOThread;
+
+  // This object is created the first time we need to dispatch to another thread
+  // to free segments. It is only freed when the nsSegmentedBufer is destroyed
+  // or when the runnable is finally handled and its refcount goes to 0.
+  RefPtr<FreeOMTPointers> mFreeOMT;
 };
 
 // NS_SEGMENTARRAY_INITIAL_SIZE: This number needs to start out as a