diff --git a/src/gpgmm/common/BuddyMemoryAllocator.cpp b/src/gpgmm/common/BuddyMemoryAllocator.cpp
index f203f0c16..047752aea 100644
--- a/src/gpgmm/common/BuddyMemoryAllocator.cpp
+++ b/src/gpgmm/common/BuddyMemoryAllocator.cpp
@@ -62,7 +62,7 @@ namespace gpgmm {
         GPGMM_TRY_ASSIGN(
             TrySubAllocateMemory(
                 &mBuddyBlockAllocator, allocationSize, request.Alignment, request.NeverAllocate,
-                [&](const auto& block) -> ResultOrError<MemoryBase*> {
+                [&](const auto& block) -> ResultOrError<std::unique_ptr<MemoryAllocationBase>> {
                     const uint64_t memoryIndex = GetMemoryIndex(block->Offset);
                     std::unique_ptr<MemoryAllocationBase> memoryAllocation =
                         mUsedPool.AcquireFromPool(memoryIndex);
@@ -72,21 +72,14 @@ namespace gpgmm {
                         MemoryAllocationRequest newRequest = request;
                         newRequest.SizeInBytes = mMemorySize;
                         newRequest.Alignment = mMemoryAlignment;
-
-                        ResultOrError<std::unique_ptr<MemoryAllocationBase>>
-                            memoryAllocationResult =
-                                GetNextInChain()->TryAllocateMemory(newRequest);
-                        if (!memoryAllocationResult.IsSuccess()) {
-                            return memoryAllocationResult.GetErrorCode();
-                        }
-
-                        memoryAllocation = memoryAllocationResult.AcquireResult();
+                        GPGMM_TRY_ASSIGN(GetNextInChain()->TryAllocateMemory(newRequest),
+                                         memoryAllocation);
                     }
 
-                    MemoryBase* memory = memoryAllocation->GetMemory();
+                    const MemoryAllocationBase buddyAllocation = *memoryAllocation;
                     mUsedPool.ReturnToPool(std::move(memoryAllocation), memoryIndex);
 
-                    return memory;
+                    return std::make_unique<MemoryAllocationBase>(buddyAllocation);
                 }),
             subAllocation);
 
diff --git a/src/gpgmm/common/MemoryAllocator.h b/src/gpgmm/common/MemoryAllocator.h
index 246247cd8..e5fe65afd 100644
--- a/src/gpgmm/common/MemoryAllocator.h
+++ b/src/gpgmm/common/MemoryAllocator.h
@@ -227,7 +227,7 @@ namespace gpgmm {
                 return {};
             }
 
-            ResultOrError<MemoryBase*> result = GetOrCreateMemory(block);
+            ResultOrError<std::unique_ptr<MemoryAllocationBase>> result = GetOrCreateMemory(block);
             if (!result.IsSuccess()) {
                 // NeverAllocate always fails, so suppress it.
                 if (!neverAllocate) {
@@ -240,16 +240,17 @@ namespace gpgmm {
                 return result.GetErrorCode();
             }
 
-            MemoryBase* memory = result.AcquireResult();
-            ASSERT(memory != nullptr);
+            std::unique_ptr<MemoryAllocationBase> memoryAllocation = result.AcquireResult();
+            ASSERT(memoryAllocation->GetMemory() != nullptr);
 
-            memory->Ref();
+            memoryAllocation->GetMemory()->Ref();
 
             // Caller is be responsible in fully initializing the memory allocation.
             // This is because TrySubAllocateMemory() does not necessarily know how to map the
             // final sub-allocated block to created memory.
             return std::make_unique<MemoryAllocationBase>(
-                nullptr, memory, kInvalidOffset, AllocationMethod::kUndefined, block, requestSize);
+                nullptr, memoryAllocation->GetMemory(), kInvalidOffset,
+                AllocationMethod::kUndefined, block, requestSize);
         }
 
         void InsertIntoChain(std::unique_ptr<MemoryAllocatorBase> next);
diff --git a/src/gpgmm/common/SlabMemoryAllocator.cpp b/src/gpgmm/common/SlabMemoryAllocator.cpp
index a493e8518..991f61e32 100644
--- a/src/gpgmm/common/SlabMemoryAllocator.cpp
+++ b/src/gpgmm/common/SlabMemoryAllocator.cpp
@@ -266,10 +266,10 @@ namespace gpgmm {
         GPGMM_TRY_ASSIGN(
             TrySubAllocateMemory(
                 &pFreeSlab->Allocator, mBlockSize, request.Alignment, request.NeverAllocate,
-                [&](const auto& block) -> ResultOrError<MemoryBase*> {
+                [&](const auto& block) -> ResultOrError<std::unique_ptr<MemoryAllocationBase>> {
                     // Re-use memory from the free slab.
                     if (pFreeSlab->Allocation.GetMemory() != nullptr) {
-                        return pFreeSlab->Allocation.GetMemory();
+                        return std::make_unique<MemoryAllocationBase>(pFreeSlab->Allocation);
                     }
 
                     // Or use pre-fetched memory if possible. Else, throw it away and create a new
@@ -277,14 +277,9 @@ namespace gpgmm {
                     if (mNextSlabAllocationEvent != nullptr) {
                         // Resolve the pending pre-fetched allocation.
                         mNextSlabAllocationEvent->Wait();
-                        ResultOrError<std::unique_ptr<MemoryAllocationBase>>
-                            memoryAllocationResult = mNextSlabAllocationEvent->AcquireAllocation();
-                        if (!memoryAllocationResult.IsSuccess()) {
-                            return memoryAllocationResult.GetErrorCode();
-                        }
-
-                        std::unique_ptr<MemoryAllocationBase> prefetchedSlabAllocation =
-                            memoryAllocationResult.AcquireResult();
+                        std::unique_ptr<MemoryAllocationBase> prefetchedSlabAllocation;
+                        GPGMM_TRY_ASSIGN(mNextSlabAllocationEvent->AcquireAllocation(),
+                                         prefetchedSlabAllocation);
                         mNextSlabAllocationEvent.reset();
 
                         // Assign pre-fetched memory to the slab.
@@ -292,7 +287,7 @@ namespace gpgmm {
                             prefetchedSlabAllocation->GetSize() == slabSize) {
                             pFreeSlab->Allocation = *prefetchedSlabAllocation;
                             mStats.PrefetchedMemoryMissesEliminated++;
-                            return pFreeSlab->Allocation.GetMemory();
+                            return std::make_unique<MemoryAllocationBase>(pFreeSlab->Allocation);
                         }
 
                         if (prefetchedSlabAllocation != nullptr) {
@@ -304,7 +299,6 @@ namespace gpgmm {
                         }
 
                         mStats.PrefetchedMemoryMisses++;
-
                         mMemoryAllocator->DeallocateMemory(std::move(prefetchedSlabAllocation));
                     }
 
@@ -313,15 +307,13 @@ namespace gpgmm {
                     newSlabRequest.SizeInBytes = slabSize;
                     newSlabRequest.Alignment = mSlabAlignment;
 
-                    ResultOrError<std::unique_ptr<MemoryAllocationBase>> slabAllocationResult =
-                        mMemoryAllocator->TryAllocateMemory(newSlabRequest);
-                    if (!slabAllocationResult.IsSuccess()) {
-                        return slabAllocationResult.GetErrorCode();
-                    }
+                    std::unique_ptr<MemoryAllocationBase> slabAllocation;
+                    GPGMM_TRY_ASSIGN(mMemoryAllocator->TryAllocateMemory(newSlabRequest),
+                                     slabAllocation);
 
-                    pFreeSlab->Allocation = *slabAllocationResult.AcquireResult();
+                    pFreeSlab->Allocation = *slabAllocation;
 
-                    return pFreeSlab->Allocation.GetMemory();
+                    return std::make_unique<MemoryAllocationBase>(pFreeSlab->Allocation);
                 }),
             subAllocation);