Updated dynamic change + test

cowpaths · Oct 24, 2024 · 5892b6a · 5892b6a
1 parent 075875a
commit 5892b6a
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 13 deletions.
diff --git a/solr/core/src/java/org/apache/solr/servlet/QueryRateLimiter.java b/solr/core/src/java/org/apache/solr/servlet/QueryRateLimiter.java
@@ -48,19 +48,11 @@ public QueryRateLimiter(RateLimiterConfig config) {
   }
 
   public static RateLimiterConfig processConfigChange(
-      RateLimiterConfig rateLimiterConfig, Map<String, Object> properties) throws IOException {
-    byte[] configInput = Utils.toJSON(properties.get(RL_CONFIG_KEY));
-
-    RateLimiterPayload rateLimiterMeta;
-    if (configInput == null || configInput.length == 0) {
-      rateLimiterMeta = null;
-    } else {
-      rateLimiterMeta = mapper.readValue(configInput, RateLimiterPayload.class);
-    }
+      RateLimiterConfig rateLimiterConfig, RateLimiterPayload rateLimiterMeta) throws IOException {
 
     // default rate limiter
     SolrRequest.SolrRequestType requestType = SolrRequest.SolrRequestType.QUERY;
-    if (rateLimiterConfig.priorityBasedEnabled) {
+    if (rateLimiterMeta.priorityBasedEnabled) {
       requestType = SolrRequest.SolrRequestType.PRIORITY_BASED;
     }
 

diff --git a/solr/core/src/java/org/apache/solr/servlet/RateLimitManager.java b/solr/core/src/java/org/apache/solr/servlet/RateLimitManager.java
@@ -19,7 +19,9 @@
 
 import static org.apache.solr.common.params.CommonParams.SOLR_REQUEST_CONTEXT_PARAM;
 import static org.apache.solr.common.params.CommonParams.SOLR_REQUEST_TYPE_PARAM;
+import static org.apache.solr.core.RateLimiterConfig.RL_CONFIG_KEY;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.lang.invoke.MethodHandles;
@@ -28,9 +30,12 @@
 import javax.servlet.http.HttpServletRequest;
 import net.jcip.annotations.ThreadSafe;
 import org.apache.solr.client.solrj.SolrRequest;
+import org.apache.solr.client.solrj.request.beans.RateLimiterPayload;
 import org.apache.solr.common.cloud.ClusterPropertiesListener;
 import org.apache.solr.common.cloud.SolrZkClient;
+import org.apache.solr.common.util.Utils;
 import org.apache.solr.core.RateLimiterConfig;
+import org.apache.solr.util.SolrJacksonAnnotationInspector;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -46,7 +51,7 @@
 @ThreadSafe
 public class RateLimitManager implements ClusterPropertiesListener {
   private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
-
+  private static final ObjectMapper mapper = SolrJacksonAnnotationInspector.createObjectMapper();
   public static final String ERROR_MESSAGE =
       "Too many requests for this request type. Please try after some time or increase the quota for this request type";
   public static final int DEFAULT_CONCURRENT_REQUESTS =
@@ -61,14 +66,29 @@ public RateLimitManager() {
   @Override
   public boolean onChange(Map<String, Object> properties) {
 
+    byte[] configInput = Utils.toJSON(properties.get(RL_CONFIG_KEY));
+
+    RateLimiterPayload rateLimiterMeta;
+    if (configInput == null || configInput.length == 0) {
+      return false;
+    } else {
+      try {
+        rateLimiterMeta = mapper.readValue(configInput, RateLimiterPayload.class);
+      } catch (IOException e) {
+        throw new UncheckedIOException(e);
+      }
+    }
+
     // Hack: We only support query rate limiting for now
     requestRateLimiterMap.compute(
-        SolrRequest.SolrRequestType.QUERY.toString(),
+        rateLimiterMeta.priorityBasedEnabled
+            ? SolrRequest.SolrRequestType.PRIORITY_BASED.name()
+            : SolrRequest.SolrRequestType.QUERY.name(),
         (k, v) -> {
           try {
             RateLimiterConfig newConfig =
                 QueryRateLimiter.processConfigChange(
-                    v == null ? null : v.getRateLimiterConfig(), properties);
+                    v == null ? null : v.getRateLimiterConfig(), rateLimiterMeta);
             if (newConfig == null) {
               return v;
             } else {

diff --git a/solr/core/src/test/org/apache/solr/servlet/TestRequestRateLimiter.java b/solr/core/src/test/org/apache/solr/servlet/TestRequestRateLimiter.java
@@ -26,7 +26,9 @@
 import java.io.Closeable;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Random;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CountDownLatch;
@@ -702,4 +704,47 @@ public void testPriorityBasedRateLimiterTimeout() throws Exception {
 
     firstRequestAllowed.close();
   }
+
+  @Test
+  public void testPriorityBasedRateLimiterDynamicChange() throws Exception {
+    RateLimitManager rateLimitManager = new RateLimitManager();
+
+    // PriorityBasedRateLimiter
+    RateLimiterConfig rateLimiterConfig =
+        new RateLimiterConfig(
+            SolrRequest.SolrRequestType.QUERY,
+            true,
+            1,
+            10,
+            1 /* allowedRequests */,
+            true /* isSlotBorrowing */,
+            false);
+
+    QueryRateLimiter requestRateLimiter = new QueryRateLimiter(rateLimiterConfig);
+
+    rateLimitManager.registerRequestRateLimiter(
+        requestRateLimiter, SolrRequest.SolrRequestType.QUERY);
+
+    RequestRateLimiter rateLimiter =
+        rateLimitManager.getRequestRateLimiter(SolrRequest.SolrRequestType.PRIORITY_BASED);
+    assertNull(rateLimiter);
+
+    Map<String, Object> properties = new HashMap<>();
+    Map<String, Object> rateLimiterProps = new HashMap<>();
+    rateLimiterProps.put("enabled", true);
+    rateLimiterProps.put("guaranteedSlots", 1);
+    rateLimiterProps.put("allowedRequests", 1);
+    rateLimiterProps.put("slotBorrowingEnabled", false);
+    rateLimiterProps.put("slotAcquisitionTimeoutInMS", 100);
+    rateLimiterProps.put("priorityBasedEnabled", true);
+    properties.put("rate-limiters", rateLimiterProps);
+
+    // updating rate limiter
+    rateLimitManager.onChange(properties);
+
+    rateLimiter =
+        rateLimitManager.getRequestRateLimiter(SolrRequest.SolrRequestType.PRIORITY_BASED);
+
+    assertEquals(true, rateLimiter.getRateLimiterConfig().priorityBasedEnabled);
+  }
 }