From b78616c5b09335fe35415c48b7b2367d810ad1db Mon Sep 17 00:00:00 2001
From: Leah Thomas <lthomas@confluent.io>
Date: Thu, 2 Sep 2021 13:39:09 -0500
Subject: [PATCH] feat: update storage utilization metrics to start when app is
 initialized

---
 .../internal/StorageUtilizationMetricsReporter.java | 13 +------------
 .../ksql/rest/server/KsqlRestApplication.java       |  8 ++++++++
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/ksqldb-engine/src/main/java/io/confluent/ksql/internal/StorageUtilizationMetricsReporter.java b/ksqldb-engine/src/main/java/io/confluent/ksql/internal/StorageUtilizationMetricsReporter.java
index 274a10fc8003..423d756ceda9 100644
--- a/ksqldb-engine/src/main/java/io/confluent/ksql/internal/StorageUtilizationMetricsReporter.java
+++ b/ksqldb-engine/src/main/java/io/confluent/ksql/internal/StorageUtilizationMetricsReporter.java
@@ -67,20 +67,9 @@ public void init(final List<KafkaMetric> list) {
 
   @Override
   public void configure(final Map<String, ?> map) {
-    final String dir;
-    if (map.containsKey(StreamsConfig.STATE_DIR_CONFIG)) {
-      dir = map.get(StreamsConfig.STATE_DIR_CONFIG).toString();
-    } else {
-      dir =  StreamsConfig
-        .configDef()
-        .defaultValues()
-        .get(StreamsConfig.STATE_DIR_CONFIG)
-        .toString();
-    }
-    configureShared(new File(dir), this.metricRegistry);
   }
 
-  private static void configureShared(final File baseDir, final Metrics metricRegistry) {
+  public static void configureShared(final File baseDir, final Metrics metricRegistry) {
     if (registeredNodeMetrics.getAndSet(true)) {
       return;
     }
diff --git a/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/KsqlRestApplication.java b/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/KsqlRestApplication.java
index 14b82dc170a1..e0d229e96477 100644
--- a/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/KsqlRestApplication.java
+++ b/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/KsqlRestApplication.java
@@ -43,6 +43,7 @@
 import io.confluent.ksql.function.MutableFunctionRegistry;
 import io.confluent.ksql.function.UserFunctionLoader;
 import io.confluent.ksql.internal.PullQueryExecutorMetrics;
+import io.confluent.ksql.internal.StorageUtilizationMetricsReporter;
 import io.confluent.ksql.logging.processing.ProcessingLogConfig;
 import io.confluent.ksql.logging.processing.ProcessingLogContext;
 import io.confluent.ksql.logging.processing.ProcessingLogServerUtils;
@@ -699,6 +700,13 @@ static KsqlRestApplication buildApplication(
 
     final SpecificQueryIdGenerator specificQueryIdGenerator =
         new SpecificQueryIdGenerator();
+    
+    final String stateDir = ksqlConfig.getKsqlStreamConfigProps().getOrDefault(
+      StreamsConfig.STATE_DIR_CONFIG,
+      StreamsConfig.configDef().defaultValues().get(StreamsConfig.STATE_DIR_CONFIG))
+      .toString();
+
+    StorageUtilizationMetricsReporter.configureShared(new File(stateDir), MetricCollectors.getMetrics());
 
     final KsqlEngine ksqlEngine = new KsqlEngine(
         serviceContext,