Fix SlabMemoryAllocator crash when near OOM.

src/gpgmm/common/SlabMemoryAllocator.cpp

@@ -103,8 +103,12 @@ namespace gpgmm {
         // Otherwise, creating a larger slab will page-out smaller slabs.
         if (availableForAllocation < slabSize) {
             const uint64_t slabSizeUnderBudget = FindNextFreeSlabOfSize(requestSize);
-            DebugEvent(GetTypename()) << "Limiting slab size due to available memory: (" << slabSize
-                                      << " vs " << availableForAllocation << " bytes).";
+            if (slabSizeUnderBudget == kInvalidSize) {
+                DebugEvent(GetTypename()) << "Slab size exceeds available memory: " << slabSize
+                                          << " vs " << availableForAllocation << " bytes.";
+                return kInvalidSize;
+            }
+
             slabSize = slabSizeUnderBudget;
         }
 
@@ -129,8 +133,7 @@ namespace gpgmm {
             }
         }
 
-        // If there are no more free slabs, use the smallest size possible.
-        return mMinSlabSize;
+        return kInvalidSize;
     }
 
     SlabMemoryAllocator::SlabCache* SlabMemoryAllocator::GetOrCreateCache(uint64_t slabSize) {
@@ -178,11 +181,15 @@ namespace gpgmm {
         if (pCache->FreeList.empty() || pFreeSlab->IsFull()) {
             // Get the next free slab.
             if (mLastUsedSlabSize > 0) {
-                uint64_t newSlabSize =
-                    std::min(ComputeSlabSize(request.SizeInBytes,
-                                             static_cast<uint64_t>(slabSize * mSlabGrowthFactor),
-                                             request.AvailableForAllocation),
-                             mMaxSlabSize);
+                uint64_t newSlabSize = ComputeSlabSize(
+                    request.SizeInBytes, static_cast<uint64_t>(slabSize * mSlabGrowthFactor),
+                    request.AvailableForAllocation);
+                GPGMM_INVALID_IF(newSlabSize == kInvalidSize);
+
+                // If the new slab size exceeds the limit, then re-use the previous, smaller size.
+                if (newSlabSize > mMaxSlabSize) {
+                    newSlabSize = slabSize;
+                }
 
                 // If the new slab size is not larger then the total size of full slabs, then re-use
                 // the previous, smaller size. Otherwise, the larger slab would likely never be
@@ -295,11 +302,14 @@ namespace gpgmm {
             // If a subsequent TryAllocateMemory() uses a request size different than the current
             // request size, memory required for the next slab could be the wrong size. If so,
             // pre-fetching did not pay off and the pre-fetched memory will be de-allocated instead.
-            uint64_t nextSlabSize = std::min(
-                ComputeSlabSize(request.SizeInBytes,
-                                static_cast<uint64_t>(mLastUsedSlabSize * mSlabGrowthFactor),
-                                request.AvailableForAllocation),
-                mMaxSlabSize);
+            uint64_t nextSlabSize = ComputeSlabSize(
+                request.SizeInBytes, static_cast<uint64_t>(mLastUsedSlabSize * mSlabGrowthFactor),
+                request.AvailableForAllocation);
+
+            // If the next slab size exceeds the limit, then re-use the previous, smaller size.
+            if (nextSlabSize > mMaxSlabSize) {
+                nextSlabSize = mLastUsedSlabSize;
+            }
 
             // If under growth phase (and accounting that the current slab will soon become
             // full), reset the slab size back to the last size. Otherwise, the pre-fetch will

src/tests/end2end/D3D12ResidencyManagerTests.cpp

@@ -136,8 +136,7 @@ TEST_F(D3D12ResidencyManagerTests, OverBudget) {
     EXPECT_LE(resourceHeaps.at(allocations.size() - 1)->GetSize(),
               resourceHeaps.at(allocations.size() - 2)->GetSize());
 
-    // But when going over budget, should evict a previous resource heap with a new one of the same
-    // size.
+    // But when going over budget, should evict some other resource heap.
     {
         RESIDENCY_INFO beforeInfo = residencyManager->GetInfo();
 
@@ -149,7 +148,7 @@ TEST_F(D3D12ResidencyManagerTests, OverBudget) {
 
         EXPECT_TRUE(allocation->IsResident());
 
-        EXPECT_EQ(afterInfo.MemoryCount - beforeInfo.MemoryCount, 0u);
-        EXPECT_EQ(afterInfo.MemoryUsage - beforeInfo.MemoryUsage, 0u);
+        EXPECT_EQ(afterInfo.MemoryCount, beforeInfo.MemoryCount);
+        EXPECT_LE(afterInfo.MemoryUsage, beforeInfo.MemoryUsage);
     }
 }

src/tests/unittests/SlabMemoryAllocatorTests.cpp

@@ -632,6 +632,7 @@ TEST_F(SlabMemoryAllocatorTests, SlabGrowthLimit) {
         // Slab B holds 1 allocation per slab.
         std::unique_ptr<MemoryAllocation> allocationAInSlabB =
             allocator.TryAllocateMemory(CreateBasicRequest(kBlockSize, 1));
+        ASSERT_NE(allocationAInSlabB, nullptr);
         EXPECT_EQ(allocationAInSlabB->GetSize(), kBlockSize);
         EXPECT_EQ(allocationAInSlabB->GetMemory()->GetSize(), kBlockSize);
 
@@ -640,26 +641,30 @@ TEST_F(SlabMemoryAllocatorTests, SlabGrowthLimit) {
         // Slab C grows 2x and holds 2 allocation per slab.
         std::unique_ptr<MemoryAllocation> allocationAInSlabC =
             allocator.TryAllocateMemory(CreateBasicRequest(kBlockSize, 1));
+        ASSERT_NE(allocationAInSlabC, nullptr);
         EXPECT_EQ(allocationAInSlabC->GetSize(), kBlockSize);
         EXPECT_EQ(allocationAInSlabC->GetMemory()->GetSize(), kBlockSize * 2);
 
         EXPECT_NE(allocationAInSlabB->GetMemory(), allocationAInSlabC->GetMemory());
 
         std::unique_ptr<MemoryAllocation> allocationBInSlabC =
             allocator.TryAllocateMemory(CreateBasicRequest(kBlockSize, 1));
+        ASSERT_NE(allocationBInSlabC, nullptr);
         EXPECT_EQ(allocationBInSlabC->GetSize(), kBlockSize);
         EXPECT_EQ(allocationBInSlabC->GetMemory()->GetSize(), kBlockSize * 2);
 
         // Slab C still holds 2 allocation per slab.
         std::unique_ptr<MemoryAllocation> allocationAInSlabD =
             allocator.TryAllocateMemory(CreateBasicRequest(kBlockSize, 1));
+        ASSERT_NE(allocationAInSlabD, nullptr);
         EXPECT_EQ(allocationAInSlabD->GetSize(), kBlockSize);
         EXPECT_EQ(allocationAInSlabD->GetMemory()->GetSize(), kBlockSize * 2);
 
         EXPECT_NE(allocationBInSlabC->GetMemory(), allocationAInSlabD->GetMemory());
 
         std::unique_ptr<MemoryAllocation> allocationBInSlabD =
             allocator.TryAllocateMemory(CreateBasicRequest(kBlockSize, 1));
+        ASSERT_NE(allocationBInSlabD, nullptr);
         EXPECT_EQ(allocationBInSlabD->GetSize(), kBlockSize);
         EXPECT_EQ(allocationBInSlabD->GetMemory()->GetSize(), kBlockSize * 2);
 
@@ -1011,3 +1016,30 @@ TEST_F(SlabCacheAllocatorTests, SlabPrefetch) {
         allocator.DeallocateMemory(std::move(allocation));
     }
 }
+
+// Verify creating more slabs than memory available fails.
+TEST_F(SlabCacheAllocatorTests, OutOfMemory) {
+    SlabCacheAllocator allocator(kDefaultSlabSize, kDefaultSlabSize, kDefaultSlabAlignment,
+                                 kDefaultSlabFragmentationLimit, kDefaultPrefetchSlab,
+                                 kNoSlabGrowthFactor, std::make_unique<DummyMemoryAllocator>());
+
+    constexpr uint64_t kTotalMemoryAvailable = 512;
+
+    MemoryAllocationRequest request = CreateBasicRequest(32, 1);
+    request.AvailableForAllocation = kTotalMemoryAvailable;
+
+    std::vector<std::unique_ptr<MemoryAllocation>> allocations = {};
+    while (true) {
+        std::unique_ptr<MemoryAllocation> allocation = allocator.TryAllocateMemory(request);
+        if (allocation == nullptr) {
+            break;
+        }
+        request.AvailableForAllocation =
+            (kTotalMemoryAvailable - allocator.GetInfo().UsedMemoryUsage);
+        allocations.push_back(std::move(allocation));
+    }
+
+    for (auto& allocation : allocations) {
+        allocator.DeallocateMemory(std::move(allocation));
+    }
+}