Allow balancing weights to be set per tier

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;