From d92f20b582ecf2f7703a3342c844c0ae64ae7614 Mon Sep 17 00:00:00 2001 From: Ajay Kumar Movva Date: Thu, 19 Oct 2023 18:01:34 +0530 Subject: [PATCH] Addressing Comments Signed-off-by: Ajay Kumar Movva --- .../common/network/NetworkModule.java | 14 +-- .../main/java/org/opensearch/node/Node.java | 6 +- .../AdmissionControlService.java | 6 +- .../controllers/AdmissionController.java | 39 ++++-- .../CPUBasedAdmissionController.java | 51 ++------ .../CPUBasedAdmissionControllerSettings.java | 10 +- .../common/network/NetworkModuleTests.java | 119 ++++++++++++++---- .../AdmissionControlServiceTests.java | 21 ++-- .../CPUBasedAdmissionControllerTests.java | 8 +- 9 files changed, 169 insertions(+), 105 deletions(-) diff --git a/server/src/main/java/org/opensearch/common/network/NetworkModule.java b/server/src/main/java/org/opensearch/common/network/NetworkModule.java index fe16867fca195..486fc55e5f58f 100644 --- a/server/src/main/java/org/opensearch/common/network/NetworkModule.java +++ b/server/src/main/java/org/opensearch/common/network/NetworkModule.java @@ -149,9 +149,13 @@ public NetworkModule( NetworkService networkService, HttpServerTransport.Dispatcher dispatcher, ClusterSettings clusterSettings, - Tracer tracer + Tracer tracer, + List coreTransportInterceptors ) { this.settings = settings; + if (coreTransportInterceptors != null) { + coreTransportInterceptors.forEach(this::registerTransportInterceptor); + } for (NetworkPlugin plugin : plugins) { Map> httpTransportFactory = plugin.getHttpTransports( settings, @@ -267,14 +271,6 @@ private void registerTransportInterceptor(TransportInterceptor interceptor) { this.transportInterceptors.add(Objects.requireNonNull(interceptor, "interceptor must not be null")); } - /** - * Registers a new {@link TransportInterceptor} - * This method used to register CoreInterceptors before the plugin interceptors - */ - public void registerCoreTransportInterceptor(TransportInterceptor interceptor) { - this.transportInterceptors.add(0, Objects.requireNonNull(interceptor, "interceptor must not be null")); - } - /** * Returns a composite {@link TransportInterceptor} containing all registered interceptors * @see #registerTransportInterceptor(TransportInterceptor) diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index b22d912b471ed..0a7968c87f9eb 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -893,6 +893,7 @@ protected Node( final RestController restController = actionModule.getRestController(); + List coreTransportInterceptors = new ArrayList<>(); final AdmissionControlService admissionControlService = new AdmissionControlService( settings, clusterService.getClusterSettings(), @@ -903,6 +904,7 @@ protected Node( admissionControlService ); + coreTransportInterceptors.add(admissionControlTransportInterceptor); final NetworkModule networkModule = new NetworkModule( settings, pluginsService.filterPlugins(NetworkPlugin.class), @@ -915,9 +917,9 @@ protected Node( networkService, restController, clusterService.getClusterSettings(), - tracer + tracer, + coreTransportInterceptors ); - networkModule.registerCoreTransportInterceptor(admissionControlTransportInterceptor); Collection>> indexTemplateMetadataUpgraders = pluginsService.filterPlugins( Plugin.class diff --git a/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/AdmissionControlService.java b/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/AdmissionControlService.java index ded50e86d2ded..2cc409b0e4465 100644 --- a/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/AdmissionControlService.java +++ b/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/AdmissionControlService.java @@ -52,7 +52,7 @@ public AdmissionControlService(Settings settings, ClusterSettings clusterSetting /** * Initialise and Register all the admissionControllers */ - public void initialise() { + private void initialise() { // Initialise different type of admission controllers registerAdmissionController(CPU_BASED_ADMISSION_CONTROLLER); } @@ -69,14 +69,14 @@ public void applyTransportAdmissionControl(String action) { * @param admissionControllerName admissionControllerName to register into the service. */ public void registerAdmissionController(String admissionControllerName) { - AdmissionController admissionController = this.getControllerImplementation(admissionControllerName); + AdmissionController admissionController = this.controllerFactory(admissionControllerName); this.ADMISSION_CONTROLLERS.put(admissionControllerName, admissionController); } /** * @return AdmissionController Instance */ - private AdmissionController getControllerImplementation(String admissionControllerName) { + private AdmissionController controllerFactory(String admissionControllerName) { switch (admissionControllerName) { case CPU_BASED_ADMISSION_CONTROLLER: return new CPUBasedAdmissionController(admissionControllerName, this.settings, this.clusterSettings); diff --git a/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/controllers/AdmissionController.java b/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/controllers/AdmissionController.java index d20d3260f1d35..00564a9967f31 100644 --- a/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/controllers/AdmissionController.java +++ b/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/controllers/AdmissionController.java @@ -8,38 +8,63 @@ package org.opensearch.ratelimitting.admissioncontrol.controllers; +import org.opensearch.ratelimitting.admissioncontrol.enums.AdmissionControlMode; + +import java.util.concurrent.atomic.AtomicLong; + /** - * Interface for Admission Controller in OpenSearch, which aims to provide resource based request admission control. + * Abstract class for Admission Controller in OpenSearch, which aims to provide resource based request admission control. * It provides methods for any tracking-object that can be incremented (such as memory size), * and admission control can be applied if configured limit has been reached */ -public interface AdmissionController { +public abstract class AdmissionController { + + private final AtomicLong rejectionCount; + private final String admissionControllerName; + + /** + * + * @param rejectionCount initialised rejectionCount value for AdmissionController + * @param admissionControllerName name of the admissionController + */ + public AdmissionController(AtomicLong rejectionCount, String admissionControllerName) { + this.rejectionCount = rejectionCount; + this.admissionControllerName = admissionControllerName; + } /** * Return the current state of the admission controller * @return true if admissionController is enabled for the transport layer else false */ - boolean isEnabledForTransportLayer(); + public boolean isEnabledForTransportLayer(AdmissionControlMode admissionControlMode) { + return admissionControlMode != AdmissionControlMode.DISABLED; + } /** * Increment the tracking-objects and apply the admission control if threshold is breached. * Mostly applicable while applying admission controller */ - void apply(String action); + public abstract void apply(String action); /** * @return name of the admission-controller */ - String getName(); + public String getName() { + return this.admissionControllerName; + } /** * Adds the rejection count for the controller. Primarily used when copying controller states. * @param count To add the value of the tracking resource object as the provided count */ - void addRejectionCount(long count); + public void addRejectionCount(long count) { + this.rejectionCount.addAndGet(count); + } /** * @return current value of the rejection count metric tracked by the admission-controller. */ - long getRejectionCount(); + public long getRejectionCount() { + return this.rejectionCount.get(); + } } diff --git a/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/controllers/CPUBasedAdmissionController.java b/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/controllers/CPUBasedAdmissionController.java index e7258d9ba145e..3a8956b2cce87 100644 --- a/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/controllers/CPUBasedAdmissionController.java +++ b/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/controllers/CPUBasedAdmissionController.java @@ -12,7 +12,6 @@ import org.apache.logging.log4j.Logger; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; -import org.opensearch.ratelimitting.admissioncontrol.enums.AdmissionControlMode; import org.opensearch.ratelimitting.admissioncontrol.settings.CPUBasedAdmissionControllerSettings; import java.util.concurrent.atomic.AtomicLong; @@ -21,29 +20,17 @@ * Class for CPU Based Admission Controller in OpenSearch, which aims to provide CPU utilisation admission control. * It provides methods to apply admission control if configured limit has been reached */ -public class CPUBasedAdmissionController implements AdmissionController { +public class CPUBasedAdmissionController extends AdmissionController { private static final Logger LOGGER = LogManager.getLogger(CPUBasedAdmissionController.class); - private final String admissionControllerName; public CPUBasedAdmissionControllerSettings settings; - private final AtomicLong rejectionCount; /** * * @param admissionControllerName State of the admission controller */ public CPUBasedAdmissionController(String admissionControllerName, Settings settings, ClusterSettings clusterSettings) { - this.admissionControllerName = admissionControllerName; + super(new AtomicLong(0), admissionControllerName); this.settings = new CPUBasedAdmissionControllerSettings(clusterSettings, settings); - this.rejectionCount = new AtomicLong(0); - } - - /** - * - * @return true if admissionController is enabled for the transport layer else false - */ - @Override - public boolean isEnabledForTransportLayer() { - return this.settings.getTransportLayerAdmissionControllerMode() != AdmissionControlMode.DISABLED; } /** @@ -53,40 +40,16 @@ public boolean isEnabledForTransportLayer() { @Override public void apply(String action) { // TODO Will extend this logic further currently just incrementing rejectionCount - if (this.isEnabledForTransportLayer()) { + if (this.isEnabledForTransportLayer(this.settings.getTransportLayerAdmissionControllerMode())) { this.applyForTransportLayer(action); } } private void applyForTransportLayer(String actionName) { - // currently incrementing counts to evaluate the controller triggering as expected and using in testing + // currently incrementing counts to evaluate the controller triggering as expected and using in testing so limiting to 10 // TODO will update rejection logic further in next PR's - this.addRejectionCount(1); - } - - /** - * @return name of the admission Controller - */ - @Override - public String getName() { - return this.admissionControllerName; - } - - /** - * Adds the rejection count for the controller. - * - * @param count the value that needs to be added to total rejection count - */ - @Override - public void addRejectionCount(long count) { - this.rejectionCount.incrementAndGet(); - } - - /** - * @return current value of the rejection count metric tracked by the admission-controller. - */ - @Override - public long getRejectionCount() { - return this.rejectionCount.get(); + if (this.getRejectionCount() < 10) { + this.addRejectionCount(1); + } } } diff --git a/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/settings/CPUBasedAdmissionControllerSettings.java b/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/settings/CPUBasedAdmissionControllerSettings.java index 39da278f3abd1..141e9b68db145 100644 --- a/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/settings/CPUBasedAdmissionControllerSettings.java +++ b/server/src/main/java/org/opensearch/ratelimitting/admissioncontrol/settings/CPUBasedAdmissionControllerSettings.java @@ -16,7 +16,6 @@ import java.util.Arrays; import java.util.List; -import java.util.function.Function; /** * Settings related to cpu based admission controller. @@ -70,20 +69,13 @@ public static class Defaults { Setting.Property.NodeScope ); - public static final Setting> CPU_BASED_ADMISSION_CONTROLLER_TRANSPORT_URI_LIST = Setting.listSetting( - "admission_control.global_cpu_usage.actions_list", - Defaults.TRANSPORT_LAYER_DEFAULT_URI_TYPE, - Function.identity(), - Setting.Property.NodeScope - ); - // currently limited to one setting will add further more settings in follow-up PR's public CPUBasedAdmissionControllerSettings(ClusterSettings clusterSettings, Settings settings) { this.transportLayerMode = CPU_BASED_ADMISSION_CONTROLLER_TRANSPORT_LAYER_MODE.get(settings); clusterSettings.addSettingsUpdateConsumer(CPU_BASED_ADMISSION_CONTROLLER_TRANSPORT_LAYER_MODE, this::setTransportLayerMode); this.searchCPULimit = SEARCH_CPU_USAGE_LIMIT.get(settings); this.indexingCPULimit = INDEXING_CPU_USAGE_LIMIT.get(settings); - this.transportActionsList = CPU_BASED_ADMISSION_CONTROLLER_TRANSPORT_URI_LIST.get(settings); + this.transportActionsList = Defaults.TRANSPORT_LAYER_DEFAULT_URI_TYPE; clusterSettings.addSettingsUpdateConsumer(INDEXING_CPU_USAGE_LIMIT, this::setIndexingCPULimit); clusterSettings.addSettingsUpdateConsumer(SEARCH_CPU_USAGE_LIMIT, this::setSearchCPULimit); } diff --git a/server/src/test/java/org/opensearch/common/network/NetworkModuleTests.java b/server/src/test/java/org/opensearch/common/network/NetworkModuleTests.java index 5f8595d0fb086..ab51cafb039c2 100644 --- a/server/src/test/java/org/opensearch/common/network/NetworkModuleTests.java +++ b/server/src/test/java/org/opensearch/common/network/NetworkModuleTests.java @@ -57,6 +57,7 @@ import org.opensearch.transport.TransportRequest; import org.opensearch.transport.TransportRequestHandler; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; @@ -124,7 +125,7 @@ public Map> getTransports( return Collections.singletonMap("custom", custom); } }; - NetworkModule module = newNetworkModule(settings, plugin); + NetworkModule module = newNetworkModule(settings, null, plugin); assertSame(custom, module.getTransportSupplier()); } @@ -135,7 +136,7 @@ public void testRegisterHttpTransport() { .build(); Supplier custom = FakeHttpTransport::new; - NetworkModule module = newNetworkModule(settings, new NetworkPlugin() { + NetworkModule module = newNetworkModule(settings, null, new NetworkPlugin() { @Override public Map> getHttpTransports( Settings settings, @@ -155,7 +156,7 @@ public Map> getHttpTransports( assertSame(custom, module.getHttpServerTransportSupplier()); settings = Settings.builder().put(NetworkModule.TRANSPORT_TYPE_KEY, "local").build(); - NetworkModule newModule = newNetworkModule(settings); + NetworkModule newModule = newNetworkModule(settings, null); expectThrows(IllegalStateException.class, () -> newModule.getHttpServerTransportSupplier()); } @@ -169,7 +170,7 @@ public void testOverrideDefault() { Supplier customTransport = () -> null; // content doesn't matter we check reference equality Supplier custom = FakeHttpTransport::new; Supplier def = FakeHttpTransport::new; - NetworkModule module = newNetworkModule(settings, new NetworkPlugin() { + NetworkModule module = newNetworkModule(settings, null, new NetworkPlugin() { @Override public Map> getTransports( Settings settings, @@ -214,7 +215,7 @@ public void testDefaultKeys() { Supplier custom = FakeHttpTransport::new; Supplier def = FakeHttpTransport::new; Supplier customTransport = () -> null; - NetworkModule module = newNetworkModule(settings, new NetworkPlugin() { + NetworkModule module = newNetworkModule(settings, null, new NetworkPlugin() { @Override public Map> getTransports( Settings settings, @@ -273,7 +274,7 @@ public TransportRequestHandler interceptHandler( return actualHandler; } }; - NetworkModule module = newNetworkModule(settings, new NetworkPlugin() { + NetworkModule module = newNetworkModule(settings, null, new NetworkPlugin() { @Override public List getTransportInterceptors( NamedWriteableRegistry namedWriteableRegistry, @@ -295,7 +296,7 @@ public List getTransportInterceptors( assertSame(((NetworkModule.CompositeTransportInterceptor) transportInterceptor).transportInterceptors.get(0), interceptor); NullPointerException nullPointerException = expectThrows(NullPointerException.class, () -> { - newNetworkModule(settings, new NetworkPlugin() { + newNetworkModule(settings, null, new NetworkPlugin() { @Override public List getTransportInterceptors( NamedWriteableRegistry namedWriteableRegistry, @@ -331,8 +332,10 @@ public TransportRequestHandler interceptHandler( } }; - NetworkModule module = newNetworkModule(settings); - module.registerCoreTransportInterceptor(interceptor); + List coreTransportInterceptors = new ArrayList<>(); + coreTransportInterceptors.add(interceptor); + + NetworkModule module = newNetworkModule(settings, coreTransportInterceptors); TransportInterceptor transportInterceptor = module.getTransportInterceptor(); assertEquals(0, called.get()); @@ -348,6 +351,7 @@ public TransportRequestHandler interceptHandler( public void testInterceptorOrder() { Settings settings = Settings.builder().put(NetworkModule.TRANSPORT_TYPE_KEY, "local").build(); AtomicInteger called = new AtomicInteger(0); + AtomicInteger called1 = new AtomicInteger(0); TransportInterceptor interceptor = new TransportInterceptor() { @Override @@ -367,7 +371,28 @@ public TransportRequestHandler interceptHandler( } }; - NetworkModule module = newNetworkModule(settings, new NetworkPlugin() { + TransportInterceptor interceptor1 = new TransportInterceptor() { + @Override + public TransportRequestHandler interceptHandler( + String action, + String executor, + boolean forceExecution, + TransportRequestHandler actualHandler + ) { + called1.incrementAndGet(); + if ("foo/bar/boom".equals(action)) { + assertTrue(forceExecution); + } else { + assertFalse(forceExecution); + } + return actualHandler; + } + }; + + List coreTransportInterceptors = new ArrayList<>(); + coreTransportInterceptors.add(interceptor1); + + NetworkModule module = newNetworkModule(settings, coreTransportInterceptors, new NetworkPlugin() { @Override public List getTransportInterceptors( NamedWriteableRegistry namedWriteableRegistry, @@ -378,9 +403,25 @@ public List getTransportInterceptors( } }); + TransportInterceptor transportInterceptor = module.getTransportInterceptor(); + assertEquals(((NetworkModule.CompositeTransportInterceptor) transportInterceptor).transportInterceptors.size(), 2); + + assertEquals(0, called.get()); + assertEquals(0, called1.get()); + transportInterceptor.interceptHandler("foo/bar/boom", null, true, null); + assertEquals(1, called.get()); + assertEquals(1, called1.get()); + transportInterceptor.interceptHandler("foo/baz/boom", null, false, null); + assertEquals(2, called.get()); + assertEquals(2, called1.get()); + } + + public void testInterceptorOrderException() { + Settings settings = Settings.builder().put(NetworkModule.TRANSPORT_TYPE_KEY, "local").build(); + AtomicInteger called = new AtomicInteger(0); AtomicInteger called1 = new AtomicInteger(0); - TransportInterceptor interceptor1 = new TransportInterceptor() { + TransportInterceptor interceptor = new TransportInterceptor() { @Override public TransportRequestHandler interceptHandler( String action, @@ -388,7 +429,7 @@ public TransportRequestHandler interceptHandler( boolean forceExecution, TransportRequestHandler actualHandler ) { - called1.incrementAndGet(); + called.incrementAndGet(); if ("foo/bar/boom".equals(action)) { assertTrue(forceExecution); } else { @@ -398,22 +439,57 @@ public TransportRequestHandler interceptHandler( } }; - module.registerCoreTransportInterceptor(interceptor1); + TransportInterceptor interceptor1 = new TransportInterceptor() { + @Override + public TransportRequestHandler interceptHandler( + String action, + String executor, + boolean forceExecution, + TransportRequestHandler actualHandler + ) { + called1.incrementAndGet(); + throw new RuntimeException("Handler Invoke Failed"); + } + }; + + List coreTransportInterceptors = new ArrayList<>(); + coreTransportInterceptors.add(interceptor1); + + NetworkModule module = newNetworkModule(settings, coreTransportInterceptors, new NetworkPlugin() { + @Override + public List getTransportInterceptors( + NamedWriteableRegistry namedWriteableRegistry, + ThreadContext threadContext + ) { + assertNotNull(threadContext); + return Collections.singletonList(interceptor); + } + }); TransportInterceptor transportInterceptor = module.getTransportInterceptor(); assertEquals(((NetworkModule.CompositeTransportInterceptor) transportInterceptor).transportInterceptors.size(), 2); assertEquals(0, called.get()); assertEquals(0, called1.get()); - transportInterceptor.interceptHandler("foo/bar/boom", null, true, null); - assertEquals(1, called.get()); - assertEquals(1, called1.get()); - transportInterceptor.interceptHandler("foo/baz/boom", null, false, null); - assertEquals(2, called.get()); - assertEquals(2, called1.get()); + try { + transportInterceptor.interceptHandler("foo/bar/boom", null, true, null); + } catch (Exception e) { + assertEquals(0, called.get()); + assertEquals(1, called1.get()); + } + try { + transportInterceptor.interceptHandler("foo/baz/boom", null, false, null); + } catch (Exception e) { + assertEquals(0, called.get()); + assertEquals(2, called1.get()); + } } - private NetworkModule newNetworkModule(Settings settings, NetworkPlugin... plugins) { + private NetworkModule newNetworkModule( + Settings settings, + List coreTransportInterceptors, + NetworkPlugin... plugins + ) { return new NetworkModule( settings, Arrays.asList(plugins), @@ -426,7 +502,8 @@ private NetworkModule newNetworkModule(Settings settings, NetworkPlugin... plugi null, new NullDispatcher(), new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS), - NoopTracer.INSTANCE + NoopTracer.INSTANCE, + coreTransportInterceptors ); } } diff --git a/server/src/test/java/org/opensearch/ratelimitting/admissioncontrol/AdmissionControlServiceTests.java b/server/src/test/java/org/opensearch/ratelimitting/admissioncontrol/AdmissionControlServiceTests.java index 7cf66113bbe45..bac4eaf3fd677 100644 --- a/server/src/test/java/org/opensearch/ratelimitting/admissioncontrol/AdmissionControlServiceTests.java +++ b/server/src/test/java/org/opensearch/ratelimitting/admissioncontrol/AdmissionControlServiceTests.java @@ -12,6 +12,7 @@ import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; import org.opensearch.ratelimitting.admissioncontrol.controllers.AdmissionController; +import org.opensearch.ratelimitting.admissioncontrol.controllers.CPUBasedAdmissionController; import org.opensearch.ratelimitting.admissioncontrol.enums.AdmissionControlMode; import org.opensearch.ratelimitting.admissioncontrol.settings.CPUBasedAdmissionControllerSettings; import org.opensearch.test.OpenSearchTestCase; @@ -65,12 +66,13 @@ public void testAdmissionControllerSettings() { AdmissionControlSettings admissionControlSettings = admissionControlService.admissionControlSettings; List admissionControllerList = admissionControlService.getAdmissionControllers(); assertEquals(admissionControllerList.size(), 1); - AdmissionController cpuBasedAdmissionController = admissionControlService.getAdmissionController( - CPUBasedAdmissionControllerSettings.CPU_BASED_ADMISSION_CONTROLLER - ); + CPUBasedAdmissionController cpuBasedAdmissionController = (CPUBasedAdmissionController) admissionControlService + .getAdmissionController(CPUBasedAdmissionControllerSettings.CPU_BASED_ADMISSION_CONTROLLER); assertEquals( admissionControlSettings.isTransportLayerAdmissionControlEnabled(), - cpuBasedAdmissionController.isEnabledForTransportLayer() + cpuBasedAdmissionController.isEnabledForTransportLayer( + cpuBasedAdmissionController.settings.getTransportLayerAdmissionControllerMode() + ) ); Settings settings = Settings.builder() @@ -79,7 +81,9 @@ public void testAdmissionControllerSettings() { clusterService.getClusterSettings().applySettings(settings); assertEquals( admissionControlSettings.isTransportLayerAdmissionControlEnabled(), - cpuBasedAdmissionController.isEnabledForTransportLayer() + cpuBasedAdmissionController.isEnabledForTransportLayer( + cpuBasedAdmissionController.settings.getTransportLayerAdmissionControllerMode() + ) ); assertFalse(admissionControlSettings.isTransportLayerAdmissionControlEnabled()); @@ -92,7 +96,11 @@ public void testAdmissionControllerSettings() { .build(); clusterService.getClusterSettings().applySettings(newSettings); assertFalse(admissionControlSettings.isTransportLayerAdmissionControlEnabled()); - assertTrue(cpuBasedAdmissionController.isEnabledForTransportLayer()); + assertTrue( + cpuBasedAdmissionController.isEnabledForTransportLayer( + cpuBasedAdmissionController.settings.getTransportLayerAdmissionControllerMode() + ) + ); } public void testApplyAdmissionControllerDisabled() { @@ -120,7 +128,6 @@ public void testApplyAdmissionControllerEnabled() { ) .build(); clusterService.getClusterSettings().applySettings(settings); - this.action = "indices:data/write/bulk[s][p]"; admissionControlService.applyTransportAdmissionControl(this.action); List admissionControllerList = admissionControlService.getAdmissionControllers(); assertEquals(admissionControllerList.size(), 1); diff --git a/server/src/test/java/org/opensearch/ratelimitting/admissioncontrol/controllers/CPUBasedAdmissionControllerTests.java b/server/src/test/java/org/opensearch/ratelimitting/admissioncontrol/controllers/CPUBasedAdmissionControllerTests.java index 7a53ed831b6fe..af6ec0749e709 100644 --- a/server/src/test/java/org/opensearch/ratelimitting/admissioncontrol/controllers/CPUBasedAdmissionControllerTests.java +++ b/server/src/test/java/org/opensearch/ratelimitting/admissioncontrol/controllers/CPUBasedAdmissionControllerTests.java @@ -50,7 +50,9 @@ public void testCheckDefaultParameters() { assertEquals(admissionController.getName(), CPUBasedAdmissionControllerSettings.CPU_BASED_ADMISSION_CONTROLLER); assertEquals(admissionController.getRejectionCount(), 0); assertEquals(admissionController.settings.getTransportLayerAdmissionControllerMode(), AdmissionControlMode.DISABLED); - assertFalse(admissionController.isEnabledForTransportLayer()); + assertFalse( + admissionController.isEnabledForTransportLayer(admissionController.settings.getTransportLayerAdmissionControllerMode()) + ); } public void testCheckUpdateSettings() { @@ -70,7 +72,7 @@ public void testCheckUpdateSettings() { assertEquals(admissionController.getName(), CPUBasedAdmissionControllerSettings.CPU_BASED_ADMISSION_CONTROLLER); assertEquals(admissionController.getRejectionCount(), 0); assertEquals(admissionController.settings.getTransportLayerAdmissionControllerMode(), AdmissionControlMode.ENFORCED); - assertTrue(admissionController.isEnabledForTransportLayer()); + assertTrue(admissionController.isEnabledForTransportLayer(admissionController.settings.getTransportLayerAdmissionControllerMode())); } public void testApplyControllerWithDefaultSettings() { @@ -98,7 +100,7 @@ public void testApplyControllerWhenSettingsEnabled() { settings, clusterService.getClusterSettings() ); - assertTrue(admissionController.isEnabledForTransportLayer()); + assertTrue(admissionController.isEnabledForTransportLayer(admissionController.settings.getTransportLayerAdmissionControllerMode())); assertEquals(admissionController.getRejectionCount(), 0); action = "indices:data/write/bulk[s][p]"; admissionController.apply(action);