Make MemoryAllocator ref-counted.

src/gpgmm/common/BuddyMemoryAllocator.cpp

@@ -24,7 +24,7 @@ namespace gpgmm {
     BuddyMemoryAllocator::BuddyMemoryAllocator(uint64_t systemSize,
                                                uint64_t memorySize,
                                                uint64_t memoryAlignment,
-                                               std::unique_ptr<MemoryAllocatorBase> memoryAllocator)
+                                               ScopedRef<MemoryAllocatorBase> memoryAllocator)
         : MemoryAllocatorBase(std::move(memoryAllocator)),
           mMemorySize(memorySize),
           mMemoryAlignment(memoryAlignment),

src/gpgmm/common/BuddyMemoryAllocator.h

@@ -40,7 +40,7 @@ namespace gpgmm {
         BuddyMemoryAllocator(uint64_t systemSize,
                              uint64_t memorySize,
                              uint64_t memoryAlignment,
-                             std::unique_ptr<MemoryAllocatorBase> memoryAllocator);
+                             ScopedRef<MemoryAllocatorBase> memoryAllocator);
 
         // MemoryAllocatorBase interface
         ResultOrError<std::unique_ptr<MemoryAllocationBase>> TryAllocateMemory(

src/gpgmm/common/DedicatedMemoryAllocator.cpp

@@ -22,7 +22,7 @@
 namespace gpgmm {
 
     DedicatedMemoryAllocator::DedicatedMemoryAllocator(
-        std::unique_ptr<MemoryAllocatorBase> memoryAllocator,
+        ScopedRef<MemoryAllocatorBase> memoryAllocator,
         uint64_t memoryAlignment)
         : MemoryAllocatorBase(std::move(memoryAllocator)), mMemoryAlignment(memoryAlignment) {
     }

src/gpgmm/common/DedicatedMemoryAllocator.h

@@ -25,7 +25,7 @@ namespace gpgmm {
     // memory to be tracked.
     class DedicatedMemoryAllocator final : public MemoryAllocatorBase {
       public:
-        DedicatedMemoryAllocator(std::unique_ptr<MemoryAllocatorBase> memoryAllocator,
+        DedicatedMemoryAllocator(ScopedRef<MemoryAllocatorBase> memoryAllocator,
                                  uint64_t memoryAlignment);
 
         // MemoryAllocatorBase interface

src/gpgmm/common/MemoryAllocator.cpp

@@ -90,11 +90,11 @@ namespace gpgmm {
 
     // MemoryAllocatorBase
 
-    MemoryAllocatorBase::MemoryAllocatorBase() {
+    MemoryAllocatorBase::MemoryAllocatorBase() : RefCounted(0) {
     }
 
-    MemoryAllocatorBase::MemoryAllocatorBase(std::unique_ptr<MemoryAllocatorBase> next) {
-        InsertIntoChain(std::move(next));
+    MemoryAllocatorBase::MemoryAllocatorBase(ScopedRef<MemoryAllocatorBase> next) : RefCounted(0) {
+        InsertIntoChain(next);
     }
 
     MemoryAllocatorBase::~MemoryAllocatorBase() {
@@ -110,7 +110,7 @@ namespace gpgmm {
 
         // Deletes adjacent node recursively (post-order).
         if (mNext != nullptr) {
-            SafeDelete(mNext);
+            mNext = nullptr;
         }
 
         if (IsInList()) {
@@ -195,17 +195,17 @@ namespace gpgmm {
     }
 
     MemoryAllocatorBase* MemoryAllocatorBase::GetNextInChain() const {
-        return mNext;
+        return mNext.Get();
     }
 
     MemoryAllocatorBase* MemoryAllocatorBase::GetParent() const {
         return mParent;
     }
 
-    void MemoryAllocatorBase::InsertIntoChain(std::unique_ptr<MemoryAllocatorBase> next) {
+    void MemoryAllocatorBase::InsertIntoChain(ScopedRef<MemoryAllocatorBase> next) {
         ASSERT(next != nullptr);
         next->mParent = this->value();
-        mNext = next.release();
+        mNext = next;
     }
 
 }  // namespace gpgmm

src/gpgmm/common/MemoryAllocator.h

@@ -25,6 +25,7 @@
 #include "gpgmm/utils/Limits.h"
 #include "gpgmm/utils/LinkedList.h"
 #include "gpgmm/utils/Log.h"
+#include "gpgmm/utils/Refcount.h"
 
 #include <memory>
 #include <mutex>
@@ -140,13 +141,15 @@ namespace gpgmm {
     // parent) and the next MemoryAllocatorBase (or child) form a one-way edge. This allows the
     // first-order MemoryAllocatorBase to sub-allocate from larger blocks provided by the
     // second-order MemoryAllocatorBase and so on.
-    class MemoryAllocatorBase : public ObjectBase, public LinkNode<MemoryAllocatorBase> {
+    class MemoryAllocatorBase : public ObjectBase,
+                                public LinkNode<MemoryAllocatorBase>,
+                                public RefCounted {
       public:
         // Constructs a standalone MemoryAllocatorBase.
         MemoryAllocatorBase();
 
         // Constructs a MemoryAllocatorBase that also owns a (child) allocator.
-        explicit MemoryAllocatorBase(std::unique_ptr<MemoryAllocatorBase> next);
+        explicit MemoryAllocatorBase(ScopedRef<MemoryAllocatorBase> next);
 
         virtual ~MemoryAllocatorBase() override;
 
@@ -252,14 +255,14 @@ namespace gpgmm {
                 nullptr, memory, kInvalidOffset, AllocationMethod::kUndefined, block, requestSize);
         }
 
-        void InsertIntoChain(std::unique_ptr<MemoryAllocatorBase> next);
+        void InsertIntoChain(ScopedRef<MemoryAllocatorBase> next);
 
         MemoryAllocatorStats mStats = {};
 
         mutable std::mutex mMutex;
 
       private:
-        MemoryAllocatorBase* mNext = nullptr;
+        ScopedRef<MemoryAllocatorBase> mNext;
         MemoryAllocatorBase* mParent = nullptr;
     };
 

src/gpgmm/common/PooledMemoryAllocator.cpp

@@ -21,10 +21,9 @@
 
 namespace gpgmm {
 
-    PooledMemoryAllocator::PooledMemoryAllocator(
-        uint64_t memorySize,
-        uint64_t memoryAlignment,
-        std::unique_ptr<MemoryAllocatorBase> memoryAllocator)
+    PooledMemoryAllocator::PooledMemoryAllocator(uint64_t memorySize,
+                                                 uint64_t memoryAlignment,
+                                                 ScopedRef<MemoryAllocatorBase> memoryAllocator)
         : MemoryAllocatorBase(std::move(memoryAllocator)),
           mPool(new LIFOMemoryPool(memorySize)),
           mMemoryAlignment(memoryAlignment) {

src/gpgmm/common/PooledMemoryAllocator.h

@@ -26,7 +26,7 @@ namespace gpgmm {
       public:
         PooledMemoryAllocator(uint64_t memorySize,
                               uint64_t memoryAlignment,
-                              std::unique_ptr<MemoryAllocatorBase> memoryAllocator);
+                              ScopedRef<MemoryAllocatorBase> memoryAllocator);
         ~PooledMemoryAllocator() override;
 
         // MemoryAllocatorBase interface

src/gpgmm/common/SegmentedMemoryAllocator.cpp

@@ -79,7 +79,7 @@ namespace gpgmm {
     // SegmentedMemoryAllocator
 
     SegmentedMemoryAllocator::SegmentedMemoryAllocator(
-        std::unique_ptr<MemoryAllocatorBase> memoryAllocator,
+        ScopedRef<MemoryAllocatorBase> memoryAllocator,
         uint64_t memoryAlignment)
         : MemoryAllocatorBase(std::move(memoryAllocator)), mMemoryAlignment(memoryAlignment) {
     }

src/gpgmm/common/SegmentedMemoryAllocator.h

@@ -34,7 +34,7 @@ namespace gpgmm {
     // variable-size memory blocks.
     class SegmentedMemoryAllocator : public MemoryAllocatorBase {
       public:
-        SegmentedMemoryAllocator(std::unique_ptr<MemoryAllocatorBase> memoryAllocator,
+        SegmentedMemoryAllocator(ScopedRef<MemoryAllocatorBase> memoryAllocator,
                                  uint64_t memoryAlignment);
         ~SegmentedMemoryAllocator() override;
 

src/gpgmm/common/SlabMemoryAllocator.cpp

@@ -492,7 +492,7 @@ namespace gpgmm {
                                            float slabFragmentationLimit,
                                            bool allowPrefetchSlab,
                                            float slabGrowthFactor,
-                                           std::unique_ptr<MemoryAllocatorBase> memoryAllocator)
+                                           ScopedRef<MemoryAllocatorBase> memoryAllocator)
         : MemoryAllocatorBase(std::move(memoryAllocator)),
           mMaxSlabSize(maxSlabSize),
           mMinSlabSize(minSlabSize),

src/gpgmm/common/SlabMemoryAllocator.h

@@ -117,7 +117,7 @@ namespace gpgmm {
                            float slabFragmentationLimit,
                            bool allowSlabPrefetch,
                            float slabGrowthFactor,
-                           std::unique_ptr<MemoryAllocatorBase> memoryAllocator);
+                           ScopedRef<MemoryAllocatorBase> memoryAllocator);
 
         ~SlabCacheAllocator() override;