[QoS] Fix/qos system table rate limiting

atlasdb-cassandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/qos/QosCassandraClient.java

@@ -19,6 +19,7 @@
 import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 
 import org.apache.cassandra.thrift.CASResult;
 import org.apache.cassandra.thrift.Cassandra;
@@ -43,12 +44,16 @@
 import com.palantir.atlasdb.keyvalue.api.TableReference;
 import com.palantir.atlasdb.keyvalue.cassandra.CassandraClient;
 import com.palantir.atlasdb.keyvalue.cassandra.CqlQuery;
+import com.palantir.atlasdb.keyvalue.cassandra.HiddenTables;
 import com.palantir.atlasdb.qos.QosClient;
+import com.palantir.atlasdb.transaction.impl.TransactionConstants;
 
 @SuppressWarnings({"all"}) // thrift variable names.
 public class QosCassandraClient implements CassandraClient {
 
     private static final Logger log = LoggerFactory.getLogger(CassandraClient.class);
+    private static final Function<TableReference, Boolean> ZERO_ESTIMATE_DETERMINING_FUNCTION = tRef ->
+            tRef.equals(TransactionConstants.TRANSACTION_TABLE) || new HiddenTables().isHidden(tRef);
 
     private final CassandraClient client;
     private final QosClient qosClient;
@@ -69,7 +74,7 @@ public Map<ByteBuffer, List<ColumnOrSuperColumn>> multiget_slice(String kvsMetho
             throws InvalidRequestException, UnavailableException, TimedOutException, TException {
         return qosClient.executeRead(
                 () -> client.multiget_slice(kvsMethodName, tableRef, keys, predicate, consistency_level),
-                ThriftQueryWeighers.multigetSlice(keys));
+                ThriftQueryWeighers.multigetSlice(keys, ZERO_ESTIMATE_DETERMINING_FUNCTION.apply(tableRef)));
     }
 
     @Override
@@ -78,7 +83,7 @@ public List<KeySlice> get_range_slices(String kvsMethodName, TableReference tabl
             throws InvalidRequestException, UnavailableException, TimedOutException, TException {
         return qosClient.executeRead(
                 () -> client.get_range_slices(kvsMethodName, tableRef, predicate, range, consistency_level),
-                ThriftQueryWeighers.getRangeSlices(range));
+                ThriftQueryWeighers.getRangeSlices(range, ZERO_ESTIMATE_DETERMINING_FUNCTION.apply(tableRef)));
     }
 
     @Override
@@ -99,7 +104,7 @@ public ColumnOrSuperColumn get(TableReference tableReference, ByteBuffer key, by
             throws InvalidRequestException, NotFoundException, UnavailableException, TimedOutException, TException {
         return qosClient.executeRead(
                 () -> client.get(tableReference, key, column, consistency_level),
-                ThriftQueryWeighers.GET);
+                ThriftQueryWeighers.get(ZERO_ESTIMATE_DETERMINING_FUNCTION.apply(tableReference)));
     }
 
     @Override

atlasdb-cassandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/qos/ThriftQueryWeighers.java

@@ -50,31 +50,79 @@ public final class ThriftQueryWeighers {
             .timeTakenNanos(TimeUnit.MILLISECONDS.toNanos(2))
             .build();
 
+    static final QueryWeight ZERO_ESTIMATED_WEIGHT = ImmutableQueryWeight.builder()
+            .numBytes(0)
+            .numDistinctRows(1)
+            .timeTakenNanos(TimeUnit.MILLISECONDS.toNanos(1))
+            .build();
+
     private ThriftQueryWeighers() { }
 
-    static QosClient.QueryWeigher<Map<ByteBuffer, List<ColumnOrSuperColumn>>> multigetSlice(List<ByteBuffer> keys) {
-        return readWeigher(ThriftObjectSizeUtils::getApproximateSizeOfColsByKey, Map::size, keys.size());
+    static QosClient.QueryWeigher<Map<ByteBuffer, List<ColumnOrSuperColumn>>> multigetSlice(List<ByteBuffer> keys,
+            boolean zeroEstimate) {
+        return zeroEstimate
+                ? readWeigherWithZeroEstimate(ThriftObjectSizeUtils::getApproximateSizeOfColsByKey, Map::size,
+                keys.size())
+                : readWeigher(ThriftObjectSizeUtils::getApproximateSizeOfColsByKey, Map::size, keys.size());
     }
 
-    static QosClient.QueryWeigher<List<KeySlice>> getRangeSlices(KeyRange keyRange) {
-        return readWeigher(ThriftObjectSizeUtils::getApproximateSizeOfKeySlices, List::size, keyRange.count);
+    static QosClient.QueryWeigher<List<KeySlice>> getRangeSlices(KeyRange keyRange, boolean zeroEstimate) {
+        return zeroEstimate
+                ? readWeigherWithZeroEstimate(ThriftObjectSizeUtils::getApproximateSizeOfKeySlices, List::size,
+                keyRange.count)
+                : readWeigher(ThriftObjectSizeUtils::getApproximateSizeOfKeySlices, List::size, keyRange.count);
     }
 
-    static final QosClient.QueryWeigher<ColumnOrSuperColumn> GET =
-            readWeigher(ThriftObjectSizeUtils::getColumnOrSuperColumnSize, ignored -> 1, 1);
+    static QosClient.QueryWeigher<ColumnOrSuperColumn> get(boolean zeroEstimate) {
+        return zeroEstimate
+                ? readWeigherWithZeroEstimate(ThriftObjectSizeUtils::getColumnOrSuperColumnSize, ignored -> 1, 1)
+                : readWeigher(ThriftObjectSizeUtils::getColumnOrSuperColumnSize, ignored -> 1, 1);
+    }
 
     static final QosClient.QueryWeigher<CqlResult> EXECUTE_CQL3_QUERY =
             // TODO(nziebart): we need to inspect the schema to see how many rows there are - a CQL row is NOT a
             // partition. rows here will depend on the type of query executed in CqlExecutor: either (column, ts) pairs,
             // or (key, column, ts) triplets
+            // Currently, transaction or metadata table queries dont use the CQL executor,
+            // but we should provide a way to estimate zero based on the tableRef if they do start using it.
             readWeigher(ThriftObjectSizeUtils::getCqlResultSize, ignored -> 1, 1);
 
     static QosClient.QueryWeigher<Void> batchMutate(Map<ByteBuffer, Map<String, List<Mutation>>> mutationMap) {
         long numRows = mutationMap.size();
         return writeWeigher(numRows, () -> ThriftObjectSizeUtils.getApproximateSizeOfMutationMap(mutationMap));
     }
 
-    private static <T> QosClient.QueryWeigher<T> readWeigher(Function<T, Long> bytesRead, Function<T, Integer> numRows,
+    private static <T> QosClient.QueryWeigher<T> readWeigherWithZeroEstimate(Function<T, Long> bytesRead,
+            Function<T, Integer> numRows,
+            int numberOfQueriedRows) {
+        return new QosClient.QueryWeigher<T>() {
+            @Override
+            public QueryWeight estimate() {
+                return ZERO_ESTIMATED_WEIGHT;
+            }
+
+            @Override
+            public QueryWeight weighSuccess(T result, long timeTakenNanos) {
+                return ImmutableQueryWeight.builder()
+                        .numBytes(safeGetNumBytesOrDefault(() -> bytesRead.apply(result)))
+                        .timeTakenNanos(timeTakenNanos)
+                        .numDistinctRows(numRows.apply(result))
+                        .build();
+            }
+
+            @Override
+            public QueryWeight weighFailure(Exception error, long timeTakenNanos) {
+                return ImmutableQueryWeight.builder()
+                        .from(DEFAULT_ESTIMATED_WEIGHT)
+                        .numBytes(ESTIMATED_NUM_BYTES_PER_ROW * numberOfQueriedRows)
+                        .timeTakenNanos(timeTakenNanos)
+                        .build();
+            }
+        };
+    }
+
+    private static <T> QosClient.QueryWeigher<T> readWeigher(Function<T, Long> bytesRead,
+            Function<T, Integer> numRows,
             int numberOfQueriedRows) {
         return new QosClient.QueryWeigher<T>() {
             @Override

atlasdb-cassandra/src/test/java/com/palantir/atlasdb/keyvalue/cassandra/qos/ThriftQueryWeighersTest.java

@@ -22,7 +22,6 @@
 import java.util.List;
 import java.util.Map;
 
-import org.apache.cassandra.thrift.Column;
 import org.apache.cassandra.thrift.ColumnOrSuperColumn;
 import org.apache.cassandra.thrift.CqlResult;
 import org.apache.cassandra.thrift.KeyRange;
@@ -41,7 +40,6 @@ public class ThriftQueryWeighersTest {
     private static final ByteBuffer BYTES1 = ByteBuffer.allocate(3);
     private static final ByteBuffer BYTES2 = ByteBuffer.allocate(7);
     private static final ColumnOrSuperColumn COLUMN_OR_SUPER = new ColumnOrSuperColumn();
-    private static final Column COLUMN = new Column();
     private static final KeySlice KEY_SLICE = new KeySlice();
     private static final Mutation MUTATION = new Mutation();
 
@@ -54,21 +52,33 @@ public class ThriftQueryWeighersTest {
 
     @Test
     public void multigetSliceWeigherEstimatesNumberOfBytesBasedOnNumberOfRows() {
-        long numBytesWithOneRow = ThriftQueryWeighers.multigetSlice(ImmutableList.of(BYTES1)).estimate().numBytes();
-        long numBytesWithTwoRows = ThriftQueryWeighers.multigetSlice(ImmutableList.of(BYTES1, BYTES2)).estimate()
-                .numBytes();
+        assertThatEstimatesAreCorrect(false,
+                ThriftQueryWeighers.multigetSlice(ImmutableList.of(BYTES1), false),
+                ThriftQueryWeighers.multigetSlice(ImmutableList.of(BYTES1, BYTES2), false));
+    }
 
-        assertThat(numBytesWithOneRow).isGreaterThan(0L);
-        assertThat(numBytesWithTwoRows).isEqualTo(numBytesWithOneRow * 2);
+    @Test
+    public void multigetSliceWithZeroEstimateWeigherEstimatesZeroNumberOfBytes() {
+        assertThatEstimatesAreCorrect(true,
+                ThriftQueryWeighers.multigetSlice(ImmutableList.of(BYTES1), true),
+                ThriftQueryWeighers.multigetSlice(ImmutableList.of(BYTES1, BYTES2), true));
     }
 
     @Test
     public void multigetSliceWeigherReturnsCorrectNumRows() {
+        assertThatWeightSuccessReturnsCorrectNumberOfRows(
+                ThriftQueryWeighers.multigetSlice(ImmutableList.of(BYTES1), false));
+        assertThatWeightSuccessReturnsCorrectNumberOfRows(
+                ThriftQueryWeighers.multigetSlice(ImmutableList.of(BYTES1), true));
+    }
+
+    private void assertThatWeightSuccessReturnsCorrectNumberOfRows(
+            QosClient.QueryWeigher<Map<ByteBuffer, List<ColumnOrSuperColumn>>> mapQueryWeigher) {
         Map<ByteBuffer, List<ColumnOrSuperColumn>> result = ImmutableMap.of(
                 BYTES1, ImmutableList.of(COLUMN_OR_SUPER, COLUMN_OR_SUPER),
                 BYTES2, ImmutableList.of(COLUMN_OR_SUPER));
 
-        long actualNumRows = ThriftQueryWeighers.multigetSlice(ImmutableList.of(BYTES1))
+        long actualNumRows = mapQueryWeigher
                 .weighSuccess(result, TIME_TAKEN)
                 .numDistinctRows();
 
@@ -77,26 +87,58 @@ public void multigetSliceWeigherReturnsCorrectNumRows() {
 
     @Test
     public void rangeSlicesWeigherEstimatesNumberOfBytesBasedOnNumberOfRows() {
-        long numBytesWithOneRow = ThriftQueryWeighers.getRangeSlices(new KeyRange(1)).estimate().numBytes();
-        long numBytesWithTwoRows = ThriftQueryWeighers.getRangeSlices(new KeyRange(2)).estimate().numBytes();
+        assertThatEstimatesAreCorrect(false,
+                ThriftQueryWeighers.getRangeSlices(new KeyRange(1), false),
+                ThriftQueryWeighers.getRangeSlices(new KeyRange(2), false));
+    }
+
+    @Test
+    public void rangeSlicesWeigherWithZeroEstimateEstimatesZeroNumberOfBytes() {
+        assertThatEstimatesAreCorrect(true,
+                ThriftQueryWeighers.getRangeSlices(new KeyRange(1), true),
+                ThriftQueryWeighers.getRangeSlices(new KeyRange(2), true));
+    }
 
-        assertThat(numBytesWithOneRow).isGreaterThan(0L);
-        assertThat(numBytesWithTwoRows).isEqualTo(numBytesWithOneRow * 2);
+    private void assertThatEstimatesAreCorrect(boolean zeroEstimate,
+            QosClient.QueryWeigher queryWeigherOneRow,
+            QosClient.QueryWeigher queryWeigherTwoRows) {
+        long numBytesWithOneRow = queryWeigherOneRow.estimate().numBytes();
+        long numBytesWithTwoRows = queryWeigherTwoRows.estimate().numBytes();
+
+        if (zeroEstimate) {
+            assertThat(numBytesWithOneRow).isEqualTo(0L);
+        } else {
+            assertThat(numBytesWithOneRow).isGreaterThan(0L);
+        }
+        assertThat(numBytesWithTwoRows).isEqualTo(zeroEstimate ? 0L : numBytesWithOneRow * 2);
     }
 
     @Test
     public void rangeSlicesWeigherReturnsCorrectNumRows() {
+        assertThatRangeSlicesWeigherReturnsCorrectNumRows(ThriftQueryWeighers.getRangeSlices(new KeyRange(1), false));
+        assertThatRangeSlicesWeigherReturnsCorrectNumRows(ThriftQueryWeighers.getRangeSlices(new KeyRange(1), true));
+    }
+
+    private void assertThatRangeSlicesWeigherReturnsCorrectNumRows(
+            QosClient.QueryWeigher<List<KeySlice>> weigher) {
         List<KeySlice> result = ImmutableList.of(KEY_SLICE, KEY_SLICE, KEY_SLICE);
 
-        long actualNumRows = ThriftQueryWeighers.getRangeSlices(new KeyRange(1)).weighSuccess(result, TIME_TAKEN)
+        long actualNumRows = weigher
+                .weighSuccess(result, TIME_TAKEN)
                 .numDistinctRows();
 
         assertThat(actualNumRows).isEqualTo(3);
     }
 
     @Test
     public void getWeigherReturnsCorrectNumRows() {
-        long actualNumRows = ThriftQueryWeighers.GET.weighSuccess(COLUMN_OR_SUPER, TIME_TAKEN).numDistinctRows();
+        assertThatGetWeighSuccessReturnsOneRow(ThriftQueryWeighers.get(false));
+        assertThatGetWeighSuccessReturnsOneRow(ThriftQueryWeighers.get(true));
+    }
+
+    private void assertThatGetWeighSuccessReturnsOneRow(
+            QosClient.QueryWeigher<ColumnOrSuperColumn> weighQuery) {
+        long actualNumRows = weighQuery.weighSuccess(COLUMN_OR_SUPER, TIME_TAKEN).numDistinctRows();
 
         assertThat(actualNumRows).isEqualTo(1);
     }
@@ -126,25 +168,20 @@ public void batchMutateWeigherReturnsCorrectNumRows() {
 
     @Test
     public void multigetSliceWeigherReturnsDefaultEstimateForFailure() {
-        QueryWeight weight = ThriftQueryWeighers.multigetSlice(ImmutableList.of(BYTES1))
-                .weighFailure(new RuntimeException(), TIME_TAKEN);
-
-        assertThat(weight).isEqualTo(DEFAULT_WEIGHT);
+        assertThatWeighFailureReturnsDefaultWeight(ThriftQueryWeighers.multigetSlice(ImmutableList.of(BYTES1), false));
+        assertThatWeighFailureReturnsDefaultWeight(ThriftQueryWeighers.multigetSlice(ImmutableList.of(BYTES1), true));
     }
 
     @Test
     public void getWeigherReturnsDefaultEstimateForFailure() {
-        QueryWeight weight = ThriftQueryWeighers.GET.weighFailure(new RuntimeException(), TIME_TAKEN);
-
-        assertThat(weight).isEqualTo(DEFAULT_WEIGHT);
+        assertThatWeighFailureReturnsDefaultWeight(ThriftQueryWeighers.get(false));
+        assertThatWeighFailureReturnsDefaultWeight(ThriftQueryWeighers.get(true));
     }
 
     @Test
     public void getRangeSlicesWeigherReturnsDefaultEstimateForFailure() {
-        QueryWeight weight = ThriftQueryWeighers.getRangeSlices(new KeyRange(1))
-                .weighFailure(new RuntimeException(), TIME_TAKEN);
-
-        assertThat(weight).isEqualTo(DEFAULT_WEIGHT);
+        assertThatWeighFailureReturnsDefaultWeight(ThriftQueryWeighers.getRangeSlices(new KeyRange(1), false));
+        assertThatWeighFailureReturnsDefaultWeight(ThriftQueryWeighers.getRangeSlices(new KeyRange(1), true));
     }
 
     @Test
@@ -158,17 +195,21 @@ public void batchMutateWeigherReturnsEstimateForFailure() {
                 .from(weigher.estimate())
                 .timeTakenNanos(TIME_TAKEN)
                 .build();
+
         QueryWeight actual = weigher.weighFailure(new RuntimeException(), TIME_TAKEN);
 
         assertThat(actual).isEqualTo(expected);
     }
 
     @Test
     public void cql3QueryWeigherReturnsDefaultEstimateForFailure() {
-        QueryWeight weight = ThriftQueryWeighers.EXECUTE_CQL3_QUERY.weighFailure(new RuntimeException(),
-                TIME_TAKEN);
+        assertThatWeighFailureReturnsDefaultWeight(ThriftQueryWeighers.EXECUTE_CQL3_QUERY);
+    }
+
+    private void assertThatWeighFailureReturnsDefaultWeight(
+            QosClient.QueryWeigher queryWeigher) {
+        QueryWeight weight = queryWeigher.weighFailure(new RuntimeException(), TIME_TAKEN);
 
         assertThat(weight).isEqualTo(DEFAULT_WEIGHT);
     }
-
 }

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

@@ -78,12 +78,11 @@ public void shouldBeAbleToReadLargeAmountsExceedingTheLimitSecondTimeWithSoftLim
 
     @Test
     public void shouldNotBeAbleToReadLargeAmountsIfSoftLimitSleepWillBeMoreThanConfiguredBackoffTime() {
-        assertThatThrownBy(() -> readOneBatchOfSize(200))
-                .isInstanceOf(RateLimitExceededException.class)
-                .hasMessage("Rate limited. Available capacity has been exhausted.");
+        // Have one quick limit-exceeding write, as the rate-limiter
+        // will let anything pass through until the limit is exceeded.
+        assertThat(readOneBatchOfSize(20)).hasSize(20);
 
-        // TODO(hsaraogi): This should not happen.
-        assertThatThrownBy(() -> readOneBatchOfSize(1))
+        assertThatThrownBy(() -> readOneBatchOfSize(100))
                 .isInstanceOf(RateLimitExceededException.class)
                 .hasMessage("Rate limited. Available capacity has been exhausted.");
     }

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

@@ -73,8 +73,10 @@ private <T, E extends Exception> T execute(
         estimatedWeightMetric.ifPresent(metric -> metric.accept(estimatedWeight));
 
         try {
-            Duration waitTime = rateLimiter.consumeWithBackoff(estimatedWeight.numBytes());
-            metrics.recordBackoffMicros(TimeUnit.NANOSECONDS.toMicros(waitTime.toNanos()));
+            if (estimatedWeight.numBytes() > 0) {
+                Duration waitTime = rateLimiter.consumeWithBackoff(estimatedWeight.numBytes());
+                metrics.recordBackoffMicros(TimeUnit.NANOSECONDS.toMicros(waitTime.toNanos()));
+            }
         } catch (RateLimitExceededException ex) {
             metrics.recordRateLimitedException();
             throw ex;