Allow balancer weights factory to be provided by ClusterPlugin

Author: nicktindall

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

@@ -26,7 +26,6 @@
 import org.elasticsearch.cluster.metadata.MetadataMappingService;
 import org.elasticsearch.cluster.metadata.NodesShutdownMetadata;
 import org.elasticsearch.cluster.metadata.RepositoriesMetadata;
-import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.project.ProjectResolver;
 import org.elasticsearch.cluster.routing.DelayedAllocationService;
 import org.elasticsearch.cluster.routing.ShardRouting;
@@ -44,7 +43,6 @@
 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.allocator.StatelessBalancingWeightsFactory;
 import org.elasticsearch.cluster.routing.allocation.decider.AllocationDecider;
 import org.elasticsearch.cluster.routing.allocation.decider.AllocationDeciders;
 import org.elasticsearch.cluster.routing.allocation.decider.AwarenessAllocationDecider;
@@ -150,11 +148,11 @@ public ClusterModule(
         this.deciderList = createAllocationDeciders(settings, clusterService.getClusterSettings(), clusterPlugins);
         this.allocationDeciders = new AllocationDeciders(deciderList);
         final BalancerSettings balancerSettings = new BalancerSettings(clusterService.getClusterSettings());
-        // I'm aware that the following is an anti-pattern and will implement as an SPI provider or plugin
-        // if we decide to go ahead with this.
-        final BalancingWeightsFactory balancingWeightsFactory = DiscoveryNode.isStateless(settings)
-            ? new StatelessBalancingWeightsFactory(balancerSettings, clusterService.getClusterSettings())
-            : new GlobalBalancingWeightsFactory(balancerSettings);
+        final BalancingWeightsFactory balancingWeightsFactory = getBalancingWeightsFactory(
+            clusterPlugins,
+            balancerSettings,
+            clusterService.getClusterSettings()
+        );
         var nodeAllocationStatsAndWeightsCalculator = new NodeAllocationStatsAndWeightsCalculator(
             writeLoadForecaster,
             balancingWeightsFactory
@@ -216,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);
     }

server/src/main/java/org/elasticsearch/plugins/ClusterPlugin.java

@@ -12,6 +12,8 @@
 import org.elasticsearch.cluster.routing.ShardRoutingRoleStrategy;
 import org.elasticsearch.cluster.routing.allocation.ExistingShardsAllocator;
 import org.elasticsearch.cluster.routing.allocation.WriteLoadForecaster;
+import org.elasticsearch.cluster.routing.allocation.allocator.BalancerSettings;
+import org.elasticsearch.cluster.routing.allocation.allocator.BalancingWeightsFactory;
 import org.elasticsearch.cluster.routing.allocation.allocator.ShardsAllocator;
 import org.elasticsearch.cluster.routing.allocation.decider.AllocationDecider;
 import org.elasticsearch.common.settings.ClusterSettings;
@@ -75,6 +77,10 @@ default ShardRoutingRoleStrategy getShardRoutingRoleStrategy() {
         return null;
     }
 
+    default BalancingWeightsFactory getBalancingWeightsFactory(BalancerSettings balancerSettings, ClusterSettings clusterSettings) {
+        return null;
+    }
+
     /**
      * Called when the node is started
      */