From 5ea27d6e0b2fd44f2455038f3415b18a2a86a695 Mon Sep 17 00:00:00 2001
From: Gordon Brown <gordon.brown@elastic.co>
Date: Thu, 25 Feb 2021 13:47:58 -0700
Subject: [PATCH] Implement Transform feature state reset

---
 .../xpack/transform/Transform.java            | 32 +++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/Transform.java b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/Transform.java
index 3c36c03186518..772fd52ffcc90 100644
--- a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/Transform.java
+++ b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/Transform.java
@@ -10,11 +10,14 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.util.SetOnce;
+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.client.Client;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.metadata.IndexTemplateMetadata;
+import org.elasticsearch.cluster.metadata.Metadata;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodeRole;
 import org.elasticsearch.cluster.node.DiscoveryNodes;
@@ -371,6 +374,35 @@ public Collection<SystemIndexDescriptor> getSystemIndexDescriptors(Settings sett
         }
     }
 
+    @Override
+    public void cleanUpFeature(
+        ClusterService clusterService,
+        Client client,
+        ActionListener<ResetFeatureStateResponse.ResetFeatureStateStatus> listener
+    ) {
+        ActionListener<StopTransformAction.Response> afterStoppingTransforms = ActionListener.wrap(stopTransformsResponse -> {
+            if (stopTransformsResponse.isAcknowledged()
+                && stopTransformsResponse.getTaskFailures().isEmpty()
+                && stopTransformsResponse.getNodeFailures().isEmpty()) {
+
+                SystemIndexPlugin.super.cleanUpFeature(clusterService, client, listener);
+            } else {
+                String errMsg = "Failed to reset Transform: "
+                    + (stopTransformsResponse.isAcknowledged() ? "" : "not acknowledged ")
+                    + (stopTransformsResponse.getNodeFailures().isEmpty()
+                        ? ""
+                        : "node failures: " + stopTransformsResponse.getNodeFailures() + " ")
+                    + (stopTransformsResponse.getTaskFailures().isEmpty()
+                        ? ""
+                        : "task failures: " + stopTransformsResponse.getTaskFailures());
+                listener.onResponse(new ResetFeatureStateResponse.ResetFeatureStateStatus(this.getFeatureName(), errMsg));
+            }
+        }, listener::onFailure);
+
+        StopTransformAction.Request stopTransformsRequest = new StopTransformAction.Request(Metadata.ALL, true, true, null, true, false);
+        client.execute(StopTransformAction.INSTANCE, stopTransformsRequest, afterStoppingTransforms);
+    }
+
     @Override
     public String getFeatureName() {
         return "transform";