Add RequestSizeAdmissionController

Signed-off-by: Mital Awachat <awachatm@amazon.com>

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<String, Long> 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<String, Long> 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<String, Long> getLimits() {
+        return limits;
+    }
+
+    @Override
+    public String toString() {
+        return "AdmissionControlSetting{" + "name='" + name + '\'' + ", enabled=" + enabled + ", limit=" + limits + '}';
+    }
+}

server/src/main/java/org/opensearch/common/throttle/AdmissionController.java

@@ -46,11 +46,6 @@ interface AdmissionController {
      */
     long setInitialQuota(long count);
 
-    /**
-     * @return limit defined for the tracking-object being tracked by the admission-controller.
-     */
-    long getTotalQuota();
-
     /**
      * @return currently acquired value of the tracking-object being tracked by the admission-controller.
      */
@@ -67,4 +62,26 @@ interface AdmissionController {
      * @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;
+        }
+    }
 }

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 () -> {};
+    }
+}

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<String, Long> 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());
+    }
+}