Use common code to set expected shard size (#101949)

This change uses common component to set expected shard size when initializing
shard from GatewayAllocator

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

@@ -109,7 +109,7 @@ public class RoutingAllocation {
         AllocationDeciders deciders,
         @Nullable RoutingNodes routingNodes,
         ClusterState clusterState,
-        @Nullable ClusterInfo clusterInfo,
+        ClusterInfo clusterInfo,
         SnapshotShardSizeInfo shardSizeInfo,
         long currentNanoTime
     ) {

server/src/main/java/org/elasticsearch/gateway/BaseGatewayShardAllocator.java

@@ -10,7 +10,6 @@ package org.elasticsearch.gateway;
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.elasticsearch.cluster.routing.RecoverySource;
 import org.elasticsearch.cluster.routing.RoutingNode;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.allocation.AllocateUnassignedDecision;
@@ -23,6 +22,8 @@ import org.elasticsearch.cluster.routing.allocation.decider.Decision;
 import java.util.ArrayList;
 import java.util.List;
 
+import static org.elasticsearch.cluster.routing.ExpectedShardSizeEstimator.getExpectedShardSize;
+
 /**
  * An abstract class that implements basic functionality for allocating
  * shards to nodes based on shard copies that already exist in the cluster.
@@ -58,7 +59,7 @@ public abstract class BaseGatewayShardAllocator {
             unassignedAllocationHandler.initialize(
                 allocateUnassignedDecision.getTargetNode().getId(),
                 allocateUnassignedDecision.getAllocationId(),
-                getExpectedShardSize(shardRouting, allocation),
+                getExpectedShardSize(shardRouting, ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE, allocation),
                 allocation.changes()
             );
         } else {
@@ -66,18 +67,6 @@ public abstract class BaseGatewayShardAllocator {
         }
     }
 
-    protected static long getExpectedShardSize(ShardRouting shardRouting, RoutingAllocation allocation) {
-        if (shardRouting.primary()) {
-            if (shardRouting.recoverySource().getType() == RecoverySource.Type.SNAPSHOT) {
-                return allocation.snapshotShardSizeInfo().getShardSize(shardRouting, ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE);
-            } else {
-                return ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE;
-            }
-        } else {
-            return allocation.clusterInfo().getShardSize(shardRouting, ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE);
-        }
-    }
-
     /**
      * Make a decision on the allocation of an unassigned shard.  This method is used by
      * {@link #allocateUnassigned(ShardRouting, RoutingAllocation, ExistingShardsAllocator.UnassignedAllocationHandler)} to make decisions

server/src/test/java/org/elasticsearch/gateway/PrimaryShardAllocatorTests.java

@@ -9,6 +9,7 @@
 package org.elasticsearch.gateway;
 
 import org.apache.lucene.index.CorruptIndexException;
+import org.elasticsearch.cluster.ClusterInfo;
 import org.elasticsearch.cluster.ClusterName;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.ESAllocationTestCase;
@@ -519,7 +520,14 @@ public class PrimaryShardAllocatorTests extends ESAllocationTestCase {
             .routingTable(routingTableBuilder.build())
             .nodes(DiscoveryNodes.builder().add(node1).add(node2).add(node3))
             .build();
-        return new RoutingAllocation(deciders, state.mutableRoutingNodes(), state, null, null, System.nanoTime());
+        return new RoutingAllocation(
+            deciders,
+            state.mutableRoutingNodes(),
+            state,
+            ClusterInfo.EMPTY,
+            SnapshotShardSizeInfo.EMPTY,
+            System.nanoTime()
+        );
     }
 
     private void assertClusterHealthStatus(RoutingAllocation allocation, ClusterHealthStatus expectedStatus) {