From 5892b6a536fff321bbe26e586e94f74656ee82d7 Mon Sep 17 00:00:00 2001 From: hitesh Date: Thu, 24 Oct 2024 16:20:29 -0700 Subject: [PATCH] Updated dynamic change + test --- .../apache/solr/servlet/QueryRateLimiter.java | 12 +---- .../apache/solr/servlet/RateLimitManager.java | 26 +++++++++-- .../solr/servlet/TestRequestRateLimiter.java | 45 +++++++++++++++++++ 3 files changed, 70 insertions(+), 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 index 0a4f7b082df..91145c03653 100644 --- 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 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 index d4ce299081c..7d244bc3024 100644 --- 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 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 index 67ea40dfb56..c309ee2f3ba 100644 --- 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 properties = new HashMap<>(); + Map 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); + } }