From bacf9c3c68f5785a547ce7ea65a857007b574b4d Mon Sep 17 00:00:00 2001
From: Jialiang Tan <jacob.jialiang.tan@gmail.com>
Date: Tue, 12 Mar 2024 13:27:48 -0700
Subject: [PATCH] [native] Add system-memory-pushback-by-abort-enabled config

---
 presto-native-execution/presto_cpp/main/common/Configs.cpp | 5 +++++
 presto-native-execution/presto_cpp/main/common/Configs.h   | 7 +++++++
 2 files changed, 12 insertions(+)

diff --git a/presto-native-execution/presto_cpp/main/common/Configs.cpp b/presto-native-execution/presto_cpp/main/common/Configs.cpp
index 3ce3eaf30a236..891806fc049a4 100644
--- a/presto-native-execution/presto_cpp/main/common/Configs.cpp
+++ b/presto-native-execution/presto_cpp/main/common/Configs.cpp
@@ -172,6 +172,7 @@ SystemConfig::SystemConfig() {
           NUM_PROP(kSystemMemLimitGb, 55),
           NUM_PROP(kSystemMemShrinkGb, 8),
           BOOL_PROP(kMallocMemHeapDumpEnabled, false),
+          BOOL_PROP(kSystemMemPushbackAbortEnabled, false),
           NUM_PROP(kMallocHeapDumpThresholdGb, 20),
           NUM_PROP(kMallocMemMinHeapDumpInterval, 10),
           NUM_PROP(kMallocMemMaxHeapDumpFiles, 5),
@@ -389,6 +390,10 @@ bool SystemConfig::systemMemPushbackEnabled() const {
   return optionalProperty<bool>(kSystemMemPushbackEnabled).value();
 }
 
+bool SystemConfig::systemMemPushBackAbortEnabled() const {
+  return optionalProperty<bool>(kSystemMemPushbackAbortEnabled).value();
+}
+
 bool SystemConfig::mallocMemHeapDumpEnabled() const {
   return optionalProperty<bool>(kMallocMemHeapDumpEnabled).value();
 }
diff --git a/presto-native-execution/presto_cpp/main/common/Configs.h b/presto-native-execution/presto_cpp/main/common/Configs.h
index f0a442deab3eb..20ec529c99436 100644
--- a/presto-native-execution/presto_cpp/main/common/Configs.h
+++ b/presto-native-execution/presto_cpp/main/common/Configs.h
@@ -267,6 +267,11 @@ class SystemConfig : public ConfigBase {
   /// get the server out of low memory condition. This only applies if
   /// 'system-mem-pushback-enabled' is true.
   static constexpr std::string_view kSystemMemShrinkGb{"system-mem-shrink-gb"};
+  /// If true, memory pushback will quickly abort queries with the most memory
+  /// usage under low memory condition. This only applies if
+  /// 'system-mem-pushback-enabled' is set.
+  static constexpr std::string_view kSystemMemPushbackAbortEnabled{
+      "system-mem-pushback-abort-enabled"};
 
   /// If true, memory allocated via malloc is periodically checked and a heap
   /// profile is dumped if usage exceeds 'malloc-heap-dump-gb-threshold'.
@@ -579,6 +584,8 @@ class SystemConfig : public ConfigBase {
 
   uint32_t systemMemShrinkGb() const;
 
+  bool systemMemPushBackAbortEnabled() const;
+
   bool mallocMemHeapDumpEnabled() const;
 
   uint32_t mallocHeapDumpThresholdGb() const;