Attempt

mohsaka · Mar 11, 2024 · 9e6dae7 · 9e6dae7
1 parent 1fcdfdd
commit 9e6dae7
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/velox/common/caching/AsyncDataCache.cpp b/velox/common/caching/AsyncDataCache.cpp
@@ -109,7 +109,7 @@ void AsyncDataCacheEntry::initialize(FileCacheKey key) {
     tinyData_.clear();
     tinyData_.shrink_to_fit();
     const auto sizePages = memory::AllocationTraits::numPages(size_);
-    if (cache->allocator()->allocateNonContiguous(sizePages, data_)) {
+    if (cache->allocator()->allocateNonContiguous(sizePages, data_, nullptr, 0, true)) {
       cache->incrementCachedPages(data().numPages());
     } else {
       // No memory to cover 'this'.

diff --git a/velox/common/memory/MemoryAllocator.cpp b/velox/common/memory/MemoryAllocator.cpp
@@ -160,13 +160,14 @@ bool MemoryAllocator::allocateNonContiguous(
     MachinePageCount numPages,
     Allocation& out,
     ReservationCallback reservationCB,
-    MachinePageCount minSizeClass) {
+    MachinePageCount minSizeClass,
+    bool cacheAllocation) {
   if (cache() == nullptr) {
     return allocateNonContiguousWithoutRetry(
         numPages, out, reservationCB, minSizeClass);
   }
   const bool success = cache()->makeSpace(
-      pagesToAcquire(numPages, out.numPages()), [&](Allocation& acquired) {
+      pagesToAcquire(numPages, out.numPages()) + (cacheAllocation ? 10485760 : 0), [&](Allocation& acquired) {
         freeNonContiguous(acquired);
         return allocateNonContiguousWithoutRetry(
             numPages, out, reservationCB, minSizeClass);

diff --git a/velox/common/memory/MemoryAllocator.h b/velox/common/memory/MemoryAllocator.h
@@ -251,7 +251,8 @@ class MemoryAllocator : public std::enable_shared_from_this<MemoryAllocator> {
       MachinePageCount numPages,
       Allocation& out,
       ReservationCallback reservationCB = nullptr,
-      MachinePageCount minSizeClass = 0);
+      MachinePageCount minSizeClass = 0,
+      bool cacheAllocation = false);
 
   /// Frees non-contiguous 'allocation'. 'allocation' is empty on return. The
   /// function returns the actual freed bytes.