Merge branch 'main' into data-stream-settings-transport-action

Author: masseyke

docs/changelog/126091.yaml

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

docs/changelog/127225.yaml

@@ -0,0 +1,6 @@
+pr: 127225
+summary: Fix count optimization with pushable union types
+area: ES|QL
+type: bug
+issues:
+ - 127200

muted-tests.yml

@@ -408,9 +408,6 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.multi_node.EsqlSpecIT
   method: test {rerank.Reranker before a limit ASYNC}
   issue: https://github.com/elastic/elasticsearch/issues/127051
-- class: org.elasticsearch.xpack.esql.qa.single_node.EsqlSpecIT
-  method: test {rrf.SimpleRrf ASYNC}
-  issue: https://github.com/elastic/elasticsearch/issues/127063
 - class: org.elasticsearch.packaging.test.DockerTests
   method: test026InstallBundledRepositoryPlugins
   issue: https://github.com/elastic/elasticsearch/issues/127081
@@ -426,9 +423,6 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeIT
   method: test
   issue: https://github.com/elastic/elasticsearch/issues/127157
-- class: org.elasticsearch.xpack.esql.qa.single_node.EsqlSpecIT
-  method: test {fork.ForkWithWhereSortDescAndLimit SYNC}
-  issue: https://github.com/elastic/elasticsearch/issues/127326
 - class: org.elasticsearch.geometry.utils.SpatialEnvelopeVisitorTests
   method: testVisitGeoPointsWrapping
   issue: https://github.com/elastic/elasticsearch/issues/123425
@@ -441,6 +435,18 @@ tests:
 - class: org.elasticsearch.xpack.esql.heap_attack.HeapAttackIT
   method: testLookupExplosionNoFetch
   issue: https://github.com/elastic/elasticsearch/issues/127365
+- class: org.elasticsearch.action.admin.cluster.state.TransportClusterStateActionDisruptionIT
+  method: testNonLocalRequestAlwaysFindsMasterAndWaitsForMetadata
+  issue: https://github.com/elastic/elasticsearch/issues/127422
+- class: org.elasticsearch.xpack.esql.qa.single_node.PushQueriesIT
+  method: testPushCaseInsensitiveEqualityOnDefaults
+  issue: https://github.com/elastic/elasticsearch/issues/127431
+- class: org.elasticsearch.action.admin.cluster.state.TransportClusterStateActionDisruptionIT
+  method: testLocalRequestAlwaysSucceeds
+  issue: https://github.com/elastic/elasticsearch/issues/127423
+- class: org.elasticsearch.action.admin.cluster.state.TransportClusterStateActionDisruptionIT
+  method: testLocalRequestWaitsForMetadata
+  issue: https://github.com/elastic/elasticsearch/issues/127466
 
 # Examples:
 #

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/SnapshotsInProgress.java

@@ -135,14 +135,16 @@ private SnapshotsInProgress(Map<ProjectRepo, ByRepo> entries, Set<String> nodesI
 
     @FixForMultiProject
     @Deprecated(forRemoval = true)
-    public SnapshotsInProgress withUpdatedEntriesForRepo(String repository, List<Entry> updatedEntries) {
-        return withUpdatedEntriesForRepo(Metadata.DEFAULT_PROJECT_ID, repository, updatedEntries);
+    public SnapshotsInProgress createCopyWithUpdatedEntriesForRepo(String repository, List<Entry> updatedEntries) {
+        return createCopyWithUpdatedEntriesForRepo(Metadata.DEFAULT_PROJECT_ID, repository, updatedEntries);
     }
 
-    public SnapshotsInProgress withUpdatedEntriesForRepo(ProjectId projectId, String repository, List<Entry> updatedEntries) {
+    public SnapshotsInProgress createCopyWithUpdatedEntriesForRepo(ProjectId projectId, String repository, List<Entry> updatedEntries) {
         if (updatedEntries.equals(forRepo(projectId, repository))) {
+            // No changes to apply, return the current object.
             return this;
         }
+
         final Map<ProjectRepo, ByRepo> copy = new HashMap<>(this.entries);
         final var projectRepo = new ProjectRepo(projectId, repository);
         if (updatedEntries.isEmpty()) {
@@ -153,13 +155,14 @@ public SnapshotsInProgress withUpdatedEntriesForRepo(ProjectId projectId, String
         } else {
             copy.put(projectRepo, new ByRepo(updatedEntries));
         }
+
         return new SnapshotsInProgress(copy, nodesIdsForRemoval);
     }
 
     public SnapshotsInProgress withAddedEntry(Entry entry) {
         final List<Entry> forRepo = new ArrayList<>(forRepo(entry.projectId(), entry.repository()));
         forRepo.add(entry);
-        return withUpdatedEntriesForRepo(entry.projectId(), entry.repository(), forRepo);
+        return createCopyWithUpdatedEntriesForRepo(entry.projectId(), entry.repository(), forRepo);
     }
 
     /**

server/src/main/java/org/elasticsearch/cluster/action/shard/ShardStateAction.java

@@ -33,6 +33,7 @@
 import org.elasticsearch.cluster.routing.IndexRoutingTable;
 import org.elasticsearch.cluster.routing.RerouteService;
 import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.cluster.routing.ShardRoutingState;
 import org.elasticsearch.cluster.routing.allocation.AllocationService;
 import org.elasticsearch.cluster.routing.allocation.FailedShard;
 import org.elasticsearch.cluster.routing.allocation.StaleShard;
@@ -634,10 +635,11 @@ public ClusterState execute(BatchExecutionContext<StartedShardUpdateTask> batchE
             List<ShardRouting> shardRoutingsToBeApplied = new ArrayList<>(batchExecutionContext.taskContexts().size());
             Set<ShardRouting> seenShardRoutings = new HashSet<>(); // to prevent duplicates
             final Map<Index, ClusterStateTimeRanges> updatedTimestampRanges = new HashMap<>();
+
             final ClusterState initialState = batchExecutionContext.initialState();
             for (var taskContext : batchExecutionContext.taskContexts()) {
-                final var task = taskContext.getTask();
-                final StartedShardEntry startedShardEntry = task.getEntry();
+                final ShardStateAction.StartedShardUpdateTask task = taskContext.getTask();
+                final StartedShardEntry startedShardEntry = task.getStartedShardEntry();
                 final Optional<ProjectMetadata> project = initialState.metadata().lookupProject(startedShardEntry.shardId.getIndex());
                 final ShardRouting matched = project.map(ProjectMetadata::id)
                     .map(id -> initialState.routingTable(id).getByAllocationId(startedShardEntry.shardId, startedShardEntry.allocationId))
@@ -917,9 +919,16 @@ public int hashCode() {
         }
     }
 
+    /**
+     * Task that runs on the master node. Handles responding to the request listener with the result of the update request.
+     * Task is created when the master node receives a data node request to mark a shard as {@link ShardRoutingState#STARTED}.
+     *
+     * @param entry Information about the newly sharted shard.
+     * @param listener Channel listener with which to respond to the data node.
+     */
     public record StartedShardUpdateTask(StartedShardEntry entry, ActionListener<Void> listener) implements ClusterStateTaskListener {
 
-        public StartedShardEntry getEntry() {
+        public StartedShardEntry getStartedShardEntry() {
             return entry;
         }
 

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;