diff --git a/src/gpgmm/common/BuddyMemoryAllocator.cpp b/src/gpgmm/common/BuddyMemoryAllocator.cpp index 7150667c..4dd5a1a1 100644 --- 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> { const uint64_t memoryIndex = GetMemoryIndex(block->Offset); - std::unique_ptr memoryAllocation = - mUsedPool.AcquireFromPool(memoryIndex); + + std::unique_ptr 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 memoryAllocation = - mUsedPool.AcquireFromPool(memoryIndex); + auto result = mUsedPool.AcquireFromPool(memoryIndex); + ASSERT(result.IsSuccess()); + + std::unique_ptr 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 index 6ef66218..b8c6b6cc 100644 --- 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 IndexedMemoryPool::AcquireFromPool(uint64_t indexInPool) { + ResultOrError> IndexedMemoryPool::AcquireFromPool( + uint64_t indexInPool) { if (indexInPool >= mPool.size()) { mPool.resize(indexInPool + 1); } return std::unique_ptr(mPool[indexInPool].release()); } - void IndexedMemoryPool::ReturnToPool(std::unique_ptr allocation, - uint64_t indexInPool) { - ASSERT(indexInPool < mPool.size()); + MaybeError IndexedMemoryPool::ReturnToPool(std::unique_ptr 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 index 4badfb58..647b147b 100644 --- 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 AcquireFromPool(uint64_t indexInPool) override; - void ReturnToPool(std::unique_ptr allocation, - uint64_t indexInPool) override; + ResultOrError> AcquireFromPool( + uint64_t indexInPool) override; + MaybeError ReturnToPool(std::unique_ptr 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 index 7032cf56..5e9f592a 100644 --- 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 LIFOMemoryPool::AcquireFromPool(uint64_t indexInPool) { - ASSERT(indexInPool == kInvalidIndex); + ResultOrError> LIFOMemoryPool::AcquireFromPool( + uint64_t indexInPool) { + GPGMM_RETURN_ERROR_IF(this, indexInPool != kInvalidIndex, + "Index was specified but not allowed", ErrorCode::kBadOperation); std::unique_ptr allocation; if (!mPool.empty()) { allocation = std::move(mPool.front()); mPool.pop_front(); } + return allocation; } - void LIFOMemoryPool::ReturnToPool(std::unique_ptr allocation, - uint64_t indexInPool) { - ASSERT(indexInPool == kInvalidIndex); + MaybeError LIFOMemoryPool::ReturnToPool(std::unique_ptr 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 index 68227162..412d26b1 100644 --- 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 AcquireFromPool( + ResultOrError> AcquireFromPool( uint64_t indexInPool = kInvalidIndex) override; - void ReturnToPool(std::unique_ptr allocation, - uint64_t indexInPool = kInvalidIndex) override; + MaybeError ReturnToPool(std::unique_ptr 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 index 74d53f85..9adda363 100644 --- 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 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 index a959d53c..f8359ce0 100644 --- 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 AcquireFromPool(uint64_t indexInPool) = 0; + virtual ResultOrError> AcquireFromPool( + uint64_t indexInPool) = 0; + + std::unique_ptr 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 allocation, - uint64_t indexInPool) = 0; + virtual MaybeError ReturnToPool(std::unique_ptr 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 index 1dac9df5..25329114 100644 --- 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 allocation = mPool->AcquireFromPool(kInvalidIndex); + std::unique_ptr 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 index ff9bb46f..c85c7825 100644 --- 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 allocation = segment->AcquireFromPool(); + std::unique_ptr 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 index 263550f4..43a16ffa 100644 --- 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);