Register database custom metrics using a multi gauge

docstore-metrics/src/main/java/org/hypertrace/core/serviceframework/docstore/metrics/DocStoreMetricsRegistry.java

@@ -3,6 +3,8 @@
 import static java.util.Collections.emptyList;
 import static java.util.concurrent.TimeUnit.MINUTES;
 import static java.util.concurrent.TimeUnit.SECONDS;
+import static java.util.stream.Collectors.toUnmodifiableList;
+import static org.hypertrace.core.serviceframework.metrics.PlatformMetricsRegistry.registerResizeableGauge;
 
 import io.micrometer.common.lang.Nullable;
 import java.time.Duration;
@@ -11,13 +13,16 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.atomic.AtomicLong;
 import lombok.NonNull;
+import lombok.extern.slf4j.Slf4j;
 import org.hypertrace.core.documentstore.Datastore;
 import org.hypertrace.core.documentstore.metric.DocStoreMetric;
 import org.hypertrace.core.documentstore.metric.DocStoreMetricProvider;
-import org.hypertrace.core.documentstore.model.config.CustomMetricConfig;
+import org.hypertrace.core.serviceframework.metrics.Measurement;
 import org.hypertrace.core.serviceframework.metrics.PlatformMetricsRegistry;
+import org.hypertrace.core.serviceframework.metrics.ResizeableGauge;
 import org.hypertrace.core.serviceframework.spi.PlatformServiceLifecycle;
 
+@Slf4j
 @SuppressWarnings("unused")
 public class DocStoreMetricsRegistry {
   private static final long INITIAL_DELAY_SECONDS = MINUTES.toSeconds(5);
@@ -93,11 +98,6 @@ public void monitor() {
     monitorCustomMetrics();
   }
 
-  /** Instantly query the datastore and report the custom metric once */
-  public void report(final CustomMetricConfig customMetricConfig) {
-    metricProvider.getCustomMetrics(customMetricConfig).forEach(this::report);
-  }
-
   /** Stop monitoring the database */
   public void shutdown() {
     if (executor != null) {
@@ -112,17 +112,40 @@ private void addShutdownHook() {
   }
 
   private void monitorCustomMetrics() {
-    customMetricConfigs.forEach(
-        reportingConfig ->
-            executor.scheduleAtFixedRate(
-                () -> report(reportingConfig.config()),
-                INITIAL_DELAY_SECONDS,
-                reportingConfig.reportingInterval().toSeconds(),
-                SECONDS));
+    customMetricConfigs.forEach(this::monitorCustomMetric);
   }
 
-  private void report(final DocStoreMetric metric) {
-    PlatformMetricsRegistry.registerGauge(metric.name(), metric.labels(), metric.value());
+  private void monitorCustomMetric(final DocStoreCustomMetricReportingConfig reportingConfig) {
+    final ResizeableGauge resizeableGauge =
+        registerResizeableGauge(reportingConfig.config().metricName());
+    executor.scheduleAtFixedRate(
+        () -> report(reportingConfig, resizeableGauge),
+        INITIAL_DELAY_SECONDS,
+        reportingConfig.reportingInterval().toSeconds(),
+        SECONDS);
+  }
+
+  private void report(
+      final DocStoreCustomMetricReportingConfig reportingConfig,
+      final ResizeableGauge resizeableGauge) {
+    try {
+      final List<DocStoreMetric> customMetrics =
+          metricProvider.getCustomMetrics(reportingConfig.config());
+
+      log.debug(
+          "Reporting custom database metrics {} for configuration {}",
+          customMetrics,
+          reportingConfig);
+
+      final List<Measurement> measurements =
+          customMetrics.stream()
+              .map(metric -> new Measurement(metric.value(), metric.labels()))
+              .collect(toUnmodifiableList());
+
+      resizeableGauge.report(measurements);
+    } catch (final Exception e) {
+      log.warn("Unable to report custom database metric for config: {}", reportingConfig, e);
+    }
   }
 
   private class StandardDocStoreMetricsRegistry {

platform-metrics/src/main/java/org/hypertrace/core/serviceframework/metrics/Measurement.java

@@ -0,0 +1,21 @@
+package org.hypertrace.core.serviceframework.metrics;
+
+import java.util.Map;
+
+public class Measurement {
+  private final double value;
+  private final Map<String, String> labels;
+
+  public Measurement(final double value, final Map<String, String> labels) {
+    this.value = value;
+    this.labels = labels;
+  }
+
+  double value() {
+    return value;
+  }
+
+  Map<String, String> labels() {
+    return labels;
+  }
+}

platform-metrics/src/main/java/org/hypertrace/core/serviceframework/metrics/PlatformMetricsRegistry.java

@@ -16,6 +16,7 @@
 import io.micrometer.core.instrument.Gauge;
 import io.micrometer.core.instrument.ImmutableTag;
 import io.micrometer.core.instrument.MeterRegistry;
+import io.micrometer.core.instrument.MultiGauge;
 import io.micrometer.core.instrument.Tag;
 import io.micrometer.core.instrument.Timer;
 import io.micrometer.core.instrument.binder.cache.GuavaCacheMetrics;
@@ -508,16 +509,6 @@ public static void monitorExecutorService(
     new ExecutorServiceMetrics(executorService, name, toIterable(tags)).bindTo(meterRegistry);
   }
 
-  private static Iterable<Tag> toIterable(Map<String, String> tags) {
-    List<Tag> newTags = new ArrayList<>();
-
-    if (tags != null) {
-      tags.forEach((k, v) -> newTags.add(new ImmutableTag(k, v)));
-    }
-
-    return newTags;
-  }
-
   public static MetricRegistry getMetricRegistry() {
     return METRIC_REGISTRY;
   }
@@ -542,6 +533,20 @@ public static synchronized void stop() {
     isInit = false;
   }
 
+  public static ResizeableGauge registerResizeableGauge(final String name) {
+    return new ResizeableGauge(MultiGauge.builder(name).register(meterRegistry));
+  }
+
+  static Iterable<Tag> toIterable(Map<String, String> tags) {
+    List<Tag> newTags = new ArrayList<>();
+
+    if (tags != null) {
+      tags.forEach((k, v) -> newTags.add(new ImmutableTag(k, v)));
+    }
+
+    return newTags;
+  }
+
   /*
    * This is needed because ConsoleMetricReporter.stop() doesn't call report for the last time
    * before closing the scheduled thread

platform-metrics/src/main/java/org/hypertrace/core/serviceframework/metrics/ResizeableGauge.java

@@ -0,0 +1,29 @@
+package org.hypertrace.core.serviceframework.metrics;
+
+import static java.util.stream.Collectors.toUnmodifiableList;
+import static org.hypertrace.core.serviceframework.metrics.PlatformMetricsRegistry.toIterable;
+
+import io.micrometer.core.instrument.MultiGauge;
+import io.micrometer.core.instrument.MultiGauge.Row;
+import io.micrometer.core.instrument.Tags;
+import java.util.Collection;
+import java.util.List;
+
+public class ResizeableGauge {
+  private final MultiGauge multiGauge;
+
+  ResizeableGauge(final MultiGauge multiGauge) {
+    this.multiGauge = multiGauge;
+  }
+
+  public void report(final Collection<Measurement> measurements) {
+    final List<Row<?>> rows =
+        measurements.stream()
+            .map(
+                measurement ->
+                    Row.of(Tags.of(toIterable(measurement.labels())), measurement::value))
+            .collect(toUnmodifiableList());
+
+    multiGauge.register(rows);
+  }
+}