Support node-local rates

1 files changed, 30 insertions, 15 deletions

diff --git a/container-search/src/test/java/com/yahoo/search/searchers/test/RateLimitingSearcherTestCase.java b/container-search/src/test/java/com/yahoo/search/searchers/test/RateLimitingSearcherTestCase.java
index 22b8d8f4d76..1682bd37fb6 100755
--- a/container-search/src/test/java/com/yahoo/search/searchers/test/RateLimitingSearcherTestCase.java
+++ b/container-search/src/test/java/com/yahoo/search/searchers/test/RateLimitingSearcherTestCase.java
@@ -31,23 +31,10 @@ public class RateLimitingSearcherTestCase {
 
     @Test
     void testRateLimiting() {
-        RateLimitingConfig.Builder rateLimitingConfig = new RateLimitingConfig.Builder();
-        rateLimitingConfig.maxAvailableCapacity(4);
-        rateLimitingConfig.capacityIncrement(2);
-        rateLimitingConfig.recheckForCapacityProbability(1.0);
-
-        ClusterInfoConfig.Builder clusterInfoConfig = new ClusterInfoConfig.Builder();
-        clusterInfoConfig.clusterId("testCluster");
-        clusterInfoConfig.nodeCount(4);
-
         ManualClock clock = new ManualClock();
         MetricReceiver.MockReceiver metric = new MetricReceiver.MockReceiver();
-
-        Chain<Searcher> chain = new Chain<>("test", new RateLimitingSearcher(new RateLimitingConfig(rateLimitingConfig),
-                        new ClusterInfoConfig(clusterInfoConfig),
-                        metric, clock),
-                new CostSettingSearcher());
-        assertEquals(2, tryRequests(chain, "id1"), "'rate' request are available initially");
+        var chain = createChain(false, clock, metric);
+        assertEquals(2, tryRequests(chain, "id1"), "'rate/nodes' request are available initially");
         assertTrue(executeWasAllowed(chain, "id1", true), "However, don't reject if we dryRun");
         clock.advance(Duration.ofMillis(1500)); // causes 2 new requests to become available
         assertEquals(2, tryRequests(chain, "id1"), "'rate' new requests became available");
@@ -76,6 +63,34 @@ public class RateLimitingSearcherTestCase {
         assertEquals(requestsToTry - 2 + requestsToTry - 4,         map.get(metric.point("id", "id2")).getCount());
     }
 
+    @Test
+    void testLocalRateLimiting() {
+        ManualClock clock = new ManualClock();
+        MetricReceiver.MockReceiver metric = new MetricReceiver.MockReceiver();
+        var chain = createChain(true, clock, metric);
+
+        assertEquals(9, tryRequests(chain, "id1"), "'rate' request are available initially");
+    }
+
+    private Chain<Searcher> createChain(boolean localRate, ManualClock clock, MetricReceiver.MockReceiver metric) {
+        RateLimitingConfig.Builder rateLimitingConfig = new RateLimitingConfig.Builder();
+        rateLimitingConfig.maxAvailableCapacity(4);
+        rateLimitingConfig.capacityIncrement(2);
+        rateLimitingConfig.recheckForCapacityProbability(1.0);
+        rateLimitingConfig.localRate(localRate);
+
+        ClusterInfoConfig.Builder clusterInfoConfig = new ClusterInfoConfig.Builder();
+        clusterInfoConfig.clusterId("testCluster");
+        clusterInfoConfig.nodeCount(4);
+
+
+        return new Chain<>("test", new RateLimitingSearcher(new RateLimitingConfig(rateLimitingConfig),
+                                                                new ClusterInfoConfig(clusterInfoConfig),
+                                                                metric,
+                                                                clock),
+                                       new CostSettingSearcher());
+    }
+
     private int requestsToTry = 50;
 
     /**