And a callback-accepting simple batching executor

joegallo · joegallo · commit ee83da0e41b8 · 2022-02-03T10:07:03.000-05:00
diff --git a/server/src/main/java/org/elasticsearch/cluster/ClusterStateTaskExecutor.java b/server/src/main/java/org/elasticsearch/cluster/ClusterStateTaskExecutor.java
@@ -13,6 +13,7 @@
 import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.BiFunction;
 
 public interface ClusterStateTaskExecutor<T extends ClusterStateTaskListener> {
     /**
@@ -151,18 +152,25 @@ public String describeTasks(List<T> tasks) {
     }
 
     static <T extends ClusterStateUpdateTask> ClusterStateTaskExecutor<T> simpleBatched() {
-        return (currentState, tasks) -> {
+        return simpleBatched((originalState, updatedState) -> updatedState);
+    }
+
+    static <T extends ClusterStateUpdateTask> ClusterStateTaskExecutor<T> simpleBatched(
+        BiFunction<ClusterState, ClusterState, ClusterState> clusterStateCallback
+    ) {
+        return (originalState, tasks) -> {
             ClusterTasksResult.Builder<T> builder = ClusterTasksResult.builder();
-            ClusterState state = currentState;
+            ClusterState updatedState = originalState;
             for (T task : tasks) {
                 try {
-                    state = task.execute(state);
+                    updatedState = task.execute(updatedState);
                     builder.success(task);
                 } catch (Exception e) {
                     builder.failure(task, e);
                 }
             }
-            return builder.build(state);
+            updatedState = clusterStateCallback.apply(originalState, updatedState);
+            return builder.build(updatedState);
         };
     }
 }
diff --git a/server/src/main/java/org/elasticsearch/cluster/metadata/MetadataUpdateSettingsService.java b/server/src/main/java/org/elasticsearch/cluster/metadata/MetadataUpdateSettingsService.java
@@ -17,7 +17,6 @@
 import org.elasticsearch.cluster.AckedClusterStateUpdateTask;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.ClusterStateTaskExecutor;
-import org.elasticsearch.cluster.ClusterStateTaskExecutor.ClusterTasksResult;
 import org.elasticsearch.cluster.block.ClusterBlock;
 import org.elasticsearch.cluster.block.ClusterBlocks;
 import org.elasticsearch.cluster.routing.RoutingTable;
@@ -72,23 +71,13 @@ public MetadataUpdateSettingsService(
         this.indicesService = indicesService;
         this.shardLimitValidator = shardLimitValidator;
         this.threadPool = threadPool;
-        this.executor = (currentState, tasks) -> {
-            ClusterTasksResult.Builder<AckedClusterStateUpdateTask> builder = ClusterTasksResult.builder();
-            ClusterState state = currentState;
-            for (AckedClusterStateUpdateTask task : tasks) {
-                try {
-                    state = task.execute(state);
-                    builder.success(task);
-                } catch (Exception e) {
-                    builder.failure(task, e);
-                }
-            }
-            if (state != currentState) {
+        this.executor = ClusterStateTaskExecutor.simpleBatched((originalState, updatedState) -> {
+            if (updatedState != originalState) {
                 // reroute in case things change that require it (like number of replicas)
-                state = allocationService.reroute(state, "settings update");
+                updatedState = allocationService.reroute(updatedState, "settings update");
             }
-            return builder.build(state);
-        };
+            return updatedState;
+        });
     }
 
     public void updateSettings(final UpdateSettingsClusterStateUpdateRequest request, final ActionListener<AcknowledgedResponse> listener) {
