MetricsPacketsHandler adds hostname dimension if absent

3 files changed, 43 insertions, 11 deletions

diff --git a/container-core/src/main/java/com/yahoo/container/jdisc/state/MetricsPacketsHandler.java b/container-core/src/main/java/com/yahoo/container/jdisc/state/MetricsPacketsHandler.java
index b949edefb31..8fa658bf7fc 100644
--- a/container-core/src/main/java/com/yahoo/container/jdisc/state/MetricsPacketsHandler.java
+++ b/container-core/src/main/java/com/yahoo/container/jdisc/state/MetricsPacketsHandler.java
@@ -22,9 +22,9 @@ import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
 import static com.yahoo.container.jdisc.state.JsonUtil.sanitizeDouble;
@@ -61,6 +61,7 @@ public class MetricsPacketsHandler extends AbstractRequestHandler {
     private final Timer timer;
     private final SnapshotProvider snapshotProvider;
     private final String applicationName;
+    private final String hostDimension;
 
     @Inject
     public MetricsPacketsHandler(StateMonitor monitor,
@@ -71,6 +72,7 @@ public class MetricsPacketsHandler extends AbstractRequestHandler {
         this.timer = timer;
         snapshotProvider = getSnapshotProviderOrThrow(snapshotProviders);
         applicationName = config.application();
+        hostDimension = config.hostname();
     }
 
 
@@ -173,16 +175,16 @@ public class MetricsPacketsHandler extends AbstractRequestHandler {
     }
 
     private void addDimensions(MetricDimensions metricDimensions, ObjectNode packet) {
-        if (metricDimensions == null) return;
-
-        Iterator<Map.Entry<String, String>> dimensionsIterator = metricDimensions.iterator();
-        if (dimensionsIterator.hasNext()) {
-            ObjectNode jsonDim = jsonMapper.createObjectNode();
-            packet.set(DIMENSIONS_KEY, jsonDim);
-            for (Map.Entry<String, String> dimensionEntry : metricDimensions) {
-                jsonDim.put(dimensionEntry.getKey(), dimensionEntry.getValue());
-            }
+        if (metricDimensions == null && hostDimension.isEmpty()) return;
+
+        ObjectNode jsonDim = jsonMapper.createObjectNode();
+        packet.set(DIMENSIONS_KEY, jsonDim);
+        Iterable<Map.Entry<String, String>> dimensionIterator = metricDimensions == null ? Set.of() : metricDimensions;
+        for (Map.Entry<String, String> dimensionEntry : dimensionIterator) {
+            jsonDim.put(dimensionEntry.getKey(), dimensionEntry.getValue());
         }
+        if (!hostDimension.isEmpty() && !jsonDim.has("host"))
+            jsonDim.put("host", hostDimension);
     }
 
     private void addMetrics(MetricSet metricSet, ObjectNode packet) {
diff --git a/container-core/src/main/resources/configdefinitions/container.jdisc.state.metrics-packets-handler.def b/container-core/src/main/resources/configdefinitions/container.jdisc.state.metrics-packets-handler.def
index 9ec81b7db1b..ab0362c3995 100644
--- a/container-core/src/main/resources/configdefinitions/container.jdisc.state.metrics-packets-handler.def
+++ b/container-core/src/main/resources/configdefinitions/container.jdisc.state.metrics-packets-handler.def
@@ -4,3 +4,5 @@ namespace=container.jdisc.state
 
 # The name of the application that is reporting metrics.
 application string
+# Optional hostname to add as dimension
+hostname string default=""
\ No newline at end of file
diff --git a/container-core/src/test/java/com/yahoo/container/jdisc/state/MetricsPacketsHandlerTest.java b/container-core/src/test/java/com/yahoo/container/jdisc/state/MetricsPacketsHandlerTest.java
index 99a31640429..6c05af95289 100644
--- a/container-core/src/test/java/com/yahoo/container/jdisc/state/MetricsPacketsHandlerTest.java
+++ b/container-core/src/test/java/com/yahoo/container/jdisc/state/MetricsPacketsHandlerTest.java
@@ -27,13 +27,14 @@ import static org.junit.Assert.assertTrue;
 public class MetricsPacketsHandlerTest extends StateHandlerTestBase {
 
     private static final String APPLICATION_NAME = "state-handler-test-base";
+    private static final String HOST_DIMENSION = "some-hostname";
 
     private static MetricsPacketsHandler metricsPacketsHandler;
 
     @Before
     public void setupHandler() {
         metricsPacketsHandlerConfig = new MetricsPacketsHandlerConfig(new MetricsPacketsHandlerConfig.Builder()
-                                                                              .application(APPLICATION_NAME));
+                                                                              .application(APPLICATION_NAME).hostname(HOST_DIMENSION));
         metricsPacketsHandler = new MetricsPacketsHandler(monitor, timer, snapshotProviderRegistry, metricsPacketsHandlerConfig);
         testDriver = new RequestHandlerTestDriver(metricsPacketsHandler);
     }
@@ -138,6 +139,26 @@ public class MetricsPacketsHandlerTest extends StateHandlerTestBase {
         List<JsonNode> packets = incrementTimeAndGetJsonPackets();
         assertEquals(3, packets.size());
     }
+
+    @Test
+    public void host_dimension_only_created_if_absent() throws Exception {
+        var context1 = StateMetricContext.newInstance(Map.of("dim1", "value1", "host", "foo.bar"));
+        var context2 = StateMetricContext.newInstance(Map.of("dim2", "value2"));
+        var snapshot = new MetricSnapshot();
+        snapshot.add(context1, "counter1", 1);
+        snapshot.add(context2, "counter2", 2);
+        snapshotProvider.setSnapshot(snapshot);
+
+        var packets = incrementTimeAndGetJsonPackets();
+        assertEquals(3, packets.size());
+
+        packets.forEach(packet -> {
+            if (!packet.has(DIMENSIONS_KEY)) return;
+            var dimensions = packet.get(DIMENSIONS_KEY);
+            if (dimensions.has("dim1")) assertDimension(packet, "host", "foo.bar");
+            if (dimensions.has("dim2")) assertDimension(packet, "host", HOST_DIMENSION);
+        });
+    }
     
     private List<JsonNode> incrementTimeAndGetJsonPackets() throws Exception {
         advanceToNextSnapshot();
@@ -163,6 +184,13 @@ public class MetricsPacketsHandlerTest extends StateHandlerTestBase {
         assertEquals(expected, counterMetrics.get(metricName).asLong());
     }
 
+    private void assertDimension(JsonNode metricsPacket, String dimensionName, String expectedDimensionValue) {
+        assertTrue(metricsPacket.has(DIMENSIONS_KEY));
+        var dimensions = metricsPacket.get(DIMENSIONS_KEY);
+        assertTrue(dimensions.has(dimensionName));
+        assertEquals(expectedDimensionValue, dimensions.get(dimensionName).asText());
+    }
+
     private void createSnapshotWithCountMetric(String name, Number value, MetricDimensions context) {
         var snapshot = new MetricSnapshot();
         snapshot.add(context, name, value);