From 6c610b91efad789db7d5b7d8dbe7503f63427de9 Mon Sep 17 00:00:00 2001
From: Prashant Wason <pwason@uber.com>
Date: Sun, 4 Oct 2020 09:30:04 -0700
Subject: [PATCH] [HUDI-1305] Added an API to shutdown and remove the metrics
 reporter. (#2132)

This helps in removing reporter once the test has complete. Prevents log pollution from un-necessary metric logs.

- Added an API to shutdown the metrics reporter after tests.
---
 .../hudi/metrics/ConsoleMetricsReporter.java  |  3 ++
 .../java/org/apache/hudi/metrics/Metrics.java | 30 +++++++++++++------
 .../metrics/TestHoodieConsoleMetrics.java     |  6 ++++
 3 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/ConsoleMetricsReporter.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/ConsoleMetricsReporter.java
index ca961096c..b65c4ade8 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/ConsoleMetricsReporter.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/ConsoleMetricsReporter.java
@@ -68,5 +68,8 @@ public class ConsoleMetricsReporter extends MetricsReporter {
 
   @Override
   public void stop() {
+    if (consoleReporter != null) {
+      consoleReporter.stop();
+    }
   }
 }
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/Metrics.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/Metrics.java
index a9cf14e69..c4107ce02 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/Metrics.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metrics/Metrics.java
@@ -52,18 +52,22 @@ public class Metrics {
     reporter.start();
 
     Runtime.getRuntime().addShutdownHook(new Thread(() -> {
-      try {
-        registerHoodieCommonMetrics();
-        reporter.report();
-        if (getReporter() != null) {
-          getReporter().close();
-        }
-      } catch (Exception e) {
-        LOG.warn("Error while closing reporter", e);
-      }
+      reportAndCloseReporter();
     }));
   }
 
+  private void reportAndCloseReporter() {
+    try {
+      registerHoodieCommonMetrics();
+      reporter.report();
+      if (getReporter() != null) {
+        getReporter().close();
+      }
+    } catch (Exception e) {
+      LOG.warn("Error while closing reporter", e);
+    }
+  }
+
   private void registerHoodieCommonMetrics() {
     registerGauges(Registry.getAllMetrics(true, true), Option.empty());
   }
@@ -85,6 +89,14 @@ public class Metrics {
     initialized = true;
   }
 
+  public static synchronized void shutdown() {
+    if (!initialized) {
+      return;
+    }
+    metrics.reportAndCloseReporter();
+    initialized = false;
+  }
+
   public static void registerGauges(Map<String, Long> metricsMap, Option<String> prefix) {
     String metricPrefix = prefix.isPresent() ? prefix.get() + "." : "";
     metricsMap.forEach((k, v) -> registerGauge(metricPrefix + k, v));
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/metrics/TestHoodieConsoleMetrics.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/metrics/TestHoodieConsoleMetrics.java
index e64473070..7424d0b92 100644
--- a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/metrics/TestHoodieConsoleMetrics.java
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/metrics/TestHoodieConsoleMetrics.java
@@ -19,6 +19,7 @@
 package org.apache.hudi.metrics;
 
 import org.apache.hudi.config.HoodieWriteConfig;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -38,6 +39,11 @@ public class TestHoodieConsoleMetrics {
     new HoodieMetrics(config, "raw_table");
   }
 
+  @AfterEach
+  public void stop() {
+    Metrics.shutdown();
+  }
+
   @Test
   public void testRegisterGauge() {
     registerGauge("metric1", 123L);