Let operator not shrink pool after reclaim

tanjialiang · May 22, 2024 · 8a3bf7a · 8a3bf7a
1 parent 4a1a391
commit 8a3bf7a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/velox/common/memory/SharedArbitrator.cpp b/velox/common/memory/SharedArbitrator.cpp
@@ -758,10 +758,8 @@ uint64_t SharedArbitrator::reclaim(
       VELOX_MEM_LOG(ERROR) << "Failed to reclaim from memory pool "
                            << pool->name() << ", aborting it: " << e.what();
       abort(pool, std::current_exception());
-      // Free up all the free capacity from the aborted pool as the associated
-      // query has failed at this point.
-      pool->shrink();
     }
+    pool->shrink();
     const uint64_t newCapacity = pool->capacity();
     VELOX_CHECK_GE(oldCapacity, newCapacity);
     reclaimedBytes = oldCapacity - newCapacity;

diff --git a/velox/exec/Operator.cpp b/velox/exec/Operator.cpp
@@ -650,8 +650,11 @@ uint64_t Operator::MemoryReclaimer::reclaim(
 
   auto reclaimBytes = memory::MemoryReclaimer::run(
       [&]() {
+        const auto oldReservedBytes = pool->reservedBytes();
         op_->reclaim(targetBytes, stats);
-        return pool->shrink(targetBytes);
+        const auto newReservedBytes = pool->reservedBytes();
+        VELOX_CHECK_GE(oldReservedBytes, newReservedBytes);
+        return oldReservedBytes - newReservedBytes;
       },
       stats);