feat: Add priority based memory reclaim framework (#11598)

Summary: X-link: facebookincubator/nimble#111 * Adds priority base reclaiming to memory reclaim framework. The priority determines which memory pool to reclaim first on the same level. This would help to make reclaim more application logic aware. * Make join node reclaim priority lower than others. This is because cost of reclaiming (spilling) on join node is high compared to other nodes. Pull Request resolved: #11598 Reviewed By: xiaoxmeng Differential Revision: D66261340 Pulled By: tanjialiang fbshipit-source-id: c03b2ef25b39dc8771f66321731d5a88da26638e
facebookincubator · Dec 6, 2024 · 939c102 · 939c102
1 parent 7672dd5
commit 939c102
Show file tree

Hide file tree

Showing 22 changed files with 385 additions and 119 deletions.
diff --git a/velox/common/memory/MemoryArbitrator.cpp b/velox/common/memory/MemoryArbitrator.cpp
@@ -174,8 +174,8 @@ void MemoryArbitrator::unregisterFactory(const std::string& kind) {
   return pool->shrink(targetBytes);
 }
 
-std::unique_ptr<MemoryReclaimer> MemoryReclaimer::create() {
-  return std::unique_ptr<MemoryReclaimer>(new MemoryReclaimer());
+std::unique_ptr<MemoryReclaimer> MemoryReclaimer::create(int32_t priority) {
+  return std::unique_ptr<MemoryReclaimer>(new MemoryReclaimer(priority));
 }
 
 // static
@@ -234,8 +234,9 @@ uint64_t MemoryReclaimer::reclaim(
     return 0;
   }
 
-  // Sort the child pools based on their reserved memory and reclaim from the
-  // child pool with most reservation first.
+  // Sort the child pools based on their reclaimer priority and reserved memory.
+  // Reclaim from the child pool with highest priority and most reservation
+  // first.
   struct Candidate {
     std::shared_ptr<memory::MemoryPool> pool;
     int64_t reclaimableBytes;
@@ -247,8 +248,13 @@ uint64_t MemoryReclaimer::reclaim(
     for (auto& entry : pool->children_) {
       auto child = entry.second.lock();
       if (child != nullptr) {
-        const int64_t reclaimableBytes = child->reclaimableBytes().value_or(0);
-        candidates.push_back(Candidate{std::move(child), reclaimableBytes});
+        const auto reclaimableBytesOpt = child->reclaimableBytes();
+        if (!reclaimableBytesOpt.has_value()) {
+          continue;
+        }
+        candidates.push_back(Candidate{
+            std::move(child),
+            static_cast<int64_t>(reclaimableBytesOpt.value())});
       }
     }
   }
@@ -257,7 +263,12 @@ uint64_t MemoryReclaimer::reclaim(
       candidates.begin(),
       candidates.end(),
       [](const auto& lhs, const auto& rhs) {
-        return lhs.reclaimableBytes > rhs.reclaimableBytes;
+        const auto lhsPrio = lhs.pool->reclaimer()->priority();
+        const auto rhsPrio = rhs.pool->reclaimer()->priority();
+        if (lhsPrio == rhsPrio) {
+          return lhs.reclaimableBytes > rhs.reclaimableBytes;
+        }
+        return lhsPrio < rhsPrio;
       });
 
   uint64_t reclaimedBytes{0};

diff --git a/velox/common/memory/MemoryArbitrator.h b/velox/common/memory/MemoryArbitrator.h
@@ -287,7 +287,7 @@ class MemoryReclaimer {
 
   virtual ~MemoryReclaimer() = default;
 
-  static std::unique_ptr<MemoryReclaimer> create();
+  static std::unique_ptr<MemoryReclaimer> create(int32_t priority = 0);
 
   /// Invoked memory reclaim function from 'pool' and record execution 'stats'.
   static uint64_t run(const std::function<int64_t()>& func, Stats& stats);
@@ -309,6 +309,26 @@ class MemoryReclaimer {
   /// enterArbitration has been called.
   virtual void leaveArbitration() noexcept {}
 
+  /// Invoked by upper layer reclaimer, to return the priority of this
+  /// reclaimer. The priority determines the reclaiming order of self among all
+  /// same level reclaimers. The smaller the number, the higher the priority.
+  /// Consider the following memory pool & reclaimer structure:
+  ///
+  ///                 rec1(pri 1)
+  ///                /            \
+  ///               /              \
+  ///              /                \
+  ///      rec2(pri 1)           rec3(pri 3)
+  ///      /        \             /        \
+  ///     /          \           /          \
+  /// rec4(pri 1) rec5(pri 0)  rec6(pri 0) rec7(pri 1)
+  ///
+  /// The reclaiming traversing order will be rec1 -> rec2 -> rec5 -> rec4 ->
+  /// rec3 -> rec6 -> rec7
+  virtual int32_t priority() const {
+    return priority_;
+  };
+
   /// Invoked by the memory arbitrator to get the amount of memory bytes that
   /// can be reclaimed from 'pool'. The function returns true if 'pool' is
   /// reclaimable and returns the estimated reclaimable bytes in
@@ -343,7 +363,10 @@ class MemoryReclaimer {
   virtual void abort(MemoryPool* pool, const std::exception_ptr& error);
 
  protected:
-  MemoryReclaimer() = default;
+  explicit MemoryReclaimer(int32_t priority) : priority_(priority){};
+
+ private:
+  const int32_t priority_;
 };
 
 /// Helper class used to measure the memory bytes reclaimed from a memory pool

diff --git a/velox/common/memory/tests/ArbitrationParticipantTest.cpp b/velox/common/memory/tests/ArbitrationParticipantTest.cpp
@@ -131,7 +131,8 @@ class MockTask : public std::enable_shared_from_this<MockTask> {
 
   class RootMemoryReclaimer : public memory::MemoryReclaimer {
    public:
-    RootMemoryReclaimer(const std::shared_ptr<MockTask>& task) : task_(task) {}
+    RootMemoryReclaimer(const std::shared_ptr<MockTask>& task)
+        : memory::MemoryReclaimer(0), task_(task) {}
 
     static std::unique_ptr<MemoryReclaimer> create(
         const std::shared_ptr<MockTask>& task) {
@@ -179,7 +180,8 @@ class MockTask : public std::enable_shared_from_this<MockTask> {
         bool reclaimable,
         ReclaimInjectionCallback reclaimInjectCb = nullptr,
         ArbitrationInjectionCallback arbitrationInjectCb = nullptr)
-        : task_(task),
+        : memory::MemoryReclaimer(0),
+          task_(task),
           reclaimable_(reclaimable),
           reclaimInjectCb_(std::move(reclaimInjectCb)),
           arbitrationInjectCb_(std::move(arbitrationInjectCb)) {}

diff --git a/velox/common/memory/tests/MemoryArbitratorTest.cpp b/velox/common/memory/tests/MemoryArbitratorTest.cpp
@@ -505,7 +505,8 @@ class MockLeafMemoryReclaimer : public MemoryReclaimer {
       std::atomic<uint64_t>& totalUsedBytes,
       bool reclaimable = true,
       bool* underArbitration = nullptr)
-      : reclaimable_(reclaimable),
+      : MemoryReclaimer(0),
+        reclaimable_(reclaimable),
         underArbitration_(underArbitration),
         totalUsedBytes_(totalUsedBytes) {}
 

diff --git a/velox/common/memory/tests/MemoryPoolTest.cpp b/velox/common/memory/tests/MemoryPoolTest.cpp
@@ -3572,7 +3572,8 @@ class MockMemoryReclaimer : public MemoryReclaimer {
   }
 
  private:
-  explicit MockMemoryReclaimer(bool doThrow) : doThrow_(doThrow) {}
+  explicit MockMemoryReclaimer(bool doThrow)
+      : MemoryReclaimer(0), doThrow_(doThrow) {}
 
   const bool doThrow_;
 };

diff --git a/velox/common/memory/tests/MockSharedArbitratorTest.cpp b/velox/common/memory/tests/MockSharedArbitratorTest.cpp
@@ -92,7 +92,8 @@ class MockTask : public std::enable_shared_from_this<MockTask> {
 
   class MemoryReclaimer : public memory::MemoryReclaimer {
    public:
-    MemoryReclaimer(const std::shared_ptr<MockTask>& task) : task_(task) {}
+    MemoryReclaimer(const std::shared_ptr<MockTask>& task)
+        : memory::MemoryReclaimer(0), task_(task) {}
 
     static std::unique_ptr<MemoryReclaimer> create(
         const std::shared_ptr<MockTask>& task) {
@@ -178,7 +179,8 @@ class MockMemoryOperator {
         bool reclaimable,
         ReclaimInjectionCallback reclaimInjectCb = nullptr,
         ArbitrationInjectionCallback arbitrationInjectCb = nullptr)
-        : op_(op),
+        : memory::MemoryReclaimer(0),
+          op_(op),
           reclaimable_(reclaimable),
           reclaimInjectCb_(std::move(reclaimInjectCb)),
           arbitrationInjectCb_(std::move(arbitrationInjectCb)) {}

diff --git a/velox/connectors/hive/HiveDataSink.h b/velox/connectors/hive/HiveDataSink.h
@@ -494,7 +494,7 @@ class HiveDataSink : public DataSink {
         HiveDataSink* dataSink,
         HiveWriterInfo* writerInfo,
         io::IoStatistics* ioStats)
-        : exec::MemoryReclaimer(),
+        : exec::MemoryReclaimer(0),
           dataSink_(dataSink),
           writerInfo_(writerInfo),
           ioStats_(ioStats) {

diff --git a/velox/core/QueryCtx.h b/velox/core/QueryCtx.h
@@ -170,8 +170,9 @@ class QueryCtx : public std::enable_shared_from_this<QueryCtx> {
    protected:
     MemoryReclaimer(
         const std::shared_ptr<QueryCtx>& queryCtx,
-        memory::MemoryPool* pool)
-        : queryCtx_(queryCtx), pool_(pool) {
+        memory::MemoryPool* pool,
+        int32_t priority = 0)
+        : memory::MemoryReclaimer(priority), queryCtx_(queryCtx), pool_(pool) {
       VELOX_CHECK_NOT_NULL(pool_);
     }
 

diff --git a/velox/dwio/common/SortingWriter.h b/velox/dwio/common/SortingWriter.h
@@ -63,8 +63,8 @@ class SortingWriter : public Writer {
         memory::MemoryReclaimer::Stats& stats) override;
 
    private:
-    explicit MemoryReclaimer(SortingWriter* writer)
-        : exec::MemoryReclaimer(),
+    MemoryReclaimer(SortingWriter* writer)
+        : exec::MemoryReclaimer(0),
           writer_(writer),
           canReclaim_(writer_->sortBuffer_->canSpill()) {}
 

diff --git a/velox/dwio/dwrf/writer/Writer.h b/velox/dwio/dwrf/writer/Writer.h
@@ -163,7 +163,8 @@ class Writer : public dwio::common::Writer {
         memory::MemoryReclaimer::Stats& stats) override;
 
    private:
-    explicit MemoryReclaimer(Writer* writer) : writer_(writer) {
+    MemoryReclaimer(Writer* writer)
+        : exec::MemoryReclaimer(0), writer_(writer) {
       VELOX_CHECK_NOT_NULL(writer_);
     }
 

diff --git a/velox/exec/HashJoinBridge.h b/velox/exec/HashJoinBridge.h
@@ -174,9 +174,10 @@ bool isLeftNullAwareJoinWithFilter(
 class HashJoinMemoryReclaimer final : public MemoryReclaimer {
  public:
   static std::unique_ptr<memory::MemoryReclaimer> create(
-      std::shared_ptr<HashJoinBridge> joinBridge) {
+      std::shared_ptr<HashJoinBridge> joinBridge,
+      int32_t priority = 0) {
     return std::unique_ptr<memory::MemoryReclaimer>(
-        new HashJoinMemoryReclaimer(joinBridge));
+        new HashJoinMemoryReclaimer(joinBridge, priority));
   }
 
   uint64_t reclaim(
@@ -186,9 +187,10 @@ class HashJoinMemoryReclaimer final : public MemoryReclaimer {
       memory::MemoryReclaimer::Stats& stats) final;
 
  private:
-  explicit HashJoinMemoryReclaimer(
-      const std::shared_ptr<HashJoinBridge>& joinBridge)
-      : MemoryReclaimer(), joinBridge_(joinBridge) {}
+  HashJoinMemoryReclaimer(
+      const std::shared_ptr<HashJoinBridge>& joinBridge,
+      int32_t priority)
+      : MemoryReclaimer(priority), joinBridge_(joinBridge) {}
   std::weak_ptr<HashJoinBridge> joinBridge_;
 };
 

diff --git a/velox/exec/MemoryReclaimer.cpp b/velox/exec/MemoryReclaimer.cpp
@@ -20,8 +20,13 @@
 #include "velox/exec/Task.h"
 
 namespace facebook::velox::exec {
-std::unique_ptr<memory::MemoryReclaimer> MemoryReclaimer::create() {
-  return std::unique_ptr<memory::MemoryReclaimer>(new MemoryReclaimer());
+MemoryReclaimer::MemoryReclaimer(int32_t priority)
+    : memory::MemoryReclaimer(priority) {}
+
+std::unique_ptr<memory::MemoryReclaimer> MemoryReclaimer::create(
+    int32_t priority) {
+  return std::unique_ptr<memory::MemoryReclaimer>(
+      new MemoryReclaimer(priority));
 }
 
 void MemoryReclaimer::enterArbitration() {
@@ -67,13 +72,15 @@ void MemoryReclaimer::abort(
 }
 
 /*static*/ std::unique_ptr<memory::MemoryReclaimer>
-ParallelMemoryReclaimer::create(folly::Executor* executor) {
+ParallelMemoryReclaimer::create(folly::Executor* executor, int32_t priority) {
   return std::unique_ptr<memory::MemoryReclaimer>(
-      new ParallelMemoryReclaimer(executor));
+      new ParallelMemoryReclaimer(executor, priority));
 }
 
-ParallelMemoryReclaimer::ParallelMemoryReclaimer(folly::Executor* executor)
-    : executor_(executor) {}
+ParallelMemoryReclaimer::ParallelMemoryReclaimer(
+    folly::Executor* executor,
+    int32_t priority)
+    : MemoryReclaimer(priority), executor_(executor) {}
 
 uint64_t ParallelMemoryReclaimer::reclaim(
     memory::MemoryPool* pool,
@@ -85,8 +92,7 @@ uint64_t ParallelMemoryReclaimer::reclaim(
         pool, targetBytes, maxWaitMs, stats);
   }
 
-  // Sort the child pools based on their reserved memory and reclaim from the
-  // child pool with most reservation first.
+  // Sort candidates based on priority.
   struct Candidate {
     std::shared_ptr<memory::MemoryPool> pool;
     int64_t reclaimableBytes;
@@ -98,11 +104,17 @@ uint64_t ParallelMemoryReclaimer::reclaim(
     for (auto& entry : pool->children_) {
       auto child = entry.second.lock();
       if (child != nullptr) {
-        const int64_t reclaimableBytes = child->reclaimableBytes().value_or(0);
-        candidates.push_back(Candidate{std::move(child), reclaimableBytes});
+        const auto reclaimableBytesOpt = child->reclaimableBytes();
+        if (!reclaimableBytesOpt.has_value()) {
+          continue;
+        }
+        candidates.push_back(Candidate{
+            std::move(child),
+            static_cast<int64_t>(reclaimableBytesOpt.value())});
       }
     }
   }
+
   struct ReclaimResult {
     const uint64_t reclaimedBytes{0};
     const Stats stats;

diff --git a/velox/exec/MemoryReclaimer.h b/velox/exec/MemoryReclaimer.h
@@ -25,7 +25,7 @@ class MemoryReclaimer : public memory::MemoryReclaimer {
  public:
   virtual ~MemoryReclaimer() = default;
 
-  static std::unique_ptr<memory::MemoryReclaimer> create();
+  static std::unique_ptr<memory::MemoryReclaimer> create(int32_t priority = 0);
 
   void enterArbitration() override;
 
@@ -35,7 +35,7 @@ class MemoryReclaimer : public memory::MemoryReclaimer {
       override;
 
  protected:
-  MemoryReclaimer() = default;
+  explicit MemoryReclaimer(int32_t priortity);
 };
 
 /// Provides the parallel memory reclaimer implementation for velox task
@@ -46,7 +46,8 @@ class ParallelMemoryReclaimer : public memory::MemoryReclaimer {
   virtual ~ParallelMemoryReclaimer() = default;
 
   static std::unique_ptr<memory::MemoryReclaimer> create(
-      folly::Executor* executor);
+      folly::Executor* executor,
+      int32_t priority = 0);
 
   uint64_t reclaim(
       memory::MemoryPool* pool,
@@ -55,7 +56,7 @@ class ParallelMemoryReclaimer : public memory::MemoryReclaimer {
       Stats& stats) override;
 
  protected:
-  explicit ParallelMemoryReclaimer(folly::Executor* executor);
+  ParallelMemoryReclaimer(folly::Executor* executor, int32_t priority);
 
   folly::Executor* const executor_{nullptr};
 };

diff --git a/velox/exec/Operator.h b/velox/exec/Operator.h
@@ -708,7 +708,7 @@ class Operator : public BaseRuntimeStatWriter {
 
    protected:
     MemoryReclaimer(const std::shared_ptr<Driver>& driver, Operator* op)
-        : driver_(driver), op_(op) {
+        : memory::MemoryReclaimer(0), driver_(driver), op_(op) {
       VELOX_CHECK_NOT_NULL(op_);
     }
 

diff --git a/velox/exec/TableWriter.h b/velox/exec/TableWriter.h
@@ -178,7 +178,8 @@ class TableWriter : public Operator {
         const std::shared_ptr<Driver>& driver,
         Operator* op)
         : ParallelMemoryReclaimer(
-              spillConfig.has_value() ? spillConfig.value().executor : nullptr),
+              spillConfig.has_value() ? spillConfig.value().executor : nullptr,
+              0),
           canReclaim_(spillConfig.has_value()),
           driver_(driver),
           op_(op) {}