diff --git a/server/src/main/java/org/opensearch/common/throttle/AdmissionControlSetting.java b/server/src/main/java/org/opensearch/common/throttle/AdmissionControlSetting.java new file mode 100644 index 0000000000000..02ac384cf4405 --- /dev/null +++ b/server/src/main/java/org/opensearch/common/throttle/AdmissionControlSetting.java @@ -0,0 +1,59 @@ +/* + * 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.common.throttle; + +import java.util.Map; + +/** + * Settings for a {@link AdmissionController} + * + * @opensearch.internal + */ +class AdmissionControlSetting { + private final String name; + private final boolean enabled; + private final Map limits; + + /** + * @param name of admission controller + * @param enabled status of admission controller + * @param limits of admission controller per endpoint (key) + */ + public AdmissionControlSetting(String name, boolean enabled, Map limits) { + this.name = name; + this.enabled = enabled; + this.limits = limits; + } + + /** + * @return name of admission controller + */ + public String getName() { + return name; + } + + /** + * @return if admission controller is enabled or not + */ + public boolean isEnabled() { + return enabled; + } + + /** + * @return limits of admission controller + */ + public Map getLimits() { + return limits; + } + + @Override + public String toString() { + return "AdmissionControlSetting{" + "name='" + name + '\'' + ", enabled=" + enabled + ", limit=" + limits + '}'; + } +} diff --git a/server/src/main/java/org/opensearch/common/throttle/AdmissionController.java b/server/src/main/java/org/opensearch/common/throttle/AdmissionController.java new file mode 100644 index 0000000000000..f766bb14425ef --- /dev/null +++ b/server/src/main/java/org/opensearch/common/throttle/AdmissionController.java @@ -0,0 +1,87 @@ +/* + * 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.common.throttle; + +import org.opensearch.common.lease.Releasable; + +/** + * Aims to provide resource based request admission control. + * Provides methods for any tracking-object that can be incremented (such as memory size), + * Admission control can be applied if configured limit has been reached. + * + * @opensearch.internal + */ +interface AdmissionController { + + /** + * @return name of the admission-controller + */ + String getName(); + + /** + * @return admission-controller enabled status + */ + boolean isEnabled(); + + /** + * Increment the tracking-object with provided value. + * Apply the admission control if threshold is breached. + * Mostly applicable while acquiring the quota. + * Later Releasable is used to decrement the tracking-object with previously acquired value. + * + * @param count value to incrementation the resource racking-object with. + * @return Releasable for tokens acquired from the resource tracking object. + */ + Releasable addBytesAndMaybeBreak(long count); + + /** + * Sets the initial used quota for the controller. Primarily used when copying controller states. + * @param count To set the value of the tracking resource object as the provided count + * @return count/value by which the resource tracking object is updated with. + */ + long setInitialQuota(long count); + + /** + * @return currently acquired value of the tracking-object being tracked by the admission-controller. + */ + long getUsedQuota(); + + /** + * @return current value of the rejection count metric tracked by the admission-controller. + */ + long getRejectionCount(); + + /** + * 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 + * @return count/value by which the resource tracking object is updated with. + */ + long addRejectionCount(long count); + + /** + * Admission Controllers + */ + public enum Controllers { + REQUEST_SIZE("RequestSize"); + + private final String name; + + Controllers(String name) { + this.name = name; + } + + public String value() { + return this.name; + } + + @Override + public String toString() { + return name; + } + } +} diff --git a/server/src/main/java/org/opensearch/common/throttle/RequestSizeAdmissionController.java b/server/src/main/java/org/opensearch/common/throttle/RequestSizeAdmissionController.java new file mode 100644 index 0000000000000..64c548519210c --- /dev/null +++ b/server/src/main/java/org/opensearch/common/throttle/RequestSizeAdmissionController.java @@ -0,0 +1,103 @@ +/* + * 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.common.throttle; + +import org.opensearch.common.lease.Releasable; + +import java.util.concurrent.atomic.AtomicLong; + +/** + * RequestSizeAdmissionController tracks the memory quota in bytes allocated for a request object based upon its + * request-size (content-length) + * + * @opensearch.internal + */ +class RequestSizeAdmissionController implements AdmissionController { + private final AdmissionControlSetting admissionControlSetting; + private final AtomicLong usedBytes; + private final AtomicLong rejectionCount; + + /** + * @param admissionControlSetting setting for request-size admission-controller + */ + public RequestSizeAdmissionController(AdmissionControlSetting admissionControlSetting) { + this.admissionControlSetting = admissionControlSetting; + this.usedBytes = new AtomicLong(0); + this.rejectionCount = new AtomicLong(0); + } + + /** + * @return name of the admission-controller + */ + @Override + public String getName() { + return this.admissionControlSetting.getName(); + } + + /** + * @return admission-controller enabled status + */ + @Override + public boolean isEnabled() { + return this.admissionControlSetting.isEnabled(); + } + + /** + * Sets the initial used quota for the controller. Primarily used when copying controller states. + * @param count To set the value of the tracking resource object as the provided count + * @return count/value by which the resource tracking object is updated with. + */ + @Override + public long setInitialQuota(long count) { + usedBytes.set(count); + return usedBytes.get(); + } + + /** + * @return currently acquired value of the tracking-object being tracked by the admission-controller. + */ + @Override + public long getUsedQuota() { + return usedBytes.get(); + } + + /** + * @return current value of the rejection count metric tracked by the admission-controller. + */ + @Override + public long getRejectionCount() { + return this.rejectionCount.get(); + } + + /** + * 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 + * @return count/value by which the resource tracking object is updated with. + */ + @Override + public long addRejectionCount(long count) { + return this.rejectionCount.addAndGet(count); + } + + /** + * Increments the memory-tracking object for request-size quota with the provided bytes; and apply the admission + * control, if threshold is breached. + * Expected to be used while acquiring the quota. + * + * @param bytes byte size value to add to the current memory-tracking object and verify the threshold. + * @return byte size value used for comparison, if controller is enabled; zero otherwise + * @throws IllegalStateException if current memory usage quota along with requested bytes is greater than + * pre-defined threshold. + */ + @Override + public Releasable addBytesAndMaybeBreak(long bytes) { + // todo: how the request should be allowed or throttled + return () -> {}; + } +} diff --git a/server/src/main/java/org/opensearch/common/throttle/package-info.java b/server/src/main/java/org/opensearch/common/throttle/package-info.java new file mode 100644 index 0000000000000..c06f89cc0e0f9 --- /dev/null +++ b/server/src/main/java/org/opensearch/common/throttle/package-info.java @@ -0,0 +1,10 @@ +/* + * 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. + */ + +/** Base admission control package. */ +package org.opensearch.common.throttle; diff --git a/server/src/test/java/org/opensearch/common/throttle/RequestSizeAdmissionControllerTests.java b/server/src/test/java/org/opensearch/common/throttle/RequestSizeAdmissionControllerTests.java new file mode 100644 index 0000000000000..87a7a84e41d50 --- /dev/null +++ b/server/src/test/java/org/opensearch/common/throttle/RequestSizeAdmissionControllerTests.java @@ -0,0 +1,43 @@ +/* + * 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.common.throttle; + +import org.junit.After; +import org.junit.Before; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Map; + +import static org.opensearch.common.throttle.AdmissionController.Controllers.REQUEST_SIZE; + +public class RequestSizeAdmissionControllerTests extends OpenSearchTestCase { + + private AdmissionControlSetting setting; + private RequestSizeAdmissionController requestSizeController; + + @Override + @Before + public void setUp() throws Exception { + super.setUp(); + Map limits = Map.of("_search", 1000L, "_bulk", 1000L); // 1000 bytes + setting = new AdmissionControlSetting(REQUEST_SIZE.value(), true, limits); + requestSizeController = new RequestSizeAdmissionController(setting); + } + + @Override + @After + public void tearDown() throws Exception { + super.tearDown(); + } + + public void testInitialState() { + assertEquals(REQUEST_SIZE.value(), requestSizeController.getName()); + assertEquals(0, requestSizeController.getUsedQuota()); + } +}