From 07685573254ea315fbeaa0bbabb627c44f40921d Mon Sep 17 00:00:00 2001
From: Ievgen Degtiarenko <ievgen.degtiarenko@elastic.co>
Date: Wed, 13 Jul 2022 08:27:41 +0200
Subject: [PATCH] Simplify map copying (#88432)

This commit introduces a method that simplifies creating a map deep copy
---
 .../cluster/routing/RoutingNode.java          |  5 +----
 .../cluster/routing/RoutingNodes.java         | 21 ++++---------------
 .../org/elasticsearch/common/util/Maps.java   | 10 +++++++++
 3 files changed, 15 insertions(+), 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
index ffe535990042f..9987facfba598 100644
--- 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
index 747a9c2ff9c7a..33872c66603c7 100644
--- 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
index 5b02f9af64a06..b41d3ff6ff159 100644
--- 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;
+    }
 }