Skip to content

Commit a08ac91

Browse files
bharath-techiebharath-techie
committed
Fixes for tests
Signed-off-by: Bharathwaj G <bharath78910@gmail.com>
1 parent d9d07eb commit a08ac91

14 files changed

Lines changed: 335 additions & 35 deletions

File tree

‎server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/AdmissionControlService.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
import java.util.concurrent.ConcurrentHashMap;
2626
import java.util.concurrent.ConcurrentMap;
2727

28-
import static org.opensearch.ratelimitting.admissioncontrol.settings.CPUBasedAdmissionControllerSettings.CPU_BASED_ADMISSION_CONTROLLER;
28+
import static org.opensearch.ratelimitting.admissioncontrol.controllers.CPUBasedAdmissionController.CPU_BASED_ADMISSION_CONTROLLER;
2929

3030
/**
3131
* Admission control Service that bootstraps and manages all the Admission Controllers in OpenSearch.

‎server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/controllers/AdmissionController.java‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,16 +34,14 @@ public abstract class AdmissionController {
3434
/**
3535
* @param admissionControllerName name of the admissionController
3636
* @param resourceUsageCollectorService instance used to get resource usage stats of the node
37-
* @param rejectionCount initialised rejectionCount value for AdmissionController
3837
* @param clusterService
3938
*/
4039
public AdmissionController(
4140
String admissionControllerName,
4241
ResourceUsageCollectorService resourceUsageCollectorService,
43-
AtomicLong rejectionCount,
4442
ClusterService clusterService
4543
) {
46-
this.rejectionCount = rejectionCount;
44+
this.rejectionCount = new AtomicLong(0);
4745
this.admissionControllerName = admissionControllerName;
4846
this.resourceUsageCollectorService = resourceUsageCollectorService;
4947
this.clusterService = clusterService;

‎server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/controllers/CPUBasedAdmissionController.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,13 @@
1919
import org.opensearch.ratelimitting.admissioncontrol.settings.CPUBasedAdmissionControllerSettings;
2020

2121
import java.util.Optional;
22-
import java.util.concurrent.atomic.AtomicLong;
2322

2423
/**
2524
* Class for CPU Based Admission Controller in OpenSearch, which aims to provide CPU utilisation admission control.
2625
* It provides methods to apply admission control if configured limit has been reached
2726
*/
2827
public class CPUBasedAdmissionController extends AdmissionController {
28+
public static final String CPU_BASED_ADMISSION_CONTROLLER = "global_cpu_usage";
2929
private static final Logger LOGGER = LogManager.getLogger(CPUBasedAdmissionController.class);
3030
public CPUBasedAdmissionControllerSettings settings;
3131

@@ -41,7 +41,7 @@ public CPUBasedAdmissionController(
4141
ClusterService clusterService,
4242
Settings settings
4343
) {
44-
super(admissionControllerName, resourceUsageCollectorService, new AtomicLong(0), clusterService);
44+
super(admissionControllerName, resourceUsageCollectorService, clusterService);
4545
this.settings = new CPUBasedAdmissionControllerSettings(clusterService.getClusterSettings(), settings);
4646
}
4747

‎server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/settings/CPUBasedAdmissionControllerSettings.java‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
* @opensearch.internal
2020
*/
2121
public class CPUBasedAdmissionControllerSettings {
22-
public static final String CPU_BASED_ADMISSION_CONTROLLER = "global_cpu_usage";
2322

2423
/**
2524
* Default parameters for the CPUBasedAdmissionControllerSettings

‎server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/stats/AdmissionControlStats.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ public AdmissionControlStats(List<AdmissionControllerStats> admissionControllerS
3838
* @throws IOException if an I/O error occurs
3939
*/
4040
public AdmissionControlStats(StreamInput in) throws IOException {
41-
this.admissionControllerStatsList = in.readNamedWriteableList(AdmissionControllerStats.class);
41+
this.admissionControllerStatsList = in.readList(AdmissionControllerStats::new);
4242
}
4343

4444
/**
@@ -51,6 +51,10 @@ public void writeTo(StreamOutput out) throws IOException {
5151
out.writeList(this.admissionControllerStatsList);
5252
}
5353

54+
public List<AdmissionControllerStats> getAdmissionControllerStatsList() {
55+
return admissionControllerStatsList;
56+
}
57+
5458
/**
5559
* @param builder
5660
* @param params
@@ -62,7 +66,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
6266
builder.startObject("admission_control");
6367
this.admissionControllerStatsList.forEach(stats -> {
6468
try {
65-
builder.field(stats.getWriteableName(), stats);
69+
builder.field(stats.getAdmissionControllerName(), stats);
6670
} catch (IOException e) {
6771
throw new RuntimeException(e);
6872
}

‎server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/stats/AdmissionControllerStats.java‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,9 @@
88

99
package org.opensearch.ratelimitting.admissioncontrol.stats;
1010

11-
import org.opensearch.core.common.io.stream.NamedWriteable;
1211
import org.opensearch.core.common.io.stream.StreamInput;
1312
import org.opensearch.core.common.io.stream.StreamOutput;
13+
import org.opensearch.core.common.io.stream.Writeable;
1414
import org.opensearch.core.xcontent.ToXContentFragment;
1515
import org.opensearch.core.xcontent.XContentBuilder;
1616
import org.opensearch.ratelimitting.admissioncontrol.controllers.AdmissionController;
@@ -21,7 +21,7 @@
2121
/**
2222
* Class for admission controller ( such as CPU ) stats which includes rejection count for each action type
2323
*/
24-
public class AdmissionControllerStats implements NamedWriteable, ToXContentFragment {
24+
public class AdmissionControllerStats implements Writeable, ToXContentFragment {
2525
public Map<String, Long> rejectionCount;
2626
public String admissionControllerName;
2727

@@ -35,11 +35,14 @@ public AdmissionControllerStats(StreamInput in) throws IOException {
3535
this.admissionControllerName = in.readString();
3636
}
3737

38-
@Override
39-
public String getWriteableName() {
38+
public String getAdmissionControllerName() {
4039
return admissionControllerName;
4140
}
4241

42+
public Map<String, Long> getRejectionCount() {
43+
return rejectionCount;
44+
}
45+
4346
/**
4447
* Write this into the {@linkplain StreamOutput}.
4548
*

‎server/src/test/java/org/opensearch/action/admin/cluster/node/stats/NodeStatsTests.java‎

Lines changed: 61 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,11 @@
6363
import org.opensearch.node.NodeResourceUsageStats;
6464
import org.opensearch.node.NodesResourceUsageStats;
6565
import org.opensearch.node.ResponseCollectorService;
66+
import org.opensearch.ratelimitting.admissioncontrol.controllers.AdmissionController;
67+
import org.opensearch.ratelimitting.admissioncontrol.controllers.CPUBasedAdmissionController;
68+
import org.opensearch.ratelimitting.admissioncontrol.enums.AdmissionControlActionType;
69+
import org.opensearch.ratelimitting.admissioncontrol.stats.AdmissionControlStats;
70+
import org.opensearch.ratelimitting.admissioncontrol.stats.AdmissionControllerStats;
6671
import org.opensearch.script.ScriptCacheStats;
6772
import org.opensearch.script.ScriptStats;
6873
import org.opensearch.test.OpenSearchTestCase;
@@ -516,15 +521,44 @@ public void testSerialization() throws IOException {
516521
assertEquals(replicationStats.getTotalBytesBehind(), deserializedReplicationStats.getTotalBytesBehind());
517522
assertEquals(replicationStats.getMaxReplicationLag(), deserializedReplicationStats.getMaxReplicationLag());
518523
}
524+
AdmissionControlStats admissionControlStats = nodeStats.getAdmissionControlStats();
525+
AdmissionControlStats deserializedAdmissionControlStats = deserializedNodeStats.getAdmissionControlStats();
526+
if (admissionControlStats == null) {
527+
assertNull(deserializedAdmissionControlStats);
528+
} else {
529+
assertEquals(
530+
admissionControlStats.getAdmissionControllerStatsList().size(),
531+
deserializedAdmissionControlStats.getAdmissionControllerStatsList().size()
532+
);
533+
AdmissionControllerStats admissionControllerStats = admissionControlStats.getAdmissionControllerStatsList().get(0);
534+
AdmissionControllerStats deserializedAdmissionControllerStats = deserializedAdmissionControlStats
535+
.getAdmissionControllerStatsList()
536+
.get(0);
537+
assertEquals(
538+
admissionControllerStats.getAdmissionControllerName(),
539+
deserializedAdmissionControllerStats.getAdmissionControllerName()
540+
);
541+
assertEquals(1, (long) admissionControllerStats.getRejectionCount().get(AdmissionControlActionType.SEARCH.getType()));
542+
assertEquals(
543+
admissionControllerStats.getRejectionCount().get(AdmissionControlActionType.SEARCH.getType()),
544+
deserializedAdmissionControllerStats.getRejectionCount().get(AdmissionControlActionType.SEARCH.getType())
545+
);
546+
547+
assertEquals(2, (long) admissionControllerStats.getRejectionCount().get(AdmissionControlActionType.INDEXING.getType()));
548+
assertEquals(
549+
admissionControllerStats.getRejectionCount().get(AdmissionControlActionType.INDEXING.getType()),
550+
deserializedAdmissionControllerStats.getRejectionCount().get(AdmissionControlActionType.INDEXING.getType())
551+
);
552+
}
519553
}
520554
}
521555
}
522556

523-
public static NodeStats createNodeStats() {
557+
public static NodeStats createNodeStats() throws IOException {
524558
return createNodeStats(false);
525559
}
526560

527-
public static NodeStats createNodeStats(boolean remoteStoreStats) {
561+
public static NodeStats createNodeStats(boolean remoteStoreStats) throws IOException {
528562
DiscoveryNode node = new DiscoveryNode(
529563
"test_node",
530564
buildNewFakeTransportAddress(),
@@ -834,6 +868,29 @@ public static NodeStats createNodeStats(boolean remoteStoreStats) {
834868
clusterManagerThrottlingStats = new ClusterManagerThrottlingStats();
835869
clusterManagerThrottlingStats.onThrottle("test-task", randomInt());
836870
}
871+
872+
AdmissionControlStats admissionControlStats = null;
873+
if (frequently()) {
874+
AdmissionController admissionController = new AdmissionController(
875+
CPUBasedAdmissionController.CPU_BASED_ADMISSION_CONTROLLER,
876+
null,
877+
null
878+
) {
879+
@Override
880+
public void apply(String action, AdmissionControlActionType admissionControlActionType) {
881+
return;
882+
}
883+
};
884+
admissionController.addRejectionCount(AdmissionControlActionType.SEARCH.getType(), 1);
885+
admissionController.addRejectionCount(AdmissionControlActionType.INDEXING.getType(), 2);
886+
AdmissionControllerStats stats = new AdmissionControllerStats(
887+
admissionController,
888+
CPUBasedAdmissionController.CPU_BASED_ADMISSION_CONTROLLER
889+
);
890+
List<AdmissionControllerStats> statsList = new ArrayList();
891+
statsList.add(stats);
892+
admissionControlStats = new AdmissionControlStats(statsList);
893+
}
837894
ScriptCacheStats scriptCacheStats = scriptStats != null ? scriptStats.toScriptCacheStats() : null;
838895

839896
WeightedRoutingStats weightedRoutingStats = null;
@@ -871,7 +928,8 @@ public static NodeStats createNodeStats(boolean remoteStoreStats) {
871928
null,
872929
null,
873930
segmentReplicationRejectionStats,
874-
null
931+
null,
932+
admissionControlStats
875933
);
876934
}
877935

‎server/src/test/java/org/opensearch/action/admin/cluster/stats/ClusterStatsNodesTests.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
import org.opensearch.core.xcontent.MediaTypeRegistry;
4242
import org.opensearch.test.OpenSearchTestCase;
4343

44+
import java.io.IOException;
4445
import java.util.Arrays;
4546
import java.util.Collections;
4647
import java.util.Iterator;
@@ -87,7 +88,13 @@ public void testNetworkTypesToXContent() throws Exception {
8788
}
8889

8990
public void testIngestStats() throws Exception {
90-
NodeStats nodeStats = randomValueOtherThanMany(n -> n.getIngestStats() == null, NodeStatsTests::createNodeStats);
91+
NodeStats nodeStats = randomValueOtherThanMany(n -> n.getIngestStats() == null, () -> {
92+
try {
93+
return NodeStatsTests.createNodeStats();
94+
} catch (IOException e) {
95+
throw new RuntimeException(e);
96+
}
97+
});
9198

9299
SortedMap<String, long[]> processorStats = new TreeMap<>();
93100
nodeStats.getIngestStats().getProcessorStats().values().forEach(stats -> {

‎server/src/test/java/org/opensearch/cluster/DiskUsageTests.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,7 @@ public void testFillDiskUsage() {
193193
null,
194194
null,
195195
null,
196+
null,
196197
null
197198
),
198199
new NodeStats(
@@ -222,6 +223,7 @@ public void testFillDiskUsage() {
222223
null,
223224
null,
224225
null,
226+
null,
225227
null
226228
),
227229
new NodeStats(
@@ -251,6 +253,7 @@ public void testFillDiskUsage() {
251253
null,
252254
null,
253255
null,
256+
null,
254257
null
255258
)
256259
);
@@ -311,6 +314,7 @@ public void testFillDiskUsageSomeInvalidValues() {
311314
null,
312315
null,
313316
null,
317+
null,
314318
null
315319
),
316320
new NodeStats(
@@ -340,6 +344,7 @@ public void testFillDiskUsageSomeInvalidValues() {
340344
null,
341345
null,
342346
null,
347+
null,
343348
null
344349
),
345350
new NodeStats(
@@ -369,6 +374,7 @@ public void testFillDiskUsageSomeInvalidValues() {
369374
null,
370375
null,
371376
null,
377+
null,
372378
null
373379
)
374380
);

‎server/src/test/java/org/opensearch/ratelimitting/admissioncontrol/AdmissionControlServiceTests.java‎

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ public void testAdmissionControllerSettings() {
6868
List<AdmissionController> admissionControllerList = admissionControlService.getAdmissionControllers();
6969
assertEquals(admissionControllerList.size(), 1);
7070
CPUBasedAdmissionController cpuBasedAdmissionController = (CPUBasedAdmissionController) admissionControlService
71-
.getAdmissionController(CPUBasedAdmissionControllerSettings.CPU_BASED_ADMISSION_CONTROLLER);
71+
.getAdmissionController(CPUBasedAdmissionController.CPU_BASED_ADMISSION_CONTROLLER);
7272
assertEquals(
7373
admissionControlSettings.isTransportLayerAdmissionControlEnabled(),
7474
cpuBasedAdmissionController.isEnabledForTransportLayer(
@@ -119,7 +119,7 @@ public void testApplyAdmissionControllerEnabled() {
119119
admissionControlService = new AdmissionControlService(Settings.EMPTY, clusterService, threadPool, null);
120120
admissionControlService.applyTransportAdmissionControl(this.action, null);
121121
assertEquals(
122-
admissionControlService.getAdmissionController(CPUBasedAdmissionControllerSettings.CPU_BASED_ADMISSION_CONTROLLER)
122+
admissionControlService.getAdmissionController(CPUBasedAdmissionController.CPU_BASED_ADMISSION_CONTROLLER)
123123
.getRejectionCount(AdmissionControlActionType.INDEXING.getType()),
124124
0
125125
);
@@ -131,13 +131,28 @@ public void testApplyAdmissionControllerEnabled() {
131131
)
132132
.build();
133133
clusterService.getClusterSettings().applySettings(settings);
134-
admissionControlService.applyTransportAdmissionControl(this.action, null);
135134
List<AdmissionController> admissionControllerList = admissionControlService.getAdmissionControllers();
136135
assertEquals(admissionControllerList.size(), 1);
136+
}
137+
138+
public void testApplyAdmissionControllerEnforced() {
139+
this.action = "indices:data/write/bulk[s][p]";
140+
admissionControlService = new AdmissionControlService(Settings.EMPTY, clusterService, threadPool, null);
141+
admissionControlService.applyTransportAdmissionControl(this.action, null);
137142
assertEquals(
138-
admissionControlService.getAdmissionController(CPUBasedAdmissionControllerSettings.CPU_BASED_ADMISSION_CONTROLLER)
143+
admissionControlService.getAdmissionController(CPUBasedAdmissionController.CPU_BASED_ADMISSION_CONTROLLER)
139144
.getRejectionCount(AdmissionControlActionType.INDEXING.getType()),
140-
1
145+
0
141146
);
147+
148+
Settings settings = Settings.builder()
149+
.put(
150+
CPUBasedAdmissionControllerSettings.CPU_BASED_ADMISSION_CONTROLLER_TRANSPORT_LAYER_MODE.getKey(),
151+
AdmissionControlMode.MONITOR.getMode()
152+
)
153+
.build();
154+
clusterService.getClusterSettings().applySettings(settings);
155+
List<AdmissionController> admissionControllerList = admissionControlService.getAdmissionControllers();
156+
assertEquals(admissionControllerList.size(), 1);
142157
}
143158
}

0 commit comments

Comments
 (0)