Allow balancing weights to be set per tier

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

@@ -37,8 +37,11 @@
 import org.elasticsearch.cluster.routing.allocation.NodeAllocationStatsAndWeightsCalculator;
 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.DesiredBalanceShardsAllocator;
 import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceShardsAllocator.DesiredBalanceReconcilerAction;
+import org.elasticsearch.cluster.routing.allocation.allocator.GlobalPartitionedClusterFactory;
+import org.elasticsearch.cluster.routing.allocation.allocator.PartitionedClusterFactory;
 import org.elasticsearch.cluster.routing.allocation.allocator.ShardsAllocator;
 import org.elasticsearch.cluster.routing.allocation.decider.AllocationDecider;
 import org.elasticsearch.cluster.routing.allocation.decider.AllocationDeciders;
@@ -145,13 +148,15 @@ public ClusterModule(
         this.clusterPlugins = clusterPlugins;
         this.deciderList = createAllocationDeciders(settings, clusterService.getClusterSettings(), clusterPlugins);
         this.allocationDeciders = new AllocationDeciders(deciderList);
+        final BalancerSettings balancerSettings = new BalancerSettings(clusterService.getClusterSettings());
         var nodeAllocationStatsAndWeightsCalculator = new NodeAllocationStatsAndWeightsCalculator(
             writeLoadForecaster,
-            clusterService.getClusterSettings()
+            new GlobalPartitionedClusterFactory(balancerSettings)
         );
         this.shardsAllocator = createShardsAllocator(
             settings,
             clusterService.getClusterSettings(),
+            balancerSettings,
             threadPool,
             clusterPlugins,
             clusterService,
@@ -440,6 +445,7 @@ private static void addAllocationDecider(Map<Class<?>, AllocationDecider> decide
     private static ShardsAllocator createShardsAllocator(
         Settings settings,
         ClusterSettings clusterSettings,
+        BalancerSettings balancerSettings,
         ThreadPool threadPool,
         List<ClusterPlugin> clusterPlugins,
         ClusterService clusterService,
@@ -449,12 +455,16 @@ private static ShardsAllocator createShardsAllocator(
         NodeAllocationStatsAndWeightsCalculator nodeAllocationStatsAndWeightsCalculator
     ) {
         Map<String, Supplier<ShardsAllocator>> allocators = new HashMap<>();
-        allocators.put(BALANCED_ALLOCATOR, () -> new BalancedShardsAllocator(clusterSettings, writeLoadForecaster));
+        final PartitionedClusterFactory partitionedClusterFactory = new GlobalPartitionedClusterFactory(balancerSettings);
+        allocators.put(
+            BALANCED_ALLOCATOR,
+            () -> new BalancedShardsAllocator(balancerSettings, writeLoadForecaster, partitionedClusterFactory)
+        );
         allocators.put(
             DESIRED_BALANCE_ALLOCATOR,
             () -> new DesiredBalanceShardsAllocator(
                 clusterSettings,
-                new BalancedShardsAllocator(clusterSettings, writeLoadForecaster),
+                new BalancedShardsAllocator(balancerSettings, writeLoadForecaster, partitionedClusterFactory),
                 threadPool,
                 clusterService,
                 reconciler,

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

@@ -15,10 +15,10 @@
 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.BalancedShardsAllocator;
 import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalance;
+import org.elasticsearch.cluster.routing.allocation.allocator.PartitionedCluster;
+import org.elasticsearch.cluster.routing.allocation.allocator.PartitionedClusterFactory;
 import org.elasticsearch.cluster.routing.allocation.allocator.WeightFunction;
-import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.core.Nullable;
 
@@ -29,11 +29,7 @@
  */
 public class NodeAllocationStatsAndWeightsCalculator {
     private final WriteLoadForecaster writeLoadForecaster;
-
-    private volatile float indexBalanceFactor;
-    private volatile float shardBalanceFactor;
-    private volatile float writeLoadBalanceFactor;
-    private volatile float diskUsageBalanceFactor;
+    private final PartitionedClusterFactory partitionedClusterFactory;
 
     /**
      * Node shard allocation stats and the total node weight.
@@ -47,18 +43,12 @@ public record NodeAllocationStatsAndWeight(
         float currentNodeWeight
     ) {}
 
-    public NodeAllocationStatsAndWeightsCalculator(WriteLoadForecaster writeLoadForecaster, ClusterSettings clusterSettings) {
+    public NodeAllocationStatsAndWeightsCalculator(
+        WriteLoadForecaster writeLoadForecaster,
+        PartitionedClusterFactory partitionedClusterFactory
+    ) {
         this.writeLoadForecaster = writeLoadForecaster;
-        clusterSettings.initializeAndWatch(BalancedShardsAllocator.SHARD_BALANCE_FACTOR_SETTING, value -> this.shardBalanceFactor = value);
-        clusterSettings.initializeAndWatch(BalancedShardsAllocator.INDEX_BALANCE_FACTOR_SETTING, value -> this.indexBalanceFactor = value);
-        clusterSettings.initializeAndWatch(
-            BalancedShardsAllocator.WRITE_LOAD_BALANCE_FACTOR_SETTING,
-            value -> this.writeLoadBalanceFactor = value
-        );
-        clusterSettings.initializeAndWatch(
-            BalancedShardsAllocator.DISK_USAGE_BALANCE_FACTOR_SETTING,
-            value -> this.diskUsageBalanceFactor = value
-        );
+        this.partitionedClusterFactory = partitionedClusterFactory;
     }
 
     /**
@@ -74,13 +64,14 @@ public Map<String, NodeAllocationStatsAndWeight> nodesAllocationStatsAndWeights(
             // must not use licensed features when just starting up
             writeLoadForecaster.refreshLicense();
         }
-        var weightFunction = new WeightFunction(shardBalanceFactor, indexBalanceFactor, writeLoadBalanceFactor, diskUsageBalanceFactor);
+        final PartitionedCluster partitionedCluster = partitionedClusterFactory.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 = partitionedCluster.weightFunctionForNode(node);
             int shards = 0;
             int undesiredShards = 0;
             double forecastedWriteLoad = 0.0;