Align snapshot restore shard size estimation condition (#102931)

This pr aligns snapshot restore shard size estimation condition with internal
one in SnapshotShardSizeInfo::getShardSize

server/src/main/java/org/elasticsearch/cluster/routing/ExpectedShardSizeEstimator.java

@@ -47,8 +47,10 @@ public class ExpectedShardSizeEstimator {
         if (indexMetadata.getResizeSourceIndex() != null
             && shard.active() == false
             && shard.recoverySource().getType() == RecoverySource.Type.LOCAL_SHARDS) {
+            assert shard.primary() : "All replica shards are recovering from " + RecoverySource.Type.PEER;
             return getExpectedSizeOfResizedShard(shard, defaultValue, indexMetadata, clusterInfo, metadata, routingTable);
-        } else if (shard.unassigned() && shard.recoverySource().getType() == RecoverySource.Type.SNAPSHOT) {
+        } else if (shard.active() == false && shard.recoverySource().getType() == RecoverySource.Type.SNAPSHOT) {
+            assert shard.primary() : "All replica shards are recovering from " + RecoverySource.Type.PEER;
             return snapshotShardSizeInfo.getShardSize(shard, defaultValue);
         } else {
             return clusterInfo.getShardSize(shard, defaultValue);

server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/DiskThresholdDecider.java

@@ -27,6 +27,7 @@ import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.settings.SettingsException;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.core.UpdateForV9;
+import org.elasticsearch.snapshots.SnapshotShardSizeInfo;
 
 import java.util.Map;
 
@@ -146,7 +147,7 @@ public class DiskThresholdDecider extends AllocationDecider {
                     routing,
                     Math.max(routing.getExpectedShardSize(), 0L),
                     clusterInfo,
-                    null,
+                    SnapshotShardSizeInfo.EMPTY,
                     metadata,
                     routingTable
                 );
@@ -158,7 +159,7 @@ public class DiskThresholdDecider extends AllocationDecider {
         if (subtractShardsMovingAway) {
             for (ShardRouting routing : node.relocating()) {
                 if (dataPath.equals(clusterInfo.getDataPath(routing))) {
-                    totalSize -= getExpectedShardSize(routing, 0L, clusterInfo, null, metadata, routingTable);
+                    totalSize -= getExpectedShardSize(routing, 0L, clusterInfo, SnapshotShardSizeInfo.EMPTY, metadata, routingTable);
                 }
             }
         }

server/src/main/java/org/elasticsearch/snapshots/SnapshotShardSizeInfo.java

@@ -9,6 +9,7 @@
 package org.elasticsearch.snapshots;
 
 import org.elasticsearch.cluster.routing.RecoverySource;
+import org.elasticsearch.cluster.routing.RecoverySource.SnapshotRecoverySource;
 import org.elasticsearch.cluster.routing.ShardRouting;
 
 import java.util.Map;
@@ -24,11 +25,8 @@ public class SnapshotShardSizeInfo {
     }
 
     public Long getShardSize(ShardRouting shardRouting) {
-        if (shardRouting.primary()
-            && shardRouting.active() == false
-            && shardRouting.recoverySource().getType() == RecoverySource.Type.SNAPSHOT) {
-            final RecoverySource.SnapshotRecoverySource snapshotRecoverySource = (RecoverySource.SnapshotRecoverySource) shardRouting
-                .recoverySource();
+        if (shardRouting.active() == false && shardRouting.recoverySource().getType() == RecoverySource.Type.SNAPSHOT) {
+            final SnapshotRecoverySource snapshotRecoverySource = (SnapshotRecoverySource) shardRouting.recoverySource();
             return snapshotShardSizes.get(
                 new InternalSnapshotsInfoService.SnapshotShard(
                     snapshotRecoverySource.snapshot(),