diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/QueryGroupStatsRequest.java b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/QueryGroupStatsRequest.java index d5a398af3ff85..66bd7df5a1adc 100644 --- a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/QueryGroupStatsRequest.java +++ b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/QueryGroupStatsRequest.java @@ -26,7 +26,7 @@ public class QueryGroupStatsRequest extends BaseNodesRequest queryGroupIds; private final Boolean breach; - protected QueryGroupStatsRequest(StreamInput in) throws IOException { + public QueryGroupStatsRequest(StreamInput in) throws IOException { super(in); this.queryGroupIds = new HashSet<>(Set.of(in.readStringArray())); this.breach = in.readOptionalBoolean(); diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/TransportQueryGroupStatsAction.java b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/TransportQueryGroupStatsAction.java index cb23312b50df9..583f06d08967f 100644 --- a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/TransportQueryGroupStatsAction.java +++ b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/TransportQueryGroupStatsAction.java @@ -11,13 +11,10 @@ import org.opensearch.action.FailedNodeException; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.nodes.TransportNodesAction; -import org.opensearch.cluster.node.DiscoveryNode; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; import org.opensearch.core.common.io.stream.StreamInput; -import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.threadpool.ThreadPool; -import org.opensearch.transport.TransportRequest; import org.opensearch.transport.TransportService; import org.opensearch.wlm.QueryGroupService; import org.opensearch.wlm.stats.QueryGroupStats; @@ -33,7 +30,7 @@ public class TransportQueryGroupStatsAction extends TransportNodesAction< QueryGroupStatsRequest, QueryGroupStatsResponse, - TransportQueryGroupStatsAction.NodeQueryGroupStatsRequest, + QueryGroupStatsRequest, QueryGroupStats> { QueryGroupService queryGroupService; @@ -53,7 +50,7 @@ public TransportQueryGroupStatsAction( transportService, actionFilters, QueryGroupStatsRequest::new, - NodeQueryGroupStatsRequest::new, + QueryGroupStatsRequest::new, ThreadPool.Names.MANAGEMENT, QueryGroupStats.class ); @@ -70,8 +67,8 @@ protected QueryGroupStatsResponse newResponse( } @Override - protected NodeQueryGroupStatsRequest newNodeRequest(QueryGroupStatsRequest request) { - return new NodeQueryGroupStatsRequest(request); + protected QueryGroupStatsRequest newNodeRequest(QueryGroupStatsRequest request) { + return request; } @Override @@ -80,37 +77,7 @@ protected QueryGroupStats newNodeResponse(StreamInput in) throws IOException { } @Override - protected QueryGroupStats nodeOperation(NodeQueryGroupStatsRequest nodeQueryGroupStatsRequest) { - QueryGroupStatsRequest request = nodeQueryGroupStatsRequest.request; - return queryGroupService.nodeStats(request.getQueryGroupIds(), request.isBreach()); - } - - /** - * Inner QueryGroupStatsRequest - * - * @opensearch.experimental - */ - public static class NodeQueryGroupStatsRequest extends TransportRequest { - - protected QueryGroupStatsRequest request; - - public NodeQueryGroupStatsRequest(StreamInput in) throws IOException { - super(in); - request = new QueryGroupStatsRequest(in); - } - - NodeQueryGroupStatsRequest(QueryGroupStatsRequest request) { - this.request = request; - } - - @Override - public void writeTo(StreamOutput out) throws IOException { - super.writeTo(out); - request.writeTo(out); - } - - public DiscoveryNode[] getDiscoveryNodes() { - return this.request.concreteNodes(); - } + protected QueryGroupStats nodeOperation(QueryGroupStatsRequest queryGroupStatsRequest) { + return queryGroupService.nodeStats(queryGroupStatsRequest.getQueryGroupIds(), queryGroupStatsRequest.isBreach()); } } diff --git a/server/src/test/java/org/opensearch/action/support/nodes/TransportQueryGroupStatsActionTests.java b/server/src/test/java/org/opensearch/action/support/nodes/TransportQueryGroupStatsActionTests.java index f27bf7f6023ee..9079410983dfd 100644 --- a/server/src/test/java/org/opensearch/action/support/nodes/TransportQueryGroupStatsActionTests.java +++ b/server/src/test/java/org/opensearch/action/support/nodes/TransportQueryGroupStatsActionTests.java @@ -10,7 +10,6 @@ import org.opensearch.action.admin.cluster.wlm.QueryGroupStatsRequest; import org.opensearch.action.admin.cluster.wlm.TransportQueryGroupStatsAction; -import org.opensearch.action.admin.cluster.wlm.TransportQueryGroupStatsAction.NodeQueryGroupStatsRequest; import org.opensearch.action.support.ActionFilters; import org.opensearch.action.support.PlainActionFuture; import org.opensearch.common.io.stream.BytesStreamOutput; @@ -35,16 +34,16 @@ public class TransportQueryGroupStatsActionTests extends TransportNodesActionTes */ public void testQueryGroupStatsActionWithRetentionOfDiscoveryNodesList() { QueryGroupStatsRequest request = new QueryGroupStatsRequest(); - Map> combinedSentRequest = performQueryGroupStatsAction(request); + Map> combinedSentRequest = performQueryGroupStatsAction(request); assertNotNull(combinedSentRequest); combinedSentRequest.forEach((node, capturedRequestList) -> { assertNotNull(capturedRequestList); - capturedRequestList.forEach(sentRequest -> { assertNull(sentRequest.getDiscoveryNodes()); }); + capturedRequestList.forEach(sentRequest -> { assertNull(sentRequest.concreteNodes()); }); }); } - private Map> performQueryGroupStatsAction(QueryGroupStatsRequest request) { + private Map> performQueryGroupStatsAction(QueryGroupStatsRequest request) { TransportNodesAction action = new TransportQueryGroupStatsAction( THREAD_POOL, clusterService, @@ -55,19 +54,18 @@ private Map> performQueryGroupStatsActi PlainActionFuture listener = new PlainActionFuture<>(); action.new AsyncAction(null, request, listener).start(); Map> capturedRequests = transport.getCapturedRequestsByTargetNodeAndClear(); - Map> combinedSentRequest = new HashMap<>(); + Map> combinedSentRequest = new HashMap<>(); capturedRequests.forEach((node, capturedRequestList) -> { - List sentRequestList = new ArrayList<>(); + List sentRequestList = new ArrayList<>(); capturedRequestList.forEach(preSentRequest -> { BytesStreamOutput out = new BytesStreamOutput(); try { - TransportQueryGroupStatsAction.NodeQueryGroupStatsRequest QueryGroupStatsRequestFromCoordinator = - (TransportQueryGroupStatsAction.NodeQueryGroupStatsRequest) preSentRequest.request; + QueryGroupStatsRequest QueryGroupStatsRequestFromCoordinator = (QueryGroupStatsRequest) preSentRequest.request; QueryGroupStatsRequestFromCoordinator.writeTo(out); StreamInput in = out.bytes().streamInput(); - NodeQueryGroupStatsRequest QueryGroupStatsRequest = new NodeQueryGroupStatsRequest(in); + QueryGroupStatsRequest QueryGroupStatsRequest = new QueryGroupStatsRequest(in); sentRequestList.add(QueryGroupStatsRequest); } catch (IOException e) { throw new RuntimeException(e);