Fix resource leak (#68)

* Fixes potential resourceleak, enhance shutdown logging * Remove empty function body
teragrep · Jul 25, 2023 · 6d2526a · 6d2526a
1 parent 637eda3
commit 6d2526a
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 36 deletions.
diff --git a/src/main/java/com/teragrep/k8s_01/KubernetesLogReader.java b/src/main/java/com/teragrep/k8s_01/KubernetesLogReader.java
@@ -111,25 +111,6 @@ public static void main(String[] args) throws IOException {
             }
         }
 
-        // Graceful shutdown so Relp sessions are gracefully terminated
-        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
-            LOGGER.info("Shutting down.");
-            for(int i=1; i <= appConfig.getRelp().getOutputThreads(); i++) {
-                LOGGER.info(
-                        "Disconnecting relp thread #{}/{}",
-                        i,
-                        appConfig.getRelp().getOutputThreads()
-                );
-                RelpOutput output;
-                try {
-                    output = relpOutputPool.take();
-                } catch (InterruptedException e) {
-                    throw new RuntimeException(e);
-                }
-                output.disconnect();
-            }
-        }));
-
         // consumer supplier, returns always the same instance
         K8SConsumerSupplier consumerSupplier = new K8SConsumerSupplier(appConfig, cacheClient, relpOutputPool);
         String[] logfiles = appConfig.getKubernetes().getLogfiles();
@@ -151,6 +132,28 @@ public static void main(String[] args) throws IOException {
             consumerSupplier
         );
 
+        // Graceful shutdown so Relp sessions are gracefully terminated
+        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
+            LOGGER.info("Shutting down.");
+            for(int i=1; i <= appConfig.getRelp().getOutputThreads(); i++) {
+                LOGGER.info(
+                        "Disconnecting relp thread #{}/{}",
+                        i,
+                        appConfig.getRelp().getOutputThreads()
+                );
+                RelpOutput output;
+                try {
+                    output = relpOutputPool.take();
+                } catch (InterruptedException e) {
+                    throw new RuntimeException(e);
+                }
+                output.disconnect();
+            }
+            prometheusMetrics.close();
+            statefulFileReader.close();
+            LOGGER.info("Goodbye");
+        }, "ShutdownHook"));
+
         // Start a new thread for all logfile watchers
         for (String logfile : logfiles) {
             Thread thread = new Thread(() -> {
@@ -199,7 +202,5 @@ public static void main(String[] args) throws IOException {
                 throw new RuntimeException(e);
             }
         }
-        prometheusMetrics.close();
-        statefulFileReader.close();
     }
 }
diff --git a/src/main/java/com/teragrep/k8s_01/PrometheusMetrics.java b/src/main/java/com/teragrep/k8s_01/PrometheusMetrics.java
@@ -35,8 +35,10 @@
 public class PrometheusMetrics {
     private static final Logger LOGGER = LoggerFactory.getLogger(PrometheusMetrics.class);
     private final Server jettyServer;
+    private final JmxReporter jmxReporter;
     public PrometheusMetrics(int port) {
         LOGGER.info("Starting prometheus metrics server on port {}", port);
+
         // prometheus-exporter
         jettyServer = new Server(port);
         ServletContextHandler context = new ServletContextHandler();
@@ -46,18 +48,8 @@ public PrometheusMetrics(int port) {
         MetricsServlet metricsServlet = new MetricsServlet();
         ServletHolder servletHolder = new ServletHolder(metricsServlet);
         context.addServlet(servletHolder, "/metrics");
-        setupDropWizard();
-        // Add metrics about CPU, JVM memory etc.
-        DefaultExports.initialize();
-        // Start the webserver.
-        try {
-            jettyServer.start();
-        } catch (Exception e) {
-            throw new RuntimeException(e);
-        }
-    }
 
-    static void setupDropWizard() {
+        // Container for all metrics
         MetricRegistry metricRegistry = new MetricRegistry();
 
         // Totals
@@ -77,19 +69,33 @@ static void setupDropWizard() {
         SharedMetricRegistries.add("default", metricRegistry);
 
         // Add to Prometheus metrics
-        CollectorRegistry.defaultRegistry.register(new DropwizardExports(metricRegistry));
+        CollectorRegistry.defaultRegistry.register(
+                new DropwizardExports(metricRegistry)
+        );
 
         // Enable JMX listener
-        JmxReporter jmxReporter = JmxReporter.forRegistry(metricRegistry).build();
+        jmxReporter = JmxReporter.forRegistry(metricRegistry).build();
         jmxReporter.start();
+
+        // Add metrics about CPU, JVM memory etc.
+        DefaultExports.initialize();
+
+        // Start the webserver.
+        try {
+            jettyServer.start();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
     }
 
     public void close() {
-        LOGGER.info("Closing prometheus metrics server");
         try {
+            LOGGER.info("Closing jettyserver");
             jettyServer.stop();
+            LOGGER.info("Closing jmxReporter");
+            jmxReporter.stop();
         } catch (Exception e) {
-            LOGGER.error("Failed to stop jettyServer:", e);
+            LOGGER.error("Failed to stop jettyServer and/or jmxReporter:", e);
         }
     }
 }