From 0647b4e757840b234acd112e3e92ad01ffbecc3e Mon Sep 17 00:00:00 2001
From: Niels Bauman <33722607+nielsbauman@users.noreply.github.com>
Date: Thu, 30 Jan 2025 09:09:00 +1000
Subject: [PATCH] Avoid unnecessarily copying enrich policies map (#121127)

Instead of always copying the map of enrich policies, we should return
the (already read-only) map straight from the `EnrichMetadata` and make
a modifiable copy only when necessary.
---
 .../xpack/core/enrich/EnrichMetadata.java     |  2 ++
 .../xpack/enrich/EnrichStore.java             | 35 ++++++++-----------
 .../esql/enrich/EnrichPolicyResolver.java     |  4 +--
 3 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/enrich/EnrichMetadata.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/enrich/EnrichMetadata.java
index b949e44ef036a..06bae0f8183bc 100644
--- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/enrich/EnrichMetadata.java
+++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/enrich/EnrichMetadata.java
@@ -36,6 +36,8 @@ public final class EnrichMetadata extends AbstractNamedDiffable<Metadata.Custom>
 
     static final ParseField POLICIES = new ParseField("policies");
 
+    public static final EnrichMetadata EMPTY = new EnrichMetadata(Collections.emptyMap());
+
     @SuppressWarnings("unchecked")
     private static final ConstructingObjectParser<EnrichMetadata, Void> PARSER = new ConstructingObjectParser<>(
         "enrich_metadata",
diff --git a/x-pack/plugin/enrich/src/main/java/org/elasticsearch/xpack/enrich/EnrichStore.java b/x-pack/plugin/enrich/src/main/java/org/elasticsearch/xpack/enrich/EnrichStore.java
index 82f9877826a5c..6e7f3846963ca 100644
--- a/x-pack/plugin/enrich/src/main/java/org/elasticsearch/xpack/enrich/EnrichStore.java
+++ b/x-pack/plugin/enrich/src/main/java/org/elasticsearch/xpack/enrich/EnrichStore.java
@@ -81,6 +81,10 @@ public static void putPolicy(
         }
 
         updateClusterState(clusterService, handler, current -> {
+            final Map<String, EnrichPolicy> originalPolicies = getPolicies(current);
+            if (originalPolicies.containsKey(name)) {
+                throw new ResourceAlreadyExistsException("policy [{}] already exists", name);
+            }
             for (String indexExpression : policy.getIndices()) {
                 // indices field in policy can contain wildcards, aliases etc.
                 String[] concreteIndices = indexNameExpressionResolver.concreteIndexNames(
@@ -101,12 +105,9 @@ public static void putPolicy(
                 }
             }
 
-            final Map<String, EnrichPolicy> policies = getPolicies(current);
-            EnrichPolicy existing = policies.putIfAbsent(name, policy);
-            if (existing != null) {
-                throw new ResourceAlreadyExistsException("policy [{}] already exists", name);
-            }
-            return policies;
+            final Map<String, EnrichPolicy> updatedPolicies = new HashMap<>(originalPolicies);
+            updatedPolicies.put(name, policy);
+            return updatedPolicies;
         });
     }
 
@@ -125,13 +126,14 @@ public static void deletePolicy(String name, ClusterService clusterService, Cons
         }
 
         updateClusterState(clusterService, handler, current -> {
-            final Map<String, EnrichPolicy> policies = getPolicies(current);
-            if (policies.containsKey(name) == false) {
+            final Map<String, EnrichPolicy> originalPolicies = getPolicies(current);
+            if (originalPolicies.containsKey(name) == false) {
                 throw new ResourceNotFoundException("policy [{}] not found", name);
             }
 
-            policies.remove(name);
-            return policies;
+            final Map<String, EnrichPolicy> updatedPolicies = new HashMap<>(originalPolicies);
+            updatedPolicies.remove(name);
+            return updatedPolicies;
         });
     }
 
@@ -153,18 +155,11 @@ public static EnrichPolicy getPolicy(String name, ClusterState state) {
      * Gets all policies in the cluster.
      *
      * @param state the cluster state
-     * @return a Map of <code>policyName, EnrichPolicy</code> of the policies
+     * @return a read-only Map of <code>policyName, EnrichPolicy</code> of the policies
      */
     public static Map<String, EnrichPolicy> getPolicies(ClusterState state) {
-        final Map<String, EnrichPolicy> policies;
-        final EnrichMetadata enrichMetadata = state.metadata().custom(EnrichMetadata.TYPE);
-        if (enrichMetadata != null) {
-            // Make a copy, because policies map inside custom metadata is read only:
-            policies = new HashMap<>(enrichMetadata.getPolicies());
-        } else {
-            policies = new HashMap<>();
-        }
-        return policies;
+        final EnrichMetadata metadata = state.metadata().custom(EnrichMetadata.TYPE, EnrichMetadata.EMPTY);
+        return metadata.getPolicies();
     }
 
     private static void updateClusterState(
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/EnrichPolicyResolver.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/EnrichPolicyResolver.java
index c8e993b7dbf0b..cd571ebb676ac 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/EnrichPolicyResolver.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/EnrichPolicyResolver.java
@@ -434,8 +434,8 @@ public void messageReceived(LookupRequest request, TransportChannel channel, Tas
     }
 
     protected Map<String, EnrichPolicy> availablePolicies() {
-        final EnrichMetadata metadata = clusterService.state().metadata().custom(EnrichMetadata.TYPE);
-        return metadata == null ? Map.of() : metadata.getPolicies();
+        final EnrichMetadata metadata = clusterService.state().metadata().custom(EnrichMetadata.TYPE, EnrichMetadata.EMPTY);
+        return metadata.getPolicies();
     }
 
     protected void getRemoteConnection(String cluster, ActionListener<Transport.Connection> listener) {