From a28c4ecdb223c9539b4718b49f9cea685e9f1817 Mon Sep 17 00:00:00 2001
From: Gagan Juneja <gjjuneja@amazon.com>
Date: Wed, 30 Oct 2024 22:38:06 +0530
Subject: [PATCH] test changes

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>
---
 .../telemetry/RTFHeapMetricsCollector.java     | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/main/java/org/opensearch/performanceanalyzer/collectors/telemetry/RTFHeapMetricsCollector.java b/src/main/java/org/opensearch/performanceanalyzer/collectors/telemetry/RTFHeapMetricsCollector.java
index 81c1a1c0..93929e87 100644
--- a/src/main/java/org/opensearch/performanceanalyzer/collectors/telemetry/RTFHeapMetricsCollector.java
+++ b/src/main/java/org/opensearch/performanceanalyzer/collectors/telemetry/RTFHeapMetricsCollector.java
@@ -140,25 +140,31 @@ private void recordMetrics() {
             heapUsedMetrics.record(
                     memoryUsage.getUsed(),
                     Tags.create().addTag(memTypeAttributeKey, entry.getKey()));
-            createGaugeInstanceIfNotAvailable(entry.getKey(), memoryUsage);
+            createGaugeInstanceIfNotAvailable(entry.getKey());
         }
     }
 
-    private void createGaugeInstanceIfNotAvailable(String key, MemoryUsage memoryUsage) {
+    private void createGaugeInstanceIfNotAvailable(String key) {
         if (!memTypeToGaugeObservableMap.containsKey(key)) {
+            LOG.info("Gauge doesn't exist for the mem type {}", key);
             Closeable observableGauge =
                     metricsRegistry.createGauge(
                             RTFMetrics.HeapValue.Constants.MAX_VALUE,
                             "Heap Max PA metrics",
                             "",
-                            () -> (double) memoryUsage.getMax(),
+                            () -> getValue(key),
                             Tags.create().addTag(memTypeAttributeKey, key));
             memTypeToGaugeObservableMap.put(key, observableGauge);
         }
     }
 
-    private double getMaxValueForMemTypeAttribute(
-            String memTypeAttributeKey, MemoryUsage memoryUsage) {
-        return memoryUsage.getMax();
+    private double getValue(String key) {
+        Map<String, Supplier<MemoryUsage>> memoryUsageSuppliers = HeapMetrics.getMemoryUsageSuppliers();
+        if(memoryUsageSuppliers.get(key) != null){
+            MemoryUsage memoryUsage = memoryUsageSuppliers.get(key).get();
+            return memoryUsage.getMax();
+        }else{
+            return 0.0;
+        }
     }
 }