diff --git a/velox/common/memory/Allocation.cpp b/velox/common/memory/Allocation.cpp index 043ff30def7d..f24ef76a359b 100644 --- a/velox/common/memory/Allocation.cpp +++ b/velox/common/memory/Allocation.cpp @@ -30,17 +30,12 @@ Allocation::~Allocation() { } } -void Allocation::append(uint8_t* address, uint32_t numPages) { +void Allocation::append(uint8_t* address, int32_t numPages) { numPages_ += numPages; VELOX_CHECK( runs_.empty() || address != runs_.back().data(), "Appending a duplicate address into a PageRun"); - while (numPages > 0) { - const auto numPagesInRun = std::min(numPages, PageRun::kMaxPagesInRun); - runs_.emplace_back(address, numPagesInRun); - address += AllocationTraits::pageBytes(numPagesInRun); - numPages -= numPagesInRun; - } + runs_.emplace_back(address, numPages); } void Allocation::appendMove(Allocation& other) { for (auto& run : other.runs_) { diff --git a/velox/common/memory/Allocation.h b/velox/common/memory/Allocation.h index 4aae76579ffd..94c930deb35b 100644 --- a/velox/common/memory/Allocation.h +++ b/velox/common/memory/Allocation.h @@ -171,7 +171,7 @@ class Allocation { VELOX_CHECK(numPages_ != 0 || pool_ == nullptr); } - void append(uint8_t* address, uint32_t numPages); + void append(uint8_t* address, int32_t numPages); void clear() { runs_.clear(); @@ -193,7 +193,6 @@ class Allocation { VELOX_FRIEND_TEST(MemoryAllocatorTest, allocationClass2); VELOX_FRIEND_TEST(AllocationTest, append); VELOX_FRIEND_TEST(AllocationTest, appendMove); - VELOX_FRIEND_TEST(AllocationTest, multiplePageRuns); }; /// Represents a run of contiguous pages that do not belong to any size class. diff --git a/velox/common/memory/MallocAllocator.cpp b/velox/common/memory/MallocAllocator.cpp index c4b2014b20ca..83e9af6772d8 100644 --- a/velox/common/memory/MallocAllocator.cpp +++ b/velox/common/memory/MallocAllocator.cpp @@ -227,19 +227,6 @@ int64_t MallocAllocator::freeNonContiguous(Allocation& allocation) { Allocation::PageRun run = allocation.runAt(i); numFreed += run.numPages(); void* ptr = run.data(); - // Check if ptr was split into multiple PageRuns. - while (i + 1 < allocation.numRuns()) { - Allocation::PageRun nextRun = allocation.runAt(i + 1); - void* nextPtr = nextRun.data(); - // NOTE: std::malloc will not return two allocated buffers which are - // contiguous in memory space. - if (static_cast(nextPtr) - static_cast(ptr) != - AllocationTraits::pageBytes(numFreed)) { - break; - } - numFreed += nextRun.numPages(); - ++i; - } { std::lock_guard l(mallocsMutex_); const auto ret = mallocs_.erase(ptr); diff --git a/velox/common/memory/MemoryAllocator.cpp b/velox/common/memory/MemoryAllocator.cpp index df0952e7a86d..6f17d883057e 100644 --- a/velox/common/memory/MemoryAllocator.cpp +++ b/velox/common/memory/MemoryAllocator.cpp @@ -102,6 +102,14 @@ MemoryAllocator::SizeMix MemoryAllocator::allocationSize( ++numUnits; needed -= size; } + if (FOLLY_UNLIKELY(numUnits * size > Allocation::PageRun::kMaxPagesInRun)) { + VELOX_MEM_ALLOC_ERROR(fmt::format( + "Too many pages {} to allocate, the number of units {} at size class of {} exceeds the PageRun limit {}", + numPages, + numUnits, + size, + Allocation::PageRun::kMaxPagesInRun)); + } mix.sizeCounts[mix.numSizes] = numUnits; pagesToAlloc += numUnits * size; mix.sizeIndices[mix.numSizes++] = sizeIndex; diff --git a/velox/common/memory/tests/AllocationTest.cpp b/velox/common/memory/tests/AllocationTest.cpp index bb9e7dc5d8b9..f27618303fa1 100644 --- a/velox/common/memory/tests/AllocationTest.cpp +++ b/velox/common/memory/tests/AllocationTest.cpp @@ -85,26 +85,4 @@ TEST_F(AllocationTest, appendMove) { allocation.clear(); } -TEST_F(AllocationTest, multiplePageRuns) { - Allocation allocation; - const uint64_t startBufAddrValue = 4096; - uint8_t* const firstBufAddr = reinterpret_cast(startBufAddrValue); - allocation.append(firstBufAddr, Allocation::PageRun::kMaxPagesInRun + 100); - ASSERT_EQ(allocation.numPages(), Allocation::PageRun::kMaxPagesInRun + 100); - ASSERT_EQ(allocation.numRuns(), 2); - - uint8_t* const secondBufAddr = reinterpret_cast( - startBufAddrValue + AllocationTraits::pageBytes(allocation.numPages())); - allocation.append(secondBufAddr, Allocation::PageRun::kMaxPagesInRun - 100); - ASSERT_EQ(allocation.numPages(), Allocation::PageRun::kMaxPagesInRun * 2); - ASSERT_EQ(allocation.numRuns(), 3); - - uint8_t* const thirdBufAddr = reinterpret_cast( - firstBufAddr + 2 * AllocationTraits::pageBytes(allocation.numPages())); - allocation.append(thirdBufAddr, Allocation::PageRun::kMaxPagesInRun * 2); - ASSERT_EQ(allocation.numPages(), Allocation::PageRun::kMaxPagesInRun * 4); - ASSERT_EQ(allocation.numRuns(), 5); - allocation.clear(); -} - } // namespace facebook::velox::memory diff --git a/velox/common/memory/tests/MemoryAllocatorTest.cpp b/velox/common/memory/tests/MemoryAllocatorTest.cpp index a2c339a3d1b6..dcf9ede9b061 100644 --- a/velox/common/memory/tests/MemoryAllocatorTest.cpp +++ b/velox/common/memory/tests/MemoryAllocatorTest.cpp @@ -1347,13 +1347,14 @@ TEST_P(MemoryAllocatorTest, StlMemoryAllocator) { } } -TEST_P(MemoryAllocatorTest, nonContiguousAllocationBounds) { +TEST_P(MemoryAllocatorTest, badNonContiguousAllocation) { // Set the num of pages to allocate exceeds one PageRun limit. constexpr MachinePageCount kNumPages = Allocation::PageRun::kMaxPagesInRun + 1; std::unique_ptr allocation(new Allocation()); - ASSERT_TRUE(instance_->allocateNonContiguous(kNumPages, *allocation)); - instance_->freeNonContiguous(*allocation); + ASSERT_THROW( + instance_->allocateNonContiguous(kNumPages, *allocation), + VeloxRuntimeError); ASSERT_TRUE(instance_->allocateNonContiguous(kNumPages - 1, *allocation)); instance_->freeNonContiguous(*allocation); } diff --git a/velox/common/memory/tests/MemoryPoolTest.cpp b/velox/common/memory/tests/MemoryPoolTest.cpp index 16b6fb6eccb8..8c4e46adf13c 100644 --- a/velox/common/memory/tests/MemoryPoolTest.cpp +++ b/velox/common/memory/tests/MemoryPoolTest.cpp @@ -1794,9 +1794,9 @@ TEST_P(MemoryPoolTest, contiguousAllocateGrowExceedMemoryPoolLimit) { ASSERT_EQ(allocation.numPages(), kMaxNumPages / 2); } -TEST_P(MemoryPoolTest, nonContiguousAllocationBounds) { +TEST_P(MemoryPoolTest, badNonContiguousAllocation) { auto manager = getMemoryManager(); - auto pool = manager->addLeafPool("nonContiguousAllocationBounds"); + auto pool = manager->addLeafPool("badNonContiguousAllocation"); Allocation allocation; // Bad zero page allocation size. ASSERT_THROW(pool->allocateNonContiguous(0, allocation), VeloxRuntimeError); @@ -1804,9 +1804,8 @@ TEST_P(MemoryPoolTest, nonContiguousAllocationBounds) { // Set the num of pages to allocate exceeds one PageRun limit. constexpr MachinePageCount kNumPages = Allocation::PageRun::kMaxPagesInRun + 1; - pool->allocateNonContiguous(kNumPages, allocation); - ASSERT_GE(allocation.numPages(), kNumPages); - pool->freeNonContiguous(allocation); + ASSERT_THROW( + pool->allocateNonContiguous(kNumPages, allocation), VeloxRuntimeError); pool->allocateNonContiguous(kNumPages - 1, allocation); ASSERT_GE(allocation.numPages(), kNumPages - 1); pool->freeNonContiguous(allocation);