From bb5ff353355f598eb4e92b9c2711d270bb9d4caa Mon Sep 17 00:00:00 2001
From: Kinshuk Bairagi <hi@kinsh.uk>
Date: Wed, 11 Dec 2024 23:49:05 +0530
Subject: [PATCH] Update DropwizardExports.java

Signed-off-by: Kinshuk Bairagi <hi@kinsh.uk>
---
 .../dropwizard/DropwizardExports.java         | 37 +++++++++++--------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/prometheus-metrics-instrumentation-dropwizard/src/main/java/io/prometheus/metrics/instrumentation/dropwizard/DropwizardExports.java b/prometheus-metrics-instrumentation-dropwizard/src/main/java/io/prometheus/metrics/instrumentation/dropwizard/DropwizardExports.java
index 7c3378d1c..a5d44b763 100644
--- a/prometheus-metrics-instrumentation-dropwizard/src/main/java/io/prometheus/metrics/instrumentation/dropwizard/DropwizardExports.java
+++ b/prometheus-metrics-instrumentation-dropwizard/src/main/java/io/prometheus/metrics/instrumentation/dropwizard/DropwizardExports.java
@@ -79,7 +79,7 @@ MetricSnapshot fromCounter(String dropwizardName, Counter counter) {
     /**
      * Export gauge as a prometheus gauge.
      */
-    MetricSnapshot fromGauge(String dropwizardName, Gauge gauge) {
+    MetricSnapshot fromGauge(String dropwizardName, Gauge<?> gauge) {
         Object obj = gauge.getValue();
         double value;
         if (obj instanceof Number) {
@@ -151,21 +151,26 @@ MetricSnapshot fromMeter(String dropwizardName, Meter meter) {
     public MetricSnapshots collect() {
         MetricSnapshots.Builder metricSnapshots = MetricSnapshots.builder();
 
-        for (SortedMap.Entry<String, Gauge> entry : registry.getGauges(metricFilter).entrySet()) {
-            Optional.ofNullable(fromGauge(entry.getKey(), entry.getValue())).map(metricSnapshots::metricSnapshot);
-        }
-        for (SortedMap.Entry<String, Counter> entry : registry.getCounters(metricFilter).entrySet()) {
-            metricSnapshots.metricSnapshot(fromCounter(entry.getKey(), entry.getValue()));
-        }
-        for (SortedMap.Entry<String, Histogram> entry : registry.getHistograms(metricFilter).entrySet()) {
-            metricSnapshots.metricSnapshot(fromHistogram(entry.getKey(), entry.getValue()));
-        }
-        for (SortedMap.Entry<String, Timer> entry : registry.getTimers(metricFilter).entrySet()) {
-            metricSnapshots.metricSnapshot(fromTimer(entry.getKey(), entry.getValue()));
-        }
-        for (SortedMap.Entry<String, Meter> entry : registry.getMeters(metricFilter).entrySet()) {
-            metricSnapshots.metricSnapshot(fromMeter(entry.getKey(), entry.getValue()));
-        }
+        registry.getGauges(metricFilter).forEach((name, gauge) -> {
+            MetricSnapshot snapshot = fromGauge(name, gauge);
+            if (snapshot != null) {
+                metricSnapshots.metricSnapshot(snapshot);
+            }
+        });
+
+        registry.getCounters(metricFilter).forEach((name, counter) ->
+                metricSnapshots.metricSnapshot(fromCounter(name, counter))
+        );
+        registry.getHistograms(metricFilter).forEach((name, histogram) ->
+                metricSnapshots.metricSnapshot(fromHistogram(name, histogram))
+        );
+        registry.getTimers(metricFilter).forEach((name, timer) ->
+                metricSnapshots.metricSnapshot(fromTimer(name, timer))
+        );
+        registry.getMeters(metricFilter).forEach((name, meter) ->
+                metricSnapshots.metricSnapshot(fromMeter(name, meter))
+        );
+
         return metricSnapshots.build();
     }