Simplify map copying (elastic#88432)

This commit introduces a method that simplifies creating a map deep copy
weizijun · Jul 13, 2022 · 0768557 · 0768557
1 parent 5d23480
commit 0768557
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 21 deletions.
diff --git a/server/src/main/java/org/elasticsearch/cluster/routing/RoutingNode.java b/server/src/main/java/org/elasticsearch/cluster/routing/RoutingNode.java
@@ -69,10 +69,7 @@ private RoutingNode(RoutingNode original) {
         this.shards = new LinkedHashMap<>(original.shards);
         this.relocatingShards = new LinkedHashSet<>(original.relocatingShards);
         this.initializingShards = new LinkedHashSet<>(original.initializingShards);
-        this.shardsByIndex = Maps.newMapWithExpectedSize(original.shardsByIndex.size());
-        for (Map.Entry<Index, Set<ShardRouting>> entry : original.shardsByIndex.entrySet()) {
-            shardsByIndex.put(entry.getKey(), new HashSet<>(entry.getValue()));
-        }
+        this.shardsByIndex = Maps.copyOf(original.shardsByIndex, HashSet::new);
         assert invariant();
     }
 

diff --git a/server/src/main/java/org/elasticsearch/cluster/routing/RoutingNodes.java b/server/src/main/java/org/elasticsearch/cluster/routing/RoutingNodes.java
@@ -152,27 +152,14 @@ private RoutingNodes(RoutingNodes routingNodes) {
         // instance
         assert routingNodes.readOnly : "tried to create a mutable copy from a mutable instance";
         this.readOnly = false;
-        this.nodesToShards = Maps.newMapWithExpectedSize(routingNodes.nodesToShards.size());
-        for (Map.Entry<String, RoutingNode> entry : routingNodes.nodesToShards.entrySet()) {
-            this.nodesToShards.put(entry.getKey(), entry.getValue().copy());
-        }
-        this.assignedShards = Maps.newMapWithExpectedSize(routingNodes.assignedShards.size());
-        for (Map.Entry<ShardId, List<ShardRouting>> entry : routingNodes.assignedShards.entrySet()) {
-            this.assignedShards.put(entry.getKey(), new ArrayList<>(entry.getValue()));
-        }
+        this.nodesToShards = Maps.copyOf(routingNodes.nodesToShards, RoutingNode::copy);
+        this.assignedShards = Maps.copyOf(routingNodes.assignedShards, ArrayList::new);
         this.unassignedShards = routingNodes.unassignedShards.copyFor(this);
-
         this.inactivePrimaryCount = routingNodes.inactivePrimaryCount;
         this.inactiveShardCount = routingNodes.inactiveShardCount;
         this.relocatingShards = routingNodes.relocatingShards;
-        this.attributeValuesByAttribute = Maps.newMapWithExpectedSize(routingNodes.attributeValuesByAttribute.size());
-        for (Map.Entry<String, Set<String>> entry : routingNodes.attributeValuesByAttribute.entrySet()) {
-            this.attributeValuesByAttribute.put(entry.getKey(), new HashSet<>(entry.getValue()));
-        }
-        this.recoveriesPerNode = Maps.newMapWithExpectedSize(routingNodes.recoveriesPerNode.size());
-        for (Map.Entry<String, Recoveries> entry : routingNodes.recoveriesPerNode.entrySet()) {
-            this.recoveriesPerNode.put(entry.getKey(), entry.getValue().copy());
-        }
+        this.attributeValuesByAttribute = Maps.copyOf(routingNodes.attributeValuesByAttribute, HashSet::new);
+        this.recoveriesPerNode = Maps.copyOf(routingNodes.recoveriesPerNode, Recoveries::copy);
     }
 
     /**

diff --git a/server/src/main/java/org/elasticsearch/common/util/Maps.java b/server/src/main/java/org/elasticsearch/common/util/Maps.java
@@ -298,4 +298,14 @@ static int capacity(int expectedSize) {
         throw new IllegalStateException("Unsupported input format");
     }
 
+    /**
+     * This method creates a copy of the {@code source} map using {@code copyValueFunction} to create a defensive copy of each value.
+     */
+    public static <K, V> Map<K, V> copyOf(Map<K, V> source, Function<V, V> copyValueFunction) {
+        var copy = Maps.<K, V>newHashMapWithExpectedSize(source.size());
+        for (var entry : source.entrySet()) {
+            copy.put(entry.getKey(), copyValueFunction.apply(entry.getValue()));
+        }
+        return copy;
+    }
 }