Add changes for graceful node decommission

CHANGELOG.md

@@ -32,7 +32,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Added precommit support for MacOS ([#4682](https://github.com/opensearch-project/OpenSearch/pull/4682))
 - Recommission API changes for service layer ([#4320](https://github.com/opensearch-project/OpenSearch/pull/4320))
 - Update GeoGrid base class access modifier to support extensibility ([#4572](https://github.com/opensearch-project/OpenSearch/pull/4572))
-
+- Added changes for graceful node decommission ([#4586](https://github.com/opensearch-project/OpenSearch/pull/4586))
 ### Dependencies
 - Bumps `log4j-core` from 2.18.0 to 2.19.0
 - Bumps `reactor-netty-http` from 1.0.18 to 1.0.23

server/src/main/java/org/opensearch/action/admin/cluster/decommission/awareness/put/DecommissionRequest.java

@@ -14,6 +14,7 @@
 import org.opensearch.common.Strings;
 import org.opensearch.common.io.stream.StreamInput;
 import org.opensearch.common.io.stream.StreamOutput;
+import org.opensearch.common.unit.TimeValue;
 
 import java.io.IOException;
 
@@ -28,23 +29,34 @@
  */
 public class DecommissionRequest extends ClusterManagerNodeRequest<DecommissionRequest> {
 
+    public static final TimeValue DEFAULT_NODE_DRAINING_TIMEOUT = TimeValue.timeValueSeconds(300);
+
     private DecommissionAttribute decommissionAttribute;
 
+    private TimeValue drainingTimeout = DEFAULT_NODE_DRAINING_TIMEOUT;
+
     public DecommissionRequest() {}
 
     public DecommissionRequest(DecommissionAttribute decommissionAttribute) {
         this.decommissionAttribute = decommissionAttribute;
     }
 
+    public DecommissionRequest(DecommissionAttribute decommissionAttribute, TimeValue drainingTimeout) {
+        this.decommissionAttribute = decommissionAttribute;
+        this.drainingTimeout = drainingTimeout;
+    }
+
     public DecommissionRequest(StreamInput in) throws IOException {
         super(in);
         decommissionAttribute = new DecommissionAttribute(in);
+        this.drainingTimeout = in.readTimeValue();
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         decommissionAttribute.writeTo(out);
+        out.writeTimeValue(drainingTimeout);
     }
 
     /**
@@ -58,13 +70,28 @@ public DecommissionRequest setDecommissionAttribute(DecommissionAttribute decomm
         return this;
     }
 
+    /**
+     * Sets the timeout for the request
+     *
+     * @param timeout time out for the request
+     * @return this request
+     */
+    public DecommissionRequest setDrainingTimeout(TimeValue timeout) {
+        this.drainingTimeout = timeout;
+        return this;
+    }
+
     /**
      * @return Returns the decommission attribute key-value
      */
     public DecommissionAttribute getDecommissionAttribute() {
         return this.decommissionAttribute;
     }
 
+    public TimeValue getDrainingTimeout() {
+        return this.drainingTimeout;
+    }
+
     @Override
     public ActionRequestValidationException validate() {
         ActionRequestValidationException validationException = null;
@@ -74,6 +101,9 @@ public ActionRequestValidationException validate() {
         if (decommissionAttribute.attributeValue() == null || Strings.isEmpty(decommissionAttribute.attributeValue())) {
             validationException = addValidationError("attribute value is missing", validationException);
         }
+        if (drainingTimeout.getSeconds() < 0) {
+            validationException = addValidationError("Invalid draining timeout - " + drainingTimeout.getSeconds(), validationException);
+        }
         return validationException;
     }
 

server/src/main/java/org/opensearch/action/admin/cluster/decommission/awareness/put/TransportDecommissionAction.java

@@ -76,6 +76,6 @@ protected ClusterBlockException checkBlock(DecommissionRequest request, ClusterS
     protected void clusterManagerOperation(DecommissionRequest request, ClusterState state, ActionListener<DecommissionResponse> listener)
         throws Exception {
         logger.info("starting awareness attribute [{}] decommissioning", request.getDecommissionAttribute().toString());
-        decommissionService.startDecommissionAction(request.getDecommissionAttribute(), listener);
+        decommissionService.startDecommissionAction(request, listener);
     }
 }

server/src/main/java/org/opensearch/cluster/decommission/DecommissionController.java

@@ -18,6 +18,13 @@
 import org.opensearch.action.admin.cluster.configuration.ClearVotingConfigExclusionsAction;
 import org.opensearch.action.admin.cluster.configuration.ClearVotingConfigExclusionsRequest;
 import org.opensearch.action.admin.cluster.configuration.ClearVotingConfigExclusionsResponse;
+import org.opensearch.action.admin.cluster.node.stats.NodeStats;
+import org.opensearch.action.admin.cluster.node.stats.NodesStatsAction;
+import org.opensearch.action.admin.cluster.node.stats.NodesStatsRequest;
+import org.opensearch.action.admin.cluster.node.stats.NodesStatsResponse;
+import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterAddWeightedRoutingAction;
+import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingRequest;
+import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingResponse;
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.ClusterStateObserver;
 import org.opensearch.cluster.ClusterStateTaskConfig;
@@ -26,20 +33,24 @@
 import org.opensearch.cluster.coordination.NodeRemovalClusterStateTaskExecutor;
 import org.opensearch.cluster.metadata.Metadata;
 import org.opensearch.cluster.node.DiscoveryNode;
+import org.opensearch.cluster.routing.WeightedRouting;
 import org.opensearch.cluster.routing.allocation.AllocationService;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.Priority;
 import org.opensearch.common.Strings;
 import org.opensearch.common.io.stream.StreamInput;
 import org.opensearch.common.unit.TimeValue;
+import org.opensearch.http.HttpStats;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.TransportException;
 import org.opensearch.transport.TransportResponseHandler;
 import org.opensearch.transport.TransportService;
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Predicate;
@@ -267,4 +278,104 @@ public void clusterStateProcessed(String source, ClusterState oldState, ClusterS
             }
         });
     }
+
+    void setRoutingWeights(
+        String awarenessAttributeName,
+        Map<String, Double> awarenessAttributeToWeightMap,
+        ActionListener<ClusterPutWeightedRoutingResponse> listener
+    ) {
+        // WRR transport action will validate invalid weights
+        final ClusterPutWeightedRoutingRequest clusterPutRoutingWeightRequest = new ClusterPutWeightedRoutingRequest(
+            awarenessAttributeName
+        );
+        clusterPutRoutingWeightRequest.setWeightedRouting(new WeightedRouting(awarenessAttributeName, awarenessAttributeToWeightMap));
+
+        transportService.sendRequest(
+            transportService.getLocalNode(),
+            ClusterAddWeightedRoutingAction.NAME,
+            clusterPutRoutingWeightRequest,
+            new TransportResponseHandler<ClusterPutWeightedRoutingResponse>() {
+                @Override
+                public void handleResponse(ClusterPutWeightedRoutingResponse response) {
+                    logger.info("Weights are successfully set. [{}]", clusterPutRoutingWeightRequest.getWeightedRouting().weights());
+                    listener.onResponse(response);
+                }
+
+                @Override
+                public void handleException(TransportException exp) {
+                    // Logging message on failure. Should we do Retry? If weights are not set should we fail?
+                    logger.error("Failed to set routing weights. ", exp);
+                    listener.onFailure(exp);
+                }
+
+                @Override
+                public String executor() {
+                    return ThreadPool.Names.SAME;
+                }
+
+                @Override
+                public ClusterPutWeightedRoutingResponse read(StreamInput in) throws IOException {
+                    return new ClusterPutWeightedRoutingResponse(in);
+                }
+            }
+        );
+    }
+
+    private void logActiveConnections(NodesStatsResponse nodesStatsResponse) {
+        if (nodesStatsResponse == null || nodesStatsResponse.getNodes() == null) {
+            logger.info("Node stats response received is null/empty.");
+            return;
+        }
+
+        Map<String, Long> nodeActiveConnectionMap = new HashMap<>();
+        List<NodeStats> responseNodes = nodesStatsResponse.getNodes();
+        for (int i = 0; i < responseNodes.size(); i++) {
+            HttpStats httpStats = responseNodes.get(i).getHttp();
+            DiscoveryNode node = responseNodes.get(i).getNode();
+            nodeActiveConnectionMap.put(node.getId(), httpStats.getServerOpen());
+        }
+        logger.info("Decommissioning node with connections : [{}]", nodeActiveConnectionMap);
+    }
+
+    void getActiveRequestCountOnDecommissionedNodes(Set<DiscoveryNode> decommissionedNodes) {
+        if (decommissionedNodes == null || decommissionedNodes.isEmpty()) {
+            return;
+        }
+
+        String[] nodes = decommissionedNodes.stream().map(DiscoveryNode::getId).toArray(String[]::new);
+        if (nodes.length == 0) {
+            return;
+        }
+
+        final NodesStatsRequest nodesStatsRequest = new NodesStatsRequest(nodes);
+        nodesStatsRequest.clear();
+        nodesStatsRequest.addMetric(NodesStatsRequest.Metric.HTTP.metricName());
+
+        transportService.sendRequest(
+            transportService.getLocalNode(),
+            NodesStatsAction.NAME,
+            nodesStatsRequest,
+            new TransportResponseHandler<NodesStatsResponse>() {
+                @Override
+                public void handleResponse(NodesStatsResponse response) {
+                    logActiveConnections(response);
+                }
+
+                @Override
+                public void handleException(TransportException exp) {
+                    logger.warn("Failure occurred while dumping connection for decommission nodes - [{}]", exp.unwrapCause().getMessage());
+                }
+
+                @Override
+                public String executor() {
+                    return ThreadPool.Names.SAME;
+                }
+
+                @Override
+                public NodesStatsResponse read(StreamInput in) throws IOException {
+                    return new NodesStatsResponse(in);
+                }
+            }
+        );
+    }
 }