ILM wait for snapshot creation to complete (#54673)

As the snapshot is not taken on the calling thread even though the request
wait_for_completion parameter is true, we can set the flag to true and make
use of the response SnapshotInfo to verify the snapshot was created successfully
or go back to the CleanSnapshotStep and retry creating the snapshot otherwise.
This branching is executed using a new step that wraps an AsyncActionStep
(keeping the wrapped's step key()) and moving to a different next step when
the wrapped action reports an incomplete response.

Co-Authored-By: Lee Hinman <[email protected]>

Author: andreidan

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ilm/AsyncActionBranchingStep.java

@@ -0,0 +1,105 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.xpack.core.ilm;
+
+import org.apache.lucene.util.SetOnce;
+import org.elasticsearch.client.Client;
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.ClusterStateObserver;
+import org.elasticsearch.cluster.metadata.IndexMetadata;
+
+import java.util.Objects;
+
+/**
+ * This step wraps an {@link AsyncActionStep} in order to be able to manipulate what the next step will be, depending on the result of the
+ * wrapped {@link AsyncActionStep}.
+ * <p>
+ * If the action response is complete, the {@link AsyncActionBranchingStep}'s nextStepKey will be the nextStepKey of the wrapped action. If
+ * the response is incomplete the nextStepKey will be the provided {@link AsyncActionBranchingStep#nextKeyOnIncompleteResponse}.
+ * Failures encountered whilst executing the wrapped action will be propagated directly.
+ */
+public class AsyncActionBranchingStep extends AsyncActionStep {
+    private final AsyncActionStep stepToExecute;
+
+    private StepKey nextKeyOnIncompleteResponse;
+    private SetOnce<Boolean> onResponseResult;
+
+    public AsyncActionBranchingStep(AsyncActionStep stepToExecute, StepKey nextKeyOnIncompleteResponse, Client client) {
+        // super.nextStepKey is set to null since it is not used by this step
+        super(stepToExecute.getKey(), null, client);
+        this.stepToExecute = stepToExecute;
+        this.nextKeyOnIncompleteResponse = nextKeyOnIncompleteResponse;
+        this.onResponseResult = new SetOnce<>();
+    }
+
+    @Override
+    public boolean isRetryable() {
+        return true;
+    }
+
+    @Override
+    public void performAction(IndexMetadata indexMetadata, ClusterState currentClusterState, ClusterStateObserver observer,
+                              Listener listener) {
+        stepToExecute.performAction(indexMetadata, currentClusterState, observer, new Listener() {
+            @Override
+            public void onResponse(boolean complete) {
+                onResponseResult.set(complete);
+                listener.onResponse(complete);
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                listener.onFailure(e);
+            }
+        });
+    }
+
+    @Override
+    public final StepKey getNextStepKey() {
+        if (onResponseResult.get() == null) {
+            throw new IllegalStateException("cannot call getNextStepKey before performAction");
+        }
+        return onResponseResult.get() ? stepToExecute.getNextStepKey() : nextKeyOnIncompleteResponse;
+    }
+
+    /**
+     * Represents the {@link AsyncActionStep} that's wrapped by this branching step.
+     */
+    AsyncActionStep getStepToExecute() {
+        return stepToExecute;
+    }
+
+    /**
+     * The step key to be reported as the {@link AsyncActionBranchingStep#getNextStepKey()} if the response of the wrapped
+     * {@link AsyncActionBranchingStep#getStepToExecute()} is incomplete.
+     */
+    StepKey getNextKeyOnIncompleteResponse() {
+        return nextKeyOnIncompleteResponse;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        if (!super.equals(o)) {
+            return false;
+        }
+        AsyncActionBranchingStep that = (AsyncActionBranchingStep) o;
+        return super.equals(o)
+            && Objects.equals(stepToExecute, that.stepToExecute)
+            && Objects.equals(nextKeyOnIncompleteResponse, that.nextKeyOnIncompleteResponse);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), stepToExecute, nextKeyOnIncompleteResponse);
+    }
+}

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ilm/CreateSnapshotStep.java

@@ -5,14 +5,15 @@
  */
 package org.elasticsearch.xpack.core.ilm;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.admin.cluster.snapshots.create.CreateSnapshotRequest;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.rest.RestStatus;
-import org.elasticsearch.snapshots.SnapshotException;
+import org.elasticsearch.snapshots.SnapshotInfo;
 
 import static org.elasticsearch.xpack.core.ilm.LifecycleExecutionState.fromIndexMetadata;
 
@@ -23,6 +24,8 @@
 public class CreateSnapshotStep extends AsyncRetryDuringSnapshotActionStep {
     public static final String NAME = "create-snapshot";
 
+    private static final Logger logger = LogManager.getLogger(CreateSnapshotStep.class);
+
     public CreateSnapshotStep(StepKey key, StepKey nextStepKey, Client client) {
         super(key, nextStepKey, client);
     }
@@ -54,19 +57,26 @@ void performDuringNoSnapshot(IndexMetadata indexMetadata, ClusterState currentCl
         }
         CreateSnapshotRequest request = new CreateSnapshotRequest(snapshotRepository, snapshotName);
         request.indices(indexName);
-        // we'll not wait for the snapshot to complete in this step as the async steps are executed from threads that shouldn't perform
-        // expensive operations (ie. clusterStateProcessed)
-        request.waitForCompletion(false);
+        // this is safe as the snapshot creation will still be async, it's just that the listener will be notified when the snapshot is
+        // complete
+        request.waitForCompletion(true);
         request.includeGlobalState(false);
         request.masterNodeTimeout(getMasterTimeout(currentClusterState));
         getClient().admin().cluster().createSnapshot(request,
             ActionListener.wrap(response -> {
-                if (response.status().equals(RestStatus.INTERNAL_SERVER_ERROR)) {
-                    listener.onFailure(new SnapshotException(snapshotRepository, snapshotName,
-                        "unable to request snapshot creation [" + snapshotName + "] for index [ " + indexName + "] as part of policy [" +
-                            policyName + "] execution due to an internal server error"));
-                } else {
+                logger.debug("create snapshot response for policy [{}] and index [{}] is: {}", policyName, indexName,
+                    Strings.toString(response));
+                final SnapshotInfo snapInfo = response.getSnapshotInfo();
+
+                // Check that there are no failed shards, since the request may not entirely
+                // fail, but may still have failures (such as in the case of an aborted snapshot)
+                if (snapInfo.failedShards() == 0) {
                     listener.onResponse(true);
+                } else {
+                    int failures = snapInfo.failedShards();
+                    int total = snapInfo.totalShards();
+                    logger.warn("failed to create snapshot successfully, {} failures  out of {} total shards failed", failures, total);
+                    listener.onResponse(false);
                 }
             }, listener::onFailure));
     }