Add memory pool debug regex

Author: tanjialiang

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

@@ -52,6 +52,7 @@ public class NativeWorkerSessionPropertyProvider
     public static final String NATIVE_DEBUG_DISABLE_COMMON_SUB_EXPRESSION = "native_debug_disable_common_sub_expressions";
     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_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";
@@ -200,6 +201,14 @@ public NativeWorkerSessionPropertyProvider(FeaturesConfig featuresConfig)
                                 "of lazy inputs unless required. Should only be used for debugging.",
                         false,
                         true),
+                stringProperty(
+                        NATIVE_DEBUG_MEMORY_POOL_NAME_REGEX,
+                        "Regex for filtering on memory pool name if not empty." +
+                                " This allows us to only track the callsites of memory allocations for" +
+                                " memory pools whose name matches the specified regular expression. Empty" +
+                                " string means no match for all.",
+                        "",
+                        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

@@ -163,9 +163,17 @@ std::shared_ptr<core::QueryCtx> QueryContextManager::findOrCreateQueryCtx(
   // though the query ctx has been evicted out of the cache. The query ctx cache
   // 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();
+  if (!debugMemoryPoolNameRegex.empty()) {
+    poolDbgOpts = memory::MemoryPool::DebugOptions{
+        .debugPoolNameRegex = debugMemoryPoolNameRegex};
+  }
   auto pool = memory::MemoryManager::getInstance()->addRootPool(
       fmt::format("{}_{}", queryId, poolId++),
-      queryConfig.queryMaxMemoryPerNode());
+      queryConfig.queryMaxMemoryPerNode(),
+      nullptr,
+      poolDbgOpts);
 
   auto queryCtx = core::QueryCtx::create(
       driverExecutor_,

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

@@ -263,6 +263,17 @@ SessionProperties::SessionProperties() {
       QueryConfig::kDebugDisableExpressionWithLazyInputs,
       boolToString(c.debugDisableExpressionsWithLazyInputs()));
 
+  addSessionProperty(
+      kDebugMemoryPoolNameRegex,
+      "Regex for filtering on memory pool name if not empty. This allows us to "
+      "only track the callsites of memory allocations for memory pools whose "
+      "name matches the specified regular expression. Empty string means no "
+      "match for all.",
+      VARCHAR(),
+      false,
+      QueryConfig::kDebugMemoryPoolNameRegex,
+      c.debugMemoryPoolNameRegex());
+
   addSessionProperty(
       kSelectiveNimbleReaderEnabled,
       "Temporary flag to control whether selective Nimble reader should be "

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

@@ -177,6 +177,13 @@ class SessionProperties {
   static constexpr const char* kDebugDisableExpressionWithLazyInputs =
       "native_debug_disable_expression_with_lazy_inputs";
 
+  /// Regex for filtering on memory pool name if not empty. This allows us to
+  /// only track the callsites of memory allocations for memory pools whose name
+  /// matches the specified regular expression. Empty string means no match for
+  /// all.
+  static constexpr const char* kDebugMemoryPoolNameRegex =
+      "native_debug_memory_pool_name_regex";
+
   /// 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.