Use ResultOrError for TryAllocateBlock.

src/gpgmm/common/BlockAllocator.h

@@ -15,6 +15,7 @@
 #ifndef SRC_GPGMM_COMMON_BLOCKALLOCATOR_H_
 #define SRC_GPGMM_COMMON_BLOCKALLOCATOR_H_
 
+#include "gpgmm/common/Error.h"
 #include "gpgmm/common/MemoryBlock.h"
 #include "gpgmm/common/Object.h"
 #include "gpgmm/utils/NonCopyable.h"
@@ -26,7 +27,8 @@ namespace gpgmm {
       public:
         ~BlockAllocator() override = default;
 
-        virtual MemoryBlock* TryAllocateBlock(uint64_t requestSize, uint64_t alignment) = 0;
+        virtual ResultOrError<MemoryBlock*> TryAllocateBlock(uint64_t requestSize,
+                                                             uint64_t alignment) = 0;
         virtual void DeallocateBlock(MemoryBlock* block) = 0;
     };
 

src/gpgmm/common/BuddyBlockAllocator.cpp

@@ -144,10 +144,11 @@ namespace gpgmm {
         }
     }
 
-    MemoryBlock* BuddyBlockAllocator::TryAllocateBlock(uint64_t requestSize, uint64_t alignment) {
+    ResultOrError<MemoryBlock*> BuddyBlockAllocator::TryAllocateBlock(uint64_t requestSize,
+                                                                      uint64_t alignment) {
         // Request cannot exceed max block size.
         if (requestSize > mMaxBlockSize) {
-            return nullptr;
+            return {};
         }
 
         // Compute the level
@@ -159,7 +160,7 @@ namespace gpgmm {
 
         // Error when no free blocks exist (allocator is full)
         if (currBlockLevel == kInvalidOffset) {
-            return nullptr;
+            return {};
         }
 
         // Split free blocks level-by-level.

src/gpgmm/common/BuddyBlockAllocator.h

@@ -42,7 +42,8 @@ namespace gpgmm {
         ~BuddyBlockAllocator() override;
 
         // BlockAllocator interface
-        MemoryBlock* TryAllocateBlock(uint64_t requestSize, uint64_t alignment) override;
+        ResultOrError<MemoryBlock*> TryAllocateBlock(uint64_t requestSize,
+                                                     uint64_t alignment) override;
         void DeallocateBlock(MemoryBlock* block) override;
 
         // For testing purposes only.

src/gpgmm/common/MemoryAllocator.h

@@ -225,13 +225,15 @@ namespace gpgmm {
             uint64_t alignment,
             bool neverAllocate,
             GetOrCreateMemoryFn&& GetOrCreateMemory) {
-            MemoryBlock* block = allocator->TryAllocateBlock(requestSize, alignment);
-            if (block == nullptr) {
-                return {};
+            ResultOrError<MemoryBlock*> allocatedBlockResult =
+                allocator->TryAllocateBlock(requestSize, alignment);
+            if (!allocatedBlockResult.IsSuccess()) {
+                return allocatedBlockResult.GetErrorCode();
             }
-
-            ResultOrError<std::unique_ptr<MemoryAllocationBase>> result = GetOrCreateMemory(block);
-            if (!result.IsSuccess()) {
+            MemoryBlock* block = allocatedBlockResult.AcquireResult();
+            ResultOrError<std::unique_ptr<MemoryAllocationBase>> allocatedMemoryResult =
+                GetOrCreateMemory(block);
+            if (!allocatedMemoryResult.IsSuccess()) {
                 // NeverAllocate always fails, so suppress it.
                 if (!neverAllocate) {
                     ErrorLog(ErrorCode::kAllocationFailed, this)
@@ -240,10 +242,11 @@ namespace gpgmm {
                         << std::to_string(block->Offset + block->Size) << ").";
                 }
                 allocator->DeallocateBlock(block);
-                return result.GetErrorCode();
+                return allocatedMemoryResult.GetErrorCode();
             }
 
-            std::unique_ptr<MemoryAllocationBase> memoryAllocation = result.AcquireResult();
+            std::unique_ptr<MemoryAllocationBase> memoryAllocation =
+                allocatedMemoryResult.AcquireResult();
             ASSERT(memoryAllocation->GetMemory() != nullptr);
 
             memoryAllocation->GetMemory()->Ref();

src/gpgmm/common/SlabBlockAllocator.cpp

@@ -65,15 +65,16 @@ namespace gpgmm {
         mFreeList.pHead = nullptr;
     }
 
-    MemoryBlock* SlabBlockAllocator::TryAllocateBlock(uint64_t requestSize, uint64_t alignment) {
+    ResultOrError<MemoryBlock*> SlabBlockAllocator::TryAllocateBlock(uint64_t requestSize,
+                                                                     uint64_t alignment) {
         // Requested cannot exceed block size.
         if (requestSize > mBlockSize) {
-            return nullptr;
+            return {};
         }
 
         // Offset must be equal to a multiple of |mBlockSize|.
         if (!IsAligned(mBlockSize, alignment)) {
-            return nullptr;
+            return {};
         }
 
         // Pop off HEAD in the free-list.

src/gpgmm/common/SlabBlockAllocator.h

@@ -47,7 +47,8 @@ namespace gpgmm {
         void ReleaseBlocks();
 
         // BlockAllocator interface
-        MemoryBlock* TryAllocateBlock(uint64_t requestSize, uint64_t alignment = 1) override;
+        ResultOrError<MemoryBlock*> TryAllocateBlock(uint64_t requestSize,
+                                                     uint64_t alignment) override;
         void DeallocateBlock(MemoryBlock* block) override;
 
         uint64_t GetBlockCount() const;

src/tests/unittests/BuddyBlockAllocatorTests.cpp

@@ -25,7 +25,7 @@ class DummyBuddyBlockAllocator {
     }
 
     MemoryBlock* TryAllocateBlock(uint64_t requestSize, uint64_t alignment = 1) {
-        return mAllocator.TryAllocateBlock(requestSize, alignment);
+        return mAllocator.TryAllocateBlock(requestSize, alignment).AcquireResult();
     }
 
     void DeallocateBlock(MemoryBlock* block) {

src/tests/unittests/SlabBlockAllocatorTests.cpp

@@ -21,11 +21,28 @@
 
 using namespace gpgmm;
 
+class DummySlabBlockAllocator {
+  public:
+    DummySlabBlockAllocator(uint64_t blockCount, uint64_t blockSize)
+        : mAllocator(blockCount, blockSize) {
+    }
+
+    MemoryBlock* TryAllocateBlock(uint64_t requestSize, uint64_t alignment = 1) {
+        return mAllocator.TryAllocateBlock(requestSize, alignment).AcquireResult();
+    }
+
+    void DeallocateBlock(MemoryBlock* block) {
+        mAllocator.DeallocateBlock(block);
+    }
+
+    SlabBlockAllocator mAllocator;
+};
+
 // Verify a single allocation in a slab.
 TEST(SlabBlockAllocatorTests, SingleBlock) {
     constexpr uint64_t blockSize = 32;
     constexpr uint64_t slabSize = 128;
-    SlabBlockAllocator allocator(slabSize / blockSize, blockSize);
+    DummySlabBlockAllocator allocator(slabSize / blockSize, blockSize);
 
     // Check that we cannot allocate a oversized block.
     EXPECT_EQ(allocator.TryAllocateBlock(slabSize * 2), nullptr);
@@ -45,7 +62,7 @@ TEST(SlabBlockAllocatorTests, SingleBlock) {
 TEST(SlabBlockAllocatorTests, SingleBlockAligned) {
     constexpr uint64_t blockSize = 32;
     constexpr uint64_t slabSize = 128;
-    SlabBlockAllocator allocator(slabSize / blockSize, blockSize);
+    DummySlabBlockAllocator allocator(slabSize / blockSize, blockSize);
 
     // Check that we cannot allocate a misaligned block.
     EXPECT_EQ(allocator.TryAllocateBlock(blockSize, 64u), nullptr);
@@ -62,7 +79,7 @@ TEST(SlabBlockAllocatorTests, MultipleBlocks) {
     // Fill entire slab in the allocator.
     constexpr uint64_t blockSize = 32;
     constexpr uint64_t slabSize = 128;
-    SlabBlockAllocator allocator(slabSize / blockSize, blockSize);
+    DummySlabBlockAllocator allocator(slabSize / blockSize, blockSize);
 
     std::unordered_set<MemoryBlock*> blocks = {};
     for (uint64_t blockIdx = 0; blockIdx < slabSize / blockSize; blockIdx++) {
@@ -88,7 +105,7 @@ TEST(SlabBlockAllocatorTests, MultipleBlocksVarious) {
     // Fill entire slab in the allocator.
     constexpr uint64_t blockSize = 32;
     constexpr uint64_t slabSize = 128;
-    SlabBlockAllocator allocator(slabSize / blockSize, blockSize);
+    DummySlabBlockAllocator allocator(slabSize / blockSize, blockSize);
 
     std::unordered_set<MemoryBlock*> blocks = {};