[Transform] Check node shutdown before fail (#107358)

Transforms continue to run even when a node is shutting down. This may
lead to a transform failing and putting itself into a failed state,
which will prevent it from restarting when the node comes back online.

The transform will now abort rather than fail, which puts itself into a
started state. When the node comes back online, or another node in the
cluster starts the transform, then the transform will pick up from its
last successful saved state and checkpoint.

Close #100891

docs/changelog/107358.yaml

@@ -0,0 +1,6 @@
+pr: 107358
+summary: Check node shutdown before fail
+area: Transform
+type: enhancement
+issues:
+ - 100891

x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/Transform.java

@@ -255,14 +255,12 @@ public Collection<?> createComponents(PluginServices services) {
             getTransformExtension().getMinFrequency()
         );
         scheduler.start();
+        var clusterStateListener = new TransformClusterStateListener(clusterService, client);
+        var transformNode = new TransformNode(clusterStateListener);
 
-        transformServices.set(new TransformServices(configManager, checkpointService, auditor, scheduler));
+        transformServices.set(new TransformServices(configManager, checkpointService, auditor, scheduler, transformNode));
 
-        return List.of(
-            transformServices.get(),
-            new TransformClusterStateListener(clusterService, client),
-            new TransformExtensionHolder(getTransformExtension())
-        );
+        return List.of(transformServices.get(), clusterStateListener, new TransformExtensionHolder(getTransformExtension()));
     }
 
     @Override

x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/TransformClusterStateListener.java

@@ -21,17 +21,21 @@
 import org.elasticsearch.gateway.GatewayService;
 import org.elasticsearch.xpack.core.transform.transforms.persistence.TransformInternalIndexConstants;
 
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 
 import static org.elasticsearch.xpack.core.ClientHelper.TRANSFORM_ORIGIN;
 import static org.elasticsearch.xpack.core.ClientHelper.executeAsyncWithOrigin;
 
-class TransformClusterStateListener implements ClusterStateListener {
+class TransformClusterStateListener implements ClusterStateListener, Supplier<Optional<ClusterState>> {
 
     private static final Logger logger = LogManager.getLogger(TransformClusterStateListener.class);
 
     private final Client client;
     private final AtomicBoolean isIndexCreationInProgress = new AtomicBoolean(false);
+    private final AtomicReference<ClusterState> clusterState = new AtomicReference<>();
 
     TransformClusterStateListener(ClusterService clusterService, Client client) {
         this.client = client;
@@ -46,6 +50,8 @@ public void clusterChanged(ClusterChangedEvent event) {
             return;
         }
 
+        clusterState.set(event.state());
+
         // The atomic flag prevents multiple simultaneous attempts to run alias creation
         // if there is a flurry of cluster state updates in quick succession
         if (event.localNodeMaster() && isIndexCreationInProgress.compareAndSet(false, true)) {
@@ -102,4 +108,12 @@ private static void createAuditAliasForDataFrameBWC(ClusterState state, Client c
         );
     }
 
+    /**
+     * Retrieves the saved cluster state from the most recent update.
+     * This differs from {@link ClusterService#state()} in that it will not throw an exception when ClusterState is null.
+     */
+    @Override
+    public Optional<ClusterState> get() {
+        return Optional.ofNullable(clusterState.get());
+    }
 }

x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/TransformNode.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.transform;
+
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.node.DiscoveryNodes;
+
+import java.util.Optional;
+import java.util.function.Supplier;
+
+/**
+ * Stateful representation of this node, relevant to the {@link org.elasticsearch.xpack.transform.transforms.TransformTask}.
+ * For stateless functions, see {@link org.elasticsearch.xpack.transform.transforms.TransformNodes}.
+ */
+public class TransformNode {
+    private final Supplier<Optional<ClusterState>> clusterState;
+
+    public TransformNode(Supplier<Optional<ClusterState>> clusterState) {
+        this.clusterState = clusterState;
+    }
+
+    /**
+     * @return an optional containing true if this node is reported as shutting down in the cluster state metadata, false if it is not
+     * reported as shutting down, or empty if the cluster state is missing or the local node has not been set yet.
+     */
+    public Optional<Boolean> isShuttingDown() {
+        return clusterState.get().map(state -> {
+            var localId = state.nodes().getLocalNodeId();
+            if (localId != null) {
+                return state.metadata().nodeShutdowns().contains(localId);
+            } else {
+                return null; // empty
+            }
+        });
+    }
+
+    /**
+     * @return the node id stored in the cluster state, or "null" if the cluster state is missing or the local node has not been set yet.
+     * This should behave exactly as {@link String#valueOf(Object)}.
+     */
+    public String nodeId() {
+        return clusterState.get().map(ClusterState::nodes).map(DiscoveryNodes::getLocalNodeId).orElse("null");
+    }
+}

x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/TransformServices.java

@@ -26,17 +26,20 @@ public final class TransformServices {
     private final TransformCheckpointService checkpointService;
     private final TransformAuditor auditor;
     private final TransformScheduler scheduler;
+    private final TransformNode transformNode;
 
     public TransformServices(
         TransformConfigManager configManager,
         TransformCheckpointService checkpointService,
         TransformAuditor auditor,
-        TransformScheduler scheduler
+        TransformScheduler scheduler,
+        TransformNode transformNode
     ) {
         this.configManager = Objects.requireNonNull(configManager);
         this.checkpointService = Objects.requireNonNull(checkpointService);
         this.auditor = Objects.requireNonNull(auditor);
         this.scheduler = Objects.requireNonNull(scheduler);
+        this.transformNode = transformNode;
     }
 
     public TransformConfigManager getConfigManager() {
@@ -54,4 +57,8 @@ public TransformAuditor getAuditor() {
     public TransformScheduler getScheduler() {
         return scheduler;
     }
+
+    public TransformNode getTransformNode() {
+        return transformNode;
+    }
 }

x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/TransformPersistentTasksExecutor.java

@@ -497,7 +497,8 @@ protected AllocatedPersistentTask createTask(
             transformServices.getScheduler(),
             auditor,
             threadPool,
-            headers
+            headers,
+            transformServices.getTransformNode()
         );
     }
 }

x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/TransformTask.java

@@ -40,6 +40,7 @@
 import org.elasticsearch.xpack.core.transform.transforms.TransformState;
 import org.elasticsearch.xpack.core.transform.transforms.TransformTaskParams;
 import org.elasticsearch.xpack.core.transform.transforms.TransformTaskState;
+import org.elasticsearch.xpack.transform.TransformNode;
 import org.elasticsearch.xpack.transform.checkpoint.TransformCheckpointService;
 import org.elasticsearch.xpack.transform.notifications.TransformAuditor;
 import org.elasticsearch.xpack.transform.transforms.scheduling.TransformScheduler;
@@ -68,6 +69,7 @@ public class TransformTask extends AllocatedPersistentTask implements TransformS
     private final TransformIndexerPosition initialPosition;
     private final IndexerState initialIndexerState;
     private final TransformContext context;
+    private final TransformNode transformNode;
     private final SetOnce<ClientTransformIndexer> indexer = new SetOnce<>();
 
     @SuppressWarnings("this-escape")
@@ -81,7 +83,8 @@ public class TransformTask extends AllocatedPersistentTask implements TransformS
         TransformScheduler transformScheduler,
         TransformAuditor auditor,
         ThreadPool threadPool,
-        Map<String, String> headers
+        Map<String, String> headers,
+        TransformNode transformNode
     ) {
         super(id, type, action, TransformField.PERSISTENT_TASK_DESCRIPTION_PREFIX + transform.getId(), parentTask, headers);
         this.transform = transform;
@@ -118,6 +121,7 @@ public class TransformTask extends AllocatedPersistentTask implements TransformS
         if (state != null) {
             this.context.setAuthState(state.getAuthState());
         }
+        this.transformNode = transformNode;
     }
 
     public String getTransformId() {
@@ -524,11 +528,26 @@ public void fail(Throwable exception, String reason, ActionListener<Void> listen
                 return;
             }
 
-            logger.atError().withThrowable(exception).log("[{}] transform has failed; experienced: [{}].", transform.getId(), reason);
-            auditor.error(transform.getId(), reason);
             // We should not keep retrying. Either the task will be stopped, or started
             // If it is started again, it is registered again.
             transformScheduler.deregisterTransform(getTransformId());
+
+            if (transformNode.isShuttingDown().orElse(false)) {
+                logger.atDebug()
+                    .withThrowable(exception)
+                    .log(
+                        "Aborting transform [{}]. Transform has failed while node [{}] is shutting down. Reason: [{}]",
+                        transform.getId(),
+                        transformNode.nodeId(),
+                        reason
+                    );
+                markAsLocallyAborted("Node is shutting down.");
+                listener.onResponse(null);
+                return;
+            }
+
+            logger.atError().withThrowable(exception).log("[{}] transform has failed; experienced: [{}].", transform.getId(), reason);
+            auditor.error(transform.getId(), reason);
             // The idea of stopping at the next checkpoint is no longer valid. Since a failed task could potentially START again,
             // we should set this flag to false.
             context.setShouldStopAtCheckpoint(false);

x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/TransformNodeTests.java

@@ -0,0 +1,112 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.transform;
+
+import org.elasticsearch.cluster.ClusterName;
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.metadata.Metadata;
+import org.elasticsearch.cluster.metadata.NodesShutdownMetadata;
+import org.elasticsearch.cluster.metadata.SingleNodeShutdownMetadata;
+import org.elasticsearch.cluster.node.DiscoveryNodeUtils;
+import org.elasticsearch.cluster.node.DiscoveryNodes;
+import org.elasticsearch.test.ESTestCase;
+
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.Supplier;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class TransformNodeTests extends ESTestCase {
+    private static final String SHUTTING_DOWN_ID = "shuttingDownNodeId";
+    private static final String NOT_SHUTTING_DOWN_ID = "notShuttingDownId";
+
+    /**
+     * When the local node is shutting down
+     * Then we return true
+     */
+    public void testIsShuttingDown() {
+        var isShuttingDown = new TransformNode(clusterState(SHUTTING_DOWN_ID)).isShuttingDown();
+        assertTrue(isShuttingDown.isPresent());
+        assertTrue(isShuttingDown.get());
+    }
+
+    /**
+     * When the local node is not shutting down
+     * Then we return false
+     */
+    public void testIsNotShuttingDown() {
+        var isShuttingDown = new TransformNode(clusterState(NOT_SHUTTING_DOWN_ID)).isShuttingDown();
+        assertTrue(isShuttingDown.isPresent());
+        assertFalse(isShuttingDown.get());
+    }
+
+    /**
+     * When the local node is null
+     * Then we return empty
+     */
+    public void testMissingLocalId() {
+        var isShuttingDown = new TransformNode(clusterState(null)).isShuttingDown();
+        assertFalse(isShuttingDown.isPresent());
+    }
+
+    /**
+     * When the cluster state is empty
+     * Then we return empty
+     */
+    public void testClusterStateMissing() {
+        var isShuttingDown = new TransformNode(Optional::empty).isShuttingDown();
+        assertFalse(isShuttingDown.isPresent());
+    }
+
+    /**
+     * When there is a local node
+     * Then return its id
+     */
+    public void testNodeId() {
+        var nodeId = new TransformNode(clusterState(SHUTTING_DOWN_ID)).nodeId();
+        assertThat(nodeId, equalTo(SHUTTING_DOWN_ID));
+    }
+
+    /**
+     * When the local node is null
+     * Then return "null"
+     */
+    public void testNodeIdMissing() {
+        var nodeId = new TransformNode(Optional::empty).nodeId();
+        assertThat(nodeId, equalTo(String.valueOf((String) null)));
+    }
+
+    private Supplier<Optional<ClusterState>> clusterState(String nodeId) {
+        var nodesShutdownMetadata = new NodesShutdownMetadata(
+            Map.of(
+                SHUTTING_DOWN_ID,
+                SingleNodeShutdownMetadata.builder()
+                    .setNodeId(SHUTTING_DOWN_ID)
+                    .setReason("shutdown for a unit test")
+                    .setType(SingleNodeShutdownMetadata.Type.RESTART)
+                    .setStartedAtMillis(randomNonNegativeLong())
+                    .setGracePeriod(null)
+                    .build()
+            )
+        );
+
+        var nodes = DiscoveryNodes.builder().add(DiscoveryNodeUtils.create(SHUTTING_DOWN_ID)).localNodeId(nodeId).masterNodeId(nodeId);
+
+        if (SHUTTING_DOWN_ID.equals(nodeId) == false && nodeId != null) {
+            nodes.add(DiscoveryNodeUtils.create(nodeId));
+        }
+
+        var state = ClusterState.builder(ClusterName.DEFAULT)
+            .metadata(Metadata.builder().putCustom(NodesShutdownMetadata.TYPE, nodesShutdownMetadata).build())
+            .nodes(nodes)
+            .build();
+
+        return () -> Optional.of(state);
+    }
+}

x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/ClientTransformIndexerTests.java

@@ -53,6 +53,7 @@
 import org.elasticsearch.xpack.core.transform.transforms.TransformProgress;
 import org.elasticsearch.xpack.core.transform.transforms.persistence.TransformInternalIndexConstants;
 import org.elasticsearch.xpack.transform.TransformExtension;
+import org.elasticsearch.xpack.transform.TransformNode;
 import org.elasticsearch.xpack.transform.TransformServices;
 import org.elasticsearch.xpack.transform.checkpoint.CheckpointProvider;
 import org.elasticsearch.xpack.transform.checkpoint.TransformCheckpointService;
@@ -140,7 +141,8 @@ public void testPitInjection() throws InterruptedException {
                     mock(IndexBasedTransformConfigManager.class),
                     mock(TransformCheckpointService.class),
                     mock(TransformAuditor.class),
-                    new TransformScheduler(Clock.systemUTC(), mock(ThreadPool.class), Settings.EMPTY, TimeValue.ZERO)
+                    new TransformScheduler(Clock.systemUTC(), mock(ThreadPool.class), Settings.EMPTY, TimeValue.ZERO),
+                    mock(TransformNode.class)
                 ),
                 mock(CheckpointProvider.class),
                 new AtomicReference<>(IndexerState.STOPPED),
@@ -237,7 +239,8 @@ public void testPitInjectionIfPitNotSupported() throws InterruptedException {
                     mock(IndexBasedTransformConfigManager.class),
                     mock(TransformCheckpointService.class),
                     mock(TransformAuditor.class),
-                    new TransformScheduler(Clock.systemUTC(), mock(ThreadPool.class), Settings.EMPTY, TimeValue.ZERO)
+                    new TransformScheduler(Clock.systemUTC(), mock(ThreadPool.class), Settings.EMPTY, TimeValue.ZERO),
+                    mock(TransformNode.class)
                 ),
                 mock(CheckpointProvider.class),
                 new AtomicReference<>(IndexerState.STOPPED),
@@ -323,7 +326,8 @@ public void testDisablePit() throws InterruptedException {
                     mock(IndexBasedTransformConfigManager.class),
                     mock(TransformCheckpointService.class),
                     mock(TransformAuditor.class),
-                    new TransformScheduler(Clock.systemUTC(), mock(ThreadPool.class), Settings.EMPTY, TimeValue.ZERO)
+                    new TransformScheduler(Clock.systemUTC(), mock(ThreadPool.class), Settings.EMPTY, TimeValue.ZERO),
+                    mock(TransformNode.class)
                 ),
                 mock(CheckpointProvider.class),
                 new AtomicReference<>(IndexerState.STOPPED),
@@ -572,7 +576,8 @@ private ClientTransformIndexer createTestIndexer(ParentTaskAssigningClient clien
                 mock(IndexBasedTransformConfigManager.class),
                 mock(TransformCheckpointService.class),
                 mock(TransformAuditor.class),
-                new TransformScheduler(Clock.systemUTC(), mock(ThreadPool.class), Settings.EMPTY, TimeValue.ZERO)
+                new TransformScheduler(Clock.systemUTC(), mock(ThreadPool.class), Settings.EMPTY, TimeValue.ZERO),
+                mock(TransformNode.class)
             ),
             mock(CheckpointProvider.class),
             new AtomicReference<>(IndexerState.STOPPED),