Fix parsing of node shutdown allocation delay

x-pack/plugin/shutdown/qa/multi-node/src/javaRestTest/java/org/elasticsearch/xpack/shutdown/NodeShutdownIT.java

@@ -15,6 +15,7 @@
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.common.xcontent.ObjectPath;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.test.rest.ESRestTestCase;
 
 import java.io.IOException;
@@ -27,6 +28,7 @@
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.equalToIgnoringCase;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.not;
@@ -35,16 +37,23 @@
 
 public class NodeShutdownIT extends ESRestTestCase {
 
+    public void testRestartCRUD() throws Exception {
+        checkCRUD(randomFrom("restart", "RESTART"), randomPositiveTimeValue());
+    }
+
+    public void testRemoveCRUD() throws Exception {
+        checkCRUD(randomFrom("remove", "REMOVE"), null);
+    }
+
     @SuppressWarnings("unchecked")
-    public void testCRUD() throws Exception {
+    public void checkCRUD(String type, String allocationDelay) 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());
         String nodeIdToShutdown = getRandomNodeId();
-        String type = randomFrom("RESTART", "REMOVE");
 
         // Ensure if we do a GET before the cluster metadata is set up, we don't get an error
         assertNoShuttingDownNodes(nodeIdToShutdown);
 
-        putNodeShutdown(nodeIdToShutdown, type);
+        putNodeShutdown(nodeIdToShutdown, type, allocationDelay);
 
         // Ensure we can read it back
         {
@@ -53,8 +62,9 @@ public void testCRUD() throws Exception {
             List<Map<String, Object>> nodesArray = (List<Map<String, Object>>) statusResponse.get("nodes");
             assertThat(nodesArray, hasSize(1));
             assertThat(nodesArray.get(0).get("node_id"), equalTo(nodeIdToShutdown));
-            assertThat(nodesArray.get(0).get("type"), equalTo(type));
+            assertThat((String) nodesArray.get(0).get("type"), equalToIgnoringCase(type));
             assertThat(nodesArray.get(0).get("reason"), equalTo(this.getTestName()));
+            assertThat(nodesArray.get(0).get("allocation_delay"), equalTo(allocationDelay));
         }
 
         // Delete it and make sure it's deleted
@@ -363,11 +373,23 @@ private void assertUnassignedShard(String nodeIdToShutdown, String indexName) th
     }
 
     private void putNodeShutdown(String nodeIdToShutdown, String type) throws IOException {
+        putNodeShutdown(nodeIdToShutdown, type, null);
+    }
+
+    private void putNodeShutdown(String nodeIdToShutdown, String type, @Nullable String allocationDelay) throws IOException {
         String reason = this.getTestName();
 
         // Put a shutdown request
         Request putShutdown = new Request("PUT", "_nodes/" + nodeIdToShutdown + "/shutdown");
-        putShutdown.setJsonEntity("{\"type\":  \"" + type + "\", \"reason\":  \"" + reason + "\"}");
+        if (type.equalsIgnoreCase("restart") && allocationDelay != null) {
+            putShutdown.setJsonEntity(
+                "{\"type\":  \"" + type + "\", \"reason\":  \"" + reason + "\", \"allocation_delay\": \"" + allocationDelay + "\"}"
+            );
+
+        } else {
+            assertNull("allocation delay parameter is only valid for RESTART-type shutdowns", allocationDelay);
+            putShutdown.setJsonEntity("{\"type\":  \"" + type + "\", \"reason\":  \"" + reason + "\"}");
+        }
         assertOK(client().performRequest(putShutdown));
     }
 

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

@@ -47,7 +47,12 @@ public static class Request extends AcknowledgedRequest<Request> {
         private static final ConstructingObjectParser<Request, String> PARSER = new ConstructingObjectParser<>(
             "put_node_shutdown_request",
             false,
-            (a, nodeId) -> new Request(nodeId, SingleNodeShutdownMetadata.Type.parse((String) a[0]), (String) a[1], (TimeValue) a[2])
+            (a, nodeId) -> new Request(
+                nodeId,
+                SingleNodeShutdownMetadata.Type.parse((String) a[0]),
+                (String) a[1],
+                a[2] == null ? null : TimeValue.parseTimeValue((String) a[2], "put-shutdown-node-request-" + nodeId)
+            )
         );
 
         static {

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

@@ -113,6 +113,12 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field(SingleNodeShutdownMetadata.NODE_ID_FIELD.getPreferredName(), metadata.getNodeId());
             builder.field(SingleNodeShutdownMetadata.TYPE_FIELD.getPreferredName(), metadata.getType());
             builder.field(SingleNodeShutdownMetadata.REASON_FIELD.getPreferredName(), metadata.getReason());
+            if (metadata.getAllocationDelay() != null) {
+                builder.field(
+                    SingleNodeShutdownMetadata.ALLOCATION_DELAY_FIELD.getPreferredName(),
+                    metadata.getAllocationDelay().getStringRep()
+                );
+            }
             builder.timeField(
                 SingleNodeShutdownMetadata.STARTED_AT_MILLIS_FIELD.getPreferredName(),
                 SingleNodeShutdownMetadata.STARTED_AT_READABLE_FIELD,