Enrich store should only update the policies via an update task. (#41944)

Author: martijnvg

x-pack/plugin/enrich/src/main/java/org/elasticsearch/xpack/enrich/EnrichStore.java

@@ -16,6 +16,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 /**
  * Helper methods for access and storage of an enrich policy.
@@ -43,9 +44,11 @@ public static void putPolicy(String name, EnrichPolicy policy, ClusterService cl
         }
         // TODO: add policy validation
 
-        final Map<String, EnrichPolicy> policies = getPolicies(clusterService.state());
-        policies.put(name, policy);
-        updateClusterState(policies, clusterService, handler);
+        updateClusterState(clusterService, handler, current -> {
+            final Map<String, EnrichPolicy> policies = getPolicies(current);
+            policies.put(name, policy);
+            return policies;
+        });
     }
 
     /**
@@ -62,13 +65,15 @@ public static void deletePolicy(String name, ClusterService clusterService, Cons
             throw new IllegalArgumentException("name is missing or empty");
         }
 
-        final Map<String, EnrichPolicy> policies = getPolicies(clusterService.state());
-        if (policies.containsKey(name) == false) {
-            throw new ResourceNotFoundException("policy [{}] not found", name);
-        }
+        updateClusterState(clusterService, handler, current -> {
+            final Map<String, EnrichPolicy> policies = getPolicies(current);
+            if (policies.containsKey(name) == false) {
+                throw new ResourceNotFoundException("policy [{}] not found", name);
+            }
 
-        policies.remove(name);
-        updateClusterState(policies, clusterService, handler);
+            policies.remove(name);
+            return policies;
+        });
     }
 
     /**
@@ -103,12 +108,14 @@ public static Map<String, EnrichPolicy> getPolicies(ClusterState state) {
         return policies;
     }
 
-    private static void updateClusterState(Map<String, EnrichPolicy> policies, ClusterService clusterService,
-                                                 Consumer<Exception> handler) {
-        clusterService.submitStateUpdateTask("update-enrich-policy", new ClusterStateUpdateTask() {
+    private static void updateClusterState(ClusterService clusterService,
+                                           Consumer<Exception> handler,
+                                           Function<ClusterState, Map<String, EnrichPolicy>> function) {
+        clusterService.submitStateUpdateTask("update-enrich-metadata", new ClusterStateUpdateTask() {
 
             @Override
             public ClusterState execute(ClusterState currentState) throws Exception {
+                Map<String, EnrichPolicy> policies = function.apply(currentState);
                 MetaData metaData = MetaData.builder(currentState.metaData())
                     .putCustom(EnrichMetadata.TYPE, new EnrichMetadata(policies))
                     .build();

x-pack/plugin/enrich/src/test/java/org/elasticsearch/xpack/enrich/EnrichStoreTests.java

@@ -36,8 +36,9 @@ public void testCrud() throws Exception {
         assertThat(listPolicies.size(), equalTo(1));
         assertThat(listPolicies.get(name), equalTo(policy));
 
-        error = deleteEnrichPolicy(name, clusterService);
-        assertThat(error.get(), nullValue());
+        deleteEnrichPolicy(name, clusterService);
+        result = EnrichStore.getPolicy(name, clusterService.state());
+        assertThat(result, nullValue());
     }
 
     public void testPutValidation() throws Exception {
@@ -110,14 +111,16 @@ private AtomicReference<Exception> saveEnrichPolicy(String name, EnrichPolicy po
         return error;
     }
 
-    private AtomicReference<Exception> deleteEnrichPolicy(String name, ClusterService clusterService) throws InterruptedException {
+    private void deleteEnrichPolicy(String name, ClusterService clusterService) throws Exception {
         CountDownLatch latch = new CountDownLatch(1);
         AtomicReference<Exception> error = new AtomicReference<>();
         EnrichStore.deletePolicy(name, clusterService, e -> {
             error.set(e);
             latch.countDown();
         });
         latch.await();
-        return error;
+        if (error.get() != null){
+            throw error.get();
+        }
     }
 }