From a09f10b9755721f90fa10a787c6ff2ec6ce5d0f5 Mon Sep 17 00:00:00 2001
From: Hendrik Muhs <hendrik.muhs@elastic.co>
Date: Thu, 14 Oct 2021 16:21:13 +0200
Subject: [PATCH] fix feature reset call lost for ESSingleNodeTestCase (#79150)

call feature reset API for all plugins inside LocalStateCompositeXPackPlugin
---
 .../core/LocalStateCompositeXPackPlugin.java  | 26 ++++++++++++++++++-
 .../xpack/ml/LocalStateMachineLearning.java   | 11 --------
 2 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/LocalStateCompositeXPackPlugin.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/LocalStateCompositeXPackPlugin.java
index 3a20c484e6457..24f6a9d39e5a3 100644
--- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/LocalStateCompositeXPackPlugin.java
+++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/LocalStateCompositeXPackPlugin.java
@@ -6,12 +6,16 @@
  */
 package org.elasticsearch.xpack.core;
 
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRequest;
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.action.RequestValidators;
+import org.elasticsearch.action.admin.cluster.snapshots.features.ResetFeatureStateResponse;
+import org.elasticsearch.action.admin.cluster.snapshots.features.ResetFeatureStateResponse.ResetFeatureStateStatus;
 import org.elasticsearch.action.admin.indices.alias.IndicesAliasesRequest;
 import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequest;
 import org.elasticsearch.action.support.ActionFilter;
+import org.elasticsearch.action.support.GroupedActionListener;
 import org.elasticsearch.bootstrap.BootstrapCheck;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.cluster.ClusterState;
@@ -37,7 +41,6 @@
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.PageCacheRecycler;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
-import org.elasticsearch.xcontent.NamedXContentRegistry;
 import org.elasticsearch.core.internal.io.IOUtils;
 import org.elasticsearch.env.Environment;
 import org.elasticsearch.env.NodeEnvironment;
@@ -85,6 +88,7 @@
 import org.elasticsearch.transport.Transport;
 import org.elasticsearch.transport.TransportInterceptor;
 import org.elasticsearch.watcher.ResourceWatcherService;
+import org.elasticsearch.xcontent.NamedXContentRegistry;
 import org.elasticsearch.xpack.core.ssl.SSLService;
 
 import java.io.IOException;
@@ -591,10 +595,30 @@ public CheckedBiConsumer<ShardSearchRequest, StreamOutput, IOException> getReque
 
     }
 
+    @Override
+    public void cleanUpFeature(
+        ClusterService clusterService,
+        Client client,
+        ActionListener<ResetFeatureStateResponse.ResetFeatureStateStatus> finalListener
+    ) {
+        List<SystemIndexPlugin> systemPlugins = filterPlugins(SystemIndexPlugin.class);
+
+        GroupedActionListener<ResetFeatureStateResponse.ResetFeatureStateStatus> allListeners = new GroupedActionListener<>(
+            ActionListener.wrap(
+                listenerResults -> finalListener.onResponse(ResetFeatureStateStatus.success(getFeatureName())),
+                finalListener::onFailure
+            ),
+            systemPlugins.size()
+        );
+        systemPlugins.forEach(plugin -> plugin.cleanUpFeature(clusterService, client, allListeners));
+    }
+
+    @Override
     public boolean safeToShutdown(String nodeId, SingleNodeShutdownMetadata.Type shutdownType) {
         return filterPlugins(ShutdownAwarePlugin.class).stream().allMatch(plugin -> plugin.safeToShutdown(nodeId, shutdownType));
     }
 
+    @Override
     public void signalShutdown(Collection<String> shutdownNodeIds) {
         filterPlugins(ShutdownAwarePlugin.class).forEach(plugin -> plugin.signalShutdown(shutdownNodeIds));
     }
diff --git a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/LocalStateMachineLearning.java b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/LocalStateMachineLearning.java
index d6f147f831df7..aba1d90ca18ba 100644
--- a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/LocalStateMachineLearning.java
+++ b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/LocalStateMachineLearning.java
@@ -9,11 +9,8 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRequest;
 import org.elasticsearch.action.ActionResponse;
-import org.elasticsearch.action.admin.cluster.snapshots.features.ResetFeatureStateResponse;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.TransportAction;
-import org.elasticsearch.client.Client;
-import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.breaker.NoopCircuitBreaker;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.settings.Settings;
@@ -80,14 +77,6 @@ protected XPackLicenseState getLicenseState() {
         plugins.add(new MockedRollupPlugin());
     }
 
-    @Override
-    public void cleanUpFeature(
-        ClusterService clusterService,
-        Client client,
-        ActionListener<ResetFeatureStateResponse.ResetFeatureStateStatus> finalListener) {
-        mlPlugin.cleanUpFeature(clusterService, client, finalListener);
-    }
-
     @Override
     public List<AggregationSpec> getAggregations() {
         return mlPlugin.getAggregations();