Break down allocation stats by tier for IngestMetricsService

Author: nicktindall

server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceMetrics.java

@@ -32,18 +32,45 @@ public class DesiredBalanceMetrics {
 
     /**
      * @param unassignedShards Shards that are not assigned to any node.
+     * @param allocationStatsByRole A breakdown of the allocations stats by {@link ShardRouting.Role}
+     */
+    public record AllocationStats(long unassignedShards, Map<ShardRouting.Role, RoleAllocationStats> allocationStatsByRole) {
+
+        public AllocationStats(long unassignedShards, long totalAllocations, long undesiredAllocationsExcludingShuttingDownNodes) {
+            this(
+                unassignedShards,
+                Map.of(ShardRouting.Role.DEFAULT, new RoleAllocationStats(totalAllocations, undesiredAllocationsExcludingShuttingDownNodes))
+            );
+        }
+
+        public long totalAllocations() {
+            return allocationStatsByRole.values().stream().mapToLong(RoleAllocationStats::totalAllocations).sum();
+        }
+
+        public long undesiredAllocationsExcludingShuttingDownNodes() {
+            return allocationStatsByRole.values()
+                .stream()
+                .mapToLong(RoleAllocationStats::undesiredAllocationsExcludingShuttingDownNodes)
+                .sum();
+        }
+    }
+
+    /**
      * @param totalAllocations Shards that are assigned to a node.
      * @param undesiredAllocationsExcludingShuttingDownNodes Shards that are assigned to a node but must move to alleviate a resource
      *                                                       constraint per the {@link AllocationDeciders}. Excludes shards that must move
      *                                                       because of a node shutting down.
-     * @param undesiredAllocationsExcludingShuttingDownNodesByRole A breakdown of the undesired allocations by {@link ShardRouting.Role}
      */
-    public record AllocationStats(
-        long unassignedShards,
-        long totalAllocations,
-        long undesiredAllocationsExcludingShuttingDownNodes,
-        Map<ShardRouting.Role, Long> undesiredAllocationsExcludingShuttingDownNodesByRole
-    ) {}
+    public record RoleAllocationStats(long totalAllocations, long undesiredAllocationsExcludingShuttingDownNodes) {
+        public static final RoleAllocationStats EMPTY = new RoleAllocationStats(0L, 0L);
+
+        public float undesiredAllocationsRatio() {
+            if (totalAllocations == 0) {
+                return 0f;
+            }
+            return undesiredAllocationsExcludingShuttingDownNodes / (float) totalAllocations;
+        }
+    }
 
     public record NodeWeightStats(long shardCount, double diskUsageInBytes, double writeLoad, double nodeWeight) {
         public static final NodeWeightStats ZERO = new NodeWeightStats(0, 0, 0, 0);
@@ -78,7 +105,7 @@ public record NodeWeightStats(long shardCount, double diskUsageInBytes, double w
     public static final String CURRENT_NODE_FORECASTED_DISK_USAGE_METRIC_NAME =
         "es.allocator.allocations.node.forecasted_disk_usage_bytes.current";
 
-    public static final AllocationStats EMPTY_ALLOCATION_STATS = new AllocationStats(-1, -1, -1, Map.of());
+    public static final AllocationStats EMPTY_ALLOCATION_STATS = new AllocationStats(-1, Map.of());
 
     private volatile boolean nodeIsMaster = false;
 
@@ -98,6 +125,12 @@ public record NodeWeightStats(long shardCount, double diskUsageInBytes, double w
      */
     private volatile long undesiredAllocationsExcludingShuttingDownNodes;
 
+    /**
+     * A breakdown of shards assigned and the undesired allocations from the last reconciliation
+     * broken down by {@link ShardRouting.Role}.
+     */
+    private volatile Map<ShardRouting.Role, RoleAllocationStats> allocationStatsByRole;
+
     private final AtomicReference<Map<DiscoveryNode, NodeWeightStats>> weightStatsPerNodeRef = new AtomicReference<>(Map.of());
     private final AtomicReference<Map<DiscoveryNode, NodeAllocationStatsAndWeight>> allocationStatsPerNodeRef = new AtomicReference<>(
         Map.of()
@@ -112,8 +145,9 @@ public void updateMetrics(
         assert weightStatsPerNode != null : "node balance weight stats cannot be null";
         if (allocationStats != EMPTY_ALLOCATION_STATS) {
             this.unassignedShards = allocationStats.unassignedShards;
-            this.totalAllocations = allocationStats.totalAllocations;
-            this.undesiredAllocationsExcludingShuttingDownNodes = allocationStats.undesiredAllocationsExcludingShuttingDownNodes;
+            this.totalAllocations = allocationStats.totalAllocations();
+            this.undesiredAllocationsExcludingShuttingDownNodes = allocationStats.undesiredAllocationsExcludingShuttingDownNodes();
+            this.allocationStatsByRole = allocationStats.allocationStatsByRole();
         }
         weightStatsPerNodeRef.set(weightStatsPerNode);
         allocationStatsPerNodeRef.set(nodeAllocationStats);
@@ -223,6 +257,10 @@ public long undesiredAllocations() {
         return undesiredAllocationsExcludingShuttingDownNodes;
     }
 
+    public Map<ShardRouting.Role, RoleAllocationStats> allocationStatsByRole() {
+        return allocationStatsByRole;
+    }
+
     private List<LongWithAttributes> getUnassignedShardsMetrics() {
         return getIfPublishing(unassignedShards);
     }

server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceReconciler.java

@@ -529,6 +529,7 @@ private DesiredBalanceMetrics.AllocationStats balance() {
             int unassignedShards = routingNodes.unassigned().size() + routingNodes.unassigned().ignored().size();
             int totalAllocations = 0;
             int undesiredAllocationsExcludingShuttingDownNodes = 0;
+            final ObjectLongMap<ShardRouting.Role> totalAllocationsByRole = new ObjectLongHashMap<>();
             final ObjectLongMap<ShardRouting.Role> undesiredAllocationsExcludingShuttingDownNodesByRole = new ObjectLongHashMap<>();
 
             // Iterate over all started shards and try to move any which are on undesired nodes. In the presence of throttling shard
@@ -538,6 +539,7 @@ private DesiredBalanceMetrics.AllocationStats balance() {
                 final var shardRouting = iterator.next();
 
                 totalAllocations++;
+                totalAllocationsByRole.addTo(shardRouting.role(), 1);
 
                 if (shardRouting.started() == false) {
                     // can only rebalance started shards
@@ -600,10 +602,16 @@ private DesiredBalanceMetrics.AllocationStats balance() {
             maybeLogUndesiredAllocationsWarning(totalAllocations, undesiredAllocationsExcludingShuttingDownNodes, routingNodes.size());
             return new DesiredBalanceMetrics.AllocationStats(
                 unassignedShards,
-                totalAllocations,
-                undesiredAllocationsExcludingShuttingDownNodes,
-                StreamSupport.stream(undesiredAllocationsExcludingShuttingDownNodesByRole.spliterator(), false)
-                    .collect(Collectors.toUnmodifiableMap(lc -> lc.key, lc -> lc.value))
+                StreamSupport.stream(totalAllocationsByRole.spliterator(), false)
+                    .collect(
+                        Collectors.toUnmodifiableMap(
+                            lc -> lc.key,
+                            lc -> new DesiredBalanceMetrics.RoleAllocationStats(
+                                totalAllocationsByRole.get(lc.key),
+                                undesiredAllocationsExcludingShuttingDownNodesByRole.get(lc.key)
+                            )
+                        )
+                    )
             );
         }
 

server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceShardsAllocator.java

@@ -438,6 +438,10 @@ public DesiredBalanceStats getStats() {
         );
     }
 
+    public Map<ShardRouting.Role, DesiredBalanceMetrics.RoleAllocationStats> getAllocationStatsByRole() {
+        return desiredBalanceMetrics.allocationStatsByRole();
+    }
+
     private void onNoLongerMaster() {
         if (indexGenerator.getAndSet(-1) != -1) {
             currentDesiredBalanceRef.set(DesiredBalance.NOT_MASTER);

server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceMetricsTests.java

@@ -27,7 +27,7 @@ public void testZeroAllMetrics() {
         long unassignedShards = randomNonNegativeLong();
         long totalAllocations = randomNonNegativeLong();
         long undesiredAllocations = randomNonNegativeLong();
-        metrics.updateMetrics(new AllocationStats(unassignedShards, totalAllocations, undesiredAllocations, Map.of()), Map.of(), Map.of());
+        metrics.updateMetrics(new AllocationStats(unassignedShards, totalAllocations, undesiredAllocations), Map.of(), Map.of());
         assertEquals(totalAllocations, metrics.totalAllocations());
         assertEquals(unassignedShards, metrics.unassignedShards());
         assertEquals(undesiredAllocations, metrics.undesiredAllocations());
@@ -44,7 +44,7 @@ public void testMetricsAreOnlyPublishedWhenNodeIsMaster() {
         long unassignedShards = randomNonNegativeLong();
         long totalAllocations = randomLongBetween(100, 10000000);
         long undesiredAllocations = randomLongBetween(0, totalAllocations);
-        metrics.updateMetrics(new AllocationStats(unassignedShards, totalAllocations, undesiredAllocations, Map.of()), Map.of(), Map.of());
+        metrics.updateMetrics(new AllocationStats(unassignedShards, totalAllocations, undesiredAllocations), Map.of(), Map.of());
 
         // Collect when not master
         meterRegistry.getRecorder().collect();
@@ -104,7 +104,7 @@ public void testUndesiredAllocationRatioIsZeroWhenTotalShardsIsZero() {
         RecordingMeterRegistry meterRegistry = new RecordingMeterRegistry();
         DesiredBalanceMetrics metrics = new DesiredBalanceMetrics(meterRegistry);
         long unassignedShards = randomNonNegativeLong();
-        metrics.updateMetrics(new AllocationStats(unassignedShards, 0, 0, Map.of()), Map.of(), Map.of());
+        metrics.updateMetrics(new AllocationStats(unassignedShards, 0, 0), Map.of(), Map.of());
 
         metrics.setNodeIsMaster(true);
         meterRegistry.getRecorder().collect();

server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceReconcilerTests.java

@@ -281,7 +281,7 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, Routing
             assertTrue(index1RoutingTable.primaryShard().unassigned());
             assertTrue(index1RoutingTable.replicaShards().stream().allMatch(ShardRouting::unassigned));
             assertNotNull(allocationStats.get());
-            assertEquals(new DesiredBalanceMetrics.AllocationStats(3, 1, 0, Map.of()), allocationStats.get());
+            assertEquals(new DesiredBalanceMetrics.AllocationStats(3, 1, 0), allocationStats.get());
         }
 
         // now relax the filter so that the replica of index-0 and the primary of index-1 can both be assigned to node-1, but the throttle
@@ -296,7 +296,7 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, Routing
             assertTrue(index1RoutingTable.primaryShard().initializing());
             assertTrue(index1RoutingTable.replicaShards().stream().allMatch(ShardRouting::unassigned));
             assertNotNull(allocationStats.get());
-            assertEquals(new DesiredBalanceMetrics.AllocationStats(2, 2, 0, Map.of()), allocationStats.get());
+            assertEquals(new DesiredBalanceMetrics.AllocationStats(2, 2, 0), allocationStats.get());
         }
 
         final var stateWithStartedPrimariesAndInitializingReplica = startInitializingShardsAndReroute(
@@ -313,7 +313,7 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, Routing
             assertTrue(index1RoutingTable.primaryShard().started());
             assertTrue(index1RoutingTable.replicaShards().stream().allMatch(ShardRouting::unassigned));
             assertNotNull(allocationStats.get());
-            assertEquals(new DesiredBalanceMetrics.AllocationStats(1, 3, 0, Map.of()), allocationStats.get());
+            assertEquals(new DesiredBalanceMetrics.AllocationStats(1, 3, 0), allocationStats.get());
         }
     }
 
@@ -910,7 +910,7 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingAllocation allocat
             assertThat(shardRouting.currentNodeId(), oneOf("node-0", "node-1"));
         }
         assertNotNull(allocationStats);
-        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 6, 0, Map.of()), allocationStats.get());
+        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 6, 0), allocationStats.get());
 
         // Only allow allocation on two of the nodes, excluding the other two nodes.
         clusterSettings.applySettings(
@@ -926,7 +926,7 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingAllocation allocat
         assertSame(clusterState, allocationService.reroute(clusterState, "test", ActionListener.noop())); // all still on desired nodes, no
                                                                                                           // movement needed
         assertNotNull(allocationStats);
-        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 6, 0, Map.of()), allocationStats.get());
+        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 6, 0), allocationStats.get());
 
         desiredBalance.set(desiredBalance(clusterState, (shardId, nodeId) -> nodeId.equals("node-2") || nodeId.equals("node-3")));
 
@@ -937,12 +937,12 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingAllocation allocat
         assertThat(reroutedState.getRoutingNodes().node("node-1").numberOfShardsWithState(ShardRoutingState.RELOCATING), equalTo(1));
         assertNotNull(allocationStats);
         // Total allocations counts relocating and intializing shards, so the two relocating shards will be counted twice.
-        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 8, 4, Map.of(ShardRouting.Role.DEFAULT, 4L)), allocationStats.get());
+        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 8, 4), allocationStats.get());
 
         // Ensuring that we check the shortcut two-param canAllocate() method up front
         canAllocateRef.set(Decision.NO);
         assertSame(clusterState, allocationService.reroute(clusterState, "test", ActionListener.noop()));
-        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 6, 6, Map.of(ShardRouting.Role.DEFAULT, 6L)), allocationStats.get());
+        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 6, 6), allocationStats.get());
         canAllocateRef.set(Decision.YES);
 
         // Restore filter to default
@@ -980,7 +980,7 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingAllocation allocat
             "test",
             ActionListener.noop()
         );
-        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 7, 3, Map.of(ShardRouting.Role.DEFAULT, 3L)), allocationStats.get());
+        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 7, 3), allocationStats.get());
 
         assertThat(shuttingDownState.getRoutingNodes().node("node-2").numberOfShardsWithState(ShardRoutingState.INITIALIZING), equalTo(1));
     }
@@ -1048,7 +1048,7 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingAllocation allocat
 
         // All still on desired nodes, no movement needed, cluster state remains the same.
         assertSame(clusterState, allocationService.reroute(clusterState, "test", ActionListener.noop()));
-        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 6, 0, Map.of()), allocationStats.get());
+        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 6, 0), allocationStats.get());
 
         desiredBalance.set(desiredBalance(clusterState, (shardId, nodeId) -> nodeId.equals("node-2") || nodeId.equals("node-3")));
 
@@ -1076,15 +1076,15 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingAllocation allocat
         assertThat(reroutedState.getRoutingNodes().node("node-0").numberOfShardsWithState(ShardRoutingState.RELOCATING), equalTo(1));
         assertThat(reroutedState.getRoutingNodes().node("node-1").numberOfShardsWithState(ShardRoutingState.RELOCATING), equalTo(1));
         assertNotNull(allocationStats.get());
-        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 6, 6, Map.of(ShardRouting.Role.DEFAULT, 6L)), allocationStats.get());
+        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 6, 6), allocationStats.get());
 
         // Test that the AllocationStats are still updated, even though throttling is active. The cluster state should remain unchanged
         // because due to throttling: the previous reroute request started relocating two shards and, since those reallocations have not
         // been completed, no additional shard relocations can begin.
         assertSame(reroutedState, allocationService.reroute(reroutedState, "test", ActionListener.noop()));
         assertNotNull(allocationStats);
         // Note: total allocations counts relocating and intializing shards, so the two relocating shards will be counted twice.
-        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 8, 4, Map.of(ShardRouting.Role.DEFAULT, 4L)), allocationStats.get());
+        assertEquals(new DesiredBalanceMetrics.AllocationStats(0, 8, 4), allocationStats.get());
     }
 
     public void testDoNotRebalanceToTheNodeThatNoLongerExists() {
@@ -1293,7 +1293,7 @@ public void testRebalanceDoesNotCauseHotSpots() {
 
             var initializing = shardsWithState(allocation.routingNodes(), ShardRoutingState.INITIALIZING);
             if (initializing.isEmpty()) {
-                assertEquals(new DesiredBalanceMetrics.AllocationStats(0, shardsPerNode * numberOfNodes, 0, Map.of()), allocationStats);
+                assertEquals(new DesiredBalanceMetrics.AllocationStats(0, shardsPerNode * numberOfNodes, 0), allocationStats);
                 break;
             }