From 8a3bf7acc7d8c144a119dcde04d8ebe0f4d47024 Mon Sep 17 00:00:00 2001
From: Jialiang Tan <jacob.jialiang.tan@gmail.com>
Date: Wed, 22 May 2024 11:57:45 -0700
Subject: [PATCH] Let operator not shrink pool after reclaim

---
 velox/common/memory/SharedArbitrator.cpp | 4 +---
 velox/exec/Operator.cpp                  | 5 ++++-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/velox/common/memory/SharedArbitrator.cpp b/velox/common/memory/SharedArbitrator.cpp
index 743c8c68eb4a8..23e9be597925a 100644
--- 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
index 932289c52f53f..3b2a51223115c 100644
--- 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);