Add MethodsCustomRateLimitResolver to enable developers to override the rate limiter's behavior (#1134)

Author: seratch

slack-api-client/src/main/java/com/slack/api/methods/MethodsConfig.java

@@ -70,6 +70,11 @@ public void setCustomThreadPoolSizes(Map<String, Integer> customThreadPoolSizes)
         public void setExecutorServiceProvider(ExecutorServiceProvider executorServiceProvider) {
             throwException();
         }
+
+        @Override
+        public void setCustomRateLimitResolver(MethodsCustomRateLimitResolver customRateLimitResolver) {
+            throwException();
+        }
     };
 
     @Builder.Default
@@ -108,4 +113,10 @@ public void setExecutorServiceProvider(ExecutorServiceProvider executorServicePr
     @Builder.Default
     private MetricsDatastore metricsDatastore = new MemoryMetricsDatastore(1);
 
+    /**
+     * Custom handlers for adjusting rate limits.
+     */
+    @Builder.Default
+    private MethodsCustomRateLimitResolver customRateLimitResolver = MethodsCustomRateLimitResolver.DEFAULT;
+
 }

slack-api-client/src/main/java/com/slack/api/methods/MethodsCustomRateLimitResolver.java

@@ -0,0 +1,48 @@
+package com.slack.api.methods;
+
+import java.util.Optional;
+
+/**
+ * Custom rate limit resolver.
+ * <p>
+ * When you need to adjust the rate limits for some reason,
+ * you can implement this interface and pass it to MethodsConfig#setCustomRateLimitResolver() method.
+ * With that, you can override the allowed requests per minute for selected methods.
+ */
+public interface MethodsCustomRateLimitResolver {
+
+    /**
+     * Return a present value only when you want to override the allowed requests per minute.
+     * Otherwise, returning Optional.empty() will result in falling back to the built-in calculation.
+     *
+     * @param teamId     the workspace ID
+     * @param methodName the method name such as "auth.test"
+     * @return the number of allowed requests per minute
+     */
+    Optional<Integer> getCustomAllowedRequestsPerMinute(String teamId, String methodName);
+
+    /**
+     * Return a present value only when you want to override the allowed requests per minute.
+     * Otherwise, returning Optional.empty() will result in falling back to the built-in calculation.
+     *
+     * @param teamId  the workspace ID
+     * @param channel either a channel ID or channel name
+     * @return the number of allowed requests per minute
+     */
+    Optional<Integer> getCustomAllowedRequestsForChatPostMessagePerMinute(String teamId, String channel);
+
+    class Default implements MethodsCustomRateLimitResolver {
+        @Override
+        public Optional<Integer> getCustomAllowedRequestsPerMinute(String teamId, String methodName) {
+            return Optional.empty();
+        }
+
+        @Override
+        public Optional<Integer> getCustomAllowedRequestsForChatPostMessagePerMinute(String teamId, String channel) {
+            return Optional.empty();
+        }
+    }
+
+    MethodsCustomRateLimitResolver DEFAULT = new Default();
+
+}

slack-api-client/src/main/java/com/slack/api/methods/impl/AsyncMethodsRateLimiter.java

@@ -1,6 +1,7 @@
 package com.slack.api.methods.impl;
 
 import com.slack.api.methods.MethodsConfig;
+import com.slack.api.methods.MethodsCustomRateLimitResolver;
 import com.slack.api.methods.MethodsRateLimits;
 import com.slack.api.rate_limits.RateLimiter;
 import com.slack.api.rate_limits.WaitTime;
@@ -10,10 +11,13 @@
 
 import java.util.Optional;
 
+import static com.slack.api.methods.MethodsRateLimitTier.SpecialTier_chat_postMessage;
+
 @Slf4j
 public class AsyncMethodsRateLimiter implements RateLimiter {
 
     private final MetricsDatastore metricsDatastore;
+    private final MethodsCustomRateLimitResolver customRateLimitResolver;
     private final WaitTimeCalculator waitTimeCalculator;
 
     public MetricsDatastore getMetricsDatastore() {
@@ -22,6 +26,7 @@ public MetricsDatastore getMetricsDatastore() {
 
     public AsyncMethodsRateLimiter(MethodsConfig config) {
         this.metricsDatastore = config.getMetricsDatastore();
+        this.customRateLimitResolver = config.getCustomRateLimitResolver();
         this.waitTimeCalculator = new MethodsWaitTimeCalculator(config);
     }
 
@@ -30,15 +35,32 @@ public WaitTime acquireWaitTime(String teamId, String methodName) {
         return waitTimeCalculator.calculateWaitTime(
                 teamId,
                 methodName,
-                waitTimeCalculator.getAllowedRequestsPerMinute(MethodsRateLimits.lookupRateLimitTier(methodName))
+                getAllowedRequestsPerMinute(teamId, methodName)
         );
     }
 
+    public int getAllowedRequestsPerMinute(String teamId, String methodName) {
+        Optional<Integer> custom = customRateLimitResolver.getCustomAllowedRequestsPerMinute(teamId, methodName);
+        if (custom.isPresent()) {
+            return custom.get();
+        }
+        return waitTimeCalculator.getAllowedRequestsPerMinute(MethodsRateLimits.lookupRateLimitTier(methodName));
+    }
+
+    public int getAllowedRequestsForChatPostMessagePerMinute(String teamId, String channel) {
+        Optional<Integer> custom = customRateLimitResolver.getCustomAllowedRequestsForChatPostMessagePerMinute(teamId, channel);
+        if (custom.isPresent()) {
+            return custom.get();
+        }
+        return waitTimeCalculator.getAllowedRequestsPerMinute(SpecialTier_chat_postMessage);
+    }
+
     @Override
     public WaitTime acquireWaitTimeForChatPostMessage(String teamId, String channel) {
         return waitTimeCalculator.calculateWaitTimeForChatPostMessage(
                 teamId,
-                channel
+                channel,
+                getAllowedRequestsForChatPostMessagePerMinute(teamId, channel)
         );
     }
 

slack-api-client/src/main/java/com/slack/api/rate_limits/WaitTimeCalculator.java

@@ -44,11 +44,19 @@ public WaitTime calculateWaitTime(String teamId, String key, int allowedRequests
         }
     }
 
+    /**
+     * @deprecated Use #calculateWaitTimeForChatPostMessage(String, String, int) instead
+     */
+    @Deprecated
     public WaitTime calculateWaitTimeForChatPostMessage(String teamId, String channel) {
+        return calculateWaitTimeForChatPostMessage(teamId, channel, getAllowedRequestsPerMinute(SpecialTier_chat_postMessage));
+    }
+
+    public WaitTime calculateWaitTimeForChatPostMessage(String teamId, String channel, int allowedRequests) {
         return calculateWaitTime(
                 teamId,
                 Methods.CHAT_POST_MESSAGE + "_" + channel,
-                getAllowedRequestsPerMinute(SpecialTier_chat_postMessage)
+                allowedRequests
         );
     }
 
@@ -76,8 +84,8 @@ private static boolean isTooFastPaced(LastMinuteRequests lastMinuteRequests, int
     }
 
     public Integer getAllowedRequestsPerMinute(MethodsRateLimitTier tier) {
-        Integer allowedRequestsForOneNode = MethodsRateLimitTier.getAllowedRequestsPerMinute(tier);
-        return allowedRequestsForOneNode / getNumberOfNodes();
+        Integer allowedRequestsPerMinute = MethodsRateLimitTier.getAllowedRequestsPerMinute(tier);
+        return allowedRequestsPerMinute / getNumberOfNodes();
     }
 
 }

slack-api-client/src/test/java/test_locally/api/methods/CustomRateLimitsTest.java

@@ -0,0 +1,72 @@
+package test_locally.api.methods;
+
+import com.slack.api.Slack;
+import com.slack.api.SlackConfig;
+import com.slack.api.methods.AsyncMethodsClient;
+import com.slack.api.methods.MethodsCustomRateLimitResolver;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import util.MockSlackApiServer;
+
+import java.util.Optional;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.greaterThan;
+import static util.MockSlackApi.ValidToken;
+
+public class CustomRateLimitsTest {
+
+    MockSlackApiServer server = new MockSlackApiServer();
+    Slack slack;
+
+    static int customLimit = 30;
+
+    static class MyCustomRateLimitResolver implements MethodsCustomRateLimitResolver {
+
+        @Override
+        public Optional<Integer> getCustomAllowedRequestsPerMinute(String teamId, String methodName) {
+            if (methodName.equals("auth.test")) {
+                return Optional.of(customLimit);
+            }
+            return Optional.empty();
+        }
+
+        @Override
+        public Optional<Integer> getCustomAllowedRequestsForChatPostMessagePerMinute(String teamId, String channel) {
+            return Optional.empty();
+        }
+    }
+
+    @Before
+    public void setup() throws Exception {
+        server.start();
+        SlackConfig config = new SlackConfig();
+        // Customize the rate limits
+        config.getMethodsConfig().setCustomRateLimitResolver(new MyCustomRateLimitResolver());
+        config.setMethodsEndpointUrlPrefix(server.getMethodsEndpointPrefix());
+        slack = Slack.getInstance(config);
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        server.stop();
+    }
+
+    @Test
+    public void burstTest() throws Exception {
+        AsyncMethodsClient client = slack.methodsAsync(ValidToken);
+        client.authTest(r -> r).get();
+        client.authTest(r -> r).get();
+        client.authTest(r -> r).get();
+        client.authTest(r -> r).get();
+
+        long startMillis = System.currentTimeMillis();
+        // With the default settings, this call is on the "safe" pace, so spentMillis will be 15L or so
+        client.authTest(r -> r).get();
+        long spentMillis = System.currentTimeMillis() - startMillis;
+        assertThat(spentMillis, is(greaterThan(500L )));
+    }
+
+}