Add request timeout param

Signed-off-by: Rishab Nahata <rnnahata@amazon.com>
opensearch-project · Dec 14, 2022 · ee90987 · ee90987
1 parent dbb404e
commit ee90987
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 9 deletions.
diff --git a/...a/org/opensearch/action/admin/cluster/decommission/awareness/put/DecommissionRequest.java b/...a/org/opensearch/action/admin/cluster/decommission/awareness/put/DecommissionRequest.java
@@ -29,11 +29,12 @@
  */
 public class DecommissionRequest extends ClusterManagerNodeRequest<DecommissionRequest> {
 
+    public static final TimeValue TIMEOUT = TimeValue.timeValueMinutes(2L);
     public static final TimeValue DEFAULT_NODE_DRAINING_TIMEOUT = TimeValue.timeValueSeconds(120);
 
     private DecommissionAttribute decommissionAttribute;
     private boolean retryOnClusterManagerChange;
-
+    private TimeValue timeout;
     private TimeValue delayTimeout = DEFAULT_NODE_DRAINING_TIMEOUT;
 
     // holder for no_delay param. To avoid draining time timeout.
@@ -42,12 +43,17 @@ public class DecommissionRequest extends ClusterManagerNodeRequest<DecommissionR
     public DecommissionRequest() {}
 
     public DecommissionRequest(DecommissionAttribute decommissionAttribute) {
-        this.decommissionAttribute = decommissionAttribute;
+        this(decommissionAttribute, false, TIMEOUT);
     }
 
-    public DecommissionRequest(DecommissionAttribute decommissionAttribute, boolean retryOnClusterManagerChange) {
+    public DecommissionRequest(DecommissionAttribute decommissionAttribute, boolean retryOnClusterManagerChange, TimeValue timeout) {
         this.decommissionAttribute = decommissionAttribute;
         this.retryOnClusterManagerChange = retryOnClusterManagerChange;
+        this.timeout = timeout;
+    }
+
+    public DecommissionRequest(DecommissionAttribute decommissionAttribute, TimeValue timeout) {
+        this(decommissionAttribute, false, timeout);
     }
 
     public DecommissionRequest(StreamInput in) throws IOException {
@@ -56,6 +62,8 @@ public DecommissionRequest(StreamInput in) throws IOException {
         this.delayTimeout = in.readTimeValue();
         this.noDelay = in.readBoolean();
         this.retryOnClusterManagerChange = in.readBoolean();
+        this.timeout = in.readTimeValue();
+
     }
 
     @Override
@@ -65,6 +73,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeTimeValue(delayTimeout);
         out.writeBoolean(noDelay);
         out.writeBoolean(retryOnClusterManagerChange);
+        out.writeTimeValue(timeout);
     }
 
     /**
@@ -122,6 +131,24 @@ public boolean retryOnClusterManagerChange() {
         return this.retryOnClusterManagerChange;
     }
 
+    /**
+     * Sets the timeout for the request
+     *
+     * @param timeout retry time out for the request
+     * @return this request
+     */
+    public DecommissionRequest setTimeout(TimeValue timeout) {
+        this.timeout = timeout;
+        return this;
+    }
+
+    /**
+     * @return timeout
+     */
+    public TimeValue timeout() {
+        return this.timeout;
+    }
+
     @Override
     public ActionRequestValidationException validate() {
         ActionRequestValidationException validationException = null;
@@ -153,12 +180,12 @@ public String toString() {
             + decommissionAttribute
             + ", retryOnClusterManagerChange="
             + retryOnClusterManagerChange
+            + ", timeout="
+            + timeout
             + ", delayTimeout="
             + delayTimeout
             + ", noDelay="
             + noDelay
-            + ", clusterManagerNodeTimeout="
-            + clusterManagerNodeTimeout
             + '}';
     }
 }
diff --git a/...pensearch/action/admin/cluster/decommission/awareness/put/DecommissionRequestBuilder.java b/...pensearch/action/admin/cluster/decommission/awareness/put/DecommissionRequestBuilder.java
@@ -57,4 +57,15 @@ public DecommissionRequestBuilder setRetryOnClusterManagerChange(boolean retryOn
         request.setRetryOnClusterManagerChange(retryOnClusterManagerChange);
         return this;
     }
+
+    /**
+     * Sets the timeout for the decommission request
+     *
+     * @param timeout time out for the request
+     * @return current object
+     */
+    public DecommissionRequestBuilder setTimeout(TimeValue timeout) {
+        request.setTimeout(timeout);
+        return this;
+    }
 }
diff --git a/server/src/main/java/org/opensearch/cluster/decommission/DecommissionController.java b/server/src/main/java/org/opensearch/cluster/decommission/DecommissionController.java
@@ -91,8 +91,7 @@ public void retryDecommissionAction(
         long startTime,
         ActionListener<DecommissionResponse> listener
     ) {
-        final long remainingTimeoutMS = decommissionRequest.clusterManagerNodeTimeout().millis() - (threadPool.relativeTimeInMillis()
-            - startTime);
+        final long remainingTimeoutMS = decommissionRequest.timeout().millis() - (threadPool.relativeTimeInMillis() - startTime);
         if (remainingTimeoutMS <= 0) {
             String errorMsg = "cluster manager node timed out before retrying ["
                 + DecommissionAction.NAME
@@ -104,7 +103,7 @@ public void retryDecommissionAction(
             return;
         }
         decommissionRequest.setRetryOnClusterManagerChange(true);
-        decommissionRequest.clusterManagerNodeTimeout(TimeValue.timeValueMillis(remainingTimeoutMS));
+        decommissionRequest.setTimeout(TimeValue.timeValueMillis(remainingTimeoutMS));
         transportService.sendRequest(
             transportService.getLocalNode(),
             DecommissionAction.NAME,

diff --git a/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestDecommissionAction.java b/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestDecommissionAction.java
@@ -21,6 +21,7 @@
 import java.util.List;
 
 import static java.util.Collections.singletonList;
+import static org.opensearch.action.admin.cluster.decommission.awareness.put.DecommissionRequest.TIMEOUT;
 import static org.opensearch.rest.RestRequest.Method.PUT;
 
 /**
@@ -59,6 +60,7 @@ DecommissionRequest createRequest(RestRequest request) throws IOException {
             decommissionRequest.setDelayTimeout(delayTimeout);
         }
         return decommissionRequest.setDecommissionAttribute(new DecommissionAttribute(attributeName, attributeValue))
-            .setRetryOnClusterManagerChange(false);
+            .setRetryOnClusterManagerChange(false)
+            .setTimeout(TimeValue.parseTimeValue(request.param("timeout"), TIMEOUT, getClass().getSimpleName() + ".timeout"));
     }
 }