Make prometheus metric updation async

Author: jaystarshot

presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.cpp

@@ -47,8 +47,9 @@ struct PrometheusStatsReporter::PrometheusImpl {
 };
 
 PrometheusStatsReporter::PrometheusStatsReporter(
-    const std::map<std::string, std::string>& labels) {
-  impl_ = std::make_shared<PrometheusImpl>(labels);
+    const std::map<std::string, std::string>& labels, int numThreads)
+    : executor_(std::make_shared<folly::CPUThreadPoolExecutor>(numThreads)),
+      impl_(std::make_shared<PrometheusImpl>(labels)) {
 }
 
 void PrometheusStatsReporter::registerMetricExportType(
@@ -161,36 +162,38 @@ void PrometheusStatsReporter::addMetricValue(
 
 void PrometheusStatsReporter::addMetricValue(const char* key, size_t value)
     const {
-  auto metricIterator = registeredMetricsMap_.find(key);
-  if (metricIterator == registeredMetricsMap_.end()) {
-    VLOG(1) << "addMetricValue called for unregistered metric " << key;
-    return;
-  }
-  auto statsInfo = metricIterator->second;
-  switch (statsInfo.statType) {
-    case velox::StatType::COUNT: {
-      auto* counter =
-          reinterpret_cast<::prometheus::Counter*>(statsInfo.metricPtr);
-      counter->Increment(static_cast<double>(value));
-      break;
-    }
-    case velox::StatType::SUM: {
-      auto* gauge = reinterpret_cast<::prometheus::Gauge*>(statsInfo.metricPtr);
-      gauge->Increment(static_cast<double>(value));
-      break;
-    }
-    case velox::StatType::AVG:
-    case velox::StatType::RATE: {
-      // Overrides the existing state.
-      auto* gauge = reinterpret_cast<::prometheus::Gauge*>(statsInfo.metricPtr);
-      gauge->Set(static_cast<double>(value));
-      break;
+  executor_->add([this, key, value]() {
+    auto metricIterator = registeredMetricsMap_.find(key);
+    if (metricIterator == registeredMetricsMap_.end()) {
+      VLOG(1) << "addMetricValue called for unregistered metric " << key;
+      return;
     }
-    default:
-      VELOX_UNSUPPORTED(
-          "Unsupported metric type {}",
-          velox::statTypeString(statsInfo.statType));
-  };
+    auto statsInfo = metricIterator->second;
+    switch (statsInfo.statType) {
+      case velox::StatType::COUNT: {
+        auto* counter =
+            reinterpret_cast<::prometheus::Counter*>(statsInfo.metricPtr);
+        counter->Increment(static_cast<double>(value));
+        break;
+      }
+      case velox::StatType::SUM: {
+        auto* gauge = reinterpret_cast<::prometheus::Gauge*>(statsInfo.metricPtr);
+        gauge->Increment(static_cast<double>(value));
+        break;
+      }
+      case velox::StatType::AVG:
+      case velox::StatType::RATE: {
+        // Overrides the existing state.
+        auto* gauge = reinterpret_cast<::prometheus::Gauge*>(statsInfo.metricPtr);
+        gauge->Set(static_cast<double>(value));
+        break;
+      }
+      default:
+        VELOX_UNSUPPORTED(
+            "Unsupported metric type {}",
+            velox::statTypeString(statsInfo.statType));
+    };
+  });
 }
 
 void PrometheusStatsReporter::addMetricValue(
@@ -208,22 +211,24 @@ void PrometheusStatsReporter::addHistogramMetricValue(
 void PrometheusStatsReporter::addHistogramMetricValue(
     const char* key,
     size_t value) const {
-  auto metricIterator = registeredMetricsMap_.find(key);
-  if (metricIterator == registeredMetricsMap_.end()) {
-    VLOG(1) << "addMetricValue for unregistered metric " << key;
-    return;
-  }
-  auto histogram = reinterpret_cast<::prometheus::Histogram*>(
-      metricIterator->second.metricPtr);
-  histogram->Observe(value);
-
-  std::string summaryKey = std::string(key).append(kSummarySuffix);
-  metricIterator = registeredMetricsMap_.find(summaryKey);
-  if (metricIterator != registeredMetricsMap_.end()) {
-    auto summary = reinterpret_cast<::prometheus::Summary*>(
+  executor_->add([this, key, value]() {
+    auto metricIterator = registeredMetricsMap_.find(key);
+    if (metricIterator == registeredMetricsMap_.end()) {
+      VLOG(1) << "addMetricValue for unregistered metric " << key;
+      return;
+    }
+    auto histogram = reinterpret_cast<::prometheus::Histogram*>(
         metricIterator->second.metricPtr);
-    summary->Observe(value);
-  }
+    histogram->Observe(value);
+
+    std::string summaryKey = std::string(key).append(kSummarySuffix);
+    metricIterator = registeredMetricsMap_.find(summaryKey);
+    if (metricIterator != registeredMetricsMap_.end()) {
+      auto summary = reinterpret_cast<::prometheus::Summary*>(
+          metricIterator->second.metricPtr);
+      summary->Observe(value);
+    }
+  });
 }
 
 void PrometheusStatsReporter::addHistogramMetricValue(
@@ -241,4 +246,8 @@ std::string PrometheusStatsReporter::fetchMetrics() {
   return serializer.Serialize(impl_->registry->Collect());
 }
 
+void PrometheusStatsReporter::waitForCompletion() const {
+  executor_->join();
+}
+
 } // namespace facebook::presto::prometheus

presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.h

@@ -16,6 +16,7 @@
 #include "velox/common/base/Exceptions.h"
 #include "velox/common/base/GTestMacros.h"
 #include "velox/common/base/StatsReporter.h"
+#include <folly/executors/CPUThreadPoolExecutor.h>
 
 namespace facebook::presto::prometheus {
 
@@ -38,8 +39,13 @@ class PrometheusStatsReporter : public facebook::velox::BaseStatsReporter {
   class PrometheusImpl;
 
  public:
+  /**
+   * @brief Constructor with optional thread count (for testing).
+   * @param labels Labels for metrics.
+   * @param numThreads Number of threads in the executor (default: 2).
+   */
   explicit PrometheusStatsReporter(
-      const std::map<std::string, std::string>& labels);
+      const std::map<std::string, std::string>& labels, int numThreads = 2);
 
   void registerMetricExportType(const char* key, velox::StatType)
       const override;
@@ -77,6 +83,12 @@ class PrometheusStatsReporter : public facebook::velox::BaseStatsReporter {
 
   std::string fetchMetrics() override;
 
+  /**
+   * Waits for all pending metric updates to complete.
+   * This is only used in tests to ensure correct timing.
+   */
+  void waitForCompletion() const;
+
   static std::unique_ptr<velox::BaseStatsReporter> createPrometheusReporter() {
     auto nodeConfig = NodeConfig::instance();
     const std::string cluster = nodeConfig->nodeEnvironment();
@@ -88,6 +100,7 @@ class PrometheusStatsReporter : public facebook::velox::BaseStatsReporter {
   }
 
  private:
+  std::shared_ptr<folly::CPUThreadPoolExecutor> executor_;
   std::shared_ptr<PrometheusImpl> impl_;
   // A map of labels assigned to each metric which helps in filtering at client
   // end.

presto-native-execution/presto_cpp/main/runtime-metrics/tests/PrometheusReporterTest.cpp

@@ -20,7 +20,7 @@ namespace facebook::presto::prometheus {
 class PrometheusReporterTest : public testing::Test {
  public:
   void SetUp() override {
-    reporter = std::make_shared<PrometheusStatsReporter>(testLabels);
+    reporter = std::make_shared<PrometheusStatsReporter>(testLabels, 1);
   }
 
   void verifySerializedResult(
@@ -76,6 +76,9 @@ TEST_F(PrometheusReporterTest, testCountAndGauge) {
   reporter->addMetricValue("test.key1");
   reporter->addMetricValue("test.key3");
 
+  // 🔹 Wait for all async updates to finish before validation
+  reporter->waitForCompletion();
+
   auto fullSerializedResult = reporter->fetchMetrics();
 
   std::vector<std::string> expected = {
@@ -114,6 +117,8 @@ TEST_F(PrometheusReporterTest, testHistogramSummary) {
     }
   }
   reporter->addHistogramMetricValue(histogramKey, 10);
+  // 🔹 Wait for all async updates to finish before validation
+  reporter->waitForCompletion();
   auto fullSerializedResult = reporter->fetchMetrics();
   std::replace(histSummaryKey.begin(), histSummaryKey.end(), '.', '_');
   std::replace(histogramKey.begin(), histogramKey.end(), '.', '_');