feat: Add session property for debugMemoryPoolWarnThresholdBytes

presto-docs/src/main/sphinx/presto_cpp/properties-session.rst

@@ -97,6 +97,29 @@ If set to ``true``, disables the optimization in expression evaluation to delay
 
 This should only be used for debugging purposes.
 
+``native_debug_memory_pool_name_regex``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``varchar``
+* **Default value:** ``""``
+
+Native Execution only. Regular expression pattern to match memory pool names for allocation callsite tracking.
+Matched pools will also perform leak checks at destruction. Empty string disables tracking.
+
+This should only be used for debugging purposes.
+
+``native_debug_memory_pool_warn_threshold_bytes``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``bigint``
+* **Default value:** ``0``
+
+Native Execution only. Warning threshold for memory pool allocations. Logs callsites when exceeded.
+Requires allocation tracking to be enabled with ``native_debug_memory_pool_name_regex``.
+Accepts B/KB/MB/GB units. Set to 0B to disable.
+
+This should only be used for debugging purposes.
+
 ``native_execution_type_rewrite_enabled``
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java

@@ -53,6 +53,7 @@ public class NativeWorkerSessionPropertyProvider
     public static final String NATIVE_DEBUG_DISABLE_EXPRESSION_WITH_MEMOIZATION = "native_debug_disable_expression_with_memoization";
     public static final String NATIVE_DEBUG_DISABLE_EXPRESSION_WITH_LAZY_INPUTS = "native_debug_disable_expression_with_lazy_inputs";
     public static final String NATIVE_DEBUG_MEMORY_POOL_NAME_REGEX = "native_debug_memory_pool_name_regex";
+    public static final String NATIVE_DEBUG_MEMORY_POOL_WARN_THRESHOLD_BYTES = "native_debug_memory_pool_warn_threshold_bytes";
     public static final String NATIVE_SELECTIVE_NIMBLE_READER_ENABLED = "native_selective_nimble_reader_enabled";
     public static final String NATIVE_MAX_PARTIAL_AGGREGATION_MEMORY = "native_max_partial_aggregation_memory";
     public static final String NATIVE_MAX_EXTENDED_PARTIAL_AGGREGATION_MEMORY = "native_max_extended_partial_aggregation_memory";
@@ -213,6 +214,15 @@ public NativeWorkerSessionPropertyProvider(FeaturesConfig featuresConfig)
                                 " string means no match for all.",
                         "",
                         true),
+                stringProperty(
+                        NATIVE_DEBUG_MEMORY_POOL_WARN_THRESHOLD_BYTES,
+                        "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.",
+                        "0B",
+                        true),
                 booleanProperty(
                         NATIVE_SELECTIVE_NIMBLE_READER_ENABLED,
                         "Temporary flag to control whether selective Nimble reader should be " +

presto-native-execution/presto_cpp/main/QueryContextManager.cpp

@@ -63,7 +63,8 @@ std::shared_ptr<core::QueryCtx> QueryContextManager::createAndCacheQueryCtx(
     QueryContextCache& cache,
     const QueryId& queryId,
     velox::core::QueryConfig&& queryConfig,
-    std::unordered_map<std::string, std::shared_ptr<config::ConfigBase>>&& connectorConfigs,
+    std::unordered_map<std::string, std::shared_ptr<config::ConfigBase>>&&
+        connectorConfigs,
     std::shared_ptr<memory::MemoryPool>&& pool) {
   auto queryCtx = core::QueryCtx::create(
       driverExecutor_,
@@ -96,10 +97,12 @@ std::shared_ptr<core::QueryCtx> QueryContextManager::findOrCreateQueryCtx(
   // is still indexed by the query id.
   static std::atomic_uint64_t poolId{0};
   std::optional<memory::MemoryPool::DebugOptions> poolDbgOpts;
-  const auto debugMemoryPoolNameRegex = queryConfig.debugMemoryPoolNameRegex();
+  auto debugMemoryPoolNameRegex = queryConfig.debugMemoryPoolNameRegex();
   if (!debugMemoryPoolNameRegex.empty()) {
     poolDbgOpts = memory::MemoryPool::DebugOptions{
-        .debugPoolNameRegex = debugMemoryPoolNameRegex};
+        .debugPoolNameRegex = std::move(debugMemoryPoolNameRegex),
+        .debugPoolWarnThresholdBytes =
+            queryConfig.debugMemoryPoolWarnThresholdBytes()};
   }
   auto pool = memory::MemoryManager::getInstance()->addRootPool(
       fmt::format("{}_{}", queryId, poolId++),

presto-native-execution/presto_cpp/main/SessionProperties.cpp

@@ -280,6 +280,19 @@ SessionProperties::SessionProperties() {
       QueryConfig::kDebugMemoryPoolNameRegex,
       c.debugMemoryPoolNameRegex());
 
+  addSessionProperty(
+      kDebugMemoryPoolWarnThresholdBytes,
+      "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. Requires allocation "
+      "tracking to be enabled with `native_debug_memory_pool_name_regex` "
+      "for the pool. A value of 0 means no warning threshold is enforced.",
+      BIGINT(),
+      false,
+      QueryConfig::kDebugMemoryPoolWarnThresholdBytes,
+      std::to_string(c.debugMemoryPoolWarnThresholdBytes()));
+
   addSessionProperty(
       kSelectiveNimbleReaderEnabled,
       "Temporary flag to control whether selective Nimble reader should be "

presto-native-execution/presto_cpp/main/SessionProperties.h

@@ -184,6 +184,12 @@ class SessionProperties {
   static constexpr const char* kDebugMemoryPoolNameRegex =
       "native_debug_memory_pool_name_regex";
 
+  /// Warning threshold in bytes for memory pool allocations. Logs callsites 
+  /// when exceeded. Requires allocation tracking to be enabled with 
+  /// `native_debug_memory_pool_name_regex` property for the pool.
+  static constexpr const char* kDebugMemoryPoolWarnThresholdBytes =
+      "native_debug_memory_pool_warn_threshold_bytes";
+
   /// Temporary flag to control whether selective Nimble reader should be used
   /// in this query or not.  Will be removed after the selective Nimble reader
   /// is fully rolled out.

presto-native-execution/presto_cpp/main/tests/SessionPropertiesTest.cpp

@@ -69,6 +69,8 @@ TEST_F(SessionPropertiesTest, validateMapping) {
        core::QueryConfig::kDebugDisableExpressionWithLazyInputs},
       {SessionProperties::kDebugMemoryPoolNameRegex,
        core::QueryConfig::kDebugMemoryPoolNameRegex},
+      {SessionProperties::kDebugMemoryPoolWarnThresholdBytes,
+       core::QueryConfig::kDebugMemoryPoolWarnThresholdBytes},
       {SessionProperties::kSelectiveNimbleReaderEnabled,
        core::QueryConfig::kSelectiveNimbleReaderEnabled},
       {SessionProperties::kQueryTraceEnabled,