Weight estimate metrics

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

@@ -16,6 +16,7 @@
 package com.palantir.atlasdb.qos.client;
 
 import java.time.Duration;
+import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
@@ -53,22 +54,26 @@ public static AtlasDbQosClient create(QosRateLimiters rateLimiters) {
 
     @Override
     public <T, E extends Exception> T executeRead(Query<T, E> query, QueryWeigher<T> weigher) throws E {
-        return execute(query, weigher, rateLimiters.read(), metrics::recordRead);
+        return execute(query, weigher, rateLimiters.read(), Optional.of(metrics::recordReadEstimate),
+                metrics::recordRead);
     }
 
     @Override
     public <T, E extends Exception> T executeWrite(Query<T, E> query, QueryWeigher<T> weigher) throws E {
-        return execute(query, weigher, rateLimiters.write(), metrics::recordWrite);
+        return execute(query, weigher, rateLimiters.write(), Optional.empty(), metrics::recordWrite);
     }
 
     private <T, E extends Exception> T execute(
             Query<T, E> query,
             QueryWeigher<T> weigher,
             QosRateLimiter rateLimiter,
+            Optional<Consumer<QueryWeight>> estimatedWeightMetric,
             Consumer<QueryWeight> weightMetric) throws E {
-        long estimatedNumBytes = weigher.estimate().numBytes();
+        QueryWeight estimatedWeight = weigher.estimate();
+        estimatedWeightMetric.ifPresent(metric -> metric.accept(estimatedWeight));
+
         try {
-            Duration waitTime = rateLimiter.consumeWithBackoff(estimatedNumBytes);
+            Duration waitTime = rateLimiter.consumeWithBackoff(estimatedWeight.numBytes());
             metrics.recordBackoffMicros(TimeUnit.NANOSECONDS.toMicros(waitTime.toNanos()));
         } catch (RateLimitExceededException ex) {
             metrics.recordRateLimitedException();
@@ -87,7 +92,7 @@ private <T, E extends Exception> T execute(
             throw ex;
         } finally {
             weightMetric.accept(actualWeight);
-            rateLimiter.recordAdjustment(actualWeight.numBytes() - estimatedNumBytes);
+            rateLimiter.recordAdjustment(actualWeight.numBytes() - estimatedWeight.numBytes());
         }
     }
 

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

@@ -36,6 +36,9 @@ public class QosMetrics {
     private final Meter readTime;
     private final Meter rowsRead;
 
+    private final Meter estimatedBytesRead;
+    private final Meter estimatedRowsRead;
+
     private final Meter writeRequestCount;
     private final Meter bytesWritten;
     private final Meter writeTime;
@@ -50,6 +53,9 @@ public QosMetrics() {
         readTime = metricsManager.registerMeter(QosMetrics.class, "readTime");
         rowsRead = metricsManager.registerMeter(QosMetrics.class, "rowsRead");
 
+        estimatedBytesRead = metricsManager.registerMeter(QosMetrics.class, "estimated.bytesRead");
+        estimatedRowsRead = metricsManager.registerMeter(QosMetrics.class, "estimated.rowsRead");
+
         writeRequestCount = metricsManager.registerMeter(QosMetrics.class, "numWriteRequests");
         bytesWritten = metricsManager.registerMeter(QosMetrics.class, "bytesWritten");
         writeTime = metricsManager.registerMeter(QosMetrics.class, "writeTime");
@@ -59,6 +65,11 @@ public QosMetrics() {
         rateLimitedExceptions = metricsManager.registerMeter(QosMetrics.class, "rateLimitedExceptions");
     }
 
+    public void recordReadEstimate(QueryWeight weight) {
+        estimatedBytesRead.mark(weight.numBytes());
+        estimatedRowsRead.mark(weight.numDistinctRows());
+    }
+
     public void recordRead(QueryWeight weight) {
         readRequestCount.mark();
         bytesRead.mark(weight.numBytes());

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

@@ -96,6 +96,7 @@ public void consumesSpecifiedNumUnitsForReads() {
     public void recordsReadMetrics() throws TestCheckedException {
         qosClient.executeRead(() -> "foo", weigher);
 
+        verify(metrics).recordReadEstimate(ESTIMATED_WEIGHT);
         verify(metrics).recordRead(ACTUAL_WEIGHT);
     }
 
@@ -120,6 +121,7 @@ public void recordsWriteMetrics() throws TestCheckedException {
         qosClient.executeWrite(() -> null, weigher);
 
         verify(metrics).recordWrite(ACTUAL_WEIGHT);
+        verify(metrics, never()).recordReadEstimate(any());
     }
 
     @Test