JAVA-1830: Surface response frame size in ExecutionInfo

olim7t · olim7t · commit 637390755ce7 · 2018-06-12T15:52:32.000-07:00
diff --git a/changelog/README.md b/changelog/README.md
@@ -4,6 +4,7 @@
 
 ### 4.0.0-alpha4 (in progress)
 
+- [improvement] JAVA-1830: Surface response frame size in ExecutionInfo
 - [improvement] JAVA-1853: Add newValue(Object...) to TupleType and UserDefinedType
 - [improvement] JAVA-1815: Reorganize configuration into basic/advanced categories
 - [improvement] JAVA-1848: Add logs to DefaultRetryPolicy
diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/cql/ExecutionInfo.java b/core/src/main/java/com/datastax/oss/driver/api/core/cql/ExecutionInfo.java
@@ -18,6 +18,7 @@
 import com.datastax.oss.driver.api.core.DefaultProtocolVersion;
 import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
 import com.datastax.oss.driver.api.core.metadata.Node;
+import com.datastax.oss.driver.api.core.retry.RetryDecision;
 import com.datastax.oss.driver.api.core.session.Request;
 import com.datastax.oss.driver.api.core.session.Session;
 import com.datastax.oss.driver.api.core.specex.SpeculativeExecutionPolicy;
@@ -147,4 +148,31 @@ default QueryTrace getQueryTrace() {
     BlockingOperation.checkNotDriverThread();
     return CompletableFutures.getUninterruptibly(getQueryTraceAsync());
   }
+
+  /**
+   * The size of the binary response in bytes.
+   *
+   * <p>This is the size of the protocol-level frame (including the frame header) before it was
+   * decoded by the driver, but after decompression (if compression is enabled).
+   *
+   * <p>If the information is not available (for example if this execution info comes from an {@link
+   * RetryDecision#IGNORE IGNORE} decision of the retry policy), this method returns -1.
+   *
+   * @see #getCompressedResponseSizeInBytes()
+   */
+  int getResponseSizeInBytes();
+
+  /**
+   * The size of the compressed binary response in bytes.
+   *
+   * <p>This is the size of the protocol-level frame (including the frame header) as it came in the
+   * TCP response, before decompression and decoding by the driver.
+   *
+   * <p>If compression is disabled, or if the information is not available (for example if this
+   * execution info comes from an {@link RetryDecision#IGNORE IGNORE} decision of the retry policy),
+   * this method returns -1.
+   *
+   * @see #getResponseSizeInBytes()
+   */
+  int getCompressedResponseSizeInBytes();
 }
diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultExecutionInfo.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultExecutionInfo.java
@@ -42,6 +42,8 @@ public class DefaultExecutionInfo implements ExecutionInfo {
   private final List<Map.Entry<Node, Throwable>> errors;
   private final ByteBuffer pagingState;
   private final UUID tracingId;
+  private final int responseSizeInBytes;
+  private final int compressedResponseSizeInBytes;
   private final List<String> warnings;
   private final Map<String, ByteBuffer> customPayload;
   private final boolean schemaInAgreement;
@@ -69,6 +71,8 @@ public DefaultExecutionInfo(
     this.pagingState = pagingState;
 
     this.tracingId = (frame == null) ? null : frame.tracingId;
+    this.responseSizeInBytes = (frame == null) ? -1 : frame.size;
+    this.compressedResponseSizeInBytes = (frame == null) ? -1 : frame.compressedSize;
     // Note: the collections returned by the protocol layer are already unmodifiable
     this.warnings = (frame == null) ? Collections.emptyList() : frame.warnings;
     this.customPayload = (frame == null) ? Collections.emptyMap() : frame.customPayload;
@@ -139,4 +143,14 @@ public CompletionStage<QueryTrace> getQueryTraceAsync() {
       return new QueryTraceFetcher(tracingId, session, context, configProfile).fetch();
     }
   }
+
+  @Override
+  public int getResponseSizeInBytes() {
+    return responseSizeInBytes;
+  }
+
+  @Override
+  public int getCompressedResponseSizeInBytes() {
+    return compressedResponseSizeInBytes;
+  }
 }
diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/api/core/compression/DirectCompressionIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/api/core/compression/DirectCompressionIT.java
@@ -19,6 +19,7 @@
 import static org.assertj.core.api.Assertions.offset;
 
 import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.cql.ExecutionInfo;
 import com.datastax.oss.driver.api.core.cql.ResultSet;
 import com.datastax.oss.driver.api.core.cql.Row;
 import com.datastax.oss.driver.api.core.cql.SimpleStatement;
@@ -90,6 +91,13 @@ private void createAndCheckCluster(String compressorOption) {
       assertThat(row.getString("t")).isEqualTo("foo");
       assertThat(row.getInt("i")).isEqualTo(42);
       assertThat(row.getFloat("f")).isEqualTo(24.03f, offset(0.1f));
+
+      ExecutionInfo executionInfo = rs.getExecutionInfo();
+      // There's not much more we can check without hard-coding sizes.
+      // We are testing with small responses, so the compressed payload is not even guaranteed to be
+      // smaller.
+      assertThat(executionInfo.getResponseSizeInBytes()).isGreaterThan(0);
+      assertThat(executionInfo.getCompressedResponseSizeInBytes()).isGreaterThan(0);
     }
   }
 }
diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/api/core/compression/HeapCompressionIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/api/core/compression/HeapCompressionIT.java
@@ -19,6 +19,7 @@
 import static org.assertj.core.api.Assertions.offset;
 
 import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.cql.ExecutionInfo;
 import com.datastax.oss.driver.api.core.cql.ResultSet;
 import com.datastax.oss.driver.api.core.cql.Row;
 import com.datastax.oss.driver.api.core.cql.SimpleStatement;
@@ -93,6 +94,13 @@ private void createAndCheckCluster(String compressorOption) {
       assertThat(row.getString("t")).isEqualTo("foo");
       assertThat(row.getInt("i")).isEqualTo(42);
       assertThat(row.getFloat("f")).isEqualTo(24.03f, offset(0.1f));
+
+      ExecutionInfo executionInfo = rs.getExecutionInfo();
+      // There's not much more we can check without hard-coding sizes.
+      // We are testing with small responses, so the compressed payload is not even guaranteed to be
+      // smaller.
+      assertThat(executionInfo.getResponseSizeInBytes()).isGreaterThan(0);
+      assertThat(executionInfo.getCompressedResponseSizeInBytes()).isGreaterThan(0);
     }
   }
 }
