From ae6a0de856f08f7575f768d8381b9552eee9c542 Mon Sep 17 00:00:00 2001
From: Amit Dutta <amit.kolorob@gmail.com>
Date: Sun, 21 Jan 2024 19:06:03 -0800
Subject: [PATCH] [native] Add config for stats reporting and counters.

---
 .../presto_cpp/main/PrestoServer.cpp                   | 10 ++++++++++
 presto-native-execution/presto_cpp/main/PrestoServer.h |  7 +++++--
 .../presto_cpp/main/common/Configs.cpp                 |  5 +++++
 .../presto_cpp/main/common/Configs.h                   |  5 +++++
 4 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.cpp b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
index 3fe7f2482b539..a3e90d35c9512 100644
--- a/presto-native-execution/presto_cpp/main/PrestoServer.cpp
+++ b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
@@ -143,6 +143,10 @@ void PrestoServer::run() {
     baseVeloxQueryConfig->initialize(
         fmt::format("{}/velox.properties", configDirectoryPath_));
 
+    if (systemConfig->enableRuntimeMetricsCollection()) {
+      enableRuntimeMetricReporting();
+    }
+
     httpPort = systemConfig->httpServerHttpPort();
     if (systemConfig->httpServerHttpsEnabled()) {
       httpsPort = systemConfig->httpServerHttpsPort();
@@ -987,6 +991,12 @@ std::string PrestoServer::getBaseSpillDirectory() const {
   return SystemConfig::instance()->spillerSpillPath().value_or("");
 }
 
+void PrestoServer::enableRuntimeMetricReporting() {
+  // This flag must be set to register the counters.
+  facebook::velox::BaseStatsReporter::registered = true;
+  registerStatsCounters();
+}
+
 void PrestoServer::populateMemAndCPUInfo() {
   auto systemConfig = SystemConfig::instance();
   const int64_t nodeMemoryGb = systemConfig->systemMemoryGb();
diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.h b/presto-native-execution/presto_cpp/main/PrestoServer.h
index 0569b7dd5e529..ed5117f5a45ad 100644
--- a/presto-native-execution/presto_cpp/main/PrestoServer.h
+++ b/presto-native-execution/presto_cpp/main/PrestoServer.h
@@ -126,8 +126,6 @@ class PrestoServer {
 
   virtual void registerMemoryArbitrators();
 
-  virtual void registerStatsCounters();
-
   /// Invoked after creating global (singleton) config objects (SystemConfig and
   /// NodeConfig) and before loading their properties from the file.
   /// In the implementation any extra config properties can be registered.
@@ -143,6 +141,9 @@ class PrestoServer {
   /// Invoked to get the spill directory.
   virtual std::string getBaseSpillDirectory() const;
 
+  /// Invoked to enable stats reporting and register counters.
+  virtual void enableRuntimeMetricReporting();
+
   /// Invoked to get the list of filters passed to the http server.
   std::vector<std::unique_ptr<proxygen::RequestHandlerFactory>>
   getHttpServerFilters();
@@ -151,6 +152,8 @@ class PrestoServer {
 
   void initializeThreadPools();
 
+  void registerStatsCounters();
+
  protected:
   void addServerPeriodicTasks();
 
diff --git a/presto-native-execution/presto_cpp/main/common/Configs.cpp b/presto-native-execution/presto_cpp/main/common/Configs.cpp
index 709778ea3a424..4bcc3f11ad888 100644
--- a/presto-native-execution/presto_cpp/main/common/Configs.cpp
+++ b/presto-native-execution/presto_cpp/main/common/Configs.cpp
@@ -215,6 +215,7 @@ SystemConfig::SystemConfig() {
           BOOL_PROP(kCacheVeloxTtlEnabled, false),
           STR_PROP(kCacheVeloxTtlThreshold, "2d"),
           STR_PROP(kCacheVeloxTtlCheckInterval, "1h"),
+          BOOL_PROP(kEnableRuntimeMetricsCollection, false),
       };
 }
 
@@ -591,6 +592,10 @@ std::chrono::duration<double> SystemConfig::cacheVeloxTtlCheckInterval() const {
       optionalProperty(kCacheVeloxTtlCheckInterval).value());
 }
 
+bool SystemConfig::enableRuntimeMetricsCollection() const {
+  return optionalProperty<bool>(kEnableRuntimeMetricsCollection).value();
+}
+
 NodeConfig::NodeConfig() {
   registeredProps_ =
       std::unordered_map<std::string, folly::Optional<std::string>>{
diff --git a/presto-native-execution/presto_cpp/main/common/Configs.h b/presto-native-execution/presto_cpp/main/common/Configs.h
index 6b2ef670d0f98..16c0d4a9da306 100644
--- a/presto-native-execution/presto_cpp/main/common/Configs.h
+++ b/presto-native-execution/presto_cpp/main/common/Configs.h
@@ -300,6 +300,9 @@ class SystemConfig : public ConfigBase {
       "cache.velox.ttl-check-interval"};
   static constexpr std::string_view kUseMmapAllocator{"use-mmap-allocator"};
 
+  static constexpr std::string_view kEnableRuntimeMetricsCollection{
+      "runtime-metrics-collection-enabled"};
+
   /// Specifies the memory arbitrator kind. If it is empty, then there is no
   /// memory arbitration.
   static constexpr std::string_view kMemoryArbitratorKind{
@@ -655,6 +658,8 @@ class SystemConfig : public ConfigBase {
   std::chrono::duration<double> cacheVeloxTtlThreshold() const;
 
   std::chrono::duration<double> cacheVeloxTtlCheckInterval() const;
+
+  bool enableRuntimeMetricsCollection() const;
 };
 
 /// Provides access to node properties defined in node.properties file.