From ff4145ce1249c7a6df3d7fa670c4c9ecf5e0f60e Mon Sep 17 00:00:00 2001 From: Ruirui Zhang Date: Tue, 8 Oct 2024 00:35:25 -0700 Subject: [PATCH] encapsulate querygroupstats in wlmstats Signed-off-by: Ruirui Zhang --- .../cluster/wlm/TransportWlmStatsAction.java | 22 +++--- .../admin/cluster/wlm/WlmStatsRequest.java | 4 +- .../admin/cluster/wlm/WlmStatsResponse.java | 16 ++-- .../org/opensearch/wlm/QueryGroupService.java | 25 ++++--- .../opensearch/wlm/stats/QueryGroupStats.java | 9 +-- .../org/opensearch/wlm/stats/WlmStats.java | 65 ++++++++++++++++ .../cluster/wlm/WlmStatsResponseTests.java | 10 ++- .../wlm/QueryGroupServiceTests.java | 11 +++ ...eryGroupRequestOperationListenerTests.java | 27 ++++--- .../wlm/stats/QueryGroupStatsTests.java | 5 +- .../opensearch/wlm/stats/WlmStatsTests.java | 75 +++++++++++++++++++ 11 files changed, 213 insertions(+), 56 deletions(-) create mode 100644 server/src/main/java/org/opensearch/wlm/stats/WlmStats.java create mode 100644 server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/TransportWlmStatsAction.java b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/TransportWlmStatsAction.java index e86e96715825f..55f2d179ec02a 100644 --- a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/TransportWlmStatsAction.java +++ b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/TransportWlmStatsAction.java @@ -17,17 +17,17 @@ import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; import org.opensearch.wlm.QueryGroupService; -import org.opensearch.wlm.stats.QueryGroupStats; +import org.opensearch.wlm.stats.WlmStats; import java.io.IOException; import java.util.List; /** - * Transport action for obtaining QueryGroupStats + * Transport action for obtaining WlmStats * * @opensearch.experimental */ -public class TransportWlmStatsAction extends TransportNodesAction { +public class TransportWlmStatsAction extends TransportNodesAction { final QueryGroupService queryGroupService; @@ -48,18 +48,14 @@ public TransportWlmStatsAction( WlmStatsRequest::new, WlmStatsRequest::new, ThreadPool.Names.MANAGEMENT, - QueryGroupStats.class + WlmStats.class ); this.queryGroupService = queryGroupService; } @Override - protected WlmStatsResponse newResponse( - WlmStatsRequest request, - List queryGroupStats, - List failures - ) { - return new WlmStatsResponse(clusterService.getClusterName(), queryGroupStats, failures); + protected WlmStatsResponse newResponse(WlmStatsRequest request, List wlmStats, List failures) { + return new WlmStatsResponse(clusterService.getClusterName(), wlmStats, failures); } @Override @@ -68,12 +64,12 @@ protected WlmStatsRequest newNodeRequest(WlmStatsRequest request) { } @Override - protected QueryGroupStats newNodeResponse(StreamInput in) throws IOException { - return new QueryGroupStats(in); + protected WlmStats newNodeResponse(StreamInput in) throws IOException { + return new WlmStats(in); } @Override - protected QueryGroupStats nodeOperation(WlmStatsRequest wlmStatsRequest) { + protected WlmStats nodeOperation(WlmStatsRequest wlmStatsRequest) { return queryGroupService.nodeStats(wlmStatsRequest.getQueryGroupIds(), wlmStatsRequest.isBreach()); } } diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsRequest.java b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsRequest.java index a09ea329d9a98..bf4f79faff478 100644 --- a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsRequest.java +++ b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsRequest.java @@ -36,7 +36,7 @@ public WlmStatsRequest(StreamInput in) throws IOException { * Get QueryGroup stats from nodes based on the nodes ids specified. If none are passed, stats * for all nodes will be returned. */ - public WlmStatsRequest(String[] nodesIds, Set queryGroupIds, boolean breach) { + public WlmStatsRequest(String[] nodesIds, Set queryGroupIds, Boolean breach) { super(false, nodesIds); this.queryGroupIds = queryGroupIds; this.breach = breach; @@ -59,7 +59,7 @@ public Set getQueryGroupIds() { return queryGroupIds; } - public boolean isBreach() { + public Boolean isBreach() { return breach; } } diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponse.java b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponse.java index 1175abc462c8e..2ce1b09a61fc6 100644 --- a/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponse.java +++ b/server/src/main/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponse.java @@ -18,6 +18,7 @@ import org.opensearch.core.xcontent.ToXContentFragment; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.wlm.stats.QueryGroupStats; +import org.opensearch.wlm.stats.WlmStats; import java.io.IOException; import java.util.List; @@ -26,30 +27,31 @@ * A response for obtaining Workload Management Stats */ @ExperimentalApi -public class WlmStatsResponse extends BaseNodesResponse implements ToXContentFragment { +public class WlmStatsResponse extends BaseNodesResponse implements ToXContentFragment { WlmStatsResponse(StreamInput in) throws IOException { super(in); } - WlmStatsResponse(ClusterName clusterName, List nodes, List failures) { + WlmStatsResponse(ClusterName clusterName, List nodes, List failures) { super(clusterName, nodes, failures); } @Override - protected List readNodesFrom(StreamInput in) throws IOException { - return in.readList(QueryGroupStats::new); + protected List readNodesFrom(StreamInput in) throws IOException { + return in.readList(WlmStats::new); } @Override - protected void writeNodesTo(StreamOutput out, List nodes) throws IOException { + protected void writeNodesTo(StreamOutput out, List nodes) throws IOException { out.writeList(nodes); } @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { - for (QueryGroupStats queryGroupStats : getNodes()) { - builder.startObject(queryGroupStats.getNode().getId()); + for (WlmStats wlmStats : getNodes()) { + builder.startObject(wlmStats.getNode().getId()); + QueryGroupStats queryGroupStats = wlmStats.getQueryGroupStats(); queryGroupStats.toXContent(builder, params); builder.endObject(); } diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java index fe81d04bca815..16be62085a4c5 100644 --- a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java @@ -33,6 +33,7 @@ import org.opensearch.wlm.stats.QueryGroupState; import org.opensearch.wlm.stats.QueryGroupStats; import org.opensearch.wlm.stats.QueryGroupStats.QueryGroupStatsHolder; +import org.opensearch.wlm.stats.WlmStats; import java.io.IOException; import java.util.HashMap; @@ -213,7 +214,7 @@ public void incrementFailuresFor(final String queryGroupId) { /** * @return node level query group stats */ - public QueryGroupStats nodeStats(Set queryGroupIds, Boolean requestedBreached) { + public WlmStats nodeStats(Set queryGroupIds, Boolean requestedBreached) { final Map statsHolderMap = new HashMap<>(); Map existingGroups = clusterService.state().metadata().queryGroups(); if (!queryGroupIds.contains("_all")) { @@ -223,16 +224,20 @@ public QueryGroupStats nodeStats(Set queryGroupIds, Boolean requestedBre } } } - queryGroupsStateAccessor.getQueryGroupStateMap().forEach((queryGroupId, currentState) -> { - boolean shouldInclude = (queryGroupIds.size() == 1 && queryGroupIds.contains("_all")) || queryGroupIds.contains(queryGroupId); - if (shouldInclude) { - boolean breached = resourceLimitBreached(existingGroups.get(queryGroupId), currentState).v1().length() != 0; - if (requestedBreached == null || requestedBreached == breached) { - statsHolderMap.put(queryGroupId, QueryGroupStatsHolder.from(currentState)); + Map stateMap = queryGroupsStateAccessor.getQueryGroupStateMap(); + if (stateMap != null) { + stateMap.forEach((queryGroupId, currentState) -> { + boolean shouldInclude = queryGroupIds.contains("_all") || queryGroupIds.contains(queryGroupId); + if (shouldInclude) { + if (requestedBreached == null + || requestedBreached == (resourceLimitBreached(existingGroups.get(queryGroupId), currentState).v1() + .length() != 0)) { + statsHolderMap.put(queryGroupId, QueryGroupStatsHolder.from(currentState)); + } } - } - }); - return new QueryGroupStats(transportService.getLocalNode(), statsHolderMap); + }); + } + return new WlmStats(transportService.getLocalNode(), new QueryGroupStats(statsHolderMap)); } /** diff --git a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java index 2222e5cc7a953..12e0f0ff96f65 100644 --- a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java +++ b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java @@ -8,8 +8,6 @@ package org.opensearch.wlm.stats; -import org.opensearch.action.support.nodes.BaseNodeResponse; -import org.opensearch.cluster.node.DiscoveryNode; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.common.io.stream.Writeable; @@ -39,22 +37,19 @@ * ... * } */ -public class QueryGroupStats extends BaseNodeResponse implements ToXContentObject, Writeable { +public class QueryGroupStats implements ToXContentObject, Writeable { private final Map stats; - public QueryGroupStats(DiscoveryNode node, Map stats) { - super(node); + public QueryGroupStats(Map stats) { this.stats = stats; } public QueryGroupStats(StreamInput in) throws IOException { - super(in); stats = in.readMap(StreamInput::readString, QueryGroupStatsHolder::new); } @Override public void writeTo(StreamOutput out) throws IOException { - super.writeTo(out); out.writeMap(stats, StreamOutput::writeString, QueryGroupStatsHolder::writeTo); } diff --git a/server/src/main/java/org/opensearch/wlm/stats/WlmStats.java b/server/src/main/java/org/opensearch/wlm/stats/WlmStats.java new file mode 100644 index 0000000000000..3313021bfae52 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/stats/WlmStats.java @@ -0,0 +1,65 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm.stats; + +import org.opensearch.action.support.nodes.BaseNodeResponse; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.common.io.stream.Writeable; +import org.opensearch.core.xcontent.ToXContentObject; +import org.opensearch.core.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.Objects; + +/** + * This class contains the stats for Workload Management + */ +public class WlmStats extends BaseNodeResponse implements ToXContentObject, Writeable { + private final QueryGroupStats queryGroupStats; + + public WlmStats(DiscoveryNode node, QueryGroupStats queryGroupStats) { + super(node); + this.queryGroupStats = queryGroupStats; + } + + public WlmStats(StreamInput in) throws IOException { + super(in); + queryGroupStats = new QueryGroupStats(in); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + queryGroupStats.writeTo(out); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + return queryGroupStats.toXContent(builder, params); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + WlmStats that = (WlmStats) o; + return Objects.equals(getQueryGroupStats(), that.getQueryGroupStats()); + } + + @Override + public int hashCode() { + return Objects.hash(queryGroupStats); + } + + public QueryGroupStats getQueryGroupStats() { + return queryGroupStats; + } +} diff --git a/server/src/test/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponseTests.java b/server/src/test/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponseTests.java index e9d5d7c31e887..86fbc341691d2 100644 --- a/server/src/test/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponseTests.java +++ b/server/src/test/java/org/opensearch/action/admin/cluster/wlm/WlmStatsResponseTests.java @@ -18,6 +18,7 @@ import org.opensearch.test.OpenSearchTestCase; import org.opensearch.wlm.ResourceType; import org.opensearch.wlm.stats.QueryGroupStats; +import org.opensearch.wlm.stats.WlmStats; import java.io.IOException; import java.util.ArrayList; @@ -38,7 +39,6 @@ public class WlmStatsResponseTests extends OpenSearchTestCase { ); Map statsHolderMap = new HashMap<>(); QueryGroupStats queryGroupStats = new QueryGroupStats( - node, Map.of( testQueryGroupId, new QueryGroupStats.QueryGroupStatsHolder( @@ -56,11 +56,12 @@ public class WlmStatsResponseTests extends OpenSearchTestCase { ) ) ); - List queryGroupStatsList = List.of(queryGroupStats); + WlmStats wlmStats = new WlmStats(node, queryGroupStats); + List wlmStatsList = List.of(wlmStats); List failedNodeExceptionList = new ArrayList<>(); public void testSerializationAndDeserialization() throws IOException { - WlmStatsResponse queryGroupStatsResponse = new WlmStatsResponse(clusterName, queryGroupStatsList, failedNodeExceptionList); + WlmStatsResponse queryGroupStatsResponse = new WlmStatsResponse(clusterName, wlmStatsList, failedNodeExceptionList); BytesStreamOutput out = new BytesStreamOutput(); queryGroupStatsResponse.writeTo(out); StreamInput in = out.bytes().streamInput(); @@ -70,7 +71,7 @@ public void testSerializationAndDeserialization() throws IOException { } public void testToString() { - WlmStatsResponse queryGroupStatsResponse = new WlmStatsResponse(clusterName, queryGroupStatsList, failedNodeExceptionList); + WlmStatsResponse queryGroupStatsResponse = new WlmStatsResponse(clusterName, wlmStatsList, failedNodeExceptionList); String responseString = queryGroupStatsResponse.toString(); assertEquals( "{\n" @@ -78,6 +79,7 @@ public void testToString() { + " \"query_groups\" : {\n" + " \"safjgagnaeekg-3r3fads\" : {\n" + " \"completions\" : 0,\n" + + " \"shard_completions\" : 0,\n" + " \"rejections\" : 0,\n" + " \"failures\" : 1,\n" + " \"total_cancellations\" : 0,\n" diff --git a/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java b/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java index c5cf0dac4f807..04d9ef7916626 100644 --- a/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java @@ -22,6 +22,7 @@ import org.opensearch.threadpool.Scheduler; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; import org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService; import org.opensearch.wlm.cancellation.TaskSelectionStrategy; import org.opensearch.wlm.stats.QueryGroupState; @@ -49,6 +50,7 @@ public class QueryGroupServiceTests extends OpenSearchTestCase { private QueryGroupService queryGroupService; + private TransportService mockTransportService; private QueryGroupTaskCancellationService mockCancellationService; private ClusterService mockClusterService; private ThreadPool mockThreadPool; @@ -61,6 +63,7 @@ public class QueryGroupServiceTests extends OpenSearchTestCase { public void setUp() throws Exception { super.setUp(); mockClusterService = Mockito.mock(ClusterService.class); + mockTransportService = Mockito.mock(TransportService.class); mockThreadPool = Mockito.mock(ThreadPool.class); mockScheduledFuture = Mockito.mock(Scheduler.Cancellable.class); mockWorkloadManagementSettings = Mockito.mock(WorkloadManagementSettings.class); @@ -71,6 +74,7 @@ public void setUp() throws Exception { queryGroupService = new QueryGroupService( mockCancellationService, + mockTransportService, mockClusterService, mockThreadPool, mockWorkloadManagementSettings, @@ -187,6 +191,7 @@ public void testRejectIfNeeded_whenQueryGroupIdIsNullOrDefaultOne() { queryGroupService = new QueryGroupService( mockCancellationService, + mockTransportService, mockClusterService, mockThreadPool, mockWorkloadManagementSettings, @@ -225,6 +230,7 @@ public void testRejectIfNeeded_whenQueryGroupIsSoftMode() { queryGroupService = new QueryGroupService( mockCancellationService, + mockTransportService, mockClusterService, mockThreadPool, mockWorkloadManagementSettings, @@ -261,6 +267,7 @@ public void testRejectIfNeeded_whenQueryGroupIsEnforcedMode_andNotBreaching() { queryGroupService = new QueryGroupService( mockCancellationService, + mockTransportService, mockClusterService, mockThreadPool, mockWorkloadManagementSettings, @@ -308,6 +315,7 @@ public void testRejectIfNeeded_whenQueryGroupIsEnforcedMode_andBreaching() { queryGroupService = new QueryGroupService( mockCancellationService, + mockTransportService, mockClusterService, mockThreadPool, mockWorkloadManagementSettings, @@ -352,6 +360,7 @@ public void testRejectIfNeeded_whenFeatureIsNotEnabled() { queryGroupService = new QueryGroupService( mockCancellationService, + mockTransportService, mockClusterService, mockThreadPool, mockWorkloadManagementSettings, @@ -375,6 +384,7 @@ public void testOnTaskCompleted() { mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); queryGroupService = new QueryGroupService( mockCancellationService, + mockTransportService, mockClusterService, mockThreadPool, mockWorkloadManagementSettings, @@ -421,6 +431,7 @@ public void testShouldSBPHandle() { mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); queryGroupService = new QueryGroupService( mockCancellationService, + mockTransportService, mockClusterService, mockThreadPool, mockWorkloadManagementSettings, diff --git a/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java b/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java index 9f7823e7085f1..580ef2e5e5999 100644 --- a/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java +++ b/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java @@ -26,6 +26,7 @@ import org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService; import org.opensearch.wlm.stats.QueryGroupState; import org.opensearch.wlm.stats.QueryGroupStats; +import org.opensearch.wlm.stats.WlmStats; import java.io.IOException; import java.util.ArrayList; @@ -60,9 +61,14 @@ public void setUp() throws Exception { mockClusterService = mock(ClusterService.class); mockWorkloadManagementSettings = mock(WorkloadManagementSettings.class); mockDiscoveryNode = mock(DiscoveryNode.class); + when(mockTransportService.getLocalNode()).thenReturn(mockDiscoveryNode); queryGroupStateMap = new HashMap<>(); testQueryGroupId = "safjgagnakg-3r3fads"; testThreadPool = new TestThreadPool("RejectionTestThreadPool"); + ClusterState mockClusterState = mock(ClusterState.class); + when(mockClusterService.state()).thenReturn(mockClusterState); + Metadata mockMetaData = mock(Metadata.class); + when(mockClusterState.metadata()).thenReturn(mockMetaData); queryGroupService = mock(QueryGroupService.class); sut = new QueryGroupRequestOperationListener(queryGroupService, testThreadPool); } @@ -90,7 +96,6 @@ public void testNonRejectionCase() { public void testValidQueryGroupRequestFailure() throws IOException { QueryGroupStats expectedStats = new QueryGroupStats( - mockDiscoveryNode, Map.of( testQueryGroupId, new QueryGroupStats.QueryGroupStatsHolder( @@ -123,7 +128,7 @@ public void testValidQueryGroupRequestFailure() throws IOException { ) ); - assertSuccess(testQueryGroupId, queryGroupStateMap, expectedStats, testQueryGroupId); + assertSuccess(testQueryGroupId, queryGroupStateMap, new WlmStats(mockDiscoveryNode, expectedStats), testQueryGroupId); } public void testMultiThreadedValidQueryGroupRequestFailures() { @@ -164,10 +169,11 @@ public void testMultiThreadedValidQueryGroupRequestFailures() { } }); - QueryGroupStats actualStats = queryGroupService.nodeStats(new HashSet<>(), null); + HashSet set = new HashSet<>(); + set.add("_all"); + WlmStats actualStats = queryGroupService.nodeStats(set, null); - QueryGroupStats expectedStats = new QueryGroupStats( - mockDiscoveryNode, + QueryGroupStats queryGroupStats = new QueryGroupStats( Map.of( testQueryGroupId, new QueryGroupStats.QueryGroupStatsHolder( @@ -200,12 +206,12 @@ public void testMultiThreadedValidQueryGroupRequestFailures() { ) ); + WlmStats expectedStats = new WlmStats(mockDiscoveryNode, queryGroupStats); assertEquals(expectedStats, actualStats); } public void testInvalidQueryGroupFailure() throws IOException { QueryGroupStats expectedStats = new QueryGroupStats( - mockDiscoveryNode, Map.of( testQueryGroupId, new QueryGroupStats.QueryGroupStatsHolder( @@ -238,14 +244,14 @@ public void testInvalidQueryGroupFailure() throws IOException { ) ); - assertSuccess(testQueryGroupId, queryGroupStateMap, expectedStats, "dummy-invalid-qg-id"); + assertSuccess(testQueryGroupId, queryGroupStateMap, new WlmStats(mockDiscoveryNode, expectedStats), "dummy-invalid-qg-id"); } private void assertSuccess( String testQueryGroupId, Map queryGroupStateMap, - QueryGroupStats expectedStats, + WlmStats expectedStats, String threadContextQG_Id ) { QueryGroupsStateAccessor stateAccessor = new QueryGroupsStateAccessor(queryGroupStateMap); @@ -266,11 +272,12 @@ private void assertSuccess( Collections.emptySet(), Collections.emptySet() ); - sut = new QueryGroupRequestOperationListener(queryGroupService, testThreadPool); sut.onRequestFailure(null, null); - QueryGroupStats actualStats = queryGroupService.nodeStats(new HashSet<>(), null); + HashSet set = new HashSet<>(); + set.add("_all"); + WlmStats actualStats = queryGroupService.nodeStats(set, null); assertEquals(expectedStats, actualStats); } diff --git a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java index e73cef0aed745..bcab0c4d124f5 100644 --- a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java +++ b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java @@ -25,7 +25,6 @@ import java.util.Map; import static java.util.Collections.emptyMap; -import static org.mockito.Mockito.mock; public class QueryGroupStatsTests extends AbstractWireSerializingTestCase { @@ -44,7 +43,7 @@ public void testToXContent() throws IOException { ) ); XContentBuilder builder = JsonXContent.contentBuilder(); - QueryGroupStats queryGroupStats = new QueryGroupStats(mock(DiscoveryNode.class), stats); + QueryGroupStats queryGroupStats = new QueryGroupStats(stats); builder.startObject(); queryGroupStats.toXContent(builder, ToXContent.EMPTY_PARAMS); builder.endObject(); @@ -87,6 +86,6 @@ protected QueryGroupStats createTestInstance() { DiscoveryNodeRole.BUILT_IN_ROLES, VersionUtils.randomCompatibleVersion(random(), Version.CURRENT) ); - return new QueryGroupStats(discoveryNode, stats); + return new QueryGroupStats(stats); } } diff --git a/server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java b/server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java new file mode 100644 index 0000000000000..604db8212e2aa --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/stats/WlmStatsTests.java @@ -0,0 +1,75 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm.stats; + +import org.opensearch.Version; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodeRole; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.io.stream.Writeable; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.test.AbstractWireSerializingTestCase; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.VersionUtils; +import org.opensearch.wlm.ResourceType; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; + +import static java.util.Collections.emptyMap; +import static org.mockito.Mockito.mock; + +class WlmStatsTests extends AbstractWireSerializingTestCase { + + public void testToXContent() throws IOException { + final Map stats = new HashMap<>(); + final String queryGroupId = "afakjklaj304041-afaka"; + stats.put( + queryGroupId, + new QueryGroupStats.QueryGroupStatsHolder( + 123456789, + 13, + 2, + 0, + 1213718, + Map.of(ResourceType.CPU, new QueryGroupStats.ResourceStats(0.3, 13, 2)) + ) + ); + XContentBuilder builder = JsonXContent.contentBuilder(); + QueryGroupStats queryGroupStats = new QueryGroupStats(stats); + WlmStats wlmStats = new WlmStats(mock(DiscoveryNode.class), queryGroupStats); + builder.startObject(); + wlmStats.toXContent(builder, ToXContent.EMPTY_PARAMS); + builder.endObject(); + assertEquals( + "{\"query_groups\":{\"afakjklaj304041-afaka\":{\"completions\":123456789,\"shard_completions\":1213718,\"rejections\":13,\"failures\":2,\"total_cancellations\":0,\"cpu\":{\"current_usage\":0.3,\"cancellations\":13,\"rejections\":2}}}}", + builder.toString() + ); + } + + @Override + protected Writeable.Reader instanceReader() { + return WlmStats::new; + } + + @Override + protected WlmStats createTestInstance() { + DiscoveryNode discoveryNode = new DiscoveryNode( + "node", + OpenSearchTestCase.buildNewFakeTransportAddress(), + emptyMap(), + DiscoveryNodeRole.BUILT_IN_ROLES, + VersionUtils.randomCompatibleVersion(random(), Version.CURRENT) + ); + QueryGroupStatsTests queryGroupStatsTests = new QueryGroupStatsTests(); + return new WlmStats(discoveryNode, queryGroupStatsTests.createTestInstance()); + } +}