[native]Add to configure total query meomry usage enforced by arbitrator

NikhilCollooru · Aug 23, 2023 · 1cdc23d · 1cdc23d
1 parent 354426d
commit 1cdc23d
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 4 deletions.
diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.cpp b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
@@ -474,9 +474,8 @@ void PrestoServer::yieldTasks() {
 }
 
 void PrestoServer::initializeVeloxMemory() {
-  auto nodeConfig = NodeConfig::instance();
-  auto systemConfig = SystemConfig::instance();
-  uint64_t memoryGb = systemConfig->systemMemoryGb();
+  auto* systemConfig = SystemConfig::instance();
+  const uint64_t memoryGb = systemConfig->systemMemoryGb();
   PRESTO_STARTUP_LOG(INFO) << "Starting with node memory " << memoryGb << "GB";
 
   const int64_t memoryBytes = memoryGb << 30;
@@ -532,6 +531,12 @@ void PrestoServer::initializeVeloxMemory() {
   options.checkUsageLeak = systemConfig->enableMemoryLeakCheck();
   if (!systemConfig->memoryArbitratorKind().empty()) {
     options.arbitratorKind = systemConfig->memoryArbitratorKind();
+    const uint64_t queryMemoryGb = systemConfig->queryMemoryGb();
+    VELOX_USER_CHECK_LE(
+        queryMemoryGb,
+        memoryGb,
+        "Query memory capacity must not be larger than system memory capacity");
+    options.queryMemoryCapacity = queryMemoryGb << 30;
     options.memoryPoolInitCapacity = systemConfig->memoryPoolInitCapacity();
     options.memoryPoolTransferCapacity =
         systemConfig->memoryPoolTransferCapacity();

diff --git a/presto-native-execution/presto_cpp/main/common/Configs.cpp b/presto-native-execution/presto_cpp/main/common/Configs.cpp
@@ -254,6 +254,7 @@ SystemConfig::SystemConfig() {
           NUM_PROP(kMmapArenaCapacityRatio, 10),
           STR_PROP(kUseMmapAllocator, "true"),
           STR_PROP(kMemoryArbitratorKind, ""),
+          NUM_PROP(kQueryMemoryGb, 38),
           STR_PROP(kEnableVeloxTaskLogging, "false"),
           STR_PROP(kEnableVeloxExprSetLogging, "false"),
           NUM_PROP(kLocalShuffleMaxPartitionBytes, 268435456),
@@ -466,6 +467,10 @@ std::string SystemConfig::memoryArbitratorKind() const {
   return optionalProperty<std::string>(kMemoryArbitratorKind).value_or("");
 }
 
+int32_t SystemConfig::queryMemoryGb() const {
+  return optionalProperty<int32_t>(kQueryMemoryGb).value();
+}
+
 uint64_t SystemConfig::memoryPoolInitCapacity() const {
   static constexpr uint64_t kMemoryPoolInitCapacityDefault = 128 << 20;
   return optionalProperty<uint64_t>(kMemoryPoolInitCapacity)

diff --git a/presto-native-execution/presto_cpp/main/common/Configs.h b/presto-native-execution/presto_cpp/main/common/Configs.h
@@ -184,6 +184,15 @@ class SystemConfig : public ConfigBase {
       "experimental.spiller-spill-path"};
   static constexpr std::string_view kShutdownOnsetSec{"shutdown-onset-sec"};
   static constexpr std::string_view kSystemMemoryGb{"system-memory-gb"};
+  /// Specifies the total memory capacity that can be used by query execution in
+  /// GB. The query memory capacity should be configured less than the system
+  /// memory capacity ('system-memory-gb') to reserve memory for system usage
+  /// such as disk spilling and cache prefetch which are not counted in query
+  /// memory usage.
+  ///
+  /// NOTE: the query memory capacity is enforced by memory arbitrator so that
+  /// this config only applies if the memory arbitration has been enabled.
+  static constexpr std::string_view kQueryMemoryGb{"query-memory-gb"};
   static constexpr std::string_view kAsyncDataCacheEnabled{
       "async-data-cache-enabled"};
   static constexpr std::string_view kAsyncCacheSsdGb{"async-cache-ssd-gb"};
@@ -213,6 +222,7 @@ class SystemConfig : public ConfigBase {
   /// NOTE: this config only applies if the memory arbitration has been enabled.
   static constexpr std::string_view kMemoryPoolInitCapacity{
       "memory-pool-init-capacity"};
+
   /// The minimal memory capacity in bytes transferred between memory pools
   /// during memory arbitration.
   ///
@@ -225,7 +235,6 @@ class SystemConfig : public ConfigBase {
   /// NOTE: this config only applies if the memory arbitration has been enabled.
   static constexpr std::string_view kReservedMemoryPoolCapacityPct{
       "reserved-memory-pool-capacity-pct"};
-
   static constexpr std::string_view kEnableVeloxTaskLogging{
       "enable_velox_task_logging"};
   static constexpr std::string_view kEnableVeloxExprSetLogging{
@@ -409,6 +418,8 @@ class SystemConfig : public ConfigBase {
 
   std::string memoryArbitratorKind() const;
 
+  int32_t queryMemoryGb() const;
+
   uint64_t memoryPoolInitCapacity() const;
 
   uint64_t memoryPoolTransferCapacity() const;