From bc3804847f433a577bb2525d03e412af03bf1831 Mon Sep 17 00:00:00 2001 From: hitesh Date: Wed, 30 Oct 2024 13:26:47 -0700 Subject: [PATCH 1/4] Added suppost to enable rate limiter in nodes --- .../solr/servlet/CoreContainerProvider.java | 4 +- .../apache/solr/servlet/RateLimitManager.java | 17 ++++-- .../solr/servlet/TestRequestRateLimiter.java | 59 ++++++++++++++++--- .../request/beans/RateLimiterPayload.java | 17 ++++++ 4 files changed, 82 insertions(+), 15 deletions(-) diff --git a/solr/core/src/java/org/apache/solr/servlet/CoreContainerProvider.java b/solr/core/src/java/org/apache/solr/servlet/CoreContainerProvider.java index 8b8bc3c927d..9ffc79e7ffc 100644 --- a/solr/core/src/java/org/apache/solr/servlet/CoreContainerProvider.java +++ b/solr/core/src/java/org/apache/solr/servlet/CoreContainerProvider.java @@ -241,7 +241,9 @@ private void init(ServletContext servletContext) { Builder builder = new Builder(zkClient); - this.rateLimitManager = builder.build(); + String hostname = zkController != null ? zkController.getHostName() : ""; + + this.rateLimitManager = builder.build(hostname); if (zkController != null) { zkController.zkStateReader.registerClusterPropertiesListener(this.rateLimitManager); 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 d339cafb5c3..7df7f6ca0a6 100644 --- a/solr/core/src/java/org/apache/solr/servlet/RateLimitManager.java +++ b/solr/core/src/java/org/apache/solr/servlet/RateLimitManager.java @@ -62,7 +62,10 @@ public class RateLimitManager implements ClusterPropertiesListener { public static final long DEFAULT_SLOT_ACQUISITION_TIMEOUT_MS = -1; private final ConcurrentHashMap requestRateLimiterMap; - public RateLimitManager() { + private final String hostname; + + public RateLimitManager(String hostname) { + this.hostname = hostname; this.requestRateLimiterMap = new ConcurrentHashMap<>(); } @@ -81,7 +84,7 @@ public boolean onChange(Map properties) { throw new UncheckedIOException(e); } } - + rateLimiterMeta.overrideNodeProperty(hostname); // Hack: We only support query rate limiting for now requestRateLimiterMap.compute( rateLimiterMeta.priorityBasedEnabled @@ -209,10 +212,10 @@ public Builder(SolrZkClient solrZkClient) { this.solrZkClient = solrZkClient; } - public RateLimitManager build() { - RateLimitManager rateLimitManager = new RateLimitManager(); + public RateLimitManager build(String hostname) { + RateLimitManager rateLimitManager = new RateLimitManager(hostname); - RateLimiterConfig rateLimiterConfig = constructQueryRateLimiterConfig(solrZkClient); + RateLimiterConfig rateLimiterConfig = constructQueryRateLimiterConfig(solrZkClient, hostname); if (rateLimiterConfig.priorityBasedEnabled) { rateLimitManager.registerRequestRateLimiter( @@ -228,7 +231,8 @@ public RateLimitManager build() { // To be used in initialization @SuppressWarnings({"unchecked"}) - private static RateLimiterConfig constructQueryRateLimiterConfig(SolrZkClient zkClient) { + private static RateLimiterConfig constructQueryRateLimiterConfig( + SolrZkClient zkClient, String hostname) { try { if (zkClient == null) { @@ -250,6 +254,7 @@ private static RateLimiterConfig constructQueryRateLimiterConfig(SolrZkClient zk RateLimiterPayload rateLimiterMeta = mapper.readValue(configInput, RateLimiterPayload.class); + rateLimiterMeta.overrideNodeProperty(hostname); return rateLimiterMeta.priorityBasedEnabled ? new RateLimiterConfig(SolrRequest.SolrRequestType.PRIORITY_BASED, rateLimiterMeta) : new RateLimiterConfig(SolrRequest.SolrRequestType.QUERY, rateLimiterMeta); 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 e4cc923787d..ef8fff60381 100644 --- a/solr/core/src/test/org/apache/solr/servlet/TestRequestRateLimiter.java +++ b/solr/core/src/test/org/apache/solr/servlet/TestRequestRateLimiter.java @@ -91,7 +91,7 @@ public void testConcurrentQueries() throws Exception { RateLimitManager.Builder builder = new MockBuilder( null /* dummy SolrZkClient */, new MockRequestRateLimiter(rateLimiterConfig)); - RateLimitManager rateLimitManager = builder.build(); + RateLimitManager rateLimitManager = builder.build("localhost"); solrDispatchFilter.replaceRateLimitManager(rateLimitManager); @@ -121,7 +121,7 @@ public void testConcurrentQueries() throws Exception { @SuppressWarnings("try") public void testSlotBorrowingAcquisitionTimeout() throws InterruptedException, IOException, ExecutionException { - RateLimitManager mgr = new RateLimitManager(); + RateLimitManager mgr = new RateLimitManager("localhost"); Random r = random(); int slotLimit = r.nextInt(20) + 1; int guaranteed = r.nextInt(slotLimit); @@ -355,7 +355,7 @@ public void testSlotBorrowing() throws Exception { null /*dummy SolrZkClient */, new MockRequestRateLimiter(queryRateLimiterConfig), new MockRequestRateLimiter(indexRateLimiterConfig)); - RateLimitManager rateLimitManager = builder.build(); + RateLimitManager rateLimitManager = builder.build("localhost"); solrDispatchFilter.replaceRateLimitManager(rateLimitManager); @@ -488,8 +488,8 @@ public MockBuilder( } @Override - public RateLimitManager build() { - RateLimitManager rateLimitManager = new RateLimitManager(); + public RateLimitManager build(String hostname) { + RateLimitManager rateLimitManager = new RateLimitManager("localhost"); rateLimitManager.registerRequestRateLimiter( queryRequestRateLimiter, SolrRequest.SolrRequestType.QUERY); @@ -632,7 +632,7 @@ public void testAdjustingConfig() throws IOException, InterruptedException { @Test public void testPriorityBasedRateLimiter() throws Exception { - RateLimitManager rateLimitManager = new RateLimitManager(); + RateLimitManager rateLimitManager = new RateLimitManager("localhost"); // PriorityBasedRateLimiter RateLimiterConfig rateLimiterConfig = @@ -683,7 +683,7 @@ public void testPriorityBasedRateLimiter() throws Exception { @Test public void testPriorityBasedRateLimiterTimeout() throws Exception { - RateLimitManager rateLimitManager = new RateLimitManager(); + RateLimitManager rateLimitManager = new RateLimitManager("localhost"); // PriorityBasedRateLimiter RateLimiterConfig rateLimiterConfig = @@ -730,7 +730,7 @@ public void testPriorityBasedRateLimiterTimeout() throws Exception { @Test public void testPriorityBasedRateLimiterDynamicChange() throws Exception { - RateLimitManager rateLimitManager = new RateLimitManager(); + RateLimitManager rateLimitManager = new RateLimitManager("localhost"); // PriorityBasedRateLimiter RateLimiterConfig rateLimiterConfig = @@ -770,4 +770,47 @@ public void testPriorityBasedRateLimiterDynamicChange() throws Exception { assertEquals(true, rateLimiter.getRateLimiterConfig().priorityBasedEnabled); } + + @Test + public void testEnableRateLimiterOnNode() throws Exception { + RateLimitManager rateLimitManager = new RateLimitManager("localhost"); + + // PriorityBasedRateLimiter + RateLimiterConfig rateLimiterConfig = + new RateLimiterConfig( + SolrRequest.SolrRequestType.QUERY, + false, + 1, + 10, + 1 /* allowedRequests */, + true /* isSlotBorrowing */, + false); + + QueryRateLimiter requestRateLimiter = new QueryRateLimiter(rateLimiterConfig); + + rateLimitManager.registerRequestRateLimiter( + requestRateLimiter, SolrRequest.SolrRequestType.QUERY); + + RequestRateLimiter rateLimiter = + rateLimitManager.getRequestRateLimiter(SolrRequest.SolrRequestType.QUERY); + assertFalse(rateLimiter.rateLimiterConfig.isEnabled); + + Map properties = new HashMap<>(); + Map rateLimiterProps = new HashMap<>(); + rateLimiterProps.put("enabled", false); + rateLimiterProps.put("guaranteedSlots", 1); + rateLimiterProps.put("allowedRequests", 1); + rateLimiterProps.put("slotBorrowingEnabled", false); + rateLimiterProps.put("slotAcquisitionTimeoutInMS", 100); + rateLimiterProps.put("priorityBasedEnabled", false); + rateLimiterProps.put("nodesEnabled", "localhost"); + properties.put("rate-limiters", rateLimiterProps); + + // updating rate limiter + rateLimitManager.onChange(properties); + + rateLimiter = rateLimitManager.getRequestRateLimiter(SolrRequest.SolrRequestType.QUERY); + + assertTrue(rateLimiter.getRateLimiterConfig().isEnabled); + } } diff --git a/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java b/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java index 6088cd00cee..5a1da664ee4 100644 --- a/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java +++ b/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java @@ -35,6 +35,8 @@ public class RateLimiterPayload implements ReflectMapWriter { @JsonProperty public Boolean priorityBasedEnabled; + @JsonProperty public String nodesEnabled; + public RateLimiterPayload copy() { RateLimiterPayload result = new RateLimiterPayload(); @@ -44,6 +46,7 @@ public RateLimiterPayload copy() { result.slotBorrowingEnabled = slotBorrowingEnabled; result.slotAcquisitionTimeoutInMS = slotAcquisitionTimeoutInMS; result.priorityBasedEnabled = priorityBasedEnabled; + result.nodesEnabled = nodesEnabled; return result; } @@ -71,4 +74,18 @@ public int hashCode() { slotAcquisitionTimeoutInMS, priorityBasedEnabled); } + + public void overrideNodeProperty(String hostname) { + if (!this.enabled && hostname != "") { + if (this.nodesEnabled != null && !this.nodesEnabled.isEmpty()) { + String[] hosts = this.nodesEnabled.split(","); + for (String host : hosts) { + if (host.trim().equals(hostname)) { + this.enabled = true; + break; + } + } + } + } + } } From 3ff9313b54b044a12f99e08d15a11e3f2f77e386 Mon Sep 17 00:00:00 2001 From: hitesh Date: Wed, 30 Oct 2024 13:36:52 -0700 Subject: [PATCH 2/4] minor change --- .../solr/client/solrj/request/beans/RateLimiterPayload.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java b/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java index 5a1da664ee4..31df69691ef 100644 --- a/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java +++ b/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java @@ -76,7 +76,7 @@ public int hashCode() { } public void overrideNodeProperty(String hostname) { - if (!this.enabled && hostname != "") { + if (!this.enabled && hostname.isEmpty()) { if (this.nodesEnabled != null && !this.nodesEnabled.isEmpty()) { String[] hosts = this.nodesEnabled.split(","); for (String host : hosts) { From b73ffa6bda5ba37b8f22a1a70e243c7a2e960044 Mon Sep 17 00:00:00 2001 From: hitesh Date: Wed, 30 Oct 2024 13:44:38 -0700 Subject: [PATCH 3/4] minor fix --- .../solr/client/solrj/request/beans/RateLimiterPayload.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java b/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java index 31df69691ef..89207636ed9 100644 --- a/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java +++ b/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java @@ -76,7 +76,7 @@ public int hashCode() { } public void overrideNodeProperty(String hostname) { - if (!this.enabled && hostname.isEmpty()) { + if (!this.enabled && !hostname.isEmpty()) { if (this.nodesEnabled != null && !this.nodesEnabled.isEmpty()) { String[] hosts = this.nodesEnabled.split(","); for (String host : hosts) { From 234275c654cd0782ae413ad682794b90052bcf9b Mon Sep 17 00:00:00 2001 From: hitesh Date: Wed, 30 Oct 2024 15:15:48 -0700 Subject: [PATCH 4/4] changed the func name --- .../src/java/org/apache/solr/servlet/RateLimitManager.java | 4 ++-- .../solr/client/solrj/request/beans/RateLimiterPayload.java | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) 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 7df7f6ca0a6..7ff03745d9b 100644 --- a/solr/core/src/java/org/apache/solr/servlet/RateLimitManager.java +++ b/solr/core/src/java/org/apache/solr/servlet/RateLimitManager.java @@ -84,7 +84,7 @@ public boolean onChange(Map properties) { throw new UncheckedIOException(e); } } - rateLimiterMeta.overrideNodeProperty(hostname); + rateLimiterMeta.maybeEnableForHost(hostname); // Hack: We only support query rate limiting for now requestRateLimiterMap.compute( rateLimiterMeta.priorityBasedEnabled @@ -254,7 +254,7 @@ private static RateLimiterConfig constructQueryRateLimiterConfig( RateLimiterPayload rateLimiterMeta = mapper.readValue(configInput, RateLimiterPayload.class); - rateLimiterMeta.overrideNodeProperty(hostname); + rateLimiterMeta.maybeEnableForHost(hostname); return rateLimiterMeta.priorityBasedEnabled ? new RateLimiterConfig(SolrRequest.SolrRequestType.PRIORITY_BASED, rateLimiterMeta) : new RateLimiterConfig(SolrRequest.SolrRequestType.QUERY, rateLimiterMeta); diff --git a/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java b/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java index 89207636ed9..ac825e1adb6 100644 --- a/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java +++ b/solr/solrj/src/java/org/apache/solr/client/solrj/request/beans/RateLimiterPayload.java @@ -75,7 +75,7 @@ public int hashCode() { priorityBasedEnabled); } - public void overrideNodeProperty(String hostname) { + public void maybeEnableForHost(String hostname) { if (!this.enabled && !hostname.isEmpty()) { if (this.nodesEnabled != null && !this.nodesEnabled.isEmpty()) { String[] hosts = this.nodesEnabled.split(",");