intel · bbernhar · Sep 11, 2023 · Sep 11, 2023
diff --git a/src/gpgmm/common/BuddyMemoryAllocator.cpp b/src/gpgmm/common/BuddyMemoryAllocator.cpp
@@ -64,8 +64,9 @@ namespace gpgmm {
                 &mBuddyBlockAllocator, allocationSize, request.Alignment, request.NeverAllocate,
                 [&](const auto& block) -> ResultOrError<std::unique_ptr<MemoryAllocationBase>> {
                     const uint64_t memoryIndex = GetMemoryIndex(block->Offset);
-                    std::unique_ptr<MemoryAllocationBase> memoryAllocation =
-                        mUsedPool.AcquireFromPool(memoryIndex);
+
+                    std::unique_ptr<MemoryAllocationBase> memoryAllocation;
+                    GPGMM_TRY_ASSIGN(mUsedPool.AcquireFromPool(memoryIndex), memoryAllocation);
 
                     // No existing, allocate new memory for the block.
                     if (memoryAllocation == nullptr) {
@@ -113,16 +114,18 @@ namespace gpgmm {
 
         mBuddyBlockAllocator.DeallocateBlock(subAllocation->GetBlock());
 
-        std::unique_ptr<MemoryAllocationBase> memoryAllocation =
-            mUsedPool.AcquireFromPool(memoryIndex);
+        auto result = mUsedPool.AcquireFromPool(memoryIndex);
+        ASSERT(result.IsSuccess());
+
+        std::unique_ptr<MemoryAllocationBase> allocation = result.AcquireResult();
 
-        MemoryBase* memory = memoryAllocation->GetMemory();
+        MemoryBase* memory = allocation->GetMemory();
         ASSERT(memory != nullptr);
 
         if (memory->Unref()) {
-            GetNextInChain()->DeallocateMemory(std::move(memoryAllocation));
+            GetNextInChain()->DeallocateMemory(std::move(allocation));
         } else {
-            mUsedPool.ReturnToPool(std::move(memoryAllocation), memoryIndex);
+            mUsedPool.ReturnToPool(std::move(allocation), memoryIndex);
         }
     }
 

diff --git a/src/gpgmm/common/IndexedMemoryPool.cpp b/src/gpgmm/common/IndexedMemoryPool.cpp
@@ -22,17 +22,20 @@ namespace gpgmm {
     IndexedMemoryPool::IndexedMemoryPool(uint64_t memorySize) : MemoryPoolBase(memorySize) {
     }
 
-    std::unique_ptr<MemoryAllocationBase> IndexedMemoryPool::AcquireFromPool(uint64_t indexInPool) {
+    ResultOrError<std::unique_ptr<MemoryAllocationBase>> IndexedMemoryPool::AcquireFromPool(
+        uint64_t indexInPool) {
         if (indexInPool >= mPool.size()) {
             mPool.resize(indexInPool + 1);
         }
         return std::unique_ptr<MemoryAllocationBase>(mPool[indexInPool].release());
     }
 
-    void IndexedMemoryPool::ReturnToPool(std::unique_ptr<MemoryAllocationBase> allocation,
-                                         uint64_t indexInPool) {
-        ASSERT(indexInPool < mPool.size());
+    MaybeError IndexedMemoryPool::ReturnToPool(std::unique_ptr<MemoryAllocationBase> allocation,
+                                               uint64_t indexInPool) {
+        GPGMM_RETURN_ERROR_IF(this, indexInPool >= mPool.size(), "Index exceeded pool size",
+                              ErrorCode::kBadOperation);
         mPool[indexInPool] = std::move(allocation);
+        return {};
     }
 
     uint64_t IndexedMemoryPool::ReleasePool(uint64_t bytesToRelease) {

diff --git a/src/gpgmm/common/IndexedMemoryPool.h b/src/gpgmm/common/IndexedMemoryPool.h
@@ -27,9 +27,10 @@ namespace gpgmm {
         ~IndexedMemoryPool() override = default;
 
         // MemoryPoolBase interface
-        std::unique_ptr<MemoryAllocationBase> AcquireFromPool(uint64_t indexInPool) override;
-        void ReturnToPool(std::unique_ptr<MemoryAllocationBase> allocation,
-                          uint64_t indexInPool) override;
+        ResultOrError<std::unique_ptr<MemoryAllocationBase>> AcquireFromPool(
+            uint64_t indexInPool) override;
+        MaybeError ReturnToPool(std::unique_ptr<MemoryAllocationBase> allocation,
+                                uint64_t indexInPool) override;
         uint64_t ReleasePool(uint64_t bytesToRelease) override;
 
         uint64_t GetPoolSize() const override;

diff --git a/src/gpgmm/common/LIFOMemoryPool.cpp b/src/gpgmm/common/LIFOMemoryPool.cpp
@@ -24,21 +24,26 @@ namespace gpgmm {
     LIFOMemoryPool::LIFOMemoryPool(uint64_t memorySize) : MemoryPoolBase(memorySize) {
     }
 
-    std::unique_ptr<MemoryAllocationBase> LIFOMemoryPool::AcquireFromPool(uint64_t indexInPool) {
-        ASSERT(indexInPool == kInvalidIndex);
+    ResultOrError<std::unique_ptr<MemoryAllocationBase>> LIFOMemoryPool::AcquireFromPool(
+        uint64_t indexInPool) {
+        GPGMM_RETURN_ERROR_IF(this, indexInPool != kInvalidIndex,
+                              "Index was specified but not allowed", ErrorCode::kBadOperation);
 
         std::unique_ptr<MemoryAllocationBase> allocation;
         if (!mPool.empty()) {
             allocation = std::move(mPool.front());
             mPool.pop_front();
         }
+
         return allocation;
     }
 
-    void LIFOMemoryPool::ReturnToPool(std::unique_ptr<MemoryAllocationBase> allocation,
-                                      uint64_t indexInPool) {
-        ASSERT(indexInPool == kInvalidIndex);
+    MaybeError LIFOMemoryPool::ReturnToPool(std::unique_ptr<MemoryAllocationBase> allocation,
+                                            uint64_t indexInPool) {
+        GPGMM_RETURN_ERROR_IF(this, indexInPool != kInvalidIndex,
+                              "Index was specified but not allowed", ErrorCode::kBadOperation);
         mPool.push_front(std::move(allocation));
+        return {};
     }
 
     uint64_t LIFOMemoryPool::ReleasePool(uint64_t bytesToRelease) {

diff --git a/src/gpgmm/common/LIFOMemoryPool.h b/src/gpgmm/common/LIFOMemoryPool.h
@@ -28,10 +28,10 @@ namespace gpgmm {
         ~LIFOMemoryPool() override = default;
 
         // MemoryPoolBase interface
-        std::unique_ptr<MemoryAllocationBase> AcquireFromPool(
+        ResultOrError<std::unique_ptr<MemoryAllocationBase>> AcquireFromPool(
             uint64_t indexInPool = kInvalidIndex) override;
-        void ReturnToPool(std::unique_ptr<MemoryAllocationBase> allocation,
-                          uint64_t indexInPool = kInvalidIndex) override;
+        MaybeError ReturnToPool(std::unique_ptr<MemoryAllocationBase> allocation,
+                                uint64_t indexInPool = kInvalidIndex) override;
         uint64_t ReleasePool(uint64_t bytesToFree = kInvalidSize) override;
 
         uint64_t GetPoolSize() const override;

diff --git a/src/gpgmm/common/MemoryPool.cpp b/src/gpgmm/common/MemoryPool.cpp
@@ -26,6 +26,11 @@ namespace gpgmm {
         GPGMM_TRACE_EVENT_OBJECT_DESTROY(this);
     }
 
+    std::unique_ptr<MemoryAllocationBase> MemoryPoolBase::AcquireFromPoolForTesting(
+        uint64_t indexInPool) {
+        return AcquireFromPool(indexInPool).AcquireResult();
+    }
+
     uint64_t MemoryPoolBase::GetMemorySize() const {
         return mMemorySize;
     }

diff --git a/src/gpgmm/common/MemoryPool.h b/src/gpgmm/common/MemoryPool.h
@@ -15,6 +15,7 @@
 #ifndef SRC_GPGMM_COMMON_MEMORYPOOL_H_
 #define SRC_GPGMM_COMMON_MEMORYPOOL_H_
 
+#include "gpgmm/common/Error.h"
 #include "gpgmm/common/MemoryAllocation.h"
 #include "gpgmm/utils/Limits.h"
 
@@ -31,12 +32,15 @@ namespace gpgmm {
 
         // Retrieves a memory allocation from the pool using an optional index.
         // Use kInvalidIndex to specify |this| pool is not indexed.
-        virtual std::unique_ptr<MemoryAllocationBase> AcquireFromPool(uint64_t indexInPool) = 0;
+        virtual ResultOrError<std::unique_ptr<MemoryAllocationBase>> AcquireFromPool(
+            uint64_t indexInPool) = 0;
+
+        std::unique_ptr<MemoryAllocationBase> AcquireFromPoolForTesting(uint64_t indexInPool);
 
         // Returns a memory allocation back to the pool using an optional index.
         // Use kInvalidIndex to specify |this| pool is not indexed.
-        virtual void ReturnToPool(std::unique_ptr<MemoryAllocationBase> allocation,
-                                  uint64_t indexInPool) = 0;
+        virtual MaybeError ReturnToPool(std::unique_ptr<MemoryAllocationBase> allocation,
+                                        uint64_t indexInPool) = 0;
 
         // Deallocate or shrink the pool.
         virtual uint64_t ReleasePool(uint64_t bytesToRelease) = 0;

diff --git a/src/gpgmm/common/PooledMemoryAllocator.cpp b/src/gpgmm/common/PooledMemoryAllocator.cpp
@@ -43,7 +43,8 @@ namespace gpgmm {
 
         GPGMM_RETURN_IF_ERROR(ValidateRequest(request));
 
-        std::unique_ptr<MemoryAllocationBase> allocation = mPool->AcquireFromPool(kInvalidIndex);
+        std::unique_ptr<MemoryAllocationBase> allocation;
+        GPGMM_TRY_ASSIGN(mPool->AcquireFromPool(kInvalidIndex), allocation);
         if (allocation == nullptr) {
             MemoryAllocationRequest memoryRequest = request;
             memoryRequest.Alignment = mMemoryAlignment;

diff --git a/src/gpgmm/common/SegmentedMemoryAllocator.cpp b/src/gpgmm/common/SegmentedMemoryAllocator.cpp
@@ -132,7 +132,8 @@ namespace gpgmm {
         MemorySegment* segment = GetOrCreateFreeSegment(memorySize);
         ASSERT(segment != nullptr);
 
-        std::unique_ptr<MemoryAllocationBase> allocation = segment->AcquireFromPool();
+        std::unique_ptr<MemoryAllocationBase> allocation;
+        GPGMM_TRY_ASSIGN(segment->AcquireFromPool(), allocation);
         if (allocation == nullptr) {
             MemoryAllocationRequest memoryRequest = request;
             memoryRequest.Alignment = mMemoryAlignment;

diff --git a/src/tests/unittests/MemoryPoolTests.cpp b/src/tests/unittests/MemoryPoolTests.cpp
@@ -46,7 +46,7 @@ TEST_F(LIFOMemoryPoolTests, SingleAllocation) {
     EXPECT_EQ(pool.GetPoolSize() * pool.GetMemorySize(), kDefaultMemorySize);
     EXPECT_EQ(pool.GetPoolSize(), 1u);
 
-    pool.ReturnToPool(pool.AcquireFromPool());
+    pool.ReturnToPool(pool.AcquireFromPoolForTesting(kInvalidIndex));
     EXPECT_EQ(pool.GetPoolSize() * pool.GetMemorySize(), kDefaultMemorySize);
     EXPECT_EQ(pool.GetPoolSize(), 1u);