Fix allocation calls where requested size is covered by the collateral (facebookincubator#9145)

Summary:

Currently calls to allocation APIs in memory allocator, namely
allocateNonContiguous and allocateContiguous where the requested size
is covered by the provided collateral allocation is incorrectly
handeled. The issue arises when the variable tracking the number of
additional pages to be allocated becomes negative. This negative value
is then passed to the MemoryPool's reserve() call, where it is
implicitly converted into an unsigned integer. Instead of failing the
reservation due to an unusually large reservation request (as a result
of the negative int64 to uint64 conversion), the reservation remains
unchanged. This happens because when MemoryPoolImpl::reserve() uses
reservationSizeLocked(int64_t size) to check if more reservation is
needed, it returns 0. This is due to another implicit conversion from
uint64 back to int64, which is then used to calculate that no
increment to the reservation is necessary. Finally, the used and
cumulative metrics also get updated correctly because they are
signed integers and the compiler implicitly converts unit64 back to
int64. Therefore, it never manifests into an error state but the
code is fragile and can cause issues with future changes if the
implementation changes.

This change fixes this silent bug by ensuring this case is properly
handled and extra memory is released with the proper APIs.

Differential Revision: D55047654

Author: Bikramjeet Vig

velox/common/memory/MemoryAllocator.cpp

@@ -177,18 +177,22 @@ bool MemoryAllocator::allocateNonContiguous(
   }
 
   const SizeMix mix = allocationSize(numPages, minSizeClass);
-  const int64_t numNeededPages = mix.totalPages - numPagesToFree;
-  // TODO: handle negative 'numNeededPages' in a follow-up.
-
   if (reservationCB != nullptr) {
-    try {
-      reservationCB(AllocationTraits::pageBytes(numNeededPages), true);
-    } catch (const std::exception&) {
-      VELOX_MEM_LOG_EVERY_MS(WARNING, 1000)
-          << "Exceeded memory reservation limit when reserve " << numNeededPages
-          << " new pages when allocate " << mix.totalPages << " pages";
-      cleanupAllocAndReleaseReservation(bytesToFree);
-      std::rethrow_exception(std::current_exception());
+    if (mix.totalPages >= numPagesToFree) {
+      const uint64_t numNeededPages = mix.totalPages - numPagesToFree;
+      try {
+        reservationCB(AllocationTraits::pageBytes(numNeededPages), true);
+      } catch (const std::exception&) {
+        VELOX_MEM_LOG_EVERY_MS(WARNING, 1000)
+            << "Exceeded memory reservation limit when reserve "
+            << numNeededPages << " new pages when allocate " << mix.totalPages
+            << " pages";
+        cleanupAllocAndReleaseReservation(bytesToFree);
+        std::rethrow_exception(std::current_exception());
+      }
+    } else {
+      const uint64_t numExtraPages = numPagesToFree - mix.totalPages;
+      reservationCB(AllocationTraits::pageBytes(numExtraPages), false);
     }
   }
 
@@ -222,7 +226,6 @@ bool MemoryAllocator::allocateContiguous(
       allocation.numPages() + (collateral ? collateral->numPages() : 0);
   const uint64_t totalCollateralBytes =
       AllocationTraits::pageBytes(numCollateralPages);
-  const int64_t newPages = numPages - numCollateralPages;
   auto cleanupCollateralAndReleaseReservation = [&](uint64_t reservationBytes) {
     if ((collateral != nullptr) && !collateral->empty()) {
       freeNonContiguous(*collateral);
@@ -239,17 +242,23 @@ bool MemoryAllocator::allocateContiguous(
     cleanupCollateralAndReleaseReservation(totalCollateralBytes);
     return true;
   }
-  // TODO: handle negative 'newPages'.
+
   if (reservationCB != nullptr) {
-    try {
-      reservationCB(AllocationTraits::pageBytes(newPages), true);
-    } catch (const std::exception& e) {
-      VELOX_MEM_LOG_EVERY_MS(WARNING, 1000)
-          << "Exceeded memory reservation limit when reserve " << newPages
-          << " new pages when allocate " << numPages
-          << " pages, error: " << e.what();
-      cleanupCollateralAndReleaseReservation(totalCollateralBytes);
-      std::rethrow_exception(std::current_exception());
+    if (numPages >= numCollateralPages) {
+      const int64_t numNeededPages = numPages - numCollateralPages;
+      try {
+        reservationCB(AllocationTraits::pageBytes(numNeededPages), true);
+      } catch (const std::exception& e) {
+        VELOX_MEM_LOG_EVERY_MS(WARNING, 1000)
+            << "Exceeded memory reservation limit when reserve "
+            << numNeededPages << " new pages when allocate " << numPages
+            << " pages, error: " << e.what();
+        cleanupCollateralAndReleaseReservation(totalCollateralBytes);
+        std::rethrow_exception(std::current_exception());
+      }
+    } else {
+      const uint64_t numExtraPages = numCollateralPages - numPages;
+      reservationCB(AllocationTraits::pageBytes(numExtraPages), false);
     }
   }
 

velox/common/memory/MemoryAllocator.h

@@ -221,7 +221,7 @@ class MemoryAllocator : public std::enable_shared_from_this<MemoryAllocator> {
   /// the same as 'this'.
   virtual void registerCache(const std::shared_ptr<Cache>& cache) = 0;
 
-  using ReservationCallback = std::function<void(int64_t, bool)>;
+  using ReservationCallback = std::function<void(uint64_t, bool)>;
 
   /// Returns the capacity of the allocator in bytes.
   virtual size_t capacity() const = 0;

velox/common/memory/MemoryPool.cpp

@@ -545,7 +545,7 @@ void MemoryPoolImpl::allocateNonContiguous(
   if (!allocator_->allocateNonContiguous(
           numPages,
           out,
-          [this](int64_t allocBytes, bool preAllocate) {
+          [this](uint64_t allocBytes, bool preAllocate) {
             if (preAllocate) {
               reserve(allocBytes);
             } else {
@@ -597,7 +597,7 @@ void MemoryPoolImpl::allocateContiguous(
           numPages,
           nullptr,
           out,
-          [this](int64_t allocBytes, bool preAlloc) {
+          [this](uint64_t allocBytes, bool preAlloc) {
             if (preAlloc) {
               reserve(allocBytes);
             } else {
@@ -632,7 +632,7 @@ void MemoryPoolImpl::growContiguous(
     MachinePageCount increment,
     ContiguousAllocation& allocation) {
   if (!allocator_->growContiguous(
-          increment, allocation, [this](int64_t allocBytes, bool preAlloc) {
+          increment, allocation, [this](uint64_t allocBytes, bool preAlloc) {
             if (preAlloc) {
               reserve(allocBytes);
             } else {

velox/common/memory/tests/MemoryPoolTest.cpp

@@ -3671,6 +3671,34 @@ TEST_P(MemoryPoolTest, overuseUnderArbitration) {
   ASSERT_EQ(child->reservedBytes(), 0);
 }
 
+TEST_P(MemoryPoolTest, allocationWithCoveredCollateral) {
+  // Verify that the memory pool's reservation is correctly updated when an
+  // allocation call is attempted with collateral that covers the allocation
+  // (that is, the collateral is larger than the requested allocation).
+  auto manager = getMemoryManager();
+  auto root = manager->addRootPool("root", kMaxMemory, nullptr);
+  ASSERT_TRUE(root->trackUsage());
+  auto pool =
+      root->addLeafChild("allocationWithCoveredCollateral", isLeafThreadSafe_);
+  ASSERT_TRUE(pool->trackUsage());
+  // Check non-contiguous allocation.
+  ASSERT_EQ(pool->reservedBytes(), 0);
+  Allocation allocation;
+  pool->allocateNonContiguous(100, allocation);
+  auto prevReservedBytes = pool->currentBytes();
+  pool->allocateNonContiguous(50, allocation);
+  ASSERT_LT(pool->currentBytes(), prevReservedBytes);
+  pool->freeNonContiguous(allocation);
+
+  // Check contiguous allocation.
+  ContiguousAllocation contiguousAllocation;
+  pool->allocateContiguous(100, contiguousAllocation);
+  prevReservedBytes = pool->currentBytes();
+  pool->allocateContiguous(50, contiguousAllocation);
+  ASSERT_LT(pool->currentBytes(), prevReservedBytes);
+  pool->freeContiguous(contiguousAllocation);
+}
+
 VELOX_INSTANTIATE_TEST_SUITE_P(
     MemoryPoolTestSuite,
     MemoryPoolTest,