Avoid needless map copy during disk usage simulation (elastic#91654)

Author: idegtiarenko

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

@@ -76,9 +76,9 @@ public ClusterInfo(
         Map<NodeAndPath, ReservedSpace> reservedSpace
     ) {
         this.leastAvailableSpaceUsage = Map.copyOf(leastAvailableSpaceUsage);
+        this.mostAvailableSpaceUsage = Map.copyOf(mostAvailableSpaceUsage);
         this.shardSizes = Map.copyOf(shardSizes);
         this.shardDataSetSizes = Map.copyOf(shardDataSetSizes);
-        this.mostAvailableSpaceUsage = Map.copyOf(mostAvailableSpaceUsage);
         this.dataPath = Map.copyOf(dataPath);
         this.reservedSpace = Map.copyOf(reservedSpace);
     }

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

@@ -9,6 +9,7 @@
 package org.elasticsearch.cluster;
 
 import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.common.util.CopyOnFirstWriteMap;
 import org.elasticsearch.index.shard.ShardId;
 
 import java.util.HashMap;
@@ -19,14 +20,14 @@ public class ClusterInfoSimulator {
 
     private final Map<String, DiskUsage> leastAvailableSpaceUsage;
     private final Map<String, DiskUsage> mostAvailableSpaceUsage;
-    private final Map<String, Long> shardSizes;
+    private final CopyOnFirstWriteMap<String, Long> shardSizes;
     private final Map<ShardId, Long> shardDataSetSizes;
     private final Map<ClusterInfo.NodeAndShard, String> dataPath;
 
     public ClusterInfoSimulator(ClusterInfo clusterInfo) {
         this.leastAvailableSpaceUsage = new HashMap<>(clusterInfo.getNodeLeastAvailableDiskUsages());
         this.mostAvailableSpaceUsage = new HashMap<>(clusterInfo.getNodeMostAvailableDiskUsages());
-        this.shardSizes = new HashMap<>(clusterInfo.shardSizes);
+        this.shardSizes = new CopyOnFirstWriteMap<>(clusterInfo.shardSizes);
         this.shardDataSetSizes = Map.copyOf(clusterInfo.shardDataSetSizes);
         this.dataPath = Map.copyOf(clusterInfo.dataPath);
     }
@@ -101,6 +102,13 @@ private static long withinRange(long min, long max, long value) {
     }
 
     public ClusterInfo getClusterInfo() {
-        return new ClusterInfo(leastAvailableSpaceUsage, mostAvailableSpaceUsage, shardSizes, shardDataSetSizes, dataPath, Map.of());
+        return new ClusterInfo(
+            leastAvailableSpaceUsage,
+            mostAvailableSpaceUsage,
+            shardSizes.toImmutableMap(),
+            shardDataSetSizes,
+            dataPath,
+            Map.of()
+        );
     }
 }

server/src/main/java/org/elasticsearch/common/util/CopyOnFirstWriteMap.java

@@ -0,0 +1,105 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.common.util;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * This map is designed to be constructed from an immutable map and be copied only if a (rare) mutation operation occurs.
+ * It should be converted back to an immutable map using `org.elasticsearch.common.util.LazyCopyOnWriteMap#toImmutableMap()`.
+ */
+public class CopyOnFirstWriteMap<K, V> implements Map<K, V> {
+
+    private Map<K, V> source;
+    private boolean wasCopied = false;
+
+    public CopyOnFirstWriteMap(Map<K, V> source) {
+        this.source = source;
+    }
+
+    private Map<K, V> getForUpdate() {
+        if (wasCopied == false) {
+            source = new HashMap<>(source);
+            wasCopied = true;
+        }
+        return source;
+    }
+
+    private Map<K, V> getForRead() {
+        return source;
+    }
+
+    public Map<K, V> toImmutableMap() {
+        return Map.copyOf(getForRead());
+    }
+
+    @Override
+    public int size() {
+        return getForRead().size();
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return getForRead().isEmpty();
+    }
+
+    @Override
+    public boolean containsKey(Object key) {
+        return getForRead().containsKey(key);
+    }
+
+    @Override
+    public boolean containsValue(Object value) {
+        return getForRead().containsValue(value);
+    }
+
+    @Override
+    public V get(Object key) {
+        return getForRead().get(key);
+    }
+
+    @Override
+    public V put(K key, V value) {
+        return getForUpdate().put(key, value);
+    }
+
+    @Override
+    public V remove(Object key) {
+        return getForUpdate().remove(key);
+    }
+
+    @Override
+    public void putAll(Map<? extends K, ? extends V> m) {
+        getForUpdate().putAll(m);
+    }
+
+    @Override
+    public void clear() {
+        getForUpdate().clear();
+    }
+
+    @Override
+    public Set<K> keySet() {
+        return Collections.unmodifiableSet(getForRead().keySet());
+    }
+
+    @Override
+    public Collection<V> values() {
+        return Collections.unmodifiableCollection(getForRead().values());
+    }
+
+    @Override
+    public Set<Entry<K, V>> entrySet() {
+        return Collections.unmodifiableSet(getForRead().entrySet());
+    }
+}

server/src/test/java/org/elasticsearch/common/util/CopyOnFirstWriteMapTests.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.common.util;
+
+import org.elasticsearch.test.ESTestCase;
+
+import java.util.Map;
+
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.not;
+import static org.hamcrest.Matchers.sameInstance;
+
+public class CopyOnFirstWriteMapTests extends ESTestCase {
+
+    public void testShouldNotCopyIfThereWereNoUpdates() {
+        var source = Map.of("key", "value");
+        var copyOnFirstWrite = new CopyOnFirstWriteMap<>(source);
+        source.get("key");
+        var copy = copyOnFirstWrite.toImmutableMap();
+
+        assertThat(copy, sameInstance(source));
+        assertThat(copy, equalTo(source));
+    }
+
+    public void testShouldBeUpdatable() {
+        var source = Map.of("key", "value");
+        var copyOnFirstWrite = new CopyOnFirstWriteMap<>(source);
+        copyOnFirstWrite.put("key", "new_value");
+        var copy = copyOnFirstWrite.toImmutableMap();
+
+        assertThat(copy, not(sameInstance(source)));
+        assertThat(copy, equalTo(Map.of("key", "new_value")));
+    }
+}