Revert "Remove ImmutableOpenMap from snapshot services" (#90287)

* Revert "Remove ImmutableOpenMap from snapshot services (#90006)"

This reverts commit a5f41a8fd3c62185bab1ea20ec4c6a801a09928a.

docs/changelog/90287.yaml

@@ -0,0 +1,5 @@
+pr: 90287
+summary: 'Revert "Remove `ImmutableOpenMap` from snapshot services"'
+area: "Infra/Core"
+type: regression
+issues: []

modules/data-streams/src/test/java/org/elasticsearch/datastreams/action/DeleteDataStreamTransportActionTests.java

@@ -17,6 +17,7 @@ import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.metadata.Metadata;
 import org.elasticsearch.cluster.metadata.MetadataDeleteIndexService;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.core.Tuple;
@@ -30,7 +31,6 @@ import org.elasticsearch.test.ESTestCase;
 
 import java.util.Collections;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 import java.util.function.Consumer;
 
@@ -123,7 +123,7 @@ public class DeleteDataStreamTransportActionTests extends ESTestCase {
             Collections.emptyList(),
             0,
             1,
-            Map.of(),
+            ImmutableOpenMap.of(),
             null,
             null,
             null

server/src/internalClusterTest/java/org/elasticsearch/cluster/ClusterStateDiffIT.java

@@ -716,7 +716,7 @@ public class ClusterStateDiffIT extends ESIntegTestCase {
                             Collections.emptyList(),
                             Math.abs(randomLong()),
                             randomIntBetween(0, 1000),
-                            Map.of(),
+                            ImmutableOpenMap.of(),
                             null,
                             SnapshotInfoTestUtils.randomUserMetadata(),
                             randomVersion(random())

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

@@ -132,12 +132,16 @@ public class RestoreInProgress extends AbstractNamedDiffable<Custom> implements
             List<String> indices,
             Map<ShardId, ShardRestoreStatus> shards
         ) {
-            this.uuid = Objects.requireNonNull(uuid);
             this.snapshot = Objects.requireNonNull(snapshot);
             this.state = Objects.requireNonNull(state);
-            this.quiet = quiet;
+            this.quiet = Objects.requireNonNull(quiet);
             this.indices = Objects.requireNonNull(indices);
-            this.shards = shards != null ? Map.copyOf(shards) : Map.of();
+            if (shards == null) {
+                this.shards = Map.of();
+            } else {
+                this.shards = shards;
+            }
+            this.uuid = Objects.requireNonNull(uuid);
         }
     }
 

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

@@ -312,11 +312,21 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
      * @return true if all shards have completed (either successfully or failed), false otherwise
      */
     public static boolean completed(Collection<ShardSnapshotStatus> shards) {
-        return shards.stream().allMatch(s -> s.state().completed());
+        for (ShardSnapshotStatus status : shards) {
+            if (status.state().completed == false) {
+                return false;
+            }
+        }
+        return true;
     }
 
     private static boolean hasFailures(Map<RepositoryShardId, ShardSnapshotStatus> clones) {
-        return clones.values().stream().anyMatch(s -> s.state().failed());
+        for (ShardSnapshotStatus value : clones.values()) {
+            if (value.state().failed()) {
+                return true;
+            }
+        }
+        return false;
     }
 
     private static boolean assertConsistentEntries(Map<String, ByRepo> entries) {
@@ -779,14 +789,14 @@ public class SnapshotsInProgress extends AbstractNamedDiffable<Custom> implement
             this.dataStreams = List.copyOf(dataStreams);
             this.featureStates = List.copyOf(featureStates);
             this.startTime = startTime;
-            this.shards = Map.copyOf(shards);
+            this.shards = shards;
             this.repositoryStateId = repositoryStateId;
             this.failure = failure;
             this.userMetadata = userMetadata == null ? null : Map.copyOf(userMetadata);
             this.version = version;
             this.source = source;
             this.shardStatusByRepoShardId = Map.copyOf(shardStatusByRepoShardId);
-            this.snapshotIndices = Map.copyOf(snapshotIndices);
+            this.snapshotIndices = snapshotIndices;
             assert assertShardsConsistent(this.source, this.state, this.indices, this.shards, this.shardStatusByRepoShardId);
         }
 

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

@@ -20,6 +20,7 @@ import org.elasticsearch.cluster.routing.RerouteService;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.Priority;
+import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Settings;
@@ -31,11 +32,9 @@ import org.elasticsearch.repositories.Repository;
 import org.elasticsearch.threadpool.ThreadPool;
 
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.LinkedList;
-import java.util.Map;
 import java.util.Objects;
 import java.util.Queue;
 import java.util.Set;
@@ -65,7 +64,7 @@ public class InternalSnapshotsInfoService implements ClusterStateListener, Snaps
     private final Supplier<RerouteService> rerouteService;
 
     /** contains the snapshot shards for which the size is known **/
-    private volatile Map<SnapshotShard, Long> knownSnapshotShards;
+    private volatile ImmutableOpenMap<SnapshotShard, Long> knownSnapshotShards;
 
     private volatile boolean isMaster;
 
@@ -92,7 +91,7 @@ public class InternalSnapshotsInfoService implements ClusterStateListener, Snaps
         this.threadPool = clusterService.getClusterApplierService().threadPool();
         this.repositoriesService = repositoriesServiceSupplier;
         this.rerouteService = rerouteServiceSupplier;
-        this.knownSnapshotShards = Map.of();
+        this.knownSnapshotShards = ImmutableOpenMap.of();
         this.unknownSnapshotShards = new LinkedHashSet<>();
         this.failedSnapshotShards = new LinkedHashSet<>();
         this.queue = new LinkedList<>();
@@ -113,15 +112,14 @@ public class InternalSnapshotsInfoService implements ClusterStateListener, Snaps
     @Override
     public SnapshotShardSizeInfo snapshotShardSizes() {
         synchronized (mutex) {
-            Map<SnapshotShard, Long> snapshotShardSizes = knownSnapshotShards;
+            final ImmutableOpenMap.Builder<SnapshotShard, Long> snapshotShardSizes = ImmutableOpenMap.builder(knownSnapshotShards);
             if (failedSnapshotShards.isEmpty() == false) {
-                snapshotShardSizes = new HashMap<>(snapshotShardSizes); // create a new map to modify
                 for (SnapshotShard snapshotShard : failedSnapshotShards) {
                     Long previous = snapshotShardSizes.put(snapshotShard, ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE);
                     assert previous == null : "snapshot shard size already known for " + snapshotShard;
                 }
             }
-            return new SnapshotShardSizeInfo(snapshotShardSizes);
+            return new SnapshotShardSizeInfo(snapshotShardSizes.build());
         }
     }
 
@@ -163,7 +161,7 @@ public class InternalSnapshotsInfoService implements ClusterStateListener, Snaps
             // have to repopulate the data over and over in an unstable master situation?
             synchronized (mutex) {
                 // information only needed on current master
-                knownSnapshotShards = Map.of();
+                knownSnapshotShards = ImmutableOpenMap.of();
                 failedSnapshotShards.clear();
                 isMaster = false;
                 SnapshotShard snapshotShard;
@@ -228,10 +226,12 @@ public class InternalSnapshotsInfoService implements ClusterStateListener, Snaps
                 removed = unknownSnapshotShards.remove(snapshotShard);
                 assert removed : "snapshot shard to remove does not exist " + snapshotShardSize;
                 if (isMaster) {
-                    final Map<SnapshotShard, Long> newSnapshotShardSizes = new HashMap<>(knownSnapshotShards);
+                    final ImmutableOpenMap.Builder<SnapshotShard, Long> newSnapshotShardSizes = ImmutableOpenMap.builder(
+                        knownSnapshotShards
+                    );
                     updated = newSnapshotShardSizes.put(snapshotShard, snapshotShardSize) == null;
                     assert updated : "snapshot shard size already exists for " + snapshotShard;
-                    knownSnapshotShards = Map.copyOf(newSnapshotShardSizes);
+                    knownSnapshotShards = newSnapshotShardSizes.build();
                 }
                 activeFetches -= 1;
                 assert invariant();
@@ -269,17 +269,17 @@ public class InternalSnapshotsInfoService implements ClusterStateListener, Snaps
 
     private void cleanUpSnapshotShardSizes(Set<SnapshotShard> requiredSnapshotShards) {
         assert Thread.holdsLock(mutex);
-        Map<SnapshotShard, Long> newSnapshotShardSizes = null;
+        ImmutableOpenMap.Builder<SnapshotShard, Long> newSnapshotShardSizes = null;
         for (SnapshotShard shard : knownSnapshotShards.keySet()) {
             if (requiredSnapshotShards.contains(shard) == false) {
                 if (newSnapshotShardSizes == null) {
-                    newSnapshotShardSizes = new HashMap<>(knownSnapshotShards);
+                    newSnapshotShardSizes = ImmutableOpenMap.builder(knownSnapshotShards);
                 }
                 newSnapshotShardSizes.remove(shard);
             }
         }
         if (newSnapshotShardSizes != null) {
-            knownSnapshotShards = Map.copyOf(newSnapshotShardSizes);
+            knownSnapshotShards = newSnapshotShardSizes.build();
         }
         failedSnapshotShards.retainAll(requiredSnapshotShards);
     }
@@ -293,11 +293,11 @@ public class InternalSnapshotsInfoService implements ClusterStateListener, Snaps
             assert failedSnapshotShards.contains(shard) == false : "cannot be known and failed at same time: " + shard;
         }
         for (SnapshotShard shard : unknownSnapshotShards) {
-            assert knownSnapshotShards.containsKey(shard) == false : "cannot be unknown and known at same time: " + shard;
+            assert knownSnapshotShards.keySet().contains(shard) == false : "cannot be unknown and known at same time: " + shard;
             assert failedSnapshotShards.contains(shard) == false : "cannot be unknown and failed at same time: " + shard;
         }
         for (SnapshotShard shard : failedSnapshotShards) {
-            assert knownSnapshotShards.containsKey(shard) == false : "cannot be failed and known at same time: " + shard;
+            assert knownSnapshotShards.keySet().contains(shard) == false : "cannot be failed and known at same time: " + shard;
             assert unknownSnapshotShards.contains(shard) == false : "cannot be failed and unknown at same time: " + shard;
         }
         return true;

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

@@ -49,6 +49,7 @@ import org.elasticsearch.cluster.service.MasterService;
 import org.elasticsearch.common.Priority;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.UUIDs;
+import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.lucene.Lucene;
 import org.elasticsearch.common.regex.Regex;
 import org.elasticsearch.common.settings.ClusterSettings;
@@ -722,17 +723,19 @@ public class RestoreService implements ClusterStateApplier {
         boolean changesMade = false;
         RestoreInProgress.Builder builder = new RestoreInProgress.Builder();
         for (RestoreInProgress.Entry entry : oldRestore) {
-            Map<ShardId, ShardRestoreStatus> shards = null;
-            for (ShardId shardId : entry.shards().keySet()) {
+            ImmutableOpenMap.Builder<ShardId, ShardRestoreStatus> shardsBuilder = null;
+            for (Map.Entry<ShardId, ShardRestoreStatus> cursor : entry.shards().entrySet()) {
+                ShardId shardId = cursor.getKey();
                 if (deletedIndices.contains(shardId.getIndex())) {
                     changesMade = true;
-                    if (shards == null) {
-                        shards = new HashMap<>(entry.shards());
+                    if (shardsBuilder == null) {
+                        shardsBuilder = ImmutableOpenMap.builder(entry.shards());
                     }
-                    shards.put(shardId, new ShardRestoreStatus(null, RestoreInProgress.State.FAILURE, "index was deleted"));
+                    shardsBuilder.put(shardId, new ShardRestoreStatus(null, RestoreInProgress.State.FAILURE, "index was deleted"));
                 }
             }
-            if (shards != null) {
+            if (shardsBuilder != null) {
+                ImmutableOpenMap<ShardId, ShardRestoreStatus> shards = shardsBuilder.build();
                 builder.add(
                     new RestoreInProgress.Entry(
                         entry.uuid(),

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

@@ -20,7 +20,7 @@ public class SnapshotShardSizeInfo {
     private final Map<InternalSnapshotsInfoService.SnapshotShard, Long> snapshotShardSizes;
 
     public SnapshotShardSizeInfo(Map<InternalSnapshotsInfoService.SnapshotShard, Long> snapshotShardSizes) {
-        this.snapshotShardSizes = Map.copyOf(snapshotShardSizes);
+        this.snapshotShardSizes = snapshotShardSizes;
     }
 
     public Long getShardSize(ShardRouting shardRouting) {

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

@@ -58,6 +58,7 @@ import org.elasticsearch.cluster.service.MasterService;
 import org.elasticsearch.common.Priority;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.UUIDs;
+import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.component.AbstractLifecycleComponent;
 import org.elasticsearch.common.regex.Regex;
 import org.elasticsearch.common.settings.Setting;
@@ -383,7 +384,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                 }
                 final Map<String, IndexId> indexIds = repositoryData.resolveNewIndices(indices, allIndices);
                 final Version version = minCompatibleVersion(currentState.nodes().getMinNodeVersion(), repositoryData, null);
-                Map<ShardId, ShardSnapshotStatus> shards = shards(
+                ImmutableOpenMap<ShardId, ShardSnapshotStatus> shards = shards(
                     snapshots,
                     deletionsInProgress,
                     currentState,
@@ -393,12 +394,12 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                     repositoryName
                 );
                 if (request.partial() == false) {
-                    Set<String> missing = shards.entrySet()
-                        .stream()
-                        .filter(e -> e.getValue().state() == ShardState.MISSING)
-                        .map(e -> e.getKey().getIndex().getName())
-                        .collect(Collectors.toSet());
-
+                    Set<String> missing = new HashSet<>();
+                    for (Map.Entry<ShardId, SnapshotsInProgress.ShardSnapshotStatus> entry : shards.entrySet()) {
+                        if (entry.getValue().state() == ShardState.MISSING) {
+                            missing.add(entry.getKey().getIndex().getName());
+                        }
+                    }
                     if (missing.isEmpty() == false) {
                         throw new SnapshotException(
                             new Snapshot(repositoryName, snapshotId),
@@ -643,7 +644,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                 final ShardGenerations shardGenerations = repoData.shardGenerations();
                 for (SnapshotsInProgress.Entry existing : existingEntries) {
                     if (cloneEntry.snapshot().getSnapshotId().equals(existing.snapshot().getSnapshotId())) {
-                        final Map<RepositoryShardId, ShardSnapshotStatus> clones = new HashMap<>();
+                        final ImmutableOpenMap.Builder<RepositoryShardId, ShardSnapshotStatus> clonesBuilder = ImmutableOpenMap.builder();
                         final boolean readyToExecute = currentState.custom(
                             SnapshotDeletionsInProgress.TYPE,
                             SnapshotDeletionsInProgress.EMPTY
@@ -660,9 +661,9 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                                 final RepositoryShardId repoShardId = new RepositoryShardId(count.v1(), shardId);
                                 final String indexName = repoShardId.indexName();
                                 if (readyToExecute == false || inFlightShardStates.isActive(indexName, shardId)) {
-                                    clones.put(repoShardId, ShardSnapshotStatus.UNASSIGNED_QUEUED);
+                                    clonesBuilder.put(repoShardId, ShardSnapshotStatus.UNASSIGNED_QUEUED);
                                 } else {
-                                    clones.put(
+                                    clonesBuilder.put(
                                         repoShardId,
                                         new ShardSnapshotStatus(
                                             localNodeId,
@@ -672,7 +673,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                                 }
                             }
                         }
-                        updatedEntry = cloneEntry.withClones(clones);
+                        updatedEntry = cloneEntry.withClones(clonesBuilder.build());
                     } else {
                         updatedEntries.add(existing);
                     }
@@ -1139,7 +1140,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                                         updatedEntriesForRepo.add(snapshot);
                                         continue;
                                     }
-                                    Map<RepositoryShardId, ShardSnapshotStatus> clones = null;
+                                    ImmutableOpenMap.Builder<RepositoryShardId, ShardSnapshotStatus> clones = null;
                                     InFlightShardSnapshotStates inFlightShardSnapshotStates = null;
                                     for (Map.Entry<RepositoryShardId, ShardSnapshotStatus> failureEntry : knownFailures.entrySet()) {
                                         final RepositoryShardId repositoryShardId = failureEntry.getKey();
@@ -1156,7 +1157,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                                                 continue;
                                             }
                                             if (clones == null) {
-                                                clones = new HashMap<>(snapshot.shardsByRepoShardId());
+                                                clones = ImmutableOpenMap.builder(snapshot.shardsByRepoShardId());
                                             }
                                             // We can use the generation from the shard failure to start the clone operation here
                                             // because #processWaitingShardsAndRemovedNodes adds generations to failure statuses that
@@ -1172,13 +1173,13 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                                     }
                                     if (clones != null) {
                                         changed = true;
-                                        updatedEntriesForRepo.add(snapshot.withClones(clones));
+                                        updatedEntriesForRepo.add(snapshot.withClones(clones.build()));
                                     } else {
                                         updatedEntriesForRepo.add(snapshot);
                                     }
                                 }
                             } else {
-                                Map<ShardId, ShardSnapshotStatus> shards = processWaitingShardsAndRemovedNodes(
+                                ImmutableOpenMap<ShardId, ShardSnapshotStatus> shards = processWaitingShardsAndRemovedNodes(
                                     snapshot,
                                     routingTable,
                                     nodes,
@@ -1268,7 +1269,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
         });
     }
 
-    private static Map<ShardId, ShardSnapshotStatus> processWaitingShardsAndRemovedNodes(
+    private static ImmutableOpenMap<ShardId, ShardSnapshotStatus> processWaitingShardsAndRemovedNodes(
         SnapshotsInProgress.Entry entry,
         RoutingTable routingTable,
         DiscoveryNodes nodes,
@@ -1276,7 +1277,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
     ) {
         assert entry.isClone() == false : "clones take a different path";
         boolean snapshotChanged = false;
-        Map<ShardId, ShardSnapshotStatus> shards = new HashMap<>();
+        ImmutableOpenMap.Builder<ShardId, ShardSnapshotStatus> shards = ImmutableOpenMap.builder();
         for (Map.Entry<RepositoryShardId, ShardSnapshotStatus> shardEntry : entry.shardsByRepoShardId().entrySet()) {
             ShardSnapshotStatus shardStatus = shardEntry.getValue();
             ShardId shardId = entry.shardId(shardEntry.getKey());
@@ -1355,7 +1356,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
             }
         }
         if (snapshotChanged) {
-            return Collections.unmodifiableMap(shards);
+            return shards.build();
         } else {
             return null;
         }
@@ -1824,7 +1825,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                 final SnapshotsInProgress.Entry previousEntry = entryList.get(i);
                 if (removedEntry.isClone()) {
                     if (previousEntry.isClone()) {
-                        Map<RepositoryShardId, ShardSnapshotStatus> updatedShardAssignments = null;
+                        ImmutableOpenMap.Builder<RepositoryShardId, ShardSnapshotStatus> updatedShardAssignments = null;
                         for (Map.Entry<RepositoryShardId, ShardSnapshotStatus> finishedShardEntry : removedEntry.shardsByRepoShardId()
                             .entrySet()) {
                             final ShardSnapshotStatus shardState = finishedShardEntry.getValue();
@@ -1839,7 +1840,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                         }
                         addCloneEntry(entries, previousEntry, updatedShardAssignments);
                     } else {
-                        Map<ShardId, ShardSnapshotStatus> updatedShardAssignments = null;
+                        ImmutableOpenMap.Builder<ShardId, ShardSnapshotStatus> updatedShardAssignments = null;
                         for (Map.Entry<RepositoryShardId, ShardSnapshotStatus> finishedShardEntry : removedEntry.shardsByRepoShardId()
                             .entrySet()) {
                             final ShardSnapshotStatus shardState = finishedShardEntry.getValue();
@@ -1860,7 +1861,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                     }
                 } else {
                     if (previousEntry.isClone()) {
-                        Map<RepositoryShardId, ShardSnapshotStatus> updatedShardAssignments = null;
+                        ImmutableOpenMap.Builder<RepositoryShardId, ShardSnapshotStatus> updatedShardAssignments = null;
                         for (Map.Entry<RepositoryShardId, ShardSnapshotStatus> finishedShardEntry : removedEntry.shardsByRepoShardId()
                             .entrySet()) {
                             final ShardSnapshotStatus shardState = finishedShardEntry.getValue();
@@ -1878,7 +1879,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                         }
                         addCloneEntry(entries, previousEntry, updatedShardAssignments);
                     } else {
-                        Map<ShardId, ShardSnapshotStatus> updatedShardAssignments = null;
+                        ImmutableOpenMap.Builder<ShardId, ShardSnapshotStatus> updatedShardAssignments = null;
                         for (Map.Entry<RepositoryShardId, ShardSnapshotStatus> finishedShardEntry : removedEntry.shardsByRepoShardId()
                             .entrySet()) {
                             final ShardSnapshotStatus shardState = finishedShardEntry.getValue();
@@ -1909,34 +1910,36 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
     private static void addSnapshotEntry(
         List<SnapshotsInProgress.Entry> entries,
         SnapshotsInProgress.Entry entryToUpdate,
-        @Nullable Map<ShardId, ShardSnapshotStatus> updatedShardAssignments
+        @Nullable ImmutableOpenMap.Builder<ShardId, ShardSnapshotStatus> updatedShardAssignments
     ) {
         if (updatedShardAssignments == null) {
             entries.add(entryToUpdate);
         } else {
-            final Map<ShardId, ShardSnapshotStatus> updatedStatus = new HashMap<>(entryToUpdate.shards());
-            updatedStatus.putAll(updatedShardAssignments);
-            entries.add(entryToUpdate.withShardStates(updatedStatus));
+            final ImmutableOpenMap.Builder<ShardId, ShardSnapshotStatus> updatedStatus = ImmutableOpenMap.builder(entryToUpdate.shards());
+            updatedStatus.putAllFromMap(updatedShardAssignments.build());
+            entries.add(entryToUpdate.withShardStates(updatedStatus.build()));
         }
     }
 
     private static void addCloneEntry(
         List<SnapshotsInProgress.Entry> entries,
         SnapshotsInProgress.Entry entryToUpdate,
-        @Nullable Map<RepositoryShardId, ShardSnapshotStatus> updatedShardAssignments
+        @Nullable ImmutableOpenMap.Builder<RepositoryShardId, ShardSnapshotStatus> updatedShardAssignments
     ) {
         if (updatedShardAssignments == null) {
             entries.add(entryToUpdate);
         } else {
-            final Map<RepositoryShardId, ShardSnapshotStatus> updatedStatus = new HashMap<>(entryToUpdate.shardsByRepoShardId());
-            updatedStatus.putAll(updatedShardAssignments);
-            entries.add(entryToUpdate.withClones(updatedStatus));
+            final ImmutableOpenMap.Builder<RepositoryShardId, ShardSnapshotStatus> updatedStatus = ImmutableOpenMap.builder(
+                entryToUpdate.shardsByRepoShardId()
+            );
+            updatedStatus.putAllFromMap(updatedShardAssignments.build());
+            entries.add(entryToUpdate.withClones(updatedStatus.build()));
         }
     }
 
     @Nullable
-    private static <T> Map<T, ShardSnapshotStatus> maybeAddUpdatedAssignment(
-        @Nullable Map<T, ShardSnapshotStatus> updatedShardAssignments,
+    private static <T> ImmutableOpenMap.Builder<T, ShardSnapshotStatus> maybeAddUpdatedAssignment(
+        @Nullable ImmutableOpenMap.Builder<T, ShardSnapshotStatus> updatedShardAssignments,
         ShardSnapshotStatus finishedShardState,
         T shardId,
         Map<T, ShardSnapshotStatus> statesToUpdate
@@ -1947,7 +1950,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
             && stateToUpdate.state() == ShardState.SUCCESS
             && Objects.equals(newGeneration, stateToUpdate.generation()) == false) {
             if (updatedShardAssignments == null) {
-                updatedShardAssignments = new HashMap<>();
+                updatedShardAssignments = ImmutableOpenMap.builder();
             }
             updatedShardAssignments.put(shardId, stateToUpdate.withUpdatedGeneration(newGeneration));
         }
@@ -2690,13 +2693,12 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                             if (inFlightShardStates == null) {
                                 inFlightShardStates = InFlightShardSnapshotStates.forEntries(snapshotsInProgress.forRepo(repoName));
                             }
-                            final Map<RepositoryShardId, ShardSnapshotStatus> updatedAssignments = new HashMap<>(
-                                entry.shardsByRepoShardId()
-                            );
+                            final ImmutableOpenMap.Builder<RepositoryShardId, ShardSnapshotStatus> updatedAssignmentsBuilder =
+                                ImmutableOpenMap.builder(entry.shardsByRepoShardId());
                             for (RepositoryShardId shardId : canBeUpdated) {
                                 if (inFlightShardStates.isActive(shardId.indexName(), shardId.shardId()) == false) {
                                     markShardReassigned(shardId, reassignedShardIds);
-                                    updatedAssignments.put(
+                                    updatedAssignmentsBuilder.put(
                                         shardId,
                                         new ShardSnapshotStatus(
                                             localNodeId,
@@ -2709,7 +2711,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                                     );
                                 }
                             }
-                            snapshotEntries.add(entry.withClones(updatedAssignments));
+                            snapshotEntries.add(entry.withClones(updatedAssignmentsBuilder.build()));
                             changed = true;
                         }
                     } else {
@@ -2729,7 +2731,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                             // No shards can be updated in this snapshot so we just add it as is again
                             snapshotEntries.add(entry);
                         } else {
-                            final Map<ShardId, ShardSnapshotStatus> shardAssignments = shards(
+                            final ImmutableOpenMap<ShardId, ShardSnapshotStatus> shardAssignments = shards(
                                 snapshotsInProgress,
                                 updatedDeletions,
                                 currentState,
@@ -2738,7 +2740,8 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                                 repositoryData,
                                 repoName
                             );
-                            final Map<ShardId, ShardSnapshotStatus> updatedAssignments = new HashMap<>(entry.shards());
+                            final ImmutableOpenMap.Builder<ShardId, ShardSnapshotStatus> updatedAssignmentsBuilder = ImmutableOpenMap
+                                .builder(entry.shards());
                             for (RepositoryShardId shardId : canBeUpdated) {
                                 final ShardId sid = entry.shardId(shardId);
                                 final ShardSnapshotStatus updated = shardAssignments.get(sid);
@@ -2746,15 +2749,15 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                                     // We don't have a new assignment for this shard because its index was concurrently deleted
                                     assert currentState.routingTable().hasIndex(sid.getIndex()) == false
                                         : "Missing assignment for [" + sid + "]";
-                                    updatedAssignments.put(sid, ShardSnapshotStatus.MISSING);
+                                    updatedAssignmentsBuilder.put(sid, ShardSnapshotStatus.MISSING);
                                 } else {
                                     if (updated.isActive()) {
                                         markShardReassigned(shardId, reassignedShardIds);
                                     }
-                                    updatedAssignments.put(sid, updated);
+                                    updatedAssignmentsBuilder.put(sid, updated);
                                 }
                             }
-                            final SnapshotsInProgress.Entry updatedEntry = entry.withShardStates(updatedAssignments);
+                            final SnapshotsInProgress.Entry updatedEntry = entry.withShardStates(updatedAssignmentsBuilder.build());
                             snapshotEntries.add(updatedEntry);
                             changed = true;
                             if (updatedEntry.state().completed()) {
@@ -2770,9 +2773,13 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                 }
             }
             if (changed && newIndexIdsToRefresh.isEmpty() == false) {
-                Map<IndexId, IndexId> updatedIndexIds = newIndexIdsToRefresh.stream()
-                    .collect(Collectors.toMap(Function.identity(), i -> new IndexId(i.getName(), UUIDs.randomBase64UUID())));
-                snapshotEntries.replaceAll(e -> e.withUpdatedIndexIds(updatedIndexIds));
+                final Map<IndexId, IndexId> updatedIndexIds = Maps.newMapWithExpectedSize(newIndexIdsToRefresh.size());
+                for (IndexId indexIdToRefresh : newIndexIdsToRefresh) {
+                    updatedIndexIds.put(indexIdToRefresh, new IndexId(indexIdToRefresh.getName(), UUIDs.randomBase64UUID()));
+                }
+                for (int i = 0; i < snapshotEntries.size(); i++) {
+                    snapshotEntries.set(i, snapshotEntries.get(i).withUpdatedIndexIds(updatedIndexIds));
+                }
             }
             return changed ? snapshotsInProgress.withUpdatedEntriesForRepo(repoName, snapshotEntries) : null;
         }
@@ -2840,7 +2847,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
      * @param useShardGenerations whether to write {@link ShardGenerations} during the snapshot
      * @return list of shard to be included into current snapshot
      */
-    private static Map<ShardId, SnapshotsInProgress.ShardSnapshotStatus> shards(
+    private static ImmutableOpenMap<ShardId, SnapshotsInProgress.ShardSnapshotStatus> shards(
         SnapshotsInProgress snapshotsInProgress,
         SnapshotDeletionsInProgress deletionsInProgress,
         ClusterState currentState,
@@ -2849,7 +2856,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
         RepositoryData repositoryData,
         String repoName
     ) {
-        Map<ShardId, SnapshotsInProgress.ShardSnapshotStatus> shards = new HashMap<>();
+        ImmutableOpenMap.Builder<ShardId, SnapshotsInProgress.ShardSnapshotStatus> builder = ImmutableOpenMap.builder();
         final ShardGenerations shardGenerations = repositoryData.shardGenerations();
         final InFlightShardSnapshotStates inFlightShardStates = InFlightShardSnapshotStates.forEntries(
             snapshotsInProgress.forRepo(repoName)
@@ -2861,7 +2868,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
             IndexMetadata indexMetadata = currentState.metadata().index(indexName);
             if (indexMetadata == null) {
                 // The index was deleted before we managed to start the snapshot - mark it as missing.
-                shards.put(new ShardId(indexName, IndexMetadata.INDEX_UUID_NA_VALUE, 0), ShardSnapshotStatus.MISSING);
+                builder.put(new ShardId(indexName, IndexMetadata.INDEX_UUID_NA_VALUE, 0), ShardSnapshotStatus.MISSING);
             } else {
                 final IndexRoutingTable indexRoutingTable = currentState.routingTable().index(indexName);
                 assert indexRoutingTable != null;
@@ -2890,12 +2897,12 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                     } else {
                         shardSnapshotStatus = initShardSnapshotStatus(shardRepoGeneration, indexRoutingTable.shard(i).primaryShard());
                     }
-                    shards.put(shardId, shardSnapshotStatus);
+                    builder.put(shardId, shardSnapshotStatus);
                 }
             }
         }
 
-        return Collections.unmodifiableMap(shards);
+        return builder.build();
     }
 
     /**
@@ -3135,11 +3142,11 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
             // iterator containing the updates yet to be applied to #entry
             private final Iterator<ShardSnapshotUpdate> iterator;
 
-            // map for updated shard snapshot status mappings if any could be computed
-            private Map<ShardId, ShardSnapshotStatus> shards = null;
+            // builder for updated shard snapshot status mappings if any could be computed
+            private ImmutableOpenMap.Builder<ShardId, ShardSnapshotStatus> shardsBuilder = null;
 
-            // map for updated shard clone status mappings if any could be computed
-            private Map<RepositoryShardId, ShardSnapshotStatus> clones = null;
+            // builder for updated shard clone status mappings if any could be computed
+            private ImmutableOpenMap.Builder<RepositoryShardId, ShardSnapshotStatus> clonesBuilder = null;
 
             EntryContext(SnapshotsInProgress.Entry entry, List<ShardSnapshotUpdate> updates) {
                 this.entry = entry;
@@ -3147,7 +3154,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
             }
 
             SnapshotsInProgress.Entry computeUpdatedEntry() {
-                assert shards == null && clones == null : "update context was already used";
+                assert shardsBuilder == null && clonesBuilder == null : "update context was already used";
 
                 // loop over all the shard updates that are potentially applicable to the current snapshot entry
                 while (iterator.hasNext()) {
@@ -3169,12 +3176,15 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                     }
                 }
 
-                if (shards != null) {
-                    assert clones == null
-                        : "Should not have updated clones when updating shard snapshots but saw " + clones + " as well as " + shards;
-                    return entry.withShardStates(shards);
-                } else if (clones != null) {
-                    return entry.withClones(clones);
+                if (shardsBuilder != null) {
+                    assert clonesBuilder == null
+                        : "Should not have updated clones when updating shard snapshots but saw "
+                            + clonesBuilder
+                            + " as well as "
+                            + shardsBuilder;
+                    return entry.withShardStates(shardsBuilder.build());
+                } else if (clonesBuilder != null) {
+                    return entry.withClones(clonesBuilder.build());
                 } else {
                     return entry;
                 }
@@ -3190,7 +3200,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
              * @param <T>         either {@link ShardId} for snapshots or {@link RepositoryShardId} for clones
              */
             private <T> void startShardOperation(
-                Map<T, ShardSnapshotStatus> newStates,
+                ImmutableOpenMap.Builder<T, ShardSnapshotStatus> newStates,
                 String nodeId,
                 ShardGeneration generation,
                 T shardId
@@ -3206,7 +3216,11 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
              * @param newState  new shard task state for operation to start
              * @param <T>       either {@link ShardId} for snapshots or {@link RepositoryShardId} for clones
              */
-            private <T> void startShardOperation(Map<T, ShardSnapshotStatus> newStates, T shardId, ShardSnapshotStatus newState) {
+            private <T> void startShardOperation(
+                ImmutableOpenMap.Builder<T, ShardSnapshotStatus> newStates,
+                T shardId,
+                ShardSnapshotStatus newState
+            ) {
                 logger.trace(
                     "[{}] Starting [{}] on [{}] with generation [{}]",
                     entry.snapshot(),
@@ -3221,7 +3235,7 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
 
             private <T> void executeShardSnapshotUpdate(
                 Map<T, ShardSnapshotStatus> existingStates,
-                Supplier<Map<T, ShardSnapshotStatus>> newStates,
+                Supplier<ImmutableOpenMap.Builder<T, ShardSnapshotStatus>> newStates,
                 ShardSnapshotUpdate updateSnapshotState,
                 T updatedShard
             ) {
@@ -3321,20 +3335,20 @@ public class SnapshotsService extends AbstractLifecycleComponent implements Clus
                 }
             }
 
-            private Map<RepositoryShardId, ShardSnapshotStatus> clonesBuilder() {
-                assert shards == null;
-                if (clones == null) {
-                    clones = new HashMap<>(entry.shardsByRepoShardId());
+            private ImmutableOpenMap.Builder<RepositoryShardId, ShardSnapshotStatus> clonesBuilder() {
+                assert shardsBuilder == null;
+                if (clonesBuilder == null) {
+                    clonesBuilder = ImmutableOpenMap.builder(entry.shardsByRepoShardId());
                 }
-                return clones;
+                return clonesBuilder;
             }
 
-            private Map<ShardId, ShardSnapshotStatus> shardsBuilder() {
-                assert clones == null;
-                if (shards == null) {
-                    shards = new HashMap<>(entry.shards());
+            private ImmutableOpenMap.Builder<ShardId, ShardSnapshotStatus> shardsBuilder() {
+                assert clonesBuilder == null;
+                if (shardsBuilder == null) {
+                    shardsBuilder = ImmutableOpenMap.builder(entry.shards());
                 }
-                return shards;
+                return shardsBuilder;
             }
         }
     }