Integrate Node Shutdown API with cluster metadata

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

@@ -46,6 +46,10 @@ public static NodeShutdownComponentStatus parse(XContentParser parser) {
         return PARSER.apply(parser, null);
     }
 
+    public NodeShutdownComponentStatus() {
+        this(SingleNodeShutdownMetadata.Status.NOT_STARTED, null, null);
+    }
+
     public NodeShutdownComponentStatus(
         SingleNodeShutdownMetadata.Status status,
         @Nullable Long startedAtMillis,

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

@@ -23,6 +23,7 @@
 import java.io.IOException;
 import java.util.Collections;
 import java.util.EnumSet;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -81,13 +82,34 @@ public void writeTo(StreamOutput out) throws IOException {
     }
 
     /**
-     * Retrieve the data about nodes which are currently in the process of shutting down.
-     * @return A map of node IDs to information about the node's shutdown status.
+     * @return A map of NodeID to shutdown metadata.
      */
-    public Map<String, SingleNodeShutdownMetadata> getPerNodeInfo() {
+    public Map<String, SingleNodeShutdownMetadata> getAllNodeMetdataMap() {
         return Collections.unmodifiableMap(nodes);
     }
 
+    /**
+     * Add or update the shutdown metadata for a single node.
+     * @param nodeShutdownMetadata The single node shutdown metadata to add or update.
+     * @return A new {@link NodesShutdownMetadata} that reflects the updated value.
+     */
+    public NodesShutdownMetadata putSingleNodeMetadata(SingleNodeShutdownMetadata nodeShutdownMetadata) {
+        HashMap<String, SingleNodeShutdownMetadata> newNodes = new HashMap<>(nodes);
+        newNodes.put(nodeShutdownMetadata.getNodeId(), nodeShutdownMetadata);
+        return new NodesShutdownMetadata(newNodes);
+    }
+
+    /**
+     * Removes all shutdown metadata for a particular node ID.
+     * @param nodeId The node ID to remove shutdown metadata for.
+     * @return A new {@link NodesShutdownMetadata} that does not contain shutdown metadata for the given node.
+     */
+    public NodesShutdownMetadata removeSingleNodeMetadata(String nodeId) {
+        HashMap<String, SingleNodeShutdownMetadata> newNodes = new HashMap<>(nodes);
+        newNodes.remove(nodeId);
+        return new NodesShutdownMetadata(newNodes);
+    }
+
     @Override
     public Diff<Metadata.Custom> diff(Metadata.Custom previousState) {
         return new NodeShutdownMetadataDiff((NodesShutdownMetadata) previousState, this);
@@ -113,12 +135,12 @@ public boolean equals(Object o) {
         if (this == o) return true;
         if ((o instanceof NodesShutdownMetadata) == false) return false;
         NodesShutdownMetadata that = (NodesShutdownMetadata) o;
-        return getPerNodeInfo().equals(that.getPerNodeInfo());
+        return nodes.equals(that.nodes);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(getPerNodeInfo());
+        return Objects.hash(nodes);
     }
 
     @Override
@@ -150,7 +172,7 @@ public NodeShutdownMetadataDiff(StreamInput in) throws IOException {
         @Override
         public Metadata.Custom apply(Metadata.Custom part) {
             TreeMap<String, SingleNodeShutdownMetadata> newNodes = new TreeMap<>(
-                nodesDiff.apply(((NodesShutdownMetadata) part).getPerNodeInfo())
+                nodesDiff.apply(((NodesShutdownMetadata) part).nodes)
             );
             return new NodesShutdownMetadata(newNodes);
         }

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

@@ -89,15 +89,26 @@ public static SingleNodeShutdownMetadata parse(XContentParser parser) {
     private final NodeShutdownComponentStatus persistentTasksStatus;
     private final NodeShutdownComponentStatus pluginsStatus;
 
-
-    public SingleNodeShutdownMetadata(
+    /**
+     * @param nodeId The node ID that this shutdown metadata refers to.
+     * @param type The type of shutdown. See {@link Type}.
+     * @param reason The reason for the shutdown, per the original shutdown request.
+     * @param status The overall status of this shutdown.
+     * @param startedAtMillis The timestamp at which this shutdown was requested.
+     * @param shardMigrationStatus The status of shard migrations away from this node.
+     * @param persistentTasksStatus The status of persistent task migration away from this node.
+     * @param pluginsStatus The status of plugin shutdown on this node.
+     */
+    private SingleNodeShutdownMetadata(
         String nodeId,
         Type type,
         String reason,
         Status status,
         long startedAtMillis,
         NodeShutdownComponentStatus shardMigrationStatus,
-        NodeShutdownComponentStatus persistentTasksStatus, NodeShutdownComponentStatus pluginsStatus) {
+        NodeShutdownComponentStatus persistentTasksStatus,
+        NodeShutdownComponentStatus pluginsStatus
+    ) {
         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");
@@ -143,7 +154,7 @@ public String getReason() {
     /**
      * @return The status of this node's shutdown.
      */
-    public Status isStatus() {
+    public Status getStatus() {
         return status;
     }
 
@@ -154,6 +165,27 @@ public long getStartedAtMillis() {
         return startedAtMillis;
     }
 
+    /**
+     * @return The status of shard migrations off of this node.
+     */
+    public NodeShutdownComponentStatus getShardMigrationStatus() {
+        return shardMigrationStatus;
+    }
+
+    /**
+     * @return The status of persistent task shutdown on this node.
+     */
+    public NodeShutdownComponentStatus getPersistentTasksStatus() {
+        return persistentTasksStatus;
+    }
+
+    /**
+     * @return The status of plugin shutdown on this node.
+     */
+    public NodeShutdownComponentStatus getPluginsStatus() {
+        return pluginsStatus;
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(nodeId);
@@ -213,11 +245,137 @@ public int hashCode() {
         );
     }
 
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    public static Builder builder(SingleNodeShutdownMetadata original) {
+        if (original == null) {
+            return builder();
+        }
+        return new Builder()
+            .setNodeId(original.getNodeId())
+            .setType(original.getType())
+            .setReason(original.getReason())
+            .setStatus(original.getStatus())
+            .setStartedAtMillis(original.getStartedAtMillis())
+            .setShardMigrationStatus(original.getShardMigrationStatus())
+            .setPersistentTasksStatus(original.getPersistentTasksStatus())
+            .setPluginsStatus(original.getPluginsStatus());
+    }
+
+    public static class Builder {
+        private String nodeId;
+        private Type type;
+        private String reason;
+        private long startedAtMillis = -1;
+        private Status status = Status.IN_PROGRESS;
+        private NodeShutdownComponentStatus shardMigrationStatus = new NodeShutdownComponentStatus();
+        private NodeShutdownComponentStatus persistentTasksStatus = new NodeShutdownComponentStatus();
+        private NodeShutdownComponentStatus pluginsStatus = new NodeShutdownComponentStatus();
+
+        private Builder() {}
+
+        /**
+         * @param nodeId The node ID this metadata refers to.
+         * @return This builder.
+         */
+        public Builder setNodeId(String nodeId) {
+            this.nodeId = nodeId;
+            return this;
+        }
+
+        /**
+         * @param type The type of shutdown.
+         * @return This builder.
+         */
+        public Builder setType(Type type) {
+            this.type = type;
+            return this;
+        }
+
+        /**
+         * @param reason The reason for the shutdown. An arbitrary string provided by the user.
+         * @return This builder.
+         */
+        public Builder setReason(String reason) {
+            this.reason = reason;
+            return this;
+        }
+
+        /**
+         * @param startedAtMillis The timestamp at which this shutdown was requested.
+         * @return This builder.
+         */
+        public Builder setStartedAtMillis(long startedAtMillis) {
+            this.startedAtMillis = startedAtMillis;
+            return this;
+        }
+
+        /**
+         * @param status The status of this shutdown.
+         * @return This builder.
+         */
+        public Builder setStatus(Status status) {
+            this.status = status;
+            return this;
+        }
+
+        /**
+         * @param shardMigrationStatus An object describing the status of shard migration away from this node.
+         * @return This builder.
+         */
+        public Builder setShardMigrationStatus(NodeShutdownComponentStatus shardMigrationStatus) {
+            this.shardMigrationStatus = shardMigrationStatus;
+            return this;
+        }
+
+        /**
+         * @param persistentTasksStatus An object describing the status of persistent task migration away from this node.
+         * @return This builder.
+         */
+        public Builder setPersistentTasksStatus(NodeShutdownComponentStatus persistentTasksStatus) {
+            this.persistentTasksStatus = persistentTasksStatus;
+            return this;
+        }
+
+        /**
+         * @param pluginsStatus An object describing the status of plugin shutdown on this node.
+         * @return
+         */
+        public Builder setPluginsStatus(NodeShutdownComponentStatus pluginsStatus) {
+            this.pluginsStatus = pluginsStatus;
+            return this;
+        }
+
+        public SingleNodeShutdownMetadata build() {
+            if (startedAtMillis == -1) {
+                throw new IllegalArgumentException("start timestamp must be set");
+            }
+            return new SingleNodeShutdownMetadata(
+                nodeId,
+                type,
+                reason,
+                status,
+                startedAtMillis,
+                shardMigrationStatus,
+                persistentTasksStatus,
+                pluginsStatus
+            );
+        }
+    }
+
+    /**
+     * Describes the type of node shutdown - permanent (REMOVE) or temporary (RESTART).
+     */
     public enum Type {
         REMOVE,
         RESTART
     }
 
+    /**
+     * Describes the status of a component of shutdown.
+     */
     public enum Status {
         NOT_STARTED,
         IN_PROGRESS,

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

@@ -16,12 +16,47 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.EnumSet;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
+import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.hasItem;
+import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.not;
+import static org.hamcrest.Matchers.nullValue;
+
 public class NodesShutdownMetadataTests extends AbstractDiffableSerializationTestCase<Metadata.Custom> {
 
+    public void testInsertNewNodeShutdownMetadata() {
+        NodesShutdownMetadata nodesShutdownMetadata = new NodesShutdownMetadata(new HashMap<>());
+        SingleNodeShutdownMetadata newNodeMetadata = randomNodeShutdownInfo();
+
+        nodesShutdownMetadata = nodesShutdownMetadata.putSingleNodeMetadata(newNodeMetadata);
+
+        assertThat(nodesShutdownMetadata.getAllNodeMetdataMap().get(newNodeMetadata.getNodeId()), equalTo(newNodeMetadata));
+        assertThat(nodesShutdownMetadata.getAllNodeMetdataMap().values(), contains(newNodeMetadata));
+    }
+
+    public void testRemoveShutdownMetadata() {
+        NodesShutdownMetadata nodesShutdownMetadata = new NodesShutdownMetadata(new HashMap<>());
+        List<SingleNodeShutdownMetadata> nodes = randomList(1, 20, this::randomNodeShutdownInfo);
+
+        for (SingleNodeShutdownMetadata node : nodes) {
+            nodesShutdownMetadata = nodesShutdownMetadata.putSingleNodeMetadata(node);
+        }
+
+        SingleNodeShutdownMetadata nodeToRemove = randomFrom(nodes);
+        nodesShutdownMetadata = nodesShutdownMetadata.removeSingleNodeMetadata(nodeToRemove.getNodeId());
+
+        assertThat(nodesShutdownMetadata.getAllNodeMetdataMap().get(nodeToRemove.getNodeId()), nullValue());
+        assertThat(nodesShutdownMetadata.getAllNodeMetdataMap().values(), hasSize(nodes.size() - 1));
+        assertThat(nodesShutdownMetadata.getAllNodeMetdataMap().values(), not(hasItem(nodeToRemove)));
+    }
+
     @Override
     protected Writeable.Reader<Diff<Metadata.Custom>> diffReader() {
         return NodesShutdownMetadata.NodeShutdownMetadataDiff::new;
@@ -45,15 +80,15 @@ protected NodesShutdownMetadata createTestInstance() {
     }
 
     private SingleNodeShutdownMetadata randomNodeShutdownInfo() {
-        return new SingleNodeShutdownMetadata(
-            randomAlphaOfLength(5),
-            randomBoolean() ? SingleNodeShutdownMetadata.Type.REMOVE : SingleNodeShutdownMetadata.Type.RESTART,
-            randomAlphaOfLength(5),
-            randomStatus(),
-            randomNonNegativeLong(),
-            randomComponentStatus(),
-            randomComponentStatus(),
-            randomComponentStatus());
+        return SingleNodeShutdownMetadata.builder().setNodeId(randomAlphaOfLength(5))
+            .setType(randomBoolean() ? SingleNodeShutdownMetadata.Type.REMOVE : SingleNodeShutdownMetadata.Type.RESTART)
+            .setReason(randomAlphaOfLength(5))
+            .setStatus(randomStatus())
+            .setStartedAtMillis(randomNonNegativeLong())
+            .setShardMigrationStatus(randomComponentStatus())
+            .setPersistentTasksStatus(randomComponentStatus())
+            .setPluginsStatus(randomComponentStatus())
+            .build();
     }
 
     private SingleNodeShutdownMetadata.Status randomStatus() {

x-pack/plugin/shutdown/qa/multi-node/build.gradle

@@ -0,0 +1,20 @@
+apply plugin: 'elasticsearch.java-rest-test'
+
+dependencies {
+  javaRestTestImplementation(testArtifact(project(xpackModule('core'))))
+}
+
+def clusterCredentials = [username: System.getProperty('tests.rest.cluster.username', 'test_admin'),
+                          password: System.getProperty('tests.rest.cluster.password', 'x-pack-test-password')]
+
+tasks.named("javaRestTest").configure {
+  systemProperty 'tests.rest.cluster.username', clusterCredentials.username
+  systemProperty 'tests.rest.cluster.password', clusterCredentials.password
+}
+
+testClusters.all {
+  testDistribution = 'DEFAULT'
+  numberOfNodes = 4
+
+  systemProperty 'es.shutdown_feature_flag_enabled', 'true'
+}