simplified way to copy and modify a cluster state (#81975)

It looks like there is a pattern in a way we modify a cluster state.
This commits adds a method to simplify that flow.

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

@@ -54,6 +54,7 @@ import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.function.Consumer;
 
 /**
  * Represents the state of the cluster, held in memory on all nodes in the cluster with updates coordinated by the elected master.
@@ -575,6 +576,16 @@ public class ClusterState implements ToXContentFragment, Diffable<ClusterState>
         return new Builder(state);
     }
 
+    public ClusterState copyAndUpdate(Consumer<Builder> updater) {
+        var builder = builder(this);
+        updater.accept(builder);
+        return builder.build();
+    }
+
+    public ClusterState copyAndUpdateMetadata(Consumer<Metadata.Builder> updater) {
+        return copyAndUpdate(builder -> builder.metadata(metadata().copyAndUpdate(updater)));
+    }
+
     public static class Builder {
 
         private ClusterState previous;

server/src/main/java/org/elasticsearch/cluster/metadata/Metadata.java

@@ -69,6 +69,7 @@ import java.util.Optional;
 import java.util.Set;
 import java.util.SortedMap;
 import java.util.TreeMap;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -1100,6 +1101,12 @@ public class Metadata implements Iterable<IndexMetadata>, Diffable<Metadata>, To
         return new Builder(metadata);
     }
 
+    public Metadata copyAndUpdate(Consumer<Builder> updater) {
+        var builder = builder(this);
+        updater.accept(builder);
+        return builder.build();
+    }
+
     public static class Builder {
 
         private String clusterUUID;

server/src/test/java/org/elasticsearch/cluster/ClusterStateTests.java

@@ -26,6 +26,7 @@ import org.elasticsearch.cluster.routing.RoutingTable;
 import org.elasticsearch.cluster.routing.ShardRoutingState;
 import org.elasticsearch.cluster.routing.TestShardRouting;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.UUIDs;
 import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.transport.TransportAddress;
@@ -34,7 +35,6 @@ import org.elasticsearch.index.Index;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.test.ESTestCase;
-import org.elasticsearch.test.TestCustomMetadata;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.json.JsonXContent;
@@ -54,6 +54,8 @@ import static java.util.Collections.singletonMap;
 import static org.elasticsearch.cluster.metadata.IndexMetadata.SETTING_VERSION_CREATED;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.not;
+import static org.hamcrest.Matchers.sameInstance;
 
 public class ClusterStateTests extends ESTestCase {
 
@@ -105,6 +107,26 @@ public class ClusterStateTests extends ESTestCase {
         assertThat(expectThrows(NullPointerException.class, () -> builder.customs(map)).getMessage(), containsString(key));
     }
 
+    public void testCopyAndUpdate() throws IOException {
+        var state = buildClusterState();
+        var newStateUuid = UUIDs.base64UUID();
+
+        var copy = state.copyAndUpdate(builder -> builder.stateUUID(newStateUuid));
+
+        assertThat(copy, not(sameInstance(state)));
+        assertThat(copy.stateUUID(), equalTo(newStateUuid));
+    }
+
+    public void testCopyAndUpdateMetadata() throws IOException {
+        var state = buildClusterState();
+        var newClusterUuid = UUIDs.base64UUID();
+
+        var copy = state.copyAndUpdateMetadata(metadata -> metadata.clusterUUID(newClusterUuid));
+
+        assertThat(copy, not(sameInstance(state)));
+        assertThat(copy.metadata().clusterUUID(), equalTo(newClusterUuid));
+    }
+
     public void testToXContent() throws IOException {
         final ClusterState clusterState = buildClusterState();
 
@@ -930,27 +952,4 @@ public class ClusterStateTests extends ESTestCase {
             )
             .build();
     }
-
-    public static class CustomMetadata extends TestCustomMetadata {
-        public static final String TYPE = "custom_md";
-
-        CustomMetadata(String data) {
-            super(data);
-        }
-
-        @Override
-        public String getWriteableName() {
-            return TYPE;
-        }
-
-        @Override
-        public Version getMinimalSupportedVersion() {
-            return Version.CURRENT;
-        }
-
-        @Override
-        public EnumSet<Metadata.XContentContext> context() {
-            return EnumSet.of(Metadata.XContentContext.GATEWAY, Metadata.XContentContext.SNAPSHOT);
-        }
-    }
 }

server/src/test/java/org/elasticsearch/cluster/metadata/MetadataTests.java

@@ -1062,6 +1062,16 @@ public class MetadataTests extends ESTestCase {
         assertThat(expectThrows(NullPointerException.class, () -> builder.customs(map)).getMessage(), containsString(key));
     }
 
+    public void testCopyAndUpdate() throws IOException {
+        var metadata = Metadata.builder().clusterUUID(UUIDs.base64UUID()).build();
+        var newClusterUuid = UUIDs.base64UUID();
+
+        var copy = metadata.copyAndUpdate(builder -> builder.clusterUUID(newClusterUuid));
+
+        assertThat(copy, not(sameInstance(metadata)));
+        assertThat(copy.clusterUUID(), equalTo(newClusterUuid));
+    }
+
     public void testBuilderRejectsDataStreamThatConflictsWithIndex() {
         final String dataStreamName = "my-data-stream";
         IndexMetadata idx = createFirstBackingIndex(dataStreamName).build();

x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportActivateAutoFollowPatternAction.java

@@ -17,7 +17,6 @@ import org.elasticsearch.cluster.ClusterStateTaskExecutor;
 import org.elasticsearch.cluster.block.ClusterBlockException;
 import org.elasticsearch.cluster.block.ClusterBlockLevel;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
-import org.elasticsearch.cluster.metadata.Metadata;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.tasks.Task;
@@ -115,19 +114,11 @@ public class TransportActivateAutoFollowPatternAction extends AcknowledgedTransp
             )
         );
 
-        return ClusterState.builder(currentState)
-            .metadata(
-                Metadata.builder(currentState.getMetadata())
-                    .putCustom(
-                        AutoFollowMetadata.TYPE,
-                        new AutoFollowMetadata(
-                            newPatterns,
-                            autoFollowMetadata.getFollowedLeaderIndexUUIDs(),
-                            autoFollowMetadata.getHeaders()
-                        )
-                    )
-                    .build()
+        return currentState.copyAndUpdateMetadata(
+            metadata -> metadata.putCustom(
+                AutoFollowMetadata.TYPE,
+                new AutoFollowMetadata(newPatterns, autoFollowMetadata.getFollowedLeaderIndexUUIDs(), autoFollowMetadata.getHeaders())
             )
-            .build();
+        );
     }
 }

x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportDeleteAutoFollowPatternAction.java

@@ -17,7 +17,6 @@ import org.elasticsearch.cluster.ClusterStateTaskExecutor;
 import org.elasticsearch.cluster.block.ClusterBlockException;
 import org.elasticsearch.cluster.block.ClusterBlockLevel;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
-import org.elasticsearch.cluster.metadata.Metadata;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.util.Maps;
@@ -76,9 +75,7 @@ public class TransportDeleteAutoFollowPatternAction extends AcknowledgedTranspor
 
         AutoFollowMetadata newAutoFollowMetadata = removePattern(currentAutoFollowMetadata, request.getName());
 
-        return ClusterState.builder(currentState)
-            .metadata(Metadata.builder(currentState.getMetadata()).putCustom(AutoFollowMetadata.TYPE, newAutoFollowMetadata))
-            .build();
+        return currentState.copyAndUpdateMetadata(metadata -> metadata.putCustom(AutoFollowMetadata.TYPE, newAutoFollowMetadata));
     }
 
     private static AutoFollowMetadata removePattern(AutoFollowMetadata metadata, String name) {

x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportPutAutoFollowPatternAction.java

@@ -203,13 +203,10 @@ public class TransportPutAutoFollowPatternAction extends AcknowledgedTransportMa
             request.getParameters().getReadPollTimeout()
         );
         patterns.put(request.getName(), autoFollowPattern);
-        ClusterState.Builder newState = ClusterState.builder(localState);
-        newState.metadata(
-            Metadata.builder(localState.getMetadata())
-                .putCustom(AutoFollowMetadata.TYPE, new AutoFollowMetadata(patterns, followedLeaderIndices, headers))
-                .build()
+
+        return localState.copyAndUpdateMetadata(
+            metadata -> metadata.putCustom(AutoFollowMetadata.TYPE, new AutoFollowMetadata(patterns, followedLeaderIndices, headers))
         );
-        return newState.build();
     }
 
     private static void markExistingIndicesAsAutoFollowedForNewPatterns(

x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportUnfollowAction.java

@@ -27,7 +27,6 @@ import org.elasticsearch.cluster.block.ClusterBlockException;
 import org.elasticsearch.cluster.block.ClusterBlockLevel;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
-import org.elasticsearch.cluster.metadata.Metadata;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.settings.Settings;
@@ -270,7 +269,6 @@ public class TransportUnfollowAction extends AcknowledgedTransportMasterNodeActi
         // Remove ccr custom metadata
         newIndexMetadata.removeCustom(Ccr.CCR_CUSTOM_METADATA_KEY);
 
-        Metadata newMetadata = Metadata.builder(current.metadata()).put(newIndexMetadata).build();
-        return ClusterState.builder(current).metadata(newMetadata).build();
+        return current.copyAndUpdateMetadata(metadata -> metadata.put(newIndexMetadata));
     }
 }