Adding threshold for logging memory pool allocations

Summary:
We have a leak check in memory pool the prints all current allocations along with call sites.
In this diff I'm expanding that functionality and adding possibility to log that information once we reach configurable threshold. It's supposed to be used for debugging, when we want to know how exactly the memory has been used in particular memory pool. Adding this only for leaf nodes for now.

Differential Revision: D79902261

Author: kletkavrubashku

velox/common/memory/MemoryPool.cpp

@@ -17,7 +17,6 @@
 #include "velox/common/memory/MemoryPool.h"
 
 #include <signal.h>
-#include <set>
 
 #include "velox/common/base/Counters.h"
 #include "velox/common/base/StatsReporter.h"
@@ -1203,10 +1202,37 @@ void MemoryPoolImpl::recordAllocDbg(const void* addr, uint64_t size) {
   if (!needRecordDbg(true)) {
     return;
   }
+  AllocationRecord allocationRecord{size, process::StackTrace()};
   std::lock_guard<std::mutex> l(debugAllocMutex_);
-  debugAllocRecords_.emplace(
-      reinterpret_cast<uint64_t>(addr),
-      AllocationRecord{size, process::StackTrace()});
+  auto [it, inserted] = debugAllocRecords_.try_emplace(
+      reinterpret_cast<uint64_t>(addr), std::move(allocationRecord));
+  VELOX_CHECK(
+      inserted,
+      fmt::format(
+          "[MemoryPool] Duplicate allocation record for {} bytes.\n"
+          "======== Allocation Stack ========\n"
+          "{}\n",
+          size,
+          allocationRecord.callStack.toString()));
+  debugAllocRecordsTotalSize_ += size;
+  if (debugOptions_->debugPoolWarnThresholdBytes > 0 &&
+      debugAllocRecordsTotalSize_ >
+          debugOptions_->debugPoolWarnThresholdBytes &&
+      !debugWarnThresholdExceeded_) {
+    debugWarnThresholdExceeded_ = true;
+    VELOX_MEM_LOG(WARNING) << fmt::format(
+        "[MemoryPool] Memory pool '{}' exceeded warning threshold of {} bytes with allocation of {} bytes, resulting in total size of {} bytes.\n"
+        "======== Allocation Stack ========\n"
+        "{}\n"
+        "======= Current Allocations ======\n"
+        "{}",
+        name_,
+        debugOptions_->debugPoolWarnThresholdBytes,
+        size,
+        debugAllocRecordsTotalSize_,
+        it->second.callStack.toString(),
+        dumpRecordsDbg());
+  }
 }
 
 void MemoryPoolImpl::recordAllocDbg(const Allocation& allocation) {
@@ -1251,6 +1277,7 @@ void MemoryPoolImpl::recordFreeDbg(const void* addr, uint64_t size) {
         freeStackTrace));
   }
   debugAllocRecords_.erase(addrUint64);
+  debugAllocRecordsTotalSize_ -= size;
 }
 
 void MemoryPoolImpl::recordFreeDbg(const Allocation& allocation) {
@@ -1280,6 +1307,8 @@ void MemoryPoolImpl::recordGrowDbg(const void* addr, uint64_t newSize) {
   if (allocResult == debugAllocRecords_.end()) {
     VELOX_FAIL("Growing of un-allocated memory. Free address {}.", addrUint64);
   }
+  debugAllocRecordsTotalSize_ =
+      debugAllocRecordsTotalSize_ - allocResult->second.size + newSize;
   allocResult->second.size = newSize;
 }
 
@@ -1288,10 +1317,16 @@ void MemoryPoolImpl::leakCheckDbg() {
   if (debugAllocRecords_.empty()) {
     return;
   }
-  std::stringbuf buf;
-  std::ostream oss(&buf);
-  oss << "[MemoryPool] : " << name_ << " - Detected total of "
-      << debugAllocRecords_.size() << " leaked allocations:\n";
+  VELOX_FAIL(fmt::format(
+      "[MemoryPool] Leak check failed for '{}' pool - {}",
+      name_,
+      dumpRecordsDbg()));
+}
+
+std::string MemoryPoolImpl::dumpRecordsDbg() {
+  VELOX_CHECK(debugEnabled());
+  std::stringstream oss;
+  oss << fmt::format("Found {} allocations:\n", debugAllocRecords_.size());
   struct AllocationStats {
     uint64_t size{0};
     uint64_t numAllocations{0};
@@ -1316,12 +1351,13 @@ void MemoryPoolImpl::leakCheckDbg() {
         return a.second.size > b.second.size;
       });
   for (const auto& pair : sortedRecords) {
-    oss << "======== Leaked memory from " << pair.second.numAllocations
-        << " total allocations of " << succinctBytes(pair.second.size)
-        << " total size ========\n"
-        << pair.first << "\n";
+    oss << fmt::format(
+        "======== {} allocations of {} total size ========\n{}\n",
+        pair.second.numAllocations,
+        succinctBytes(pair.second.size),
+        pair.first);
   }
-  VELOX_FAIL(buf.str());
+  return oss.str();
 }
 
 void MemoryPoolImpl::handleAllocationFailure(

velox/common/memory/MemoryPool.h

@@ -112,6 +112,12 @@ class MemoryPool : public std::enable_shared_from_this<MemoryPool> {
     /// memory pools whose name matches the specified regular expression. Empty
     /// string means no match for all.
     std::string debugPoolNameRegex;
+
+    /// Warning threshold in bytes for debug memory pools. When set to a
+    /// non-zero value, a warning will be logged once per memory pool when
+    /// allocations cause the pool to exceed this threshold. A value of
+    /// 0 means no warning threshold is enforced.
+    uint64_t debugPoolWarnThresholdBytes{0};
   };
 
   struct Options {
@@ -963,7 +969,11 @@ class MemoryPoolImpl : public MemoryPool {
   bool needRecordDbg(bool isAlloc);
 
   // Invoked to record the call stack of a buffer allocation if debug mode of
-  // this memory pool is enabled.
+  // this memory pool is enabled. If debugPoolWarnThresholdBytes is configured
+  // and the total debug allocation size exceeds this threshold, a warning is
+  // logged once per memory pool containing the allocation stack trace and
+  // current pool allocations summary for debugging purposes. Same for other
+  // 'recordAllocDbg(..)' calls below.
   void recordAllocDbg(const void* addr, uint64_t size);
 
   // Invoked to record the call stack of a non-contiguous allocation if debug
@@ -997,6 +1007,10 @@ class MemoryPoolImpl : public MemoryPool {
   // pool is enabled.
   void leakCheckDbg();
 
+  // Dump the recorded call sites of the memory allocations in
+  // 'debugAllocRecords_' to the string.
+  std::string dumpRecordsDbg();
+
   void handleAllocationFailure(const std::string& failureMessage);
 
   MemoryManager* const manager_;
@@ -1064,6 +1078,12 @@ class MemoryPoolImpl : public MemoryPool {
 
   // Map from address to 'AllocationRecord'.
   std::unordered_map<uint64_t, AllocationRecord> debugAllocRecords_;
+
+  // Total size of all the memory allocations recorded in 'debugAllocRecords_'.
+  uint64_t debugAllocRecordsTotalSize_{0};
+
+  // Flag to track if warning threshold has been exceeded once for this pool.
+  bool debugWarnThresholdExceeded_{false};
 };
 
 /// An Allocator backed by a memory pool for STL containers.

velox/core/QueryConfig.h

@@ -527,6 +527,14 @@ class QueryConfig {
   static constexpr const char* kDebugMemoryPoolNameRegex =
       "debug_memory_pool_name_regex";
 
+  /// Warning threshold in bytes for debug memory pools. When set to a
+  /// non-zero value, a warning will be logged once per memory pool when
+  /// allocations cause the pool to exceed this threshold. This is useful for
+  /// identifying memory usage patterns during debugging. A value of
+  /// 0 means no warning threshold is enforced.
+  static constexpr const char* kDebugMemoryPoolWarnThresholdBytes =
+      "debug_memory_pool_warn_threshold_bytes";
+
   /// Some lambda functions over arrays and maps are evaluated in batches of the
   /// underlying elements that comprise the arrays/maps. This is done to make
   /// the batch size managable as array vectors can have thousands of elements
@@ -692,6 +700,12 @@ class QueryConfig {
     return get<std::string>(kDebugMemoryPoolNameRegex, "");
   }
 
+  uint64_t debugMemoryPoolWarnThresholdBytes() const {
+    return config::toCapacity(
+        get<std::string>(kDebugMemoryPoolWarnThresholdBytes, "0B"),
+        config::CapacityUnit::BYTE);
+  }
+
   std::optional<uint32_t> debugAggregationApproxPercentileFixedRandomSeed()
       const {
     return get<uint32_t>(kDebugAggregationApproxPercentileFixedRandomSeed);