From d7c3de1fd9600c67fdfe620aa2ee675109ec1481 Mon Sep 17 00:00:00 2001 From: Jon Bratseth Date: Wed, 29 Nov 2023 16:35:56 +0100 Subject: Support node-local rates --- .../test/RateLimitingSearcherTestCase.java | 45 ++++++++++++++-------- 1 file changed, 30 insertions(+), 15 deletions(-) (limited to 'container-search/src/test/java/com/yahoo/search') 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 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 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; /** -- cgit v1.2.3