Allow specifying a consumer, and propagate it all the way.

5 files changed, 41 insertions, 23 deletions

diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/ApplicationMetricsRetriever.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/ApplicationMetricsRetriever.java
index 332c070943e..2c527c7396e 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/ApplicationMetricsRetriever.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/ApplicationMetricsRetriever.java
@@ -4,6 +4,7 @@
 
 package ai.vespa.metricsproxy.http.application;
 
+import ai.vespa.metricsproxy.metric.model.ConsumerId;
 import ai.vespa.metricsproxy.metric.model.MetricsPacket;
 import ai.vespa.util.http.VespaHttpClientBuilder;
 import com.google.inject.Inject;
@@ -23,6 +24,7 @@ import java.util.concurrent.ForkJoinPool;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
 
+import static ai.vespa.metricsproxy.http.ValuesFetcher.DEFAULT_PUBLIC_CONSUMER_ID;
 import static java.util.Collections.emptyList;
 import static java.util.stream.Collectors.toMap;
 
@@ -63,9 +65,13 @@ public class ApplicationMetricsRetriever extends AbstractComponent {
     }
 
     public Map<Node, List<MetricsPacket.Builder>> getMetrics() {
+        return getMetrics(DEFAULT_PUBLIC_CONSUMER_ID);
+    }
+
+    public Map<Node, List<MetricsPacket.Builder>> getMetrics(ConsumerId consumer) {
         log.info(() -> "Retrieving metrics from " + clients.size() + " nodes.");
         var forkJoinTask = forkJoinPool.submit(() -> clients.parallelStream()
-                .map(this::getNodeMetrics)
+                .map(client -> getNodeMetrics(client, consumer))
                 .collect(toMap(Map.Entry::getKey, Map.Entry::getValue)));
 
         try {
@@ -82,11 +88,11 @@ public class ApplicationMetricsRetriever extends AbstractComponent {
         }
     }
 
-    private Map.Entry<Node, List<MetricsPacket.Builder>> getNodeMetrics(NodeMetricsClient client) {
+    private Map.Entry<Node, List<MetricsPacket.Builder>> getNodeMetrics(NodeMetricsClient client, ConsumerId consumer) {
         try {
-            return new AbstractMap.SimpleEntry<>(client.node, client.getMetrics());
+            return new AbstractMap.SimpleEntry<>(client.node, client.getMetrics(consumer));
         } catch (Exception e) {
-            log.warning("Could not retrieve metrics from " + client.node.metricsUri);
+            log.warning("Could not retrieve metrics from " + client.node.metricsUri(consumer));
         }
         return new AbstractMap.SimpleEntry<>(client.node, emptyList());
     }
diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/Node.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/Node.java
index c22f81b42d5..ad8deacf7ab 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/Node.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/Node.java
@@ -4,6 +4,8 @@
 
 package ai.vespa.metricsproxy.http.application;
 
+import ai.vespa.metricsproxy.metric.model.ConsumerId;
+
 import java.net.URI;
 import java.util.Objects;
 
@@ -19,7 +21,7 @@ public class Node {
     final int port;
     final String path;
 
-    final URI metricsUri;
+    private final String metricsUriBase;
 
     public Node(VespaNodesConfig.Node nodeConfig) {
         this(nodeConfig.configId(), nodeConfig.hostname(), nodeConfig.port() ,nodeConfig.path());
@@ -30,11 +32,11 @@ public class Node {
         this.host = host;
         this.port = port;
         this.path = path;
-        metricsUri = getMetricsUri(host, port, path);
+        metricsUriBase = "http://" + host + ":" + port + path;
     }
 
-    private static URI getMetricsUri(String host, int port, String path) {
-        return URI.create("http://" + host + ":" + port + path);
+    URI metricsUri(ConsumerId consumer) {
+        return URI.create(metricsUriBase + "?consumer=" + consumer.id);
     }
 
     @Override
diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/NodeMetricsClient.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/NodeMetricsClient.java
index 1584a33d772..8987ca25be1 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/NodeMetricsClient.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/NodeMetricsClient.java
@@ -4,6 +4,7 @@
 
 package ai.vespa.metricsproxy.http.application;
 
+import ai.vespa.metricsproxy.metric.model.ConsumerId;
 import ai.vespa.metricsproxy.metric.model.MetricsPacket;
 import ai.vespa.metricsproxy.metric.model.json.GenericJsonUtil;
 import com.yahoo.yolean.Exceptions;
@@ -18,6 +19,7 @@ import java.time.Instant;
 import java.util.List;
 import java.util.logging.Logger;
 
+import static ai.vespa.metricsproxy.http.ValuesFetcher.DEFAULT_PUBLIC_CONSUMER_ID;
 import static com.yahoo.log.LogLevel.DEBUG;
 import static java.util.Collections.emptyList;
 
@@ -48,24 +50,29 @@ public class NodeMetricsClient {
     }
 
     public List<MetricsPacket.Builder> getMetrics() {
+        return getMetrics(DEFAULT_PUBLIC_CONSUMER_ID);
+    }
+
+    public List<MetricsPacket.Builder> getMetrics(ConsumerId consumer) {
         if (Instant.now(clock).isAfter(metricsTimestamp.plus(METRICS_TTL))) {
-            retrieveMetrics();
+            retrieveMetrics(consumer);
         }
         return metrics;
     }
 
-    private void retrieveMetrics() {
-        log.log(DEBUG, () -> "Retrieving metrics from host " + node.metricsUri);
+    private void retrieveMetrics(ConsumerId consumer) {
+        String metricsUri = node.metricsUri(consumer).toString();
+        log.log(DEBUG, () -> "Retrieving metrics from host " + metricsUri);
 
         try {
-            String metricsJson = httpClient.execute(new HttpGet(node.metricsUri), new BasicResponseHandler());
+            String metricsJson = httpClient.execute(new HttpGet(metricsUri), new BasicResponseHandler());
             metrics = GenericJsonUtil.toMetricsPackets(metricsJson);
             metricsTimestamp = Instant.now(clock);
             snapshotsRetrieved ++;
-            log.log(DEBUG, () -> "Successfully retrieved " + metrics.size() + " metrics packets from " + node.metricsUri);
+            log.log(DEBUG, () -> "Successfully retrieved " + metrics.size() + " metrics packets from " + metricsUri);
 
         } catch (IOException e) {
-            log.warning("Unable to retrieve metrics from " + node.metricsUri + ": " + Exceptions.toMessageString(e));
+            log.warning("Unable to retrieve metrics from " + metricsUri + ": " + Exceptions.toMessageString(e));
             metrics = emptyList();
         }
     }
diff --git a/metrics-proxy/src/test/java/ai/vespa/metricsproxy/http/application/ApplicationMetricsRetrieverTest.java b/metrics-proxy/src/test/java/ai/vespa/metricsproxy/http/application/ApplicationMetricsRetrieverTest.java
index e2bd5c84dc0..375795fb23e 100644
--- a/metrics-proxy/src/test/java/ai/vespa/metricsproxy/http/application/ApplicationMetricsRetrieverTest.java
+++ b/metrics-proxy/src/test/java/ai/vespa/metricsproxy/http/application/ApplicationMetricsRetrieverTest.java
@@ -16,7 +16,7 @@ import java.util.concurrent.TimeoutException;
 import static ai.vespa.metricsproxy.TestUtil.getFileContents;
 import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
 import static com.github.tomakehurst.wiremock.client.WireMock.get;
-import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
+import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo;
 import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options;
 import static java.util.stream.Collectors.toList;
 import static org.junit.Assert.assertEquals;
@@ -51,7 +51,7 @@ public class ApplicationMetricsRetrieverTest {
     }
 
     private void verifyRetrievingMetricsFromSingleNode(VespaNodesConfig config, Node node) {
-        wireMockRule.stubFor(get(urlEqualTo(config.node(0).path()))
+        wireMockRule.stubFor(get(urlPathEqualTo(config.node(0).path()))
                                      .willReturn(aResponse().withBody(RESPONSE)));
 
         ApplicationMetricsRetriever retriever = new ApplicationMetricsRetriever(config);
@@ -66,9 +66,9 @@ public class ApplicationMetricsRetrieverTest {
         Node node0 = new Node(config.node(0));
         Node node1 = new Node(config.node(1));
 
-        wireMockRule.stubFor(get(urlEqualTo(config.node(0).path()))
+        wireMockRule.stubFor(get(urlPathEqualTo(config.node(0).path()))
                                      .willReturn(aResponse().withBody(RESPONSE)));
-        wireMockRule.stubFor(get(urlEqualTo(config.node(1).path()))
+        wireMockRule.stubFor(get(urlPathEqualTo(config.node(1).path()))
                                      .willReturn(aResponse().withBody(RESPONSE)));
 
         ApplicationMetricsRetriever retriever = new ApplicationMetricsRetriever(config);
@@ -96,7 +96,7 @@ public class ApplicationMetricsRetrieverTest {
         Node node0 = new Node(config.node(0));
         Node node1 = new Node(config.node(1));
 
-        wireMockRule.stubFor(get(urlEqualTo(config.node(1).path()))
+        wireMockRule.stubFor(get(urlPathEqualTo(config.node(1).path()))
                                      .willReturn(aResponse().withBody(RESPONSE)));
 
         ApplicationMetricsRetriever retriever = new ApplicationMetricsRetriever(config);
@@ -110,7 +110,7 @@ public class ApplicationMetricsRetrieverTest {
     public void an_exception_is_thrown_when_retrieving_times_out() {
         var config = nodesConfig("/node0");
 
-        wireMockRule.stubFor(get(urlEqualTo(config.node(0).path()))
+        wireMockRule.stubFor(get(urlPathEqualTo(config.node(0).path()))
                                      .willReturn(aResponse()
                                                          .withBody(RESPONSE)
                                                          .withFixedDelay(10)));
@@ -131,7 +131,7 @@ public class ApplicationMetricsRetrieverTest {
         var config = nodesConfig("/node0");
         Node node = new Node(config.node(0));
 
-        var delayedStub = wireMockRule.stubFor(get(urlEqualTo(config.node(0).path()))
+        var delayedStub = wireMockRule.stubFor(get(urlPathEqualTo(config.node(0).path()))
                                                        .willReturn(aResponse()
                                                                            .withBody(RESPONSE)
                                                                            .withFixedDelay(10)));
diff --git a/metrics-proxy/src/test/java/ai/vespa/metricsproxy/http/application/NodeMetricsClientTest.java b/metrics-proxy/src/test/java/ai/vespa/metricsproxy/http/application/NodeMetricsClientTest.java
index f953e809f85..d028db93d43 100644
--- a/metrics-proxy/src/test/java/ai/vespa/metricsproxy/http/application/NodeMetricsClientTest.java
+++ b/metrics-proxy/src/test/java/ai/vespa/metricsproxy/http/application/NodeMetricsClientTest.java
@@ -15,12 +15,14 @@ import org.junit.BeforeClass;
 import org.junit.ClassRule;
 import org.junit.Test;
 
+import java.net.URI;
 import java.util.List;
 
 import static ai.vespa.metricsproxy.TestUtil.getFileContents;
+import static ai.vespa.metricsproxy.http.ValuesFetcher.DEFAULT_PUBLIC_CONSUMER_ID;
 import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
 import static com.github.tomakehurst.wiremock.client.WireMock.get;
-import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
+import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo;
 import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options;
 import static org.junit.Assert.assertEquals;
 
@@ -49,7 +51,8 @@ public class NodeMetricsClientTest {
     @BeforeClass
     public static void setupWireMock() {
         node = new Node("id", "localhost", wireMockRule.port(), MetricsHandler.VALUES_PATH);
-        wireMockRule.stubFor(get(urlEqualTo(node.metricsUri.getPath()))
+        URI metricsUri = node.metricsUri(DEFAULT_PUBLIC_CONSUMER_ID);
+        wireMockRule.stubFor(get(urlPathEqualTo(metricsUri.getPath()))
                                      .willReturn(aResponse().withBody(RESPONSE)));
     }