Use Map in snapshot/restore tracking (#87666)

Snapshot/restore keeps track of progress in ImmutableOpenMaps. This
commit changes these tracking structures to use HashMap.

relates #86239

server/src/main/java/org/elasticsearch/action/admin/cluster/snapshots/restore/RestoreClusterStateListener.java

@@ -15,12 +15,12 @@ import org.elasticsearch.cluster.ClusterChangedEvent;
 import org.elasticsearch.cluster.ClusterStateListener;
 import org.elasticsearch.cluster.RestoreInProgress;
 import org.elasticsearch.cluster.service.ClusterService;
-import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.snapshots.RestoreInfo;
 import org.elasticsearch.snapshots.RestoreService;
 
+import java.util.Map;
 import java.util.function.Supplier;
 
 import static org.elasticsearch.snapshots.RestoreService.restoreInProgress;
@@ -59,7 +59,7 @@ public class RestoreClusterStateListener implements ClusterStateListener {
                 listener.onResponse(new RestoreSnapshotResponse((RestoreInfo) null));
             } else if (newEntry == null) {
                 clusterService.removeListener(this);
-                ImmutableOpenMap<ShardId, RestoreInProgress.ShardRestoreStatus> shards = prevEntry.shards();
+                Map<ShardId, RestoreInProgress.ShardRestoreStatus> shards = prevEntry.shards();
                 assert prevEntry.state().completed() : "expected completed snapshot state but was " + prevEntry.state();
                 assert RestoreService.completed(shards) : "expected all restore entries to be completed";
                 RestoreInfo ri = new RestoreInfo(

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

@@ -10,10 +10,10 @@ package org.elasticsearch.cluster;
 
 import org.elasticsearch.Version;
 import org.elasticsearch.cluster.ClusterState.Custom;
-import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.snapshots.Snapshot;
 import org.elasticsearch.xcontent.ToXContent;
@@ -21,6 +21,7 @@ import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -33,16 +34,16 @@ public class RestoreInProgress extends AbstractNamedDiffable<Custom> implements
 
     public static final String TYPE = "restore";
 
-    public static final RestoreInProgress EMPTY = new RestoreInProgress(ImmutableOpenMap.of());
+    public static final RestoreInProgress EMPTY = new RestoreInProgress(Map.of());
 
-    private final ImmutableOpenMap<String, Entry> entries;
+    private final Map<String, Entry> entries;
 
     /**
      * Constructs new restore metadata
      *
      * @param entries map of currently running restore processes keyed by their restore uuid
      */
-    private RestoreInProgress(ImmutableOpenMap<String, Entry> entries) {
+    private RestoreInProgress(Map<String, Entry> entries) {
         this.entries = entries;
     }
 
@@ -82,12 +83,12 @@ public class RestoreInProgress extends AbstractNamedDiffable<Custom> implements
 
     public static final class Builder {
 
-        private final ImmutableOpenMap.Builder<String, Entry> entries = ImmutableOpenMap.builder();
+        private final Map<String, Entry> entries = new HashMap<>();
 
         public Builder() {}
 
         public Builder(RestoreInProgress restoreInProgress) {
-            entries.putAllFromMap(restoreInProgress.entries);
+            entries.putAll(restoreInProgress.entries);
         }
 
         public Builder add(Entry entry) {
@@ -96,20 +97,14 @@ public class RestoreInProgress extends AbstractNamedDiffable<Custom> implements
         }
 
         public RestoreInProgress build() {
-            return entries.isEmpty() ? EMPTY : new RestoreInProgress(entries.build());
+            return entries.isEmpty() ? EMPTY : new RestoreInProgress(Collections.unmodifiableMap(entries));
         }
     }
 
     /**
      * Restore metadata
      */
-    public record Entry(
-        String uuid,
-        Snapshot snapshot,
-        State state,
-        List<String> indices,
-        ImmutableOpenMap<ShardId, ShardRestoreStatus> shards
-    ) {
+    public record Entry(String uuid, Snapshot snapshot, State state, List<String> indices, Map<ShardId, ShardRestoreStatus> shards) {
         /**
          * Creates new restore metadata
          *
@@ -119,18 +114,12 @@ public class RestoreInProgress extends AbstractNamedDiffable<Custom> implements
          * @param indices  list of indices being restored
          * @param shards   map of shards being restored to their current restore status
          */
-        public Entry(
-            String uuid,
-            Snapshot snapshot,
-            State state,
-            List<String> indices,
-            ImmutableOpenMap<ShardId, ShardRestoreStatus> shards
-        ) {
+        public Entry(String uuid, Snapshot snapshot, State state, List<String> indices, Map<ShardId, ShardRestoreStatus> shards) {
             this.snapshot = Objects.requireNonNull(snapshot);
             this.state = Objects.requireNonNull(state);
             this.indices = Objects.requireNonNull(indices);
             if (shards == null) {
-                this.shards = ImmutableOpenMap.of();
+                this.shards = Map.of();
             } else {
                 this.shards = shards;
             }
@@ -347,25 +336,19 @@ public class RestoreInProgress extends AbstractNamedDiffable<Custom> implements
 
     public RestoreInProgress(StreamInput in) throws IOException {
         int count = in.readVInt();
-        final ImmutableOpenMap.Builder<String, Entry> entriesBuilder = ImmutableOpenMap.builder(count);
+        final Map<String, Entry> entriesBuilder = Maps.newHashMapWithExpectedSize(count);
         for (int i = 0; i < count; i++) {
             final String uuid;
             uuid = in.readString();
             Snapshot snapshot = new Snapshot(in);
             State state = State.fromValue(in.readByte());
-            List<String> indexBuilder = in.readStringList();
+            List<String> indices = in.readImmutableList(StreamInput::readString);
             entriesBuilder.put(
                 uuid,
-                new Entry(
-                    uuid,
-                    snapshot,
-                    state,
-                    Collections.unmodifiableList(indexBuilder),
-                    in.readImmutableOpenMap(ShardId::new, ShardRestoreStatus::readShardRestoreStatus)
-                )
+                new Entry(uuid, snapshot, state, indices, in.readImmutableMap(ShardId::new, ShardRestoreStatus::readShardRestoreStatus))
             );
         }
-        this.entries = entriesBuilder.build();
+        this.entries = Collections.unmodifiableMap(entriesBuilder);
     }
 
     @Override

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

@@ -11,7 +11,6 @@ package org.elasticsearch.cluster;
 import org.elasticsearch.Version;
 import org.elasticsearch.cluster.ClusterState.Custom;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
@@ -251,7 +250,7 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
         List<String> dataStreams,
         long startTime,
         long repositoryStateId,
-        ImmutableOpenMap<ShardId, ShardSnapshotStatus> shards,
+        Map<ShardId, ShardSnapshotStatus> shards,
         Map<String, Object> userMetadata,
         Version version,
         List<SnapshotFeatureInfo> featureStates
@@ -302,12 +301,12 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
             Collections.emptyList(),
             startTime,
             repositoryStateId,
-            ImmutableOpenMap.of(),
+            Map.of(),
             null,
             Collections.emptyMap(),
             version,
             source,
-            ImmutableOpenMap.of()
+            Map.of()
         );
     }
 
@@ -326,7 +325,7 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
         return true;
     }
 
-    private static boolean hasFailures(ImmutableOpenMap<RepositoryShardId, ShardSnapshotStatus> clones) {
+    private static boolean hasFailures(Map<RepositoryShardId, ShardSnapshotStatus> clones) {
         for (ShardSnapshotStatus value : clones.values()) {
             if (value.state().failed()) {
                 return true;
@@ -685,7 +684,7 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
         /**
          * Map of {@link RepositoryShardId} to {@link ShardSnapshotStatus} tracking the state of each shard operation in this entry.
          */
-        private final ImmutableOpenMap<RepositoryShardId, ShardSnapshotStatus> shardStatusByRepoShardId;
+        private final Map<RepositoryShardId, ShardSnapshotStatus> shardStatusByRepoShardId;
 
         @Nullable
         private final Map<String, Object> userMetadata;
@@ -723,7 +722,7 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
                 userMetadata,
                 version,
                 null,
-                ImmutableOpenMap.of()
+                Map.of()
             );
         }
 
@@ -742,7 +741,7 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
             Map<String, Object> userMetadata,
             Version version,
             @Nullable SnapshotId source,
-            ImmutableOpenMap<RepositoryShardId, ShardSnapshotStatus> shardStatusByRepoShardId
+            Map<RepositoryShardId, ShardSnapshotStatus> shardStatusByRepoShardId
         ) {
             this.state = state;
             this.snapshot = snapshot;
@@ -762,9 +761,7 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
                 assert shardStatusByRepoShardId == null || shardStatusByRepoShardId.isEmpty()
                     : "Provided explict repo shard id statuses [" + shardStatusByRepoShardId + "] but no source";
                 final Map<String, Index> res = Maps.newMapWithExpectedSize(indices.size());
-                final ImmutableOpenMap.Builder<RepositoryShardId, ShardSnapshotStatus> byRepoShardIdBuilder = ImmutableOpenMap.builder(
-                    shards.size()
-                );
+                final Map<RepositoryShardId, ShardSnapshotStatus> byRepoShardIdBuilder = Maps.newHashMapWithExpectedSize(shards.size());
                 for (Map.Entry<ShardId, ShardSnapshotStatus> entry : shards.entrySet()) {
                     final ShardId shardId = entry.getKey();
                     final IndexId indexId = indices.get(shardId.getIndexName());
@@ -773,12 +770,12 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
                     assert existing == null || existing.equals(index) : "Conflicting indices [" + existing + "] and [" + index + "]";
                     byRepoShardIdBuilder.put(new RepositoryShardId(indexId, shardId.id()), entry.getValue());
                 }
-                this.shardStatusByRepoShardId = byRepoShardIdBuilder.build();
-                snapshotIndices = Map.copyOf(res);
+                this.shardStatusByRepoShardId = Map.copyOf(byRepoShardIdBuilder);
+                this.snapshotIndices = Map.copyOf(res);
             } else {
                 assert shards.isEmpty();
                 this.shardStatusByRepoShardId = shardStatusByRepoShardId;
-                snapshotIndices = Map.of();
+                this.snapshotIndices = Map.of();
             }
             assert assertShardsConsistent(this.source, this.state, this.indices, this.shards, this.shardStatusByRepoShardId);
         }
@@ -801,17 +798,14 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
                 indices = Collections.unmodifiableMap(idx);
             }
             final long startTime = in.readLong();
-            final ImmutableOpenMap<ShardId, ShardSnapshotStatus> shards = in.readImmutableOpenMap(
-                ShardId::new,
-                ShardSnapshotStatus::readFrom
-            );
+            final Map<ShardId, ShardSnapshotStatus> shards = in.readImmutableMap(ShardId::new, ShardSnapshotStatus::readFrom);
             final long repositoryStateId = in.readLong();
             final String failure = in.readOptionalString();
             final Map<String, Object> userMetadata = in.readMap();
             final Version version = Version.readVersion(in);
             final List<String> dataStreams = in.readStringList();
             final SnapshotId source = in.readOptionalWriteable(SnapshotId::new);
-            final ImmutableOpenMap<RepositoryShardId, ShardSnapshotStatus> clones = in.readImmutableOpenMap(
+            final Map<RepositoryShardId, ShardSnapshotStatus> clones = in.readImmutableMap(
                 RepositoryShardId::new,
                 ShardSnapshotStatus::readFrom
             );
@@ -840,7 +834,7 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
             State state,
             Map<String, IndexId> indices,
             Map<ShardId, ShardSnapshotStatus> shards,
-            ImmutableOpenMap<RepositoryShardId, ShardSnapshotStatus> statusByRepoShardId
+            Map<RepositoryShardId, ShardSnapshotStatus> statusByRepoShardId
         ) {
             if ((state == State.INIT || state == State.ABORTED) && shards.isEmpty()) {
                 return true;
@@ -895,7 +889,7 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
                 userMetadata,
                 version,
                 source,
-                source == null ? ImmutableOpenMap.of() : shardStatusByRepoShardId
+                source == null ? Map.of() : shardStatusByRepoShardId
             );
         }
 
@@ -943,13 +937,13 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
                     userMetadata,
                     version,
                     source,
-                    ImmutableOpenMap.of()
+                    Map.of()
                 );
             }
             return this;
         }
 
-        public Entry withClones(ImmutableOpenMap<RepositoryShardId, ShardSnapshotStatus> updatedClones) {
+        public Entry withClones(Map<RepositoryShardId, ShardSnapshotStatus> updatedClones) {
             if (updatedClones.equals(shardStatusByRepoShardId)) {
                 return this;
             }
@@ -964,7 +958,7 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
                 featureStates,
                 startTime,
                 repositoryStateId,
-                ImmutableOpenMap.of(),
+                Map.of(),
                 failure,
                 userMetadata,
                 version,
@@ -985,7 +979,7 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
          */
         @Nullable
         public Entry abort() {
-            final ImmutableOpenMap.Builder<ShardId, ShardSnapshotStatus> shardsBuilder = ImmutableOpenMap.builder();
+            final Map<ShardId, ShardSnapshotStatus> shardsBuilder = new HashMap<>();
             boolean completed = true;
             boolean allQueued = true;
             for (Map.Entry<ShardId, ShardSnapshotStatus> shardEntry : shards.entrySet()) {
@@ -1016,12 +1010,12 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
                 featureStates,
                 startTime,
                 repositoryStateId,
-                shardsBuilder.build(),
+                Map.copyOf(shardsBuilder),
                 ABORTED_FAILURE_TEXT,
                 userMetadata,
                 version,
                 source,
-                ImmutableOpenMap.of()
+                Map.of()
             );
         }
 
@@ -1033,7 +1027,7 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
          * @param shards new shard snapshot states
          * @return new snapshot entry
          */
-        public Entry withShardStates(ImmutableOpenMap<ShardId, ShardSnapshotStatus> shards) {
+        public Entry withShardStates(Map<ShardId, ShardSnapshotStatus> shards) {
             if (completed(shards.values())) {
                 return new Entry(
                     snapshot,
@@ -1058,7 +1052,7 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
          * Same as {@link #withShardStates} but does not check if the snapshot completed and thus is only to be used when starting new
          * shard snapshots on data nodes for a running snapshot.
          */
-        public Entry withStartedShards(ImmutableOpenMap<ShardId, ShardSnapshotStatus> shards) {
+        public Entry withStartedShards(Map<ShardId, ShardSnapshotStatus> shards) {
             final SnapshotsInProgress.Entry updated = new Entry(
                 snapshot,
                 includeGlobalState,
@@ -1305,7 +1299,7 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
             out.writeStringCollection(dataStreams);
             out.writeOptionalWriteable(source);
             if (source == null) {
-                out.writeMap(ImmutableOpenMap.of());
+                out.writeMap(Map.of());
             } else {
                 out.writeMap(shardStatusByRepoShardId);
             }

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

@@ -859,9 +859,9 @@ public class RestoreService implements ClusterStateApplier {
                 RestoreInProgress.Builder builder = new RestoreInProgress.Builder();
                 for (RestoreInProgress.Entry entry : oldRestore) {
                     Map<ShardId, ShardRestoreStatus> updates = shardChanges.get(entry.uuid());
-                    ImmutableOpenMap<ShardId, ShardRestoreStatus> shardStates = entry.shards();
+                    Map<ShardId, ShardRestoreStatus> shardStates = entry.shards();
                     if (updates != null && updates.isEmpty() == false) {
-                        ImmutableOpenMap.Builder<ShardId, ShardRestoreStatus> shardsBuilder = ImmutableOpenMap.builder(shardStates);
+                        Map<ShardId, ShardRestoreStatus> shardsBuilder = new HashMap<>(shardStates);
                         for (Map.Entry<ShardId, ShardRestoreStatus> shard : updates.entrySet()) {
                             ShardId shardId = shard.getKey();
                             ShardRestoreStatus status = shardStates.get(shardId);
@@ -870,7 +870,7 @@ public class RestoreService implements ClusterStateApplier {
                             }
                         }
 
-                        ImmutableOpenMap<ShardId, ShardRestoreStatus> shards = shardsBuilder.build();
+                        Map<ShardId, ShardRestoreStatus> shards = Map.copyOf(shardsBuilder);
                         RestoreInProgress.State newState = overallState(RestoreInProgress.State.STARTED, shards);
                         builder.add(new RestoreInProgress.Entry(entry.uuid(), entry.snapshot(), newState, entry.indices(), shards));
                     } else {
@@ -886,7 +886,7 @@ public class RestoreService implements ClusterStateApplier {
 
     private static RestoreInProgress.State overallState(
         RestoreInProgress.State nonCompletedState,
-        ImmutableOpenMap<ShardId, RestoreInProgress.ShardRestoreStatus> shards
+        Map<ShardId, RestoreInProgress.ShardRestoreStatus> shards
     ) {
         boolean hasFailed = false;
         for (RestoreInProgress.ShardRestoreStatus status : shards.values()) {
@@ -904,7 +904,7 @@ public class RestoreService implements ClusterStateApplier {
         }
     }
 
-    public static boolean completed(ImmutableOpenMap<ShardId, RestoreInProgress.ShardRestoreStatus> shards) {
+    public static boolean completed(Map<ShardId, RestoreInProgress.ShardRestoreStatus> shards) {
         for (RestoreInProgress.ShardRestoreStatus status : shards.values()) {
             if (status.state().completed() == false) {
                 return false;
@@ -913,7 +913,7 @@ public class RestoreService implements ClusterStateApplier {
         return true;
     }
 
-    public static int failedShards(ImmutableOpenMap<ShardId, RestoreInProgress.ShardRestoreStatus> shards) {
+    public static int failedShards(Map<ShardId, RestoreInProgress.ShardRestoreStatus> shards) {
         int failedShards = 0;
         for (RestoreInProgress.ShardRestoreStatus status : shards.values()) {
             if (status.state() == RestoreInProgress.State.FAILURE) {
@@ -1266,7 +1266,7 @@ public class RestoreService implements ClusterStateApplier {
             final ClusterBlocks.Builder blocks = ClusterBlocks.builder().blocks(currentState.blocks());
             final RoutingTable.Builder rtBuilder = RoutingTable.builder(currentState.routingTable());
 
-            final ImmutableOpenMap.Builder<ShardId, ShardRestoreStatus> shardsBuilder = ImmutableOpenMap.builder();
+            final Map<ShardId, ShardRestoreStatus> shards = new HashMap<>();
 
             final Version minIndexCompatibilityVersion = currentState.getNodes().getMaxNodeVersion().minimumIndexCompatibilityVersion();
             final String localNodeId = clusterService.state().nodes().getLocalNodeId();
@@ -1350,7 +1350,7 @@ public class RestoreService implements ClusterStateApplier {
                 mdBuilder.put(updatedIndexMetadata, true);
                 final Index renamedIndex = updatedIndexMetadata.getIndex();
                 for (int shard = 0; shard < snapshotIndexMetadata.getNumberOfShards(); shard++) {
-                    shardsBuilder.put(
+                    shards.put(
                         new ShardId(renamedIndex, shard),
                         ignoreShards.contains(shard)
                             ? new ShardRestoreStatus(localNodeId, RestoreInProgress.State.FAILURE)
@@ -1364,7 +1364,6 @@ public class RestoreService implements ClusterStateApplier {
             }
 
             final ClusterState.Builder builder = ClusterState.builder(currentState);
-            final ImmutableOpenMap<ShardId, ShardRestoreStatus> shards = shardsBuilder.build();
             if (shards.isEmpty() == false) {
                 builder.putCustom(
                     RestoreInProgress.TYPE,
@@ -1374,7 +1373,7 @@ public class RestoreService implements ClusterStateApplier {
                             snapshot,
                             overallState(RestoreInProgress.State.INIT, shards),
                             List.copyOf(indicesToRestore.keySet()),
-                            shards
+                            Map.copyOf(shards)
                         )
                     ).build()
                 );

x-pack/plugin/ccr/src/test/java/org/elasticsearch/xpack/CcrIntegTestCase.java

@@ -41,7 +41,6 @@ import org.elasticsearch.common.Priority;
 import org.elasticsearch.common.Randomness;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.bytes.BytesReference;
-import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.network.NetworkModule;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.ByteSizeValue;
@@ -859,7 +858,7 @@ public abstract class CcrIntegTestCase extends ESTestCase {
                                 listener.onResponse(null);
                             } else if (newEntry == null) {
                                 clusterService.removeListener(this);
-                                ImmutableOpenMap<ShardId, RestoreInProgress.ShardRestoreStatus> shards = prevEntry.shards();
+                                Map<ShardId, RestoreInProgress.ShardRestoreStatus> shards = prevEntry.shards();
                                 RestoreInfo ri = new RestoreInfo(
                                     prevEntry.snapshot().getSnapshotId().getName(),
                                     prevEntry.indices(),