Hsaraogi/use qos exceptions when rate limiting (#2926)

* Bump remoting

* Throw QosExceptions.Throttle

* Delete RateLimitExceededException

* Bump guava (#2929)

* release note

atlasdb-cassandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/CassandraKeyValueServiceImpl.java

@@ -123,7 +123,6 @@
 import com.palantir.atlasdb.qos.FakeQosClient;
 import com.palantir.atlasdb.qos.QosClient;
 import com.palantir.atlasdb.qos.ratelimit.QosAwareThrowables;
-import com.palantir.atlasdb.qos.ratelimit.RateLimitExceededException;
 import com.palantir.atlasdb.util.AnnotatedCallable;
 import com.palantir.atlasdb.util.AnnotationType;
 import com.palantir.atlasdb.util.AtlasDbMetrics;
@@ -2517,9 +2516,6 @@ private <V> List<V> runAllTasksCancelOnFailure(List<Callable<V>> tasks) {
  try {
  //Callable<Void> returns null, so can't use immutable list
  return Collections.singletonList(tasks.get(0).call());
- } catch (RateLimitExceededException e) {
- // Prioritise over
- throw e;
  } catch (Exception e) {
  throw QosAwareThrowables.unwrapAndThrowRateLimitExceededOrAtlasDbDependencyException(e);
  }

atlasdb-client/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosAwareThrowables.java

@@ -20,6 +20,7 @@
 import java.util.concurrent.ExecutionException;
 
 import com.palantir.common.base.Throwables;
+import com.palantir.remoting.api.errors.QosException;
 
 public final class QosAwareThrowables {
  private QosAwareThrowables() {
@@ -36,8 +37,8 @@ public static RuntimeException unwrapAndThrowRateLimitExceededOrAtlasDbDependenc
  if (ex instanceof ExecutionException || ex instanceof InvocationTargetException) {
  // Needs to be this way in case you have ITE(RLE) or variants of that.
  throw unwrapAndThrowRateLimitExceededOrAtlasDbDependencyException(ex.getCause());
- } else if (ex instanceof RateLimitExceededException) {
- throw (RateLimitExceededException) ex;
+ } else if (ex instanceof QosException.Throttle) {
+ throw (QosException.Throttle) ex;
  }
  throw Throwables.unwrapAndThrowAtlasDbDependencyException(ex);
  }

atlasdb-client/src/test/java/com/palantir/atlasdb/qos/ratelimit/QosAwareThrowablesTest.java

@@ -25,17 +25,17 @@
 import org.junit.Test;
 
 import com.palantir.common.exception.AtlasDbDependencyException;
+import com.palantir.remoting.api.errors.QosException;
 
 public class QosAwareThrowablesTest {
- private static final Exception RATE_LIMIT_EXCEEDED_EXCEPTION =
- new RateLimitExceededException("Stop!");
+ private static final Exception THROTTLE_EXCEPTION = QosException.throttle();
  private static final Exception ATLASDB_DEPENDENCY_EXCEPTION =
  new AtlasDbDependencyException("The TimeLock is dead, long live the TimeLock");
 
  @Test
  public void unwrapAndThrowRateLimitExceededOrAtlasDbDependencyExceptionCanThrowRateLimitExceededException() {
  assertThatThrownBy(() -> QosAwareThrowables.unwrapAndThrowRateLimitExceededOrAtlasDbDependencyException(
- RATE_LIMIT_EXCEEDED_EXCEPTION)).isEqualTo(RATE_LIMIT_EXCEEDED_EXCEPTION);
+ THROTTLE_EXCEPTION)).isEqualTo(THROTTLE_EXCEPTION);
  }
 
  @Test
@@ -47,9 +47,9 @@ public void unwrapAndThrowRateLimitExceededOrAtlasDbDependencyExceptionCanThrowA
  @Test
  public void unwrapAndThrowRateLimitExceededOrAtlasDbDependencyExceptionThrowsWrappedRateLimitExceededExceptions() {
  assertThatThrownBy(() -> QosAwareThrowables.unwrapAndThrowRateLimitExceededOrAtlasDbDependencyException(
- new ExecutionException(RATE_LIMIT_EXCEEDED_EXCEPTION))).isEqualTo(RATE_LIMIT_EXCEEDED_EXCEPTION);
+ new ExecutionException(THROTTLE_EXCEPTION))).isEqualTo(THROTTLE_EXCEPTION);
  assertThatThrownBy(() -> QosAwareThrowables.unwrapAndThrowRateLimitExceededOrAtlasDbDependencyException(
- new InvocationTargetException(RATE_LIMIT_EXCEEDED_EXCEPTION))).isEqualTo(RATE_LIMIT_EXCEEDED_EXCEPTION);
+ new InvocationTargetException(THROTTLE_EXCEPTION))).isEqualTo(THROTTLE_EXCEPTION);
  }
 
  @Test

atlasdb-ete-tests/src/test/java/com/palantir/atlasdb/ete/QosCassandraReadEteTest.java

@@ -37,12 +37,12 @@
 import com.google.common.collect.ImmutableList;
 import com.palantir.atlasdb.keyvalue.api.RangeRequest;
 import com.palantir.atlasdb.keyvalue.api.RowResult;
-import com.palantir.atlasdb.qos.ratelimit.RateLimitExceededException;
 import com.palantir.atlasdb.table.description.ValueType;
 import com.palantir.atlasdb.todo.ImmutableTodo;
 import com.palantir.atlasdb.todo.Todo;
 import com.palantir.atlasdb.todo.TodoSchema;
 import com.palantir.common.base.BatchingVisitable;
+import com.palantir.remoting.api.errors.QosException;
 
 public class QosCassandraReadEteTest extends QosCassandraEteTestSetup {
  private static final int ONE_TODO_SIZE_IN_BYTES = 1050;
@@ -83,8 +83,8 @@ public void shouldNotBeAbleToReadLargeAmountsIfSoftLimitSleepWillBeMoreThanConfi
  assertThat(readOneBatchOfSize(20)).hasSize(20);
 
  assertThatThrownBy(() -> readOneBatchOfSize(100))
- .isInstanceOf(RateLimitExceededException.class)
- .hasMessage("Rate limited. Available capacity has been exhausted.");
+ .isInstanceOf(QosException.Throttle.class)
+ .hasMessage("Suggesting request throttling with optional retryAfter duration: Optional.empty");
  }
 
  @Test
@@ -121,7 +121,7 @@ private void assertThatAllReadsWereSuccessful(List<Future<List<Todo>>> futures,
  try {
  assertThat(future.get()).hasSize(numReadsPerThread);
  } catch (ExecutionException e) {
- if (e.getCause() instanceof RateLimitExceededException) {
+ if (e.getCause() instanceof QosException.Throttle) {
  exceptionCounter.getAndIncrement();
  }
  } catch (InterruptedException e) {

atlasdb-ete-tests/src/test/java/com/palantir/atlasdb/ete/QosCassandraWriteEteTest.java

@@ -34,7 +34,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.base.Stopwatch;
 import com.google.common.base.Throwables;
-import com.palantir.atlasdb.qos.ratelimit.RateLimitExceededException;
+import com.palantir.remoting.api.errors.QosException;
 
 public class QosCassandraWriteEteTest extends QosCassandraEteTestSetup {
 
@@ -74,13 +74,13 @@ public void shouldNotBeAbleToWriteLargeAmountsIfSoftLimitSleepWillBeMoreThanConf
  writeNTodosOfSize(1, 100_000);
 
  assertThatThrownBy(() -> writeNTodosOfSize(1, 100_000))
- .isInstanceOf(RateLimitExceededException.class)
- .hasMessage("Rate limited. Available capacity has been exhausted.");
+ .isInstanceOf(QosException.Throttle.class)
+ .hasMessage("Suggesting request throttling with optional retryAfter duration: Optional.empty");
 
  // One write smaller than the rate limit should also be rate limited.
  assertThatThrownBy(() -> writeNTodosOfSize(5, 10))
- .isInstanceOf(RateLimitExceededException.class)
- .hasMessage("Rate limited. Available capacity has been exhausted.");
+ .isInstanceOf(QosException.Throttle.class)
+ .hasMessage("Suggesting request throttling with optional retryAfter duration: Optional.empty");
  }
 
  @Test
@@ -117,7 +117,7 @@ private void assertThatAllWritesWereSuccessful(List<Future> futures) {
  try {
  future.get();
  } catch (ExecutionException e) {
- if (e.getCause() instanceof RateLimitExceededException) {
+ if (e.getCause() instanceof QosException.Throttle) {
  exceptionCounter.getAndIncrement();
  }
  } catch (InterruptedException e) {

atlasdb-tests-shared/src/test/java/com/palantir/atlasdb/transaction/impl/SnapshotTransactionTest.java

@@ -603,6 +603,7 @@ private Optional<String> runTaskWithInvalidLocks(Pair<String, LockAwareTransacti
  }
  }
 
+ @SuppressWarnings("CheckReturnValue")
  private List<Pair<String, LockAwareTransactionTask<Void, Exception>>> getThoroughTableReadTasks() {
  ImmutableList.Builder<Pair<String, LockAwareTransactionTask<Void, Exception>>> tasks = ImmutableList.builder();
  final int batchHint = 1;

docs/source/release_notes/release-notes.rst

@@ -50,6 +50,14 @@ develop
  * - Type
  - Change
 
+ * - |fixed| |devbreak|
+ - AtlasDB clients will receive a ``QosException.Throttle`` for requests that are throttled and http-remoting
+ should handle them appropriately based on the backOff strategy provided by the application. Note that this is
+ an experimental feature and we do not expect it to be enabled anywhere. This is a dev break as the exception type
+ has changed from ``RateLimitExceededException`` to ``QosException.Throttle``.
+ (`Throttle <https://github.com/palantir/http-remoting/blob/a14a0894c2f5d1a415c5ee2727e9c79d73255b7b/okhttp-clients/src/main/java/com/palantir/remoting3/okhttp/RemotingOkHttpCall.java#L221-L233>`__)
+ (`Pull Request <https://github.com/palantir/atlasdb/pull/2926>`__)
+
  * - |improved| |metrics|
  - Use tags in sweep outcome metrics instead of using each name per outcome.
  (`Pull Request <https://github.com/palantir/atlasdb/pull/2927>`__)

docs/source/services/qos_service/enabling-rate-limiting.rst

@@ -19,7 +19,7 @@ QoS also publishes client level metrics (meters) for the following:
  `com.palantir.atlasdb.qos.metrics.QosMetrics.writeTime`
  `com.palantir.atlasdb.qos.metrics.QosMetrics.rowsWritten`
  `com.palantir.atlasdb.qos.metrics.QosMetrics.backoffTime`
- `com.palantir.atlasdb.qos.metrics.QosMetrics.rateLimitedExceptions`
+ `com.palantir.atlasdb.qos.metrics.QosMetrics.throttleExceptions`
 
 A client should ideally look at the historical values of the ``bytesRead`` and ``bytesWritten`` metrics to determine the limits.
 

qos-service-impl/src/main/java/com/palantir/atlasdb/qos/client/AtlasDbQosClient.java

@@ -31,7 +31,7 @@
 import com.palantir.atlasdb.qos.metrics.QosMetrics;
 import com.palantir.atlasdb.qos.ratelimit.QosRateLimiter;
 import com.palantir.atlasdb.qos.ratelimit.QosRateLimiters;
-import com.palantir.atlasdb.qos.ratelimit.RateLimitExceededException;
+import com.palantir.remoting.api.errors.QosException;
 
 public class AtlasDbQosClient implements QosClient {
 
@@ -77,8 +77,8 @@ private <T, E extends Exception> T execute(
  Duration waitTime = rateLimiter.consumeWithBackoff(estimatedWeight.numBytes());
  metrics.recordBackoffMicros(TimeUnit.NANOSECONDS.toMicros(waitTime.toNanos()));
  }
- } catch (RateLimitExceededException ex) {
- metrics.recordRateLimitedException();
+ } catch (QosException.Throttle ex) {
+ metrics.recordThrottleExceptions();
  throw ex;
  }
 

qos-service-impl/src/main/java/com/palantir/atlasdb/qos/metrics/QosMetrics.java

@@ -45,7 +45,7 @@ public class QosMetrics {
  private final Meter rowsWritten;
 
  private final Meter backoffTime;
- private final Meter rateLimitedExceptions;
+ private final Meter throttleExceptions;
 
  public QosMetrics() {
  readRequestCount = metricsManager.registerOrGetMeter(QosMetrics.class, "numReadRequests");
@@ -62,7 +62,7 @@ public QosMetrics() {
  rowsWritten = metricsManager.registerOrGetMeter(QosMetrics.class, "rowsWritten");
 
  backoffTime = metricsManager.registerOrGetMeter(QosMetrics.class, "backoffTime");
- rateLimitedExceptions = metricsManager.registerOrGetMeter(QosMetrics.class, "rateLimitedExceptions");
+ throttleExceptions = metricsManager.registerOrGetMeter(QosMetrics.class, "throttleExceptions");
  }
 
  public void recordReadEstimate(QueryWeight weight) {
@@ -91,10 +91,10 @@ public void recordBackoffMicros(long backoffTimeMicros) {
  }
  }
 
- public void recordRateLimitedException() {
+ public void recordThrottleExceptions() {
  log.info("Rate limit exceeded and backoff time would be more than the configured maximum. "
  + "Throwing a throttling exception");
- rateLimitedExceptions.mark();
+ throttleExceptions.mark();
  }
 
 }

qos-service-impl/src/main/java/com/palantir/atlasdb/qos/ratelimit/QosRateLimiter.java

@@ -30,6 +30,7 @@
 import com.palantir.atlasdb.qos.ratelimit.guava.RateLimiter;
 import com.palantir.atlasdb.qos.ratelimit.guava.SmoothRateLimiter;
 import com.palantir.logsafe.SafeArg;
+import com.palantir.remoting.api.errors.QosException;
 
 /**
  * A rate limiter for database queries, based on "units" of expense. This limiter strives to maintain an upper limit on
@@ -107,7 +108,7 @@ public Duration consumeWithBackoff(long estimatedNumUnits) {
  TimeUnit.MILLISECONDS);
 
  if (!waitTime.isPresent()) {
- throw new RateLimitExceededException("Rate limited. Available capacity has been exhausted.");
+ throw QosException.throttle();
  }
 
  return waitTime.get();

qos-service-impl/src/test/java/com/palantir/atlasdb/qos/client/AtlasDbQosClientTest.java

@@ -42,7 +42,7 @@
 import com.palantir.atlasdb.qos.ratelimit.ImmutableQosRateLimiters;
 import com.palantir.atlasdb.qos.ratelimit.QosRateLimiter;
 import com.palantir.atlasdb.qos.ratelimit.QosRateLimiters;
-import com.palantir.atlasdb.qos.ratelimit.RateLimitExceededException;
+import com.palantir.remoting.api.errors.QosException;
 
 @RunWith(Parameterized.class)
 public final class AtlasDbQosClientTest {
@@ -181,7 +181,7 @@ public void propagatesCheckedExceptions() throws TestCheckedException {
  throw new TestCheckedException();
  }, weigher)).isInstanceOf(TestCheckedException.class);
 
- verify(metrics, never()).recordRateLimitedException();
+ verify(metrics, never()).recordThrottleExceptions();
  }
 
  @Test
@@ -194,11 +194,11 @@ public void recordsBackoffTime() {
 
  @Test
  public void recordsBackoffExceptions() {
- when(readLimiter.consumeWithBackoff(anyLong())).thenThrow(new RateLimitExceededException("rate limited"));
+ when(readLimiter.consumeWithBackoff(anyLong())).thenThrow(QosException.throttle());
  assertThatThrownBy(() -> qosClient.executeRead(() -> "foo", weigher)).isInstanceOf(
- RateLimitExceededException.class);
+ QosException.Throttle.class);
 
- verify(metrics).recordRateLimitedException();
+ verify(metrics).recordThrottleExceptions();
  }
 
  @Test
@@ -207,7 +207,7 @@ public void doesNotRecordRuntimeExceptions() {
  assertThatThrownBy(() -> qosClient.executeRead(() -> "foo", weigher)).isInstanceOf(
  RuntimeException.class);
 
- verify(metrics, never()).recordRateLimitedException();
+ verify(metrics, never()).recordThrottleExceptions();
  }
 
  static class TestCheckedException extends Exception {}