Account for reserved disk size (#103903)

Author: idegtiarenko

docs/changelog/103903.yaml

@@ -0,0 +1,5 @@
+pr: 103903
+summary: Account for reserved disk size
+area: Allocation
+type: enhancement
+issues: []

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

@@ -376,10 +376,6 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeCollection(shardIds);
         }
 
-        public long getTotal() {
-            return total;
-        }
-
         public boolean containsShardId(ShardId shardId) {
             return shardIds.contains(shardId);
         }

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

@@ -35,13 +35,48 @@ public class ClusterInfoSimulator {
 
     public ClusterInfoSimulator(RoutingAllocation allocation) {
         this.allocation = allocation;
-        this.leastAvailableSpaceUsage = new HashMap<>(allocation.clusterInfo().getNodeLeastAvailableDiskUsages());
-        this.mostAvailableSpaceUsage = new HashMap<>(allocation.clusterInfo().getNodeMostAvailableDiskUsages());
+        this.leastAvailableSpaceUsage = getAdjustedDiskSpace(allocation, allocation.clusterInfo().getNodeLeastAvailableDiskUsages());
+        this.mostAvailableSpaceUsage = getAdjustedDiskSpace(allocation, allocation.clusterInfo().getNodeMostAvailableDiskUsages());
         this.shardSizes = new CopyOnFirstWriteMap<>(allocation.clusterInfo().shardSizes);
         this.shardDataSetSizes = Map.copyOf(allocation.clusterInfo().shardDataSetSizes);
         this.dataPath = Map.copyOf(allocation.clusterInfo().dataPath);
     }
 
+    /**
+     * Cluster info contains a reserved space that is necessary to finish initializing shards (that are currently in progress).
+     * for all initializing shards sum(expected size) = reserved space + already used space
+     * This deducts already used space from disk usage as when shard start is simulated it is going to add entire expected shard size.
+     */
+    private static Map<String, DiskUsage> getAdjustedDiskSpace(RoutingAllocation allocation, Map<String, DiskUsage> diskUsage) {
+        var diskUsageCopy = new HashMap<>(diskUsage);
+        for (var entry : diskUsageCopy.entrySet()) {
+            var nodeId = entry.getKey();
+            var usage = entry.getValue();
+
+            var reserved = allocation.clusterInfo().getReservedSpace(nodeId, usage.path());
+            if (reserved.total() == 0) {
+                continue;
+            }
+            var node = allocation.routingNodes().node(nodeId);
+            if (node == null) {
+                continue;
+            }
+
+            long adjustment = 0;
+            for (ShardId shardId : reserved.shardIds()) {
+                var shard = node.getByShardId(shardId);
+                if (shard != null) {
+                    var expectedSize = getExpectedShardSize(shard, 0, allocation);
+                    adjustment += expectedSize;
+                }
+            }
+            adjustment -= reserved.total();
+
+            entry.setValue(updateWithFreeBytes(usage, adjustment));
+        }
+        return diskUsageCopy;
+    }
+
     /**
      * This method updates disk usage to reflect shard relocations and new replica initialization.
      * In case of a single data path both mostAvailableSpaceUsage and leastAvailableSpaceUsage are update to reflect the change.

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

@@ -232,7 +232,7 @@ public void onNewInfo(ClusterInfo info) {
                 }
             }
 
-            final long reservedSpace = info.getReservedSpace(usage.getNodeId(), usage.getPath()).getTotal();
+            final long reservedSpace = info.getReservedSpace(usage.getNodeId(), usage.getPath()).total();
             final DiskUsage usageWithReservedSpace = new DiskUsage(
                 usage.getNodeId(),
                 usage.getNodeName(),

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

@@ -125,7 +125,7 @@ public static long sizeOfUnaccountedShards(
     ) {
         // Account for reserved space wherever it is available
         final ClusterInfo.ReservedSpace reservedSpace = clusterInfo.getReservedSpace(node.nodeId(), dataPath);
-        long totalSize = reservedSpace.getTotal();
+        long totalSize = reservedSpace.total();
         // NB this counts all shards on the node when the ClusterInfoService retrieved the node stats, which may include shards that are
         // no longer initializing because their recovery failed or was cancelled.
 

server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/ClusterInfoSimulatorTests.java

@@ -42,6 +42,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static org.elasticsearch.cluster.ClusterInfo.shardIdentifierFromRouting;
 import static org.elasticsearch.cluster.metadata.IndexMetadata.INDEX_RESIZE_SOURCE_NAME_KEY;
@@ -158,6 +159,53 @@ public void testInitializeNewReplica() {
         );
     }
 
+    public void testInitializeNewReplicaWithReservedSpace() {
+
+        var recoveredSize = 70;
+        var remainingSize = 30;
+        var totalShardSize = recoveredSize + remainingSize;
+
+        var indexMetadata = IndexMetadata.builder("my-index").settings(indexSettings(IndexVersion.current(), 1, 1)).build();
+        var existingPrimary = newShardRouting(new ShardId(indexMetadata.getIndex(), 0), "node-0", true, STARTED);
+        var newReplica = newShardRouting(
+            new ShardId(indexMetadata.getIndex(), 0),
+            "node-1",
+            false,
+            INITIALIZING,
+            RecoverySource.PeerRecoverySource.INSTANCE
+        );
+
+        var initialClusterInfo = new ClusterInfoTestBuilder() //
+            .withNode("node-0", new DiskUsageBuilder("/data", 1000, 1000 - totalShardSize))
+            .withNode("node-1", new DiskUsageBuilder("/data", 1000, 1000 - recoveredSize))
+            .withShard(existingPrimary, totalShardSize)
+            .withReservedSpace("node-1", "/data", remainingSize, newReplica.shardId())
+            .build();
+
+        var state = ClusterState.builder(ClusterName.DEFAULT)
+            .metadata(Metadata.builder().put(indexMetadata, false))
+            .routingTable(
+                RoutingTable.builder()
+                    .add(IndexRoutingTable.builder(indexMetadata.getIndex()).addShard(existingPrimary).addShard(newReplica))
+            )
+            .build();
+        var allocation = createRoutingAllocation(state, initialClusterInfo, SnapshotShardSizeInfo.EMPTY);
+        var simulator = new ClusterInfoSimulator(allocation);
+        simulator.simulateShardStarted(newReplica);
+
+        assertThat(
+            simulator.getClusterInfo(),
+            equalTo(
+                new ClusterInfoTestBuilder() //
+                    .withNode("node-0", new DiskUsageBuilder("/data", 1000, 1000 - totalShardSize))
+                    .withNode("node-1", new DiskUsageBuilder("/data", 1000, 1000 - totalShardSize))
+                    .withShard(existingPrimary, totalShardSize)
+                    .withShard(newReplica, totalShardSize)
+                    .build()
+            )
+        );
+    }
+
     public void testRelocateShard() {
 
         var fromNodeId = "node-0";
@@ -631,6 +679,11 @@ public ClusterInfoTestBuilder withShard(ShardRouting shard, long size) {
             return this;
         }
 
+        public ClusterInfoTestBuilder withReservedSpace(String nodeId, String path, long size, ShardId... shardIds) {
+            reservedSpace.put(new NodeAndPath(nodeId, nodeId + path), new ReservedSpace(size, Set.of(shardIds)));
+            return this;
+        }
+
         public ClusterInfo build() {
             return new ClusterInfo(leastAvailableSpaceUsage, mostAvailableSpaceUsage, shardSizes, Map.of(), Map.of(), reservedSpace);
         }

server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceComputerTests.java

@@ -905,8 +905,7 @@ public void testAccountForSizeOfMisplacedShardsDuringNewComputation() {
                     IndexRoutingTable.builder(indexMetadata2.getIndex())
                         .addShard(newShardRouting(index2ShardId, "node-1", true, INITIALIZING, index2SnapshotRecoverySource))
                 );
-                // TODO enable in https://github.com/elastic/elasticsearch/pull/103903
-                if (false && randomBoolean()) {
+                if (randomBoolean()) {
                     // Shard is 75% downloaded
                     clusterInfoBuilder //
                         .withNodeUsedSpace("node-1", ByteSizeValue.ofMb(768).getBytes())
@@ -920,8 +919,7 @@ public void testAccountForSizeOfMisplacedShardsDuringNewComputation() {
                     IndexRoutingTable.builder(indexMetadata2.getIndex())
                         .addShard(newShardRouting(index2ShardId, "node-2", true, INITIALIZING, index2SnapshotRecoverySource))
                 );
-                // TODO enable in https://github.com/elastic/elasticsearch/pull/103903
-                if (false && randomBoolean()) {
+                if (randomBoolean()) {
                     // Shard is 75% downloaded
                     clusterInfoBuilder //
                         .withNodeUsedSpace("node-2", ByteSizeValue.ofMb(768).getBytes())