From 21491d74a52e63312f766cfc8af281f4ed36c086 Mon Sep 17 00:00:00 2001 From: Kevin Wilfong Date: Tue, 19 Mar 2024 20:04:17 -0700 Subject: [PATCH] Fix NPE caused by testingRunArbitration (#9137) Summary: Pull Request resolved: https://github.com/facebookincubator/velox/pull/9137 testingRunArbitration is called from within an operator to force a call to shrinkPools to simulate what would happen if arbitration was triggered by a call to growCapacity. It does not check if the pool was aborted due to an error during arbitration which leads to all sorts of memory issues as the Task will be terminated as part of aborting the pool causing operators to get destructed. Adding a check to simulate what already happens in growCapacity. Reviewed By: xiaoxmeng Differential Revision: D55038530 fbshipit-source-id: be443619235956a43bd08158dca276126cca4789 --- velox/common/memory/MemoryArbitrator.cpp | 4 ++++ velox/common/memory/MemoryPool.h | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/velox/common/memory/MemoryArbitrator.cpp b/velox/common/memory/MemoryArbitrator.cpp index 6cb5bcb2ad99..e3a8d6a054e1 100644 --- a/velox/common/memory/MemoryArbitrator.cpp +++ b/velox/common/memory/MemoryArbitrator.cpp @@ -478,5 +478,9 @@ void testingRunArbitration( static_cast(pool)->testingManager()->shrinkPools( targetBytes, allowSpill); pool->leaveArbitration(); + + // This function is simulating an arbitration triggered by growCapacity, which + // would check this. + static_cast(pool)->testingCheckIfAborted(); } } // namespace facebook::velox::memory diff --git a/velox/common/memory/MemoryPool.h b/velox/common/memory/MemoryPool.h index a7676bb08191..2e3ac5c6f110 100644 --- a/velox/common/memory/MemoryPool.h +++ b/velox/common/memory/MemoryPool.h @@ -687,6 +687,10 @@ class MemoryPoolImpl : public MemoryPool { return allocator_; } + void testingCheckIfAborted() const { + checkIfAborted(); + } + uint64_t testingMinReservationBytes() const { return minReservationBytes_; }