From c7dff29470a9247a61a1ed952f153c7f84dae3e4 Mon Sep 17 00:00:00 2001
From: Nathan Ziebart <nziebart@palantir.com>
Date: Wed, 15 Nov 2017 21:38:02 +0000
Subject: [PATCH] respect max backoff itme

---
 .../atlasdb/factory/TransactionManagers.java        |  4 +++-
 .../atlasdb/qos/ratelimit/QosRateLimiter.java       | 13 +++++++------
 .../atlasdb/qos/ratelimit/QosRateLimiters.java      |  6 +++---
 .../atlasdb/qos/ratelimit/QosRateLimiterTest.java   | 12 +++++++++++-
 4 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/atlasdb-config/src/main/java/com/palantir/atlasdb/factory/TransactionManagers.java b/atlasdb-config/src/main/java/com/palantir/atlasdb/factory/TransactionManagers.java
index fd662a8fae3..cf1450fb8eb 100644
--- a/atlasdb-config/src/main/java/com/palantir/atlasdb/factory/TransactionManagers.java
+++ b/atlasdb-config/src/main/java/com/palantir/atlasdb/factory/TransactionManagers.java
@@ -405,7 +405,9 @@ SerializableTransactionManager serializable() {
 
     private QosClient getQosClient(QosClientConfig config) {
         // TODO(nziebart): create a RefreshingRateLimiter
-        QosRateLimiters rateLimiters = QosRateLimiters.create(config.limits());
+        QosRateLimiters rateLimiters = QosRateLimiters.create(
+                config.limits(),
+                config.maxBackoffSleepTime().toMilliseconds());
         return AtlasDbQosClient.create(rateLimiters);
     }
 
diff --git a/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiter.java b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiter.java
index 4ea8fc5e227..7bb5773349b 100644
--- a/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiter.java
+++ b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiter.java
@@ -34,21 +34,22 @@ public class QosRateLimiter {
 
     private static final double MAX_BURST_SECONDS = 5;
     private static final double UNLIMITED_RATE = Double.MAX_VALUE;
-    private static final int MAX_WAIT_TIME_SECONDS = 10;
 
+    private final long maxBackoffTimeMillis;
     private RateLimiter rateLimiter;
 
-    public static QosRateLimiter create() {
-        return new QosRateLimiter(RateLimiter.SleepingStopwatch.createFromSystemTimer());
+    public static QosRateLimiter create(long maxBackoffTimeMillis) {
+        return new QosRateLimiter(RateLimiter.SleepingStopwatch.createFromSystemTimer(), maxBackoffTimeMillis);
     }
 
     @VisibleForTesting
-    QosRateLimiter(RateLimiter.SleepingStopwatch stopwatch) {
+    QosRateLimiter(RateLimiter.SleepingStopwatch stopwatch, long maxBackoffTimeMillis) {
         rateLimiter = new SmoothRateLimiter.SmoothBursty(
                 stopwatch,
                 MAX_BURST_SECONDS);
 
         rateLimiter.setRate(UNLIMITED_RATE);
+        this.maxBackoffTimeMillis = maxBackoffTimeMillis;
     }
 
     /**
@@ -67,8 +68,8 @@ public void updateRate(int unitsPerSecond) {
     public Duration consumeWithBackoff(int estimatedNumUnits) {
         Optional<Duration> waitTime = rateLimiter.tryAcquire(
                 estimatedNumUnits,
-                MAX_WAIT_TIME_SECONDS,
-                TimeUnit.SECONDS);
+                maxBackoffTimeMillis,
+                TimeUnit.MILLISECONDS);
 
         if (!waitTime.isPresent()) {
             throw new RuntimeException("rate limited");
diff --git a/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiters.java b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiters.java
index 503b5b79d39..a60696dc649 100644
--- a/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiters.java
+++ b/qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiters.java
@@ -23,11 +23,11 @@
 @Value.Immutable
 public interface QosRateLimiters {
 
-    static QosRateLimiters create(QosLimitsConfig config) {
-        QosRateLimiter readLimiter = QosRateLimiter.create();
+    static QosRateLimiters create(QosLimitsConfig config, long maxBackoffSleepTimeMillis) {
+        QosRateLimiter readLimiter = QosRateLimiter.create(maxBackoffSleepTimeMillis);
         readLimiter.updateRate(config.readBytesPerSecond());
 
-        QosRateLimiter writeLimiter = QosRateLimiter.create();
+        QosRateLimiter writeLimiter = QosRateLimiter.create(maxBackoffSleepTimeMillis);
         writeLimiter.updateRate(config.writeBytesPerSecond());
 
         return ImmutableQosRateLimiters.builder()
diff --git a/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiterTest.java b/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiterTest.java
index 8d88ad58b21..c68278e43d3 100644
--- a/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiterTest.java
+++ b/qos-service-impl/src/test/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiterTest.java
@@ -30,9 +30,10 @@
 public class QosRateLimiterTest {
 
     private static final long START_TIME_MICROS = 0L;
+    private static final long MAX_BACKOFF_TIME_MILLIS = 10_000;
 
     RateLimiter.SleepingStopwatch stopwatch = mock(RateLimiter.SleepingStopwatch.class);
-    QosRateLimiter limiter = new QosRateLimiter(stopwatch);
+    QosRateLimiter limiter = new QosRateLimiter(stopwatch, MAX_BACKOFF_TIME_MILLIS);
 
     @Before
     public void before() {
@@ -63,6 +64,15 @@ public void limitsByThrowingIfSleepTimeIsTooGreat() {
                 .hasMessageContaining("rate limited");
     }
 
+    @Test
+    public void doesNotThrowIfMaxBackoffTimeIsVeryLarge() {
+        QosRateLimiter limiterWithLargeBackoffLimit = new QosRateLimiter(stopwatch, Long.MAX_VALUE);
+        limiterWithLargeBackoffLimit.updateRate(10);
+
+        limiterWithLargeBackoffLimit.consumeWithBackoff(1_000_000_000);
+        limiterWithLargeBackoffLimit.consumeWithBackoff(1_000_000_000);
+    }
+
     @Test
     public void consumingAdditionalUnitsPenalizesFutureCallers() {
         limiter.updateRate(10);