From f0b671901f310f7c751d3137fff123bdf54a012b Mon Sep 17 00:00:00 2001 From: Bikramjeet Vig Date: Mon, 18 Mar 2024 18:07:48 -0700 Subject: [PATCH] Fix allocation calls where requested size is covered by the collateral 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 --- velox/common/memory/MemoryAllocator.cpp | 58 ++++++++++++-------- velox/common/memory/MemoryAllocator.h | 2 +- velox/common/memory/MemoryPool.cpp | 6 +- velox/common/memory/tests/MemoryPoolTest.cpp | 28 ++++++++++ 4 files changed, 66 insertions(+), 28 deletions(-) diff --git a/velox/common/memory/MemoryAllocator.cpp b/velox/common/memory/MemoryAllocator.cpp index 3001b5b839f01..6c9c1f2b2b488 100644 --- a/velox/common/memory/MemoryAllocator.cpp +++ b/velox/common/memory/MemoryAllocator.cpp @@ -173,23 +173,28 @@ bool MemoryAllocator::allocateNonContiguous( return true; } 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"; - if (bytesToFree > 0) { - freeNonContiguous(out); - reservationCB(bytesToFree, false); + 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"; + if (bytesToFree > 0) { + freeNonContiguous(out); + reservationCB(bytesToFree, false); + } + std::rethrow_exception(std::current_exception()); } - std::rethrow_exception(std::current_exception()); + } else { + const uint64_t numExtraPages = numPagesToFree - mix.totalPages; + reservationCB(AllocationTraits::pageBytes(numExtraPages), false); } } + const auto totalBytesReserved = AllocationTraits::pageBytes(mix.totalPages); bool success = false; if (cache() == nullptr) { @@ -225,7 +230,6 @@ bool MemoryAllocator::allocateContiguous( allocation.numPages() + (collateral ? collateral->numPages() : 0); const auto totalCollateralBytes = AllocationTraits::pageBytes(numCollateralPages); - const int64_t newPages = numPages - numCollateralPages; auto cleanupCollateralAndReleaseReservation = [&](uint64_t reservationBytes) { if ((collateral != nullptr) && !collateral->empty()) { freeNonContiguous(*collateral); @@ -242,17 +246,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); } } diff --git a/velox/common/memory/MemoryAllocator.h b/velox/common/memory/MemoryAllocator.h index c38ed99c4b3c4..a5c915c570918 100644 --- a/velox/common/memory/MemoryAllocator.h +++ b/velox/common/memory/MemoryAllocator.h @@ -221,7 +221,7 @@ class MemoryAllocator : public std::enable_shared_from_this { /// the same as 'this'. virtual void registerCache(const std::shared_ptr& cache) = 0; - using ReservationCallback = std::function; + using ReservationCallback = std::function; /// Returns the capacity of the allocator in bytes. virtual size_t capacity() const = 0; diff --git a/velox/common/memory/MemoryPool.cpp b/velox/common/memory/MemoryPool.cpp index 681ce471d5168..3da310558deb9 100644 --- a/velox/common/memory/MemoryPool.cpp +++ b/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 { diff --git a/velox/common/memory/tests/MemoryPoolTest.cpp b/velox/common/memory/tests/MemoryPoolTest.cpp index 918659fcecd38..abed396910297 100644 --- a/velox/common/memory/tests/MemoryPoolTest.cpp +++ b/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,