Back out "Handle allocations that exceed PageRun limits" (#7908)

Summary: Pull Request resolved: #7908 Original commit changeset: 8a568db14e51 Original Phabricator Diff: D51167059 Reviewed By: mbasmanova Differential Revision: D51919841 fbshipit-source-id: fad209763cda543af5c961218d41cb0a9630cfb4
facebookincubator · Dec 7, 2023 · dffc39e · dffc39e
1 parent fb045b0
commit dffc39e
Show file tree

Hide file tree

Showing 7 changed files with 19 additions and 52 deletions.
diff --git 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
@@ -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
@@ -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<char*>(nextPtr) - static_cast<char*>(ptr) !=
-          AllocationTraits::pageBytes(numFreed)) {
-        break;
-      }
-      numFreed += nextRun.numPages();
-      ++i;
-    }
     {
       std::lock_guard<std::mutex> l(mallocsMutex_);
       const auto ret = mallocs_.erase(ptr);

diff --git 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
@@ -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<uint8_t*>(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<uint8_t*>(
-      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<uint8_t*>(
-      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
@@ -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> 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
@@ -1794,19 +1794,18 @@ 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);
 
   // 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);