Add function to get aggregated node stats per distributor.

4 files changed, 90 insertions, 21 deletions

diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterStatsAggregator.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterStatsAggregator.java
index 3b8cb14e977..ed2ca32dab0 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterStatsAggregator.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterStatsAggregator.java
@@ -27,6 +27,7 @@ import java.util.*;
 public class ClusterStatsAggregator {
 
     private final Set<Integer> distributors;
+    private final Set<Integer> storageNodes;
     private final Set<Integer> nonUpdatedDistributors;
 
     // Maps the distributor node index to a map of content node index to the
@@ -40,14 +41,26 @@ public class ClusterStatsAggregator {
 
     ClusterStatsAggregator(Set<Integer> distributors, Set<Integer> storageNodes) {
         this.distributors = distributors;
+        this.storageNodes = storageNodes;
         nonUpdatedDistributors = new HashSet<>(distributors);
         aggregatedStats = new ContentClusterStats(storageNodes);
     }
 
-    ContentClusterStats getAggregatedStats() {
+    public ContentClusterStats getAggregatedStats() {
         return aggregatedStats;
     }
 
+    public ContentNodeStats getAggregatedStatsForDistributor(int distributorIndex) {
+        ContentNodeStats result = new ContentNodeStats(distributorIndex);
+        ContentClusterStats distributorStats = distributorToStats.get(distributorIndex);
+        if (distributorStats != null) {
+            for (Iterator<ContentNodeStats> itr = distributorStats.iterator(); itr.hasNext(); ) {
+                result.add(itr.next());
+            }
+        }
+        return result;
+    }
+
     boolean hasUpdatesFromAllDistributors() {
         return nonUpdatedDistributors.isEmpty();
     }
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterStatsAggregatorTest.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterStatsAggregatorTest.java
index ac6417d0077..16530b2ceb6 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterStatsAggregatorTest.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterStatsAggregatorTest.java
@@ -18,24 +18,50 @@ public class ClusterStatsAggregatorTest {
 
     private static class Fixture {
         private ClusterStatsAggregator aggregator;
+
         public Fixture(Set<Integer> distributorNodes,
                 Set<Integer> contentNodes) {
             aggregator = new ClusterStatsAggregator(distributorNodes, contentNodes);
         }
+
         public void update(int distributorIndex, ContentClusterStatsBuilder clusterStats) {
             aggregator.updateForDistributor(distributorIndex, clusterStats.build());
         }
+
         public void verify(ContentClusterStatsBuilder expectedStats) {
             assertEquals(expectedStats.build(), aggregator.getAggregatedStats());
         }
+
+        public void verify(int distributorIndex, ContentNodeStatsBuilder expectedStats) {
+            assertEquals(expectedStats.build(), aggregator.getAggregatedStatsForDistributor(distributorIndex));
+        }
+
         public boolean hasUpdatesFromAllDistributors() {
             return aggregator.hasUpdatesFromAllDistributors();
         }
+
         public boolean mayHaveBucketsPendingInGlobalSpace() {
             return aggregator.mayHaveBucketsPendingInGlobalSpace();
         }
     }
 
+    private static class FourNodesFixture extends Fixture {
+        public FourNodesFixture() {
+            super(distributorNodes(1, 2), contentNodes(3, 4));
+
+            update(1, new ContentClusterStatsBuilder()
+                    .add(3, "default", 10, 1)
+                    .add(3, "global", 11, 2)
+                    .add(4, "default", 12, 3)
+                    .add(4, "global", 13, 4));
+            update(2, new ContentClusterStatsBuilder()
+                    .add(3, "default", 14, 5)
+                    .add(3, "global", 15, 6)
+                    .add(4, "default", 16, 7)
+                    .add(4, "global", 17, 8));
+        }
+    }
+
     private static Set<Integer> distributorNodes(Integer... indices) {
         return Sets.newHashSet(indices);
     }
@@ -56,18 +82,8 @@ public class ClusterStatsAggregatorTest {
 
     @Test
     public void aggregator_handles_updates_to_multiple_distributors_and_content_nodes() {
-        Fixture f = new Fixture(distributorNodes(1, 2), contentNodes(3, 4));
+        Fixture f = new FourNodesFixture();
 
-        f.update(1, new ContentClusterStatsBuilder()
-                .add(3, "default", 10, 1)
-                .add(3, "global", 11, 2)
-                .add(4, "default", 12, 3)
-                .add(4, "global", 13, 4));
-        f.update(2, new ContentClusterStatsBuilder()
-                .add(3, "default", 14, 5)
-                .add(3, "global", 15, 6)
-                .add(4, "default", 16, 7)
-                .add(4, "global", 17, 8));
         f.verify(new ContentClusterStatsBuilder()
                 .add(3, "default", 10 + 14, 1 + 5)
                 .add(3, "global", 11 + 15, 2 + 6)
@@ -154,4 +170,17 @@ public class ClusterStatsAggregatorTest {
         assertFalse(f.mayHaveBucketsPendingInGlobalSpace());
     }
 
+    @Test
+    public void aggregator_can_provide_aggregated_stats_per_distributor() {
+        Fixture f = new FourNodesFixture();
+
+        f.verify(1, new ContentNodeStatsBuilder(1)
+                .add("default", 10 + 12, 1 + 3)
+                .add("global", 11 + 13, 2 + 4));
+
+        f.verify(2, new ContentNodeStatsBuilder(2)
+                .add("default", 14 + 16, 5 + 7)
+                .add("global", 15 + 17, 6 + 8));
+    }
+
 }
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ContentClusterStatsBuilder.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ContentClusterStatsBuilder.java
index 16767cafa8f..a1a5fd4a4e1 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ContentClusterStatsBuilder.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ContentClusterStatsBuilder.java
@@ -9,7 +9,7 @@ import java.util.Map;
  */
 public class ContentClusterStatsBuilder {
 
-    private final Map<Integer, Map<String, ContentNodeStats.BucketSpaceStats>> stats = new HashMap<>();
+    private final Map<Integer, ContentNodeStatsBuilder> stats = new HashMap<>();
 
     public ContentClusterStatsBuilder add(int nodeIndex, String bucketSpace, long bucketsTotal, long bucketsPending) {
         return add(nodeIndex, bucketSpace, ContentNodeStats.BucketSpaceStats.of(bucketsTotal, bucketsPending));
@@ -24,24 +24,24 @@ public class ContentClusterStatsBuilder {
     }
 
     public ContentClusterStatsBuilder add(int nodeIndex, String bucketSpace, ContentNodeStats.BucketSpaceStats bucketSpaceStats) {
-        Map<String, ContentNodeStats.BucketSpaceStats> contentNodeStats = stats.get(nodeIndex);
-        if (contentNodeStats == null) {
-            contentNodeStats = new HashMap<>();
-            stats.put(nodeIndex, contentNodeStats);
+        ContentNodeStatsBuilder nodeStatsBuilder = stats.get(nodeIndex);
+        if (nodeStatsBuilder == null) {
+            nodeStatsBuilder = new ContentNodeStatsBuilder(nodeIndex);
+            stats.put(nodeIndex, nodeStatsBuilder);
         }
-        contentNodeStats.put(bucketSpace, bucketSpaceStats);
+        nodeStatsBuilder.add(bucketSpace, bucketSpaceStats);
         return this;
     }
 
     public ContentClusterStatsBuilder add(int nodeIndex) {
-        stats.put(nodeIndex, new HashMap<>());
+        stats.put(nodeIndex, new ContentNodeStatsBuilder(nodeIndex));
         return this;
     }
 
     public ContentClusterStats build() {
         Map<Integer, ContentNodeStats> nodeToStatsMap = new HashMap<>();
-        stats.forEach((nodeIndex, bucketSpaces) ->
-                nodeToStatsMap.put(nodeIndex, new ContentNodeStats(nodeIndex, bucketSpaces)));
+        stats.forEach((nodeIndex, nodeStatsBuilder) ->
+                nodeToStatsMap.put(nodeIndex, nodeStatsBuilder.build()));
         return new ContentClusterStats(nodeToStatsMap);
     }
 }
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ContentNodeStatsBuilder.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ContentNodeStatsBuilder.java
new file mode 100644
index 00000000000..088fc018d88
--- /dev/null
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ContentNodeStatsBuilder.java
@@ -0,0 +1,27 @@
+package com.yahoo.vespa.clustercontroller.core;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class ContentNodeStatsBuilder {
+
+    private final int nodeIndex;
+    private final Map<String, ContentNodeStats.BucketSpaceStats> stats = new HashMap<>();
+
+    public ContentNodeStatsBuilder(int nodeIndex) {
+        this.nodeIndex = nodeIndex;
+    }
+
+    public ContentNodeStatsBuilder add(String bucketSpace, long bucketsTotal, long bucketsPending) {
+        return add(bucketSpace, ContentNodeStats.BucketSpaceStats.of(bucketsTotal, bucketsPending));
+    }
+
+    public ContentNodeStatsBuilder add(String bucketSpace, ContentNodeStats.BucketSpaceStats bucketSpaceStats) {
+        stats.put(bucketSpace, bucketSpaceStats);
+        return this;
+    }
+
+    public ContentNodeStats build() {
+        return new ContentNodeStats(nodeIndex, stats);
+    }
+}