Fixed AdmissionControlSingleNodeTests flaky tests

Signed-off-by: Ajay Kumar Movva <movvaam@amazon.com>
opensearch-project · Mar 18, 2024 · 4717c53 · 4717c53
1 parent 4ef6a0b
commit 4717c53
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -19,7 +19,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Remote reindex: Add support for configurable retry mechanism ([#12561](https://github.com/opensearch-project/OpenSearch/pull/12561))
 - [Admission Control] Integrate IO Usage Tracker to the Resource Usage Collector Service and Emit IO Usage Stats ([#11880](https://github.com/opensearch-project/OpenSearch/pull/11880))
 - Tracing for deep search path ([#12103](https://github.com/opensearch-project/OpenSearch/pull/12103))
-- [Admission Control] Integrated IO Based AdmissionController to AdmissionControl Framework ([#12583](https://github.com/opensearch-project/OpenSearch/pull/12583))
 
 ### Dependencies
 - Bump `log4j-core` from 2.18.0 to 2.19.0
@@ -117,6 +116,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Make search query counters dynamic to support all query types ([#12601](https://github.com/opensearch-project/OpenSearch/pull/12601))
 - [Tiered caching] Add policies controlling which values can enter pluggable caches [EXPERIMENTAL] ([#12542](https://github.com/opensearch-project/OpenSearch/pull/12542))
 - [Tiered caching] Add Stale keys Management and CacheCleaner to IndicesRequestCache ([#12625](https://github.com/opensearch-project/OpenSearch/pull/12625))
+- [Admission Control] Integrated IO Based AdmissionController to AdmissionControl Framework ([#12583](https://github.com/opensearch-project/OpenSearch/pull/12583))
 
 ### Dependencies
 - Bump `peter-evans/find-comment` from 2 to 3 ([#12288](https://github.com/opensearch-project/OpenSearch/pull/12288))

diff --git a/...t/java/org/opensearch/ratelimitting/admissioncontrol/AdmissionControlSingleNodeTests.java b/...t/java/org/opensearch/ratelimitting/admissioncontrol/AdmissionControlSingleNodeTests.java
@@ -34,9 +34,11 @@
 import java.util.Map;
 
 import static org.opensearch.ratelimitting.admissioncontrol.AdmissionControlSettings.ADMISSION_CONTROL_TRANSPORT_LAYER_MODE;
+import static org.opensearch.ratelimitting.admissioncontrol.settings.CpuBasedAdmissionControllerSettings.CPU_BASED_ADMISSION_CONTROLLER_TRANSPORT_LAYER_MODE;
 import static org.opensearch.ratelimitting.admissioncontrol.settings.CpuBasedAdmissionControllerSettings.INDEXING_CPU_USAGE_LIMIT;
 import static org.opensearch.ratelimitting.admissioncontrol.settings.CpuBasedAdmissionControllerSettings.SEARCH_CPU_USAGE_LIMIT;
 import static org.opensearch.ratelimitting.admissioncontrol.settings.IoBasedAdmissionControllerSettings.INDEXING_IO_USAGE_LIMIT;
+import static org.opensearch.ratelimitting.admissioncontrol.settings.IoBasedAdmissionControllerSettings.IO_BASED_ADMISSION_CONTROLLER_TRANSPORT_LAYER_MODE;
 import static org.opensearch.ratelimitting.admissioncontrol.settings.IoBasedAdmissionControllerSettings.SEARCH_IO_USAGE_LIMIT;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
 import static org.hamcrest.Matchers.is;
@@ -72,7 +74,7 @@ protected Settings nodeSettings() {
             .put(ResourceTrackerSettings.GLOBAL_CPU_USAGE_AC_WINDOW_DURATION_SETTING.getKey(), TimeValue.timeValueMillis(500))
             .put(ResourceTrackerSettings.GLOBAL_JVM_USAGE_AC_WINDOW_DURATION_SETTING.getKey(), TimeValue.timeValueMillis(500))
             .put(ResourceTrackerSettings.GLOBAL_IO_USAGE_AC_WINDOW_DURATION_SETTING.getKey(), TimeValue.timeValueMillis(5000))
-            .put(ADMISSION_CONTROL_TRANSPORT_LAYER_MODE.getKey(), AdmissionControlMode.ENFORCED)
+            .put(CPU_BASED_ADMISSION_CONTROLLER_TRANSPORT_LAYER_MODE.getKey(), AdmissionControlMode.ENFORCED.getMode())
             .put(SEARCH_CPU_USAGE_LIMIT.getKey(), 0)
             .put(INDEXING_CPU_USAGE_LIMIT.getKey(), 0)
             .build();
@@ -143,8 +145,8 @@ public void testAdmissionControlRejectionEnforcedMode() throws Exception {
                 .put(super.nodeSettings())
                 .put(SEARCH_IO_USAGE_LIMIT.getKey(), 0)
                 .put(INDEXING_IO_USAGE_LIMIT.getKey(), 0)
-                .put(SEARCH_CPU_USAGE_LIMIT.getKey(), 101)
-                .put(INDEXING_CPU_USAGE_LIMIT.getKey(), 101)
+                .put(CPU_BASED_ADMISSION_CONTROLLER_TRANSPORT_LAYER_MODE.getKey(), AdmissionControlMode.DISABLED.getMode())
+                .put(IO_BASED_ADMISSION_CONTROLLER_TRANSPORT_LAYER_MODE.getKey(), AdmissionControlMode.ENFORCED.getMode())
         );
         assertAcked(client().admin().cluster().updateSettings(updateSettingsRequest).actionGet());
 
@@ -211,7 +213,7 @@ public void testAdmissionControlRejectionMonitorOnlyMode() throws Exception {
         updateSettingsRequest.transientSettings(
             Settings.builder()
                 .put(super.nodeSettings())
-                .put(ADMISSION_CONTROL_TRANSPORT_LAYER_MODE.getKey(), AdmissionControlMode.MONITOR.getMode())
+                .put(CPU_BASED_ADMISSION_CONTROLLER_TRANSPORT_LAYER_MODE.getKey(), AdmissionControlMode.MONITOR.getMode())
         );
         assertAcked(client().admin().cluster().updateSettings(updateSettingsRequest).actionGet());
 
@@ -252,6 +254,8 @@ public void testAdmissionControlRejectionMonitorOnlyMode() throws Exception {
                 .put(INDEXING_IO_USAGE_LIMIT.getKey(), 0)
                 .put(SEARCH_CPU_USAGE_LIMIT.getKey(), 101)
                 .put(INDEXING_CPU_USAGE_LIMIT.getKey(), 101)
+                .put(CPU_BASED_ADMISSION_CONTROLLER_TRANSPORT_LAYER_MODE.getKey(), AdmissionControlMode.DISABLED.getMode())
+                .put(IO_BASED_ADMISSION_CONTROLLER_TRANSPORT_LAYER_MODE.getKey(), AdmissionControlMode.MONITOR.getMode())
         );
         assertAcked(client().admin().cluster().updateSettings(updateSettingsRequest).actionGet());
         bulk = client().prepareBulk();
@@ -307,7 +311,9 @@ public void testAdmissionControlRejectionDisabledMode() throws Exception {
 
         ClusterUpdateSettingsRequest updateSettingsRequest = new ClusterUpdateSettingsRequest();
         updateSettingsRequest.transientSettings(
-            Settings.builder().put(super.nodeSettings()).put(ADMISSION_CONTROL_TRANSPORT_LAYER_MODE.getKey(), AdmissionControlMode.DISABLED)
+            Settings.builder()
+                .put(super.nodeSettings())
+                .put(CPU_BASED_ADMISSION_CONTROLLER_TRANSPORT_LAYER_MODE.getKey(), AdmissionControlMode.DISABLED.getMode())
         );
         assertAcked(client().admin().cluster().updateSettings(updateSettingsRequest).actionGet());
 
@@ -334,10 +340,7 @@ public void testAdmissionControlRejectionDisabledMode() throws Exception {
         updateSettingsRequest.transientSettings(
             Settings.builder()
                 .put(super.nodeSettings())
-                .put(SEARCH_IO_USAGE_LIMIT.getKey(), 0)
-                .put(INDEXING_IO_USAGE_LIMIT.getKey(), 0)
-                .put(SEARCH_CPU_USAGE_LIMIT.getKey(), 101)
-                .put(INDEXING_CPU_USAGE_LIMIT.getKey(), 101)
+                .put(IO_BASED_ADMISSION_CONTROLLER_TRANSPORT_LAYER_MODE.getKey(), AdmissionControlMode.DISABLED.getMode())
         );
         assertAcked(client().admin().cluster().updateSettings(updateSettingsRequest).actionGet());
         bulk = client().prepareBulk();
@@ -376,7 +379,7 @@ public void testAdmissionControlWithinLimits() throws Exception {
         updateSettingsRequest.transientSettings(
             Settings.builder()
                 .put(super.nodeSettings())
-                .put(ADMISSION_CONTROL_TRANSPORT_LAYER_MODE.getKey(), AdmissionControlMode.ENFORCED)
+                .put(ADMISSION_CONTROL_TRANSPORT_LAYER_MODE.getKey(), AdmissionControlMode.ENFORCED.getMode())
                 .put(SEARCH_CPU_USAGE_LIMIT.getKey(), 101)
                 .put(INDEXING_CPU_USAGE_LIMIT.getKey(), 101)
                 .put(SEARCH_IO_USAGE_LIMIT.getKey(), 101)