Skip to content

Commit

Permalink
unit tests for resource limit group
Browse files Browse the repository at this point in the history
Signed-off-by: Ruirui Zhang <mariazrr@amazon.com>
  • Loading branch information
ruai0511 committed Apr 19, 2024
1 parent b722313 commit ea727b9
Show file tree
Hide file tree
Showing 27 changed files with 917 additions and 102 deletions.
Original file line number Diff line number Diff line change
@@ -1,3 +1,11 @@
/*
* 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.plugin.resource_limit_group;

import org.opensearch.action.ActionRequest;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,14 @@
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.core.rest.RestStatus;
import org.opensearch.core.xcontent.ToXContent;
import org.opensearch.plugin.resource_limit_group.*;
import org.opensearch.rest.*;
import org.opensearch.plugin.resource_limit_group.CreateResourceLimitGroupAction;
import org.opensearch.plugin.resource_limit_group.CreateResourceLimitGroupRequest;
import org.opensearch.plugin.resource_limit_group.CreateResourceLimitGroupResponse;
import org.opensearch.rest.BaseRestHandler;
import org.opensearch.rest.BytesRestResponse;
import org.opensearch.rest.RestChannel;
import org.opensearch.rest.RestRequest;
import org.opensearch.rest.RestResponse;
import org.opensearch.rest.action.RestResponseListener;

import java.io.IOException;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,14 @@
import org.opensearch.client.node.NodeClient;
import org.opensearch.core.rest.RestStatus;
import org.opensearch.core.xcontent.ToXContent;
import org.opensearch.plugin.resource_limit_group.*;
import org.opensearch.rest.*;
import org.opensearch.plugin.resource_limit_group.DeleteResourceLimitGroupAction;
import org.opensearch.plugin.resource_limit_group.DeleteResourceLimitGroupRequest;
import org.opensearch.plugin.resource_limit_group.DeleteResourceLimitGroupResponse;
import org.opensearch.rest.BaseRestHandler;
import org.opensearch.rest.BytesRestResponse;
import org.opensearch.rest.RestChannel;
import org.opensearch.rest.RestRequest;
import org.opensearch.rest.RestResponse;
import org.opensearch.rest.action.RestResponseListener;

import java.io.IOException;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,14 @@
import org.opensearch.client.node.NodeClient;
import org.opensearch.core.rest.RestStatus;
import org.opensearch.core.xcontent.ToXContent;
import org.opensearch.plugin.resource_limit_group.*;
import org.opensearch.rest.*;
import org.opensearch.plugin.resource_limit_group.GetResourceLimitGroupAction;
import org.opensearch.plugin.resource_limit_group.GetResourceLimitGroupRequest;
import org.opensearch.plugin.resource_limit_group.GetResourceLimitGroupResponse;
import org.opensearch.rest.BaseRestHandler;
import org.opensearch.rest.BytesRestResponse;
import org.opensearch.rest.RestChannel;
import org.opensearch.rest.RestRequest;
import org.opensearch.rest.RestResponse;
import org.opensearch.rest.action.RestResponseListener;

import java.io.IOException;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,14 @@
import org.opensearch.core.rest.RestStatus;
import org.opensearch.core.xcontent.ToXContent;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.plugin.resource_limit_group.*;
import org.opensearch.rest.*;
import org.opensearch.plugin.resource_limit_group.UpdateResourceLimitGroupAction;
import org.opensearch.plugin.resource_limit_group.UpdateResourceLimitGroupRequest;
import org.opensearch.plugin.resource_limit_group.UpdateResourceLimitGroupResponse;
import org.opensearch.rest.BaseRestHandler;
import org.opensearch.rest.BytesRestResponse;
import org.opensearch.rest.RestChannel;
import org.opensearch.rest.RestRequest;
import org.opensearch.rest.RestResponse;
import org.opensearch.rest.action.RestResponseListener;

import java.io.IOException;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,10 @@
import org.opensearch.common.settings.Settings;
import org.opensearch.core.action.ActionListener;
import org.opensearch.core.rest.RestStatus;
import org.opensearch.plugin.resource_limit_group.*;
import org.opensearch.plugin.resource_limit_group.CreateResourceLimitGroupResponse;
import org.opensearch.plugin.resource_limit_group.DeleteResourceLimitGroupResponse;
import org.opensearch.plugin.resource_limit_group.GetResourceLimitGroupResponse;
import org.opensearch.plugin.resource_limit_group.UpdateResourceLimitGroupResponse;

import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
Expand Down Expand Up @@ -103,7 +106,6 @@ public void update(ResourceLimitGroup resourceLimitGroup, String existingName, A
return;
}


if (resourceLimitGroup.getName() != null && currentGroupsMap.containsKey(resourceLimitGroup.getName())) {
logger.warn("Resource Limit Group already exists with the updated name: {}", resourceLimitGroup.getName());
Exception e = new RuntimeException("Resource Limit Group already exists with the provided name: " + resourceLimitGroup.getName());
Expand Down Expand Up @@ -171,21 +173,16 @@ public void update(ResourceLimitGroup resourceLimitGroup, String existingName, A
@Override
public void get(String name, ActionListener<GetResourceLimitGroupResponse> listener) {
ClusterState currentState = clusterService.state();
List<ResourceLimitGroup> resourceLimitGroups;
Map<String, ResourceLimitGroup> currentGroupsMap = currentState.metadata().resourceLimitGroups();
if (name == null || name.equals("")) {
resourceLimitGroups = new ArrayList<>(currentGroupsMap.values());
} else if (currentGroupsMap.containsKey(name)) {
resourceLimitGroups = List.of(currentGroupsMap.get(name));
} else {
List<ResourceLimitGroup> resultGroups = getFromClusterStateMetadata(name, currentState);
if (resultGroups.isEmpty() && name != null && !name.isEmpty()) {
logger.warn("No Resource Limit Group exists with the provided name: {}", name);
Exception e = new RuntimeException("No Resource Limit Group exists with the provided name: " + name);
GetResourceLimitGroupResponse response = new GetResourceLimitGroupResponse();
response.setRestStatus(RestStatus.NOT_FOUND);
listener.onFailure(e);
return;
}
GetResourceLimitGroupResponse response = new GetResourceLimitGroupResponse(resourceLimitGroups);
GetResourceLimitGroupResponse response = new GetResourceLimitGroupResponse(resultGroups);
response.setRestStatus(RestStatus.OK);
listener.onResponse(response);
}
Expand Down Expand Up @@ -213,11 +210,6 @@ public ThrottlingKey getClusterManagerThrottlingKey() {

@Override
public void onFailure(String source, Exception e) {
inflightCreateResourceLimitGroupRequestCount.decrementAndGet();
for (ResourceLimit rl: resourceLimitGroup.getResourceLimits()) {
String currResourceName = rl.getResourceName();
inflightResourceLimitValues.get(currResourceName).add(-getResourceLimitValue(currResourceName, resourceLimitGroup));
}
logger.warn("failed to save Resource Limit Group object due to error: {}, for source: {}", e.getMessage(), source);
CreateResourceLimitGroupResponse response = new CreateResourceLimitGroupResponse();
response.setRestStatus(RestStatus.FAILED_DEPENDENCY);
Expand All @@ -226,11 +218,6 @@ public void onFailure(String source, Exception e) {

@Override
public void clusterStateProcessed(String source, ClusterState oldState, ClusterState newState) {
inflightCreateResourceLimitGroupRequestCount.decrementAndGet();
for (ResourceLimit rl: resourceLimitGroup.getResourceLimits()) {
String currResourceName = rl.getResourceName();
inflightResourceLimitValues.get(currResourceName).add(-getResourceLimitValue(currResourceName, resourceLimitGroup));
}
CreateResourceLimitGroupResponse response = new CreateResourceLimitGroupResponse(resourceLimitGroup);
response.setRestStatus(RestStatus.OK);
listener.onResponse(response);
Expand Down Expand Up @@ -286,17 +273,32 @@ ClusterState saveResourceLimitGroupInClusterState(final ResourceLimitGroup resou
}
}
if (inflightCreateResourceLimitGroupRequestCount.incrementAndGet() + previousGroups.size() > maxResourceLimitGroupCount) {
restoreInflightValues(resourceLimitGroup);
logger.error("{} value exceeded its assigned limit of {}", RESOURCE_LIMIT_GROUP_COUNT_SETTING_NAME, maxResourceLimitGroupCount);
throw new RuntimeException("Can't create more than " + maxResourceLimitGroupCount + " Resource Limit Groups in the system");
}
if (!resourceNameWithThresholdExceeded.isEmpty()) {
restoreInflightValues(resourceLimitGroup);
logger.error("Total resource allocation for {} will go above the max limit of 1.0", resourceNameWithThresholdExceeded);
throw new RuntimeException("Total resource allocation for " + resourceNameWithThresholdExceeded+ " will go above the max limit of 1.0");
}

restoreInflightValues(resourceLimitGroup);
return ClusterState.builder(currentClusterState).metadata(Metadata.builder(metadata).put(resourceLimitGroup).build()).build();
}

/**
* This method restores the inflight values to be before the resource limit group is processed
* @param resourceLimitGroup - the resource limit group we're currently creating
*/
void restoreInflightValues(ResourceLimitGroup resourceLimitGroup) {
inflightCreateResourceLimitGroupRequestCount.decrementAndGet();
for (ResourceLimit rl: resourceLimitGroup.getResourceLimits()) {
String currResourceName = rl.getResourceName();
inflightResourceLimitValues.get(currResourceName).add(-getResourceLimitValue(currResourceName, resourceLimitGroup));
}
}

/**
* Modify cluster state to update the Resource Limit Group
* @param existingGroup {@link ResourceLimitGroup} - the existing resource limit group that we want to update
Expand Down Expand Up @@ -417,4 +419,16 @@ ClusterState deleteResourceLimitGroupInClusterState(final String name, final Clu
}
return ClusterState.builder(currentClusterState).metadata(Metadata.builder(metadata).resourceLimitGroups(resultGroups).build()).build();
}

List<ResourceLimitGroup> getFromClusterStateMetadata(String name, ClusterState currentState) {
Map<String, ResourceLimitGroup> currentGroupsMap = currentState.getMetadata().resourceLimitGroups();
List<ResourceLimitGroup> currentGroups = new ArrayList<>(currentGroupsMap.values());
List<ResourceLimitGroup> resultGroups = new ArrayList<>();
if (name == null || name.equals("")) {
resultGroups = currentGroups;
} else if (currentGroupsMap.containsKey(name)) {
resultGroups = List.of(currentGroupsMap.get(name));
}
return resultGroups;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
/*
* 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.plugin.resource_limit_group;

import org.opensearch.common.io.stream.BytesStreamOutput;
import org.opensearch.cluster.metadata.ResourceLimitGroup;
import org.opensearch.cluster.metadata.ResourceLimitGroup.ResourceLimit;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.test.OpenSearchTestCase;

import java.io.IOException;
import java.util.List;

import static org.opensearch.plugin.resource_limit_group.ResourceLimitGroupTestUtils.NAME_ONE;
import static org.opensearch.plugin.resource_limit_group.ResourceLimitGroupTestUtils.MONITOR;
import static org.opensearch.plugin.resource_limit_group.ResourceLimitGroupTestUtils.compareResourceLimits;

public class CreateResourceLimitGroupRequestTests extends OpenSearchTestCase {

public void testSerialization() throws IOException {
ResourceLimitGroup resourceLimitGroup = new ResourceLimitGroup(NAME_ONE, List.of(new ResourceLimit("jvm", 0.4)), MONITOR);
CreateResourceLimitGroupRequest request = new CreateResourceLimitGroupRequest(resourceLimitGroup);
BytesStreamOutput out = new BytesStreamOutput();
request.writeTo(out);
StreamInput streamInput = out.bytes().streamInput();
CreateResourceLimitGroupRequest otherRequest = new CreateResourceLimitGroupRequest(streamInput);
assertEquals(request.getName(), otherRequest.getName());
assertEquals(request.getResourceLimits().size(), otherRequest.getResourceLimits().size());
assertEquals(request.getEnforcement(), otherRequest.getEnforcement());
compareResourceLimits(request.getResourceLimits(), otherRequest.getResourceLimits());
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
/*
* 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.plugin.resource_limit_group;

import org.opensearch.cluster.metadata.ResourceLimitGroup;
import org.opensearch.cluster.metadata.ResourceLimitGroup.ResourceLimit;
import org.opensearch.common.io.stream.BytesStreamOutput;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.test.OpenSearchTestCase;

import java.io.IOException;
import java.util.List;

import static org.opensearch.plugin.resource_limit_group.ResourceLimitGroupTestUtils.NAME_ONE;
import static org.opensearch.plugin.resource_limit_group.ResourceLimitGroupTestUtils.MONITOR;
import static org.opensearch.plugin.resource_limit_group.ResourceLimitGroupTestUtils.compareResourceLimits;

public class CreateResourceLimitGroupResponseTests extends OpenSearchTestCase {

public void testSerialization() throws IOException {
ResourceLimitGroup resourceLimitGroup = new ResourceLimitGroup(NAME_ONE, List.of(new ResourceLimit("jvm", 0.4)), MONITOR);
CreateResourceLimitGroupResponse response = new CreateResourceLimitGroupResponse(resourceLimitGroup);
BytesStreamOutput out = new BytesStreamOutput();
response.writeTo(out);
StreamInput streamInput = out.bytes().streamInput();
CreateResourceLimitGroupResponse otherResponse = new CreateResourceLimitGroupResponse(streamInput);
assertEquals(response.getRestStatus(), otherResponse.getRestStatus());
ResourceLimitGroup responseGroup = response.getResourceLimitGroup();
ResourceLimitGroup otherResponseGroup = otherResponse.getResourceLimitGroup();
assertEquals(responseGroup.getName(), otherResponseGroup.getName());
assertEquals(responseGroup.getResourceLimits().size(), otherResponseGroup.getResourceLimits().size());
assertEquals(responseGroup.getEnforcement(), otherResponseGroup.getEnforcement());
compareResourceLimits(responseGroup.getResourceLimits(), otherResponseGroup.getResourceLimits());
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
/*
* 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.plugin.resource_limit_group;

import org.opensearch.common.io.stream.BytesStreamOutput;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.test.OpenSearchTestCase;

import java.io.IOException;

import static org.opensearch.plugin.resource_limit_group.ResourceLimitGroupTestUtils.NAME_ONE;

public class DeleteResourceLimitGroupRequestTests extends OpenSearchTestCase {

public void testSerialization() throws IOException {
DeleteResourceLimitGroupRequest request = new DeleteResourceLimitGroupRequest(NAME_ONE);
assertEquals(NAME_ONE, request.getName());
BytesStreamOutput out = new BytesStreamOutput();
request.writeTo(out);
StreamInput streamInput = out.bytes().streamInput();
DeleteResourceLimitGroupRequest otherRequest = new DeleteResourceLimitGroupRequest(streamInput);
assertEquals(request.getName(), otherRequest.getName());
}

public void testSerializationWithNull() throws IOException {
DeleteResourceLimitGroupRequest request = new DeleteResourceLimitGroupRequest((String) null);
assertNull(request.getName());
BytesStreamOutput out = new BytesStreamOutput();
request.writeTo(out);
StreamInput streamInput = out.bytes().streamInput();
DeleteResourceLimitGroupRequest otherRequest = new DeleteResourceLimitGroupRequest(streamInput);
assertEquals(request.getName(), otherRequest.getName());
}
}
Loading

0 comments on commit ea727b9

Please sign in to comment.