Revert "Revert "MetricReceiverWrapper has type indicating if tags are populated""

This reverts commit 6968ba5b89d9a0a02f08925fc338e429f6efbf1e.

2 files changed, 61 insertions, 15 deletions

diff --git a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/metrics/MetricReceiverWrapper.java b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/metrics/MetricReceiverWrapper.java
index 05d24c671ab..58126a59cbb 100644
--- a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/metrics/MetricReceiverWrapper.java
+++ b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/metrics/MetricReceiverWrapper.java
@@ -6,11 +6,13 @@ import com.yahoo.metrics.simple.MetricReceiver;
 import com.yahoo.metrics.simple.Point;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -27,7 +29,7 @@ public class MetricReceiverWrapper {
     public static final String APPLICATION_NODE = "vespa.node";
 
     private final Object monitor = new Object();
-    private final Map<String, ApplicationMetrics> applicationMetrics = new HashMap<>(); // key is application name
+    private final Map<DimensionType, Map<String, ApplicationMetrics>> metrics = new HashMap<>();
     private final MetricReceiver metricReceiver;
 
     @Inject
@@ -39,8 +41,12 @@ public class MetricReceiverWrapper {
      *  Declaring the same dimensions and name results in the same CounterWrapper instance (idempotent).
      */
     public CounterWrapper declareCounter(String application, Dimensions dimensions, String name) {
+        return declareCounter(application, dimensions, name, DimensionType.DEFAULT);
+    }
+
+    public CounterWrapper declareCounter(String application, Dimensions dimensions, String name, DimensionType type) {
         synchronized (monitor) {
-            Map<Dimensions, Map<String, MetricValue>> metricsByDimensions = getOrCreateApplicationMetrics(application);
+            Map<Dimensions, Map<String, MetricValue>> metricsByDimensions = getOrCreateApplicationMetrics(application, type);
             if (!metricsByDimensions.containsKey(dimensions)) metricsByDimensions.put(dimensions, new HashMap<>());
             if (!metricsByDimensions.get(dimensions).containsKey(name)) {
                 CounterWrapper counter = new CounterWrapper(metricReceiver.declareCounter(name, new Point(dimensions.dimensionsMap)));
@@ -55,8 +61,12 @@ public class MetricReceiverWrapper {
      *  Declaring the same dimensions and name results in the same GaugeWrapper instance (idempotent).
      */
     public GaugeWrapper declareGauge(String application, Dimensions dimensions, String name) {
+        return declareGauge(application, dimensions, name, DimensionType.DEFAULT);
+    }
+
+    public GaugeWrapper declareGauge(String application, Dimensions dimensions, String name, DimensionType type) {
         synchronized (monitor) {
-            Map<Dimensions, Map<String, MetricValue>> metricsByDimensions = getOrCreateApplicationMetrics(application);
+            Map<Dimensions, Map<String, MetricValue>> metricsByDimensions = getOrCreateApplicationMetrics(application, type);
             if (!metricsByDimensions.containsKey(dimensions))
                 metricsByDimensions.put(dimensions, new HashMap<>());
             if (!metricsByDimensions.get(dimensions).containsKey(name)) {
@@ -68,40 +78,58 @@ public class MetricReceiverWrapper {
         }
     }
 
-    public List<DimensionMetrics> getAllMetrics() {
+    public List<DimensionMetrics> getDefaultMetrics() {
+        return getMetricsByType(DimensionType.DEFAULT);
+    }
+
+    // For testing, returns same as getDefaultMetrics(), but without "timestamp"
+    public Set<Map<String, Object>> getDefaultMetricsRaw() {
         synchronized (monitor) {
-            List<DimensionMetrics> dimensionMetrics = new ArrayList<>();
-            applicationMetrics.forEach((application, applicationMetrics) -> applicationMetrics.metricsByDimensions().entrySet().stream()
+            Set<Map<String, Object>> dimensionMetrics = new HashSet<>();
+            metrics.getOrDefault(DimensionType.DEFAULT, new HashMap<>())
+                    .forEach((application, applicationMetrics) -> applicationMetrics.metricsByDimensions().entrySet().stream()
                     .map(entry -> new DimensionMetrics(application, entry.getKey(),
                             entry.getValue().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, value -> value.getValue().getValue()))))
+                    .map(DimensionMetrics::getMetrics)
                     .forEach(dimensionMetrics::add));
             return dimensionMetrics;
         }
     }
 
-    // For testing, returns same as getAllMetrics(), but without "timestamp"
-    public Set<Map<String, Object>> getAllMetricsRaw() {
+    public List<DimensionMetrics> getMetricsByType(DimensionType type) {
         synchronized (monitor) {
-            Set<Map<String, Object>> dimensionMetrics = new HashSet<>();
-            applicationMetrics.forEach((application, applicationMetrics) -> applicationMetrics.metricsByDimensions().entrySet().stream()
+            List<DimensionMetrics> dimensionMetrics = new ArrayList<>();
+            metrics.getOrDefault(type, new HashMap<>())
+                    .forEach((application, applicationMetrics) -> applicationMetrics.metricsByDimensions().entrySet().stream()
                     .map(entry -> new DimensionMetrics(application, entry.getKey(),
                             entry.getValue().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, value -> value.getValue().getValue()))))
-                    .map(DimensionMetrics::getMetrics)
                     .forEach(dimensionMetrics::add));
             return dimensionMetrics;
         }
     }
 
+    public void deleteMetricByDimension(String name,  Dimensions dimensionsToRemove, DimensionType type) {
+        synchronized (monitor) {
+            Optional.ofNullable(metrics.get(type))
+                    .map(m -> m.get(name))
+                    .map(ApplicationMetrics::metricsByDimensions)
+                    .ifPresent(m -> m.remove(dimensionsToRemove));
+        }
+    }
+
     // For testing
     Map<String, Number> getMetricsForDimension(String application, Dimensions dimensions) {
         synchronized (monitor) {
-            Map<Dimensions, Map<String, MetricValue>> metricsByDimensions = getOrCreateApplicationMetrics(application);
-            return metricsByDimensions.get(dimensions).entrySet().stream().collect(
-                    Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().getValue()));
+            Map<Dimensions, Map<String, MetricValue>> metricsByDimensions = getOrCreateApplicationMetrics(application, DimensionType.DEFAULT);
+            return metricsByDimensions.getOrDefault(dimensions, Collections.emptyMap())
+                    .entrySet()
+                    .stream()
+                    .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().getValue()));
         }
     }
 
-    private Map<Dimensions, Map<String, MetricValue>> getOrCreateApplicationMetrics(String application) {
+    private Map<Dimensions, Map<String, MetricValue>> getOrCreateApplicationMetrics(String application, DimensionType type) {
+        Map<String, ApplicationMetrics> applicationMetrics = metrics.computeIfAbsent(type, m -> new HashMap<>());
         if (! applicationMetrics.containsKey(application)) {
             ApplicationMetrics metrics = new ApplicationMetrics();
             applicationMetrics.put(application, metrics);
@@ -117,4 +145,7 @@ public class MetricReceiverWrapper {
             return metricsByDimensions;
         }
     }
+
+    // Used to distinguish whether metrics have been populated with all tag vaules
+    public enum DimensionType {DEFAULT, PRETAGGED}
 }
diff --git a/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/metrics/MetricReceiverWrapperTest.java b/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/metrics/MetricReceiverWrapperTest.java
index 08b3489fbea..c20e64d906e 100644
--- a/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/metrics/MetricReceiverWrapperTest.java
+++ b/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/metrics/MetricReceiverWrapperTest.java
@@ -78,4 +78,19 @@ public class MetricReceiverWrapperTest {
         assertEquals(metricReceiver.getMetricsForDimension(applicationDocker, hostDimension).get("test.gauge"), 42.);
         assertEquals(metricReceiver.getMetricsForDimension(applicationDocker, newDimension).get("test.gauge"), 56.);
     }
+
+    @Test
+    public void testDeletingMetric() {
+        MetricReceiverWrapper metricReceiver = new MetricReceiverWrapper(MetricReceiver.nullImplementation);
+        metricReceiver.declareGauge(applicationDocker, hostDimension, "test.gauge");
+
+        Dimensions differentDimension = new Dimensions.Builder().add("host", "abcd.yahoo.com").build();
+        metricReceiver.declareGauge(applicationDocker, differentDimension, "test.gauge");
+
+        assertEquals(2, metricReceiver.getDefaultMetricsRaw().size());
+        metricReceiver.deleteMetricByDimension(applicationDocker, differentDimension, MetricReceiverWrapper.DimensionType.DEFAULT);
+        assertEquals(1, metricReceiver.getDefaultMetricsRaw().size());
+        assertEquals(metricReceiver.getMetricsForDimension(applicationDocker, hostDimension).size(), 1);
+        assertEquals(metricReceiver.getMetricsForDimension(applicationDocker, differentDimension).size(), 0);
+    }
 }