Fix NOT_STARTED statuses appearing inappropirately during node shutdown

server/src/main/java/org/elasticsearch/cluster/metadata/SingleNodeShutdownMetadata.java

@@ -35,14 +35,16 @@ public class SingleNodeShutdownMetadata extends AbstractDiffable<SingleNodeShutd
     public static final ParseField REASON_FIELD = new ParseField("reason");
     public static final String STARTED_AT_READABLE_FIELD = "shutdown_started";
     public static final ParseField STARTED_AT_MILLIS_FIELD = new ParseField(STARTED_AT_READABLE_FIELD + "millis");
+    public static final ParseField NODE_SEEN_FIELD = new ParseField("node_seen");
 
     public static final ConstructingObjectParser<SingleNodeShutdownMetadata, Void> PARSER = new ConstructingObjectParser<>(
         "node_shutdown_info",
         a -> new SingleNodeShutdownMetadata(
             (String) a[0],
             Type.valueOf((String) a[1]),
             (String) a[2],
-            (long) a[3]
+            (long) a[3],
+            (boolean) a[4]
         )
     );
 
@@ -51,6 +53,7 @@ public class SingleNodeShutdownMetadata extends AbstractDiffable<SingleNodeShutd
         PARSER.declareString(ConstructingObjectParser.constructorArg(), TYPE_FIELD);
         PARSER.declareString(ConstructingObjectParser.constructorArg(), REASON_FIELD);
         PARSER.declareLong(ConstructingObjectParser.constructorArg(), STARTED_AT_MILLIS_FIELD);
+        PARSER.declareBoolean(ConstructingObjectParser.constructorArg(), NODE_SEEN_FIELD);
     }
 
     public static SingleNodeShutdownMetadata parse(XContentParser parser) {
@@ -61,6 +64,7 @@ public static SingleNodeShutdownMetadata parse(XContentParser parser) {
     private final Type type;
     private final String reason;
     private final long startedAtMillis;
+    private final boolean nodeSeen;
 
     /**
      * @param nodeId The node ID that this shutdown metadata refers to.
@@ -72,19 +76,22 @@ private SingleNodeShutdownMetadata(
         String nodeId,
         Type type,
         String reason,
-        long startedAtMillis
+        long startedAtMillis,
+        boolean nodeSeen
     ) {
         this.nodeId = Objects.requireNonNull(nodeId, "node ID must not be null");
         this.type = Objects.requireNonNull(type, "shutdown type must not be null");
         this.reason = Objects.requireNonNull(reason, "shutdown reason must not be null");
         this.startedAtMillis = startedAtMillis;
+        this.nodeSeen = nodeSeen;
     }
 
     public SingleNodeShutdownMetadata(StreamInput in) throws IOException {
         this.nodeId = in.readString();
         this.type = in.readEnum(Type.class);
         this.reason = in.readString();
         this.startedAtMillis = in.readVLong();
+        this.nodeSeen = in.readBoolean();
     }
 
     /**
@@ -115,12 +122,20 @@ public long getStartedAtMillis() {
         return startedAtMillis;
     }
 
+    /**
+     * @return A boolean indicated whether this node has been seen in the cluster since the shutdown was registered.
+     */
+    public boolean getNodeSeen() {
+        return nodeSeen;
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(nodeId);
         out.writeEnum(type);
         out.writeString(reason);
         out.writeVLong(startedAtMillis);
+        out.writeBoolean(nodeSeen);
     }
 
     @Override
@@ -131,6 +146,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field(TYPE_FIELD.getPreferredName(), type);
             builder.field(REASON_FIELD.getPreferredName(), reason);
             builder.timeField(STARTED_AT_MILLIS_FIELD.getPreferredName(), STARTED_AT_READABLE_FIELD, startedAtMillis);
+            builder.field(NODE_SEEN_FIELD.getPreferredName(), nodeSeen);
         }
         builder.endObject();
 
@@ -145,7 +161,8 @@ public boolean equals(Object o) {
         return getStartedAtMillis() == that.getStartedAtMillis()
             && getNodeId().equals(that.getNodeId())
             && getType() == that.getType()
-            && getReason().equals(that.getReason());
+            && getReason().equals(that.getReason())
+            && getNodeSeen() == that.getNodeSeen();
     }
 
     @Override
@@ -154,7 +171,8 @@ public int hashCode() {
             getNodeId(),
             getType(),
             getReason(),
-            getStartedAtMillis()
+            getStartedAtMillis(),
+            getNodeSeen()
         );
     }
 
@@ -170,14 +188,16 @@ public static Builder builder(SingleNodeShutdownMetadata original) {
             .setNodeId(original.getNodeId())
             .setType(original.getType())
             .setReason(original.getReason())
-            .setStartedAtMillis(original.getStartedAtMillis());
+            .setStartedAtMillis(original.getStartedAtMillis())
+            .setNodeSeen(original.getNodeSeen());
     }
 
     public static class Builder {
         private String nodeId;
         private Type type;
         private String reason;
         private long startedAtMillis = -1;
+        private boolean nodeSeen = false;
 
         private Builder() {}
 
@@ -217,6 +237,15 @@ public Builder setStartedAtMillis(long startedAtMillis) {
             return this;
         }
 
+        /**
+         * @param nodeSeen Whether or not the node has been seen since the shutdown was registered.
+         * @return This builder.
+         */
+        public Builder setNodeSeen(boolean nodeSeen) {
+            this.nodeSeen = nodeSeen;
+            return this;
+        }
+
         public SingleNodeShutdownMetadata build() {
             if (startedAtMillis == -1) {
                 throw new IllegalArgumentException("start timestamp must be set");
@@ -225,7 +254,8 @@ public SingleNodeShutdownMetadata build() {
                 nodeId,
                 type,
                 reason,
-                startedAtMillis
+                startedAtMillis,
+                nodeSeen
             );
         }
     }

server/src/test/java/org/elasticsearch/cluster/metadata/NodesShutdownMetadataTests.java

@@ -82,6 +82,7 @@ private SingleNodeShutdownMetadata randomNodeShutdownInfo() {
             .setType(randomBoolean() ? SingleNodeShutdownMetadata.Type.REMOVE : SingleNodeShutdownMetadata.Type.RESTART)
             .setReason(randomAlphaOfLength(5))
             .setStartedAtMillis(randomNonNegativeLong())
+            .setNodeSeen(randomBoolean())
             .build();
     }
 

x-pack/plugin/shutdown/src/internalClusterTest/java/org/elasticsearch/xpack/shutdown/NodeShutdownShardsIT.java

@@ -0,0 +1,94 @@
+/*
+ * 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.shutdown;
+
+import org.elasticsearch.Build;
+import org.elasticsearch.action.admin.cluster.node.info.NodeInfo;
+import org.elasticsearch.action.admin.cluster.node.info.NodesInfoResponse;
+import org.elasticsearch.action.support.master.AcknowledgedResponse;
+import org.elasticsearch.cluster.metadata.SingleNodeShutdownMetadata;
+import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.test.ESIntegTestCase;
+
+import java.util.Arrays;
+import java.util.Collection;
+
+import static org.elasticsearch.cluster.metadata.SingleNodeShutdownMetadata.Status.COMPLETE;
+import static org.hamcrest.Matchers.equalTo;
+
+@ESIntegTestCase.ClusterScope(scope = ESIntegTestCase.Scope.TEST, numDataNodes = 0, numClientNodes = 0)
+public class NodeShutdownShardsIT extends ESIntegTestCase {
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        return Arrays.asList(ShutdownPlugin.class);
+    }
+
+    public void testShardStatusStaysCompleteAfterNodeLeaves() throws Exception {
+        assumeTrue("must be on a snapshot build of ES to run in order for the feature flag to be set", Build.CURRENT.isSnapshot());
+        final String nodeToRestartName = internalCluster().startNode();
+        final String nodeToRestartId = getNodeId(nodeToRestartName);
+        internalCluster().startNode();
+
+        // Mark the node for shutdown
+        PutShutdownNodeAction.Request putShutdownRequest = new PutShutdownNodeAction.Request(
+            nodeToRestartId,
+            SingleNodeShutdownMetadata.Type.REMOVE,
+            this.getTestName()
+        );
+        AcknowledgedResponse putShutdownResponse = client().execute(PutShutdownNodeAction.INSTANCE, putShutdownRequest).get();
+        assertTrue(putShutdownResponse.isAcknowledged());
+
+        internalCluster().stopNode(nodeToRestartName);
+
+        NodesInfoResponse nodes = client().admin().cluster().prepareNodesInfo().clear().get();
+        assertThat(nodes.getNodes().size(), equalTo(1));
+
+        GetShutdownStatusAction.Response getResp = client().execute(
+            GetShutdownStatusAction.INSTANCE,
+            new GetShutdownStatusAction.Request(nodeToRestartId)
+        ).get();
+
+        assertThat(getResp.getShutdownStatuses().get(0).migrationStatus().getStatus(), equalTo(COMPLETE));
+    }
+
+    public void testShardStatusIsCompleteOnNonDataNodes() throws Exception {
+        assumeTrue("must be on a snapshot build of ES to run in order for the feature flag to be set", Build.CURRENT.isSnapshot());
+        final String nodeToShutDownName = internalCluster().startMasterOnlyNode();
+        internalCluster().startMasterOnlyNode(); // Just to have at least one other node
+        final String nodeToRestartId = getNodeId(nodeToShutDownName);
+
+        // Mark the node for shutdown
+        PutShutdownNodeAction.Request putShutdownRequest = new PutShutdownNodeAction.Request(
+            nodeToRestartId,
+            SingleNodeShutdownMetadata.Type.REMOVE,
+            this.getTestName()
+        );
+        AcknowledgedResponse putShutdownResponse = client().execute(PutShutdownNodeAction.INSTANCE, putShutdownRequest).get();
+        assertTrue(putShutdownResponse.isAcknowledged());
+
+        GetShutdownStatusAction.Response getResp = client().execute(
+            GetShutdownStatusAction.INSTANCE,
+            new GetShutdownStatusAction.Request(nodeToRestartId)
+        ).get();
+
+        assertThat(getResp.getShutdownStatuses().get(0).migrationStatus().getStatus(), equalTo(COMPLETE));
+    }
+
+    private String getNodeId(String nodeName) throws Exception {
+        NodesInfoResponse nodes = client().admin().cluster().prepareNodesInfo().clear().get();
+        return nodes.getNodes()
+            .stream()
+            .map(NodeInfo::getNode)
+            .filter(node -> node.getName().equals(nodeName))
+            .map(DiscoveryNode::getId)
+            .findFirst()
+            .orElseThrow();
+    }
+}

x-pack/plugin/shutdown/src/main/java/org/elasticsearch/xpack/shutdown/ShutdownService.java

@@ -0,0 +1,85 @@
+/*
+ * 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.shutdown;
+
+import org.apache.log4j.LogManager;
+import org.apache.log4j.Logger;
+import org.apache.logging.log4j.message.ParameterizedMessage;
+import org.elasticsearch.cluster.ClusterChangedEvent;
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.ClusterStateListener;
+import org.elasticsearch.cluster.ClusterStateUpdateTask;
+import org.elasticsearch.cluster.metadata.Metadata;
+import org.elasticsearch.cluster.metadata.NodesShutdownMetadata;
+import org.elasticsearch.cluster.metadata.SingleNodeShutdownMetadata;
+import org.elasticsearch.cluster.service.ClusterService;
+
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+public class ShutdownService implements ClusterStateListener {
+    private static final Logger logger = LogManager.getLogger(ShutdownService.class);
+
+    final ClusterService clusterService;
+
+    public ShutdownService(ClusterService clusterService) {
+        this.clusterService = clusterService;
+        clusterService.addListener(this);
+    }
+
+    @Override
+    public void clusterChanged(ClusterChangedEvent event) {
+        if (event.state().nodes().isLocalNodeElectedMaster() == false) {
+            // Only do this if we're the current master node.
+            return;
+        }
+        NodesShutdownMetadata eventShutdownMetadata = event.state().metadata().custom(NodesShutdownMetadata.TYPE);
+        final Set<String> nodesNotPreviouslySeen = eventShutdownMetadata.getAllNodeMetadataMap()
+            .values()
+            .stream()
+            .filter(singleNodeShutdownMetadata -> singleNodeShutdownMetadata.getNodeSeen() == false)
+            .map(SingleNodeShutdownMetadata::getNodeId)
+            .filter(nodeId -> event.state().nodes().nodeExists(nodeId))
+            .collect(Collectors.toUnmodifiableSet());
+
+        if (nodesNotPreviouslySeen.isEmpty() == false) {
+            clusterService.submitStateUpdateTask("shutdown-seen-nodes-updater", new ClusterStateUpdateTask() {
+                @Override
+                public ClusterState execute(ClusterState currentState) throws Exception {
+                    NodesShutdownMetadata shutdownMetadata = currentState.metadata().custom(NodesShutdownMetadata.TYPE);
+
+                    final Map<String, SingleNodeShutdownMetadata> newShutdownMetadataMap = shutdownMetadata.getAllNodeMetadataMap()
+                        .values()
+                        .stream()
+                        .map(singleNodeShutdownMetadata -> {
+                            if (nodesNotPreviouslySeen.contains(singleNodeShutdownMetadata.getNodeId())) {
+                                return SingleNodeShutdownMetadata.builder(singleNodeShutdownMetadata).setNodeSeen(true).build();
+                            }
+                            return singleNodeShutdownMetadata;
+                        })
+                        .collect(Collectors.toUnmodifiableMap(SingleNodeShutdownMetadata::getNodeId, Function.identity()));
+
+                    return ClusterState.builder(currentState)
+                        .metadata(
+                            Metadata.builder(currentState.metadata())
+                                .putCustom(NodesShutdownMetadata.TYPE, new NodesShutdownMetadata(newShutdownMetadataMap))
+                                .build()
+                        )
+                        .build();
+                }
+
+                @Override
+                public void onFailure(String source, Exception e) {
+                    logger.warn(new ParameterizedMessage("failed to mark shutting down nodes as seen: {}", nodesNotPreviouslySeen), e);
+                }
+            });
+        }
+    }
+}