Allow balancing weights to be set per tier (elastic#126091)

Instead of assuming the BalancedShardsAllocator applies to the entire cluster, I've added the concept of "partitions" into the balancing. The partitions must be mutually disjoint subsets of the shards and nodes - i.e. the set of shards in a partition are only ever allocated to the set of nodes in the same partition, as is the case in serverless. WeightFunctions and NodeSorters are scoped to partitions. a BalancingWeights implementation knows how to partition a cluster.

The status quo behaviour is defined by the GlobalBalancingWeightsFactory, it produces a single global partition.

Author: nicktindall

docs/changelog/126091.yaml

@@ -0,0 +1,5 @@
+pr: 126091
+summary: Allow balancing weights to be set per tier
+area: Allocation
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/cluster/ClusterModule.java

@@ -38,8 +38,10 @@
 import org.elasticsearch.cluster.routing.allocation.WriteLoadForecaster;
 import org.elasticsearch.cluster.routing.allocation.allocator.BalancedShardsAllocator;
 import org.elasticsearch.cluster.routing.allocation.allocator.BalancerSettings;
+import org.elasticsearch.cluster.routing.allocation.allocator.BalancingWeightsFactory;
 import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceShardsAllocator;
 import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceShardsAllocator.DesiredBalanceReconcilerAction;
+import org.elasticsearch.cluster.routing.allocation.allocator.GlobalBalancingWeightsFactory;
 import org.elasticsearch.cluster.routing.allocation.allocator.ShardsAllocator;
 import org.elasticsearch.cluster.routing.allocation.decider.AllocationDecider;
 import org.elasticsearch.cluster.routing.allocation.decider.AllocationDeciders;
@@ -146,11 +148,20 @@ public ClusterModule(
         this.deciderList = createAllocationDeciders(settings, clusterService.getClusterSettings(), clusterPlugins);
         this.allocationDeciders = new AllocationDeciders(deciderList);
         final BalancerSettings balancerSettings = new BalancerSettings(clusterService.getClusterSettings());
-        var nodeAllocationStatsAndWeightsCalculator = new NodeAllocationStatsAndWeightsCalculator(writeLoadForecaster, balancerSettings);
+        final BalancingWeightsFactory balancingWeightsFactory = getBalancingWeightsFactory(
+            clusterPlugins,
+            balancerSettings,
+            clusterService.getClusterSettings()
+        );
+        var nodeAllocationStatsAndWeightsCalculator = new NodeAllocationStatsAndWeightsCalculator(
+            writeLoadForecaster,
+            balancingWeightsFactory
+        );
         this.shardsAllocator = createShardsAllocator(
             settings,
             clusterService.getClusterSettings(),
             balancerSettings,
+            balancingWeightsFactory,
             threadPool,
             clusterPlugins,
             clusterService,
@@ -203,6 +214,22 @@ public ShardRouting.Role newEmptyRole(int copyIndex) {
         };
     }
 
+    static BalancingWeightsFactory getBalancingWeightsFactory(
+        List<ClusterPlugin> clusterPlugins,
+        BalancerSettings balancerSettings,
+        ClusterSettings clusterSettings
+    ) {
+        final var strategies = clusterPlugins.stream()
+            .map(pl -> pl.getBalancingWeightsFactory(balancerSettings, clusterSettings))
+            .filter(Objects::nonNull)
+            .toList();
+        return switch (strategies.size()) {
+            case 0 -> new GlobalBalancingWeightsFactory(balancerSettings);
+            case 1 -> strategies.getFirst();
+            default -> throw new IllegalArgumentException("multiple plugins define balancing weights factories, which is not permitted");
+        };
+    }
+
     private ClusterState reconcile(ClusterState clusterState, RerouteStrategy rerouteStrategy) {
         return allocationService.executeWithRoutingAllocation(clusterState, "reconcile-desired-balance", rerouteStrategy);
     }
@@ -439,6 +466,7 @@ private static ShardsAllocator createShardsAllocator(
         Settings settings,
         ClusterSettings clusterSettings,
         BalancerSettings balancerSettings,
+        BalancingWeightsFactory balancingWeightsFactory,
         ThreadPool threadPool,
         List<ClusterPlugin> clusterPlugins,
         ClusterService clusterService,
@@ -448,12 +476,15 @@ private static ShardsAllocator createShardsAllocator(
         NodeAllocationStatsAndWeightsCalculator nodeAllocationStatsAndWeightsCalculator
     ) {
         Map<String, Supplier<ShardsAllocator>> allocators = new HashMap<>();
-        allocators.put(BALANCED_ALLOCATOR, () -> new BalancedShardsAllocator(balancerSettings, writeLoadForecaster));
+        allocators.put(
+            BALANCED_ALLOCATOR,
+            () -> new BalancedShardsAllocator(balancerSettings, writeLoadForecaster, balancingWeightsFactory)
+        );
         allocators.put(
             DESIRED_BALANCE_ALLOCATOR,
             () -> new DesiredBalanceShardsAllocator(
                 clusterSettings,
-                new BalancedShardsAllocator(balancerSettings, writeLoadForecaster),
+                new BalancedShardsAllocator(balancerSettings, writeLoadForecaster, balancingWeightsFactory),
                 threadPool,
                 clusterService,
                 reconciler,

server/src/main/java/org/elasticsearch/cluster/routing/allocation/NodeAllocationStatsAndWeightsCalculator.java

@@ -15,7 +15,8 @@
 import org.elasticsearch.cluster.routing.RoutingNode;
 import org.elasticsearch.cluster.routing.RoutingNodes;
 import org.elasticsearch.cluster.routing.ShardRouting;
-import org.elasticsearch.cluster.routing.allocation.allocator.BalancerSettings;
+import org.elasticsearch.cluster.routing.allocation.allocator.BalancingWeights;
+import org.elasticsearch.cluster.routing.allocation.allocator.BalancingWeightsFactory;
 import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalance;
 import org.elasticsearch.cluster.routing.allocation.allocator.WeightFunction;
 import org.elasticsearch.common.util.Maps;
@@ -28,7 +29,7 @@
  */
 public class NodeAllocationStatsAndWeightsCalculator {
     private final WriteLoadForecaster writeLoadForecaster;
-    private final BalancerSettings balancerSettings;
+    private final BalancingWeightsFactory balancingWeightsFactory;
 
     /**
      * Node shard allocation stats and the total node weight.
@@ -42,9 +43,12 @@ public record NodeAllocationStatsAndWeight(
         float currentNodeWeight
     ) {}
 
-    public NodeAllocationStatsAndWeightsCalculator(WriteLoadForecaster writeLoadForecaster, BalancerSettings balancerSettings) {
+    public NodeAllocationStatsAndWeightsCalculator(
+        WriteLoadForecaster writeLoadForecaster,
+        BalancingWeightsFactory balancingWeightsFactory
+    ) {
         this.writeLoadForecaster = writeLoadForecaster;
-        this.balancerSettings = balancerSettings;
+        this.balancingWeightsFactory = balancingWeightsFactory;
     }
 
     /**
@@ -60,18 +64,14 @@ public Map<String, NodeAllocationStatsAndWeight> nodesAllocationStatsAndWeights(
             // must not use licensed features when just starting up
             writeLoadForecaster.refreshLicense();
         }
-        var weightFunction = new WeightFunction(
-            balancerSettings.getShardBalanceFactor(),
-            balancerSettings.getIndexBalanceFactor(),
-            balancerSettings.getWriteLoadBalanceFactor(),
-            balancerSettings.getDiskUsageBalanceFactor()
-        );
+        final BalancingWeights balancingWeights = balancingWeightsFactory.create();
         var avgShardsPerNode = WeightFunction.avgShardPerNode(metadata, routingNodes);
         var avgWriteLoadPerNode = WeightFunction.avgWriteLoadPerNode(writeLoadForecaster, metadata, routingNodes);
         var avgDiskUsageInBytesPerNode = WeightFunction.avgDiskUsageInBytesPerNode(clusterInfo, metadata, routingNodes);
 
         var nodeAllocationStatsAndWeights = Maps.<String, NodeAllocationStatsAndWeight>newMapWithExpectedSize(routingNodes.size());
         for (RoutingNode node : routingNodes) {
+            WeightFunction weightFunction = balancingWeights.weightFunctionForNode(node);
             int shards = 0;
             int undesiredShards = 0;
             double forecastedWriteLoad = 0.0;