Add stats collection to memory reclaimer (#6895)

Summary:
During reclaiming there are information that also needs to be collected. Adding the stats collection framework to facilitate reclaiming stats collection.

Pull Request resolved: #6895

Reviewed By: xiaoxmeng

Differential Revision: D50020233

Pulled By: tanjialiang

fbshipit-source-id: 9ebd78e9b0bfdf26652fd7fa9a452a2ffaf93869

velox/common/memory/MemoryArbitrator.cpp

@@ -184,7 +184,8 @@ bool MemoryReclaimer::reclaimableBytes(
   return reclaimable;
 }
 
-uint64_t MemoryReclaimer::reclaim(MemoryPool* pool, uint64_t targetBytes) {
+uint64_t
+MemoryReclaimer::reclaim(MemoryPool* pool, uint64_t targetBytes, Stats& stats) {
   if (pool->kind() == MemoryPool::Kind::kLeaf) {
     return 0;
   }
@@ -214,7 +215,7 @@ uint64_t MemoryReclaimer::reclaim(MemoryPool* pool, uint64_t targetBytes) {
 
   uint64_t reclaimedBytes{0};
   for (const auto& candidate : candidates) {
-    const auto bytes = candidate.pool->reclaim(targetBytes);
+    const auto bytes = candidate.pool->reclaim(targetBytes, stats);
     reclaimedBytes += bytes;
     if (targetBytes != 0) {
       if (bytes >= targetBytes) {
@@ -243,6 +244,16 @@ void MemoryReclaimer::abort(MemoryPool* pool, const std::exception_ptr& error) {
   });
 }
 
+bool MemoryReclaimer::Stats::operator==(
+    const MemoryReclaimer::Stats& other) const {
+  return numNonReclaimableAttempts == other.numNonReclaimableAttempts;
+}
+
+bool MemoryReclaimer::Stats::operator!=(
+    const MemoryReclaimer::Stats& other) const {
+  return !(*this == other);
+}
+
 MemoryArbitrator::Stats::Stats(
     uint64_t _numRequests,
     uint64_t _numSucceeded,
@@ -254,7 +265,8 @@ MemoryArbitrator::Stats::Stats(
     uint64_t _numReclaimedBytes,
     uint64_t _maxCapacityBytes,
     uint64_t _freeCapacityBytes,
-    uint64_t _reclaimTimeUs)
+    uint64_t _reclaimTimeUs,
+    uint64_t _numNonReclaimableAttempts)
     : numRequests(_numRequests),
       numSucceeded(_numSucceeded),
       numAborted(_numAborted),
@@ -265,15 +277,17 @@ MemoryArbitrator::Stats::Stats(
       numReclaimedBytes(_numReclaimedBytes),
       maxCapacityBytes(_maxCapacityBytes),
       freeCapacityBytes(_freeCapacityBytes),
-      reclaimTimeUs(_reclaimTimeUs) {}
+      reclaimTimeUs(_reclaimTimeUs),
+      numNonReclaimableAttempts(_numNonReclaimableAttempts) {}
 
 std::string MemoryArbitrator::Stats::toString() const {
   return fmt::format(
-      "STATS[numRequests {} numSucceeded {} numAborted {} numFailures {} queueTime {} arbitrationTime {} reclaimTime {} shrunkMemory {} reclaimedMemory {} maxCapacity {} freeCapacity {}]",
+      "STATS[numRequests {} numSucceeded {} numAborted {} numFailures {} numNonReclaimableAttempts {} queueTime {} arbitrationTime {} reclaimTime {} shrunkMemory {} reclaimedMemory {} maxCapacity {} freeCapacity {}]",
       numRequests,
       numSucceeded,
       numAborted,
       numFailures,
+      numNonReclaimableAttempts,
       succinctMicros(queueTimeUs),
       succinctMicros(arbitrationTimeUs),
       succinctMicros(reclaimTimeUs),
@@ -297,6 +311,8 @@ MemoryArbitrator::Stats MemoryArbitrator::Stats::operator-(
   result.maxCapacityBytes = maxCapacityBytes;
   result.freeCapacityBytes = freeCapacityBytes;
   result.reclaimTimeUs = reclaimTimeUs - other.reclaimTimeUs;
+  result.numNonReclaimableAttempts =
+      numNonReclaimableAttempts - other.numNonReclaimableAttempts;
   return result;
 }
 
@@ -312,7 +328,8 @@ bool MemoryArbitrator::Stats::operator==(const Stats& other) const {
              numReclaimedBytes,
              maxCapacityBytes,
              freeCapacityBytes,
-             reclaimTimeUs) ==
+             reclaimTimeUs,
+             numNonReclaimableAttempts) ==
       std::tie(
              other.numRequests,
              other.numSucceeded,
@@ -324,7 +341,8 @@ bool MemoryArbitrator::Stats::operator==(const Stats& other) const {
              other.numReclaimedBytes,
              other.maxCapacityBytes,
              other.freeCapacityBytes,
-             other.reclaimTimeUs);
+             other.reclaimTimeUs,
+             other.numNonReclaimableAttempts);
 }
 
 bool MemoryArbitrator::Stats::operator!=(const Stats& other) const {
@@ -355,6 +373,7 @@ bool MemoryArbitrator::Stats::operator<(const Stats& other) const {
   UPDATE_COUNTER(numShrunkBytes);
   UPDATE_COUNTER(numReclaimedBytes);
   UPDATE_COUNTER(reclaimTimeUs);
+  UPDATE_COUNTER(numNonReclaimableAttempts);
 #undef UPDATE_COUNTER
   VELOX_CHECK(
       !((gtCount > 0) && (ltCount > 0)),

velox/common/memory/MemoryArbitrator.h

@@ -176,6 +176,9 @@ class MemoryArbitrator {
     /// The sum of all reclaim operation durations during arbitration in
     /// microseconds.
     uint64_t reclaimTimeUs{0};
+    /// The total number of times of the reclaim attempts that end up failing
+    /// due to reclaiming at non-reclaimable stage.
+    uint64_t numNonReclaimableAttempts{0};
 
     Stats(
         uint64_t _numRequests,
@@ -188,7 +191,8 @@ class MemoryArbitrator {
         uint64_t _numReclaimedBytes,
         uint64_t _maxCapacityBytes,
         uint64_t _freeCapacityBytes,
-        uint64_t _reclaimTimeUs);
+        uint64_t _reclaimTimeUs,
+        uint64_t _numNonReclaimableAttempts);
 
     Stats() = default;
 
@@ -252,6 +256,17 @@ FOLLY_ALWAYS_INLINE std::ostream& operator<<(
 /// through techniques such as disks spilling.
 class MemoryReclaimer {
  public:
+  /// Used to collect memory reclaim execution stats.
+  struct Stats {
+    /// The total number of times of the reclaim attempts that end up failing
+    /// due to reclaiming at non-reclaimable stage.
+    uint64_t numNonReclaimableAttempts{0};
+
+    bool operator==(const Stats& other) const;
+
+    bool operator!=(const Stats& other) const;
+  };
+
   virtual ~MemoryReclaimer() = default;
 
   static std::unique_ptr<MemoryReclaimer> create();
@@ -286,7 +301,8 @@ class MemoryReclaimer {
   /// memory bytes but there is no guarantees. If 'targetBytes' is zero, then it
   /// reclaims all the reclaimable memory from the memory 'pool'. The function
   /// returns the actual reclaimed memory bytes.
-  virtual uint64_t reclaim(MemoryPool* pool, uint64_t targetBytes);
+  virtual uint64_t
+  reclaim(MemoryPool* pool, uint64_t targetBytes, Stats& stats);
 
   /// Invoked by the memory arbitrator to abort memory 'pool' and the associated
   /// query execution when encounters non-recoverable memory reclaim error or

velox/common/memory/MemoryPool.cpp

@@ -923,11 +923,13 @@ bool MemoryPoolImpl::reclaimableBytes(uint64_t& reclaimableBytes) const {
   return reclaimer()->reclaimableBytes(*this, reclaimableBytes);
 }
 
-uint64_t MemoryPoolImpl::reclaim(uint64_t targetBytes) {
+uint64_t MemoryPoolImpl::reclaim(
+    uint64_t targetBytes,
+    memory::MemoryReclaimer::Stats& stats) {
   if (reclaimer() == nullptr) {
     return 0;
   }
-  return reclaimer()->reclaim(this, targetBytes);
+  return reclaimer()->reclaim(this, targetBytes, stats);
 }
 
 void MemoryPoolImpl::enterArbitration() {

velox/common/memory/MemoryPool.h

@@ -408,7 +408,9 @@ class MemoryPool : public std::enable_shared_from_this<MemoryPool> {
   /// noop if the reclaimer is not set, otherwise invoke the reclaimer's
   /// corresponding method. The function returns the actually freed capacity
   /// from the root of this memory pool.
-  virtual uint64_t reclaim(uint64_t targetBytes) = 0;
+  virtual uint64_t reclaim(
+      uint64_t targetBytes,
+      memory::MemoryReclaimer::Stats& stats) = 0;
 
   /// Invoked by the memory arbitrator to abort a root memory pool. The function
   /// forwards the request to the corresponding query object to abort its
@@ -627,7 +629,8 @@ class MemoryPoolImpl : public MemoryPool {
 
   bool reclaimableBytes(uint64_t& reclaimableBytes) const override;
 
-  uint64_t reclaim(uint64_t targetBytes) override;
+  uint64_t reclaim(uint64_t targetBytes, memory::MemoryReclaimer::Stats& stats)
+      override;
 
   uint64_t shrink(uint64_t targetBytes = 0) override;
 

velox/common/memory/SharedArbitrator.cpp

@@ -398,13 +398,14 @@ uint64_t SharedArbitrator::reclaim(
   uint64_t reclaimDurationUs{0};
   uint64_t reclaimedBytes{0};
   uint64_t freedBytes{0};
+  MemoryReclaimer::Stats reclaimerStats;
   {
     MicrosecondTimer reclaimTimer(&reclaimDurationUs);
     const uint64_t oldCapacity = pool->capacity();
     try {
       freedBytes = pool->shrink(targetBytes);
       if (freedBytes < targetBytes) {
-        pool->reclaim(targetBytes - freedBytes);
+        pool->reclaim(targetBytes - freedBytes, reclaimerStats);
       }
     } catch (const std::exception& e) {
       VELOX_MEM_LOG(ERROR) << "Failed to reclaim from memory pool "
@@ -421,6 +422,7 @@ uint64_t SharedArbitrator::reclaim(
   numReclaimedBytes_ += reclaimedBytes - freedBytes;
   numShrunkBytes_ += freedBytes;
   reclaimTimeUs_ += reclaimDurationUs;
+  numNonReclaimableAttempts_ += reclaimerStats.numNonReclaimableAttempts;
   VELOX_MEM_LOG(INFO) << "Reclaimed from memory pool " << pool->name()
                       << " with target of " << succinctBytes(targetBytes)
                       << ", actually reclaimed " << succinctBytes(freedBytes)
@@ -492,6 +494,7 @@ MemoryArbitrator::Stats SharedArbitrator::statsLocked() const {
   stats.maxCapacityBytes = capacity_;
   stats.freeCapacityBytes = freeCapacity_;
   stats.reclaimTimeUs = reclaimTimeUs_;
+  stats.numNonReclaimableAttempts = numNonReclaimableAttempts_;
   return stats;
 }
 

velox/common/memory/SharedArbitrator.h

@@ -200,5 +200,6 @@ class SharedArbitrator : public MemoryArbitrator {
   tsan_atomic<uint64_t> numShrunkBytes_{0};
   tsan_atomic<uint64_t> numReclaimedBytes_{0};
   tsan_atomic<uint64_t> reclaimTimeUs_{0};
+  tsan_atomic<uint64_t> numNonReclaimableAttempts_{0};
 };
 } // namespace facebook::velox::memory