Add ExecutionInfo to RequestTracker methods

Author: Andrew Weaver

core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.java

@@ -19,11 +19,13 @@
 
 import com.datastax.dse.driver.api.core.DseProtocolVersion;
 import com.datastax.dse.driver.api.core.cql.continuous.ContinuousAsyncResultSet;
+import com.datastax.dse.driver.api.core.graph.AsyncGraphResultSet;
 import com.datastax.dse.driver.internal.core.DseProtocolFeature;
 import com.datastax.dse.driver.internal.core.cql.DseConversions;
 import com.datastax.dse.protocol.internal.request.Revise;
 import com.datastax.dse.protocol.internal.response.result.DseRowsMetadata;
 import com.datastax.oss.driver.api.core.AllNodesFailedException;
+import com.datastax.oss.driver.api.core.AsyncPagingIterable;
 import com.datastax.oss.driver.api.core.CqlIdentifier;
 import com.datastax.oss.driver.api.core.DriverTimeoutException;
 import com.datastax.oss.driver.api.core.NodeUnavailableException;
@@ -626,7 +628,7 @@ public void operationComplete(@NonNull Future<java.lang.Void> future) {
         Throwable error = future.cause();
         if (error instanceof EncoderException
             && error.getCause() instanceof FrameTooLongException) {
-          trackNodeError(node, error.getCause());
+          trackNodeError(node, error.getCause(), null);
           lock.lock();
           try {
             abort(error.getCause(), false);
@@ -643,7 +645,7 @@ public void operationComplete(@NonNull Future<java.lang.Void> future) {
               .getMetricUpdater()
               .incrementCounter(DefaultNodeMetric.UNSENT_REQUESTS, executionProfile.getName());
           recordError(node, error);
-          trackNodeError(node, error.getCause());
+          trackNodeError(node, error.getCause(), null);
           sendRequest(statement, null, executionIndex, retryCount, scheduleSpeculativeExecution);
         }
       } else {
@@ -738,7 +740,8 @@ private void onPageTimeout(int expectedPage) {
      * Invoked when a continuous paging response is received, either a successful or failed one.
      *
      * <p>Delegates further processing to appropriate methods: {@link #processResultResponse(Result,
-     * Frame)} if the response was successful, or {@link #processErrorResponse(Error)} if it wasn't.
+     * Frame)} if the response was successful, or {@link #processErrorResponse(Error, Frame)} if it
+     * wasn't.
      *
      * @param response the received {@link Frame}.
      */
@@ -759,15 +762,15 @@ public void onResponse(@NonNull Frame response) {
             processResultResponse((Result) responseMessage, response);
           } else if (responseMessage instanceof Error) {
             LOG.trace("[{}] Got error response", logPrefix);
-            processErrorResponse((Error) responseMessage);
+            processErrorResponse((Error) responseMessage, response);
           } else {
             IllegalStateException error =
                 new IllegalStateException("Unexpected response " + responseMessage);
-            trackNodeError(node, error);
+            trackNodeError(node, error, response);
             abort(error, false);
           }
         } catch (Throwable t) {
-          trackNodeError(node, t);
+          trackNodeError(node, t, response);
           abort(t, false);
         }
       } finally {
@@ -901,7 +904,7 @@ private void processResultResponse(@NonNull Result result, @Nullable Frame frame
      * @param errorMessage the error message received.
      */
     @SuppressWarnings("GuardedBy") // this method is only called with the lock held
-    private void processErrorResponse(@NonNull Error errorMessage) {
+    private void processErrorResponse(@NonNull Error errorMessage, @NonNull Frame frame) {
       assert lock.isHeldByCurrentThread();
       if (errorMessage instanceof Unprepared) {
         processUnprepared((Unprepared) errorMessage);
@@ -910,7 +913,7 @@ private void processErrorResponse(@NonNull Error errorMessage) {
         if (error instanceof BootstrappingException) {
           LOG.trace("[{}] {} is bootstrapping, trying next node", logPrefix, node);
           recordError(node, error);
-          trackNodeError(node, error);
+          trackNodeError(node, error, frame);
           sendRequest(statement, null, executionIndex, retryCount, false);
         } else if (error instanceof QueryValidationException
             || error instanceof FunctionFailureException
@@ -922,7 +925,7 @@ private void processErrorResponse(@NonNull Error errorMessage) {
           NodeMetricUpdater metricUpdater = ((DefaultNode) node).getMetricUpdater();
           metricUpdater.incrementCounter(
               DefaultNodeMetric.OTHER_ERRORS, executionProfile.getName());
-          trackNodeError(node, error);
+          trackNodeError(node, error, frame);
           abort(error, true);
         } else {
           try {
@@ -1061,7 +1064,7 @@ private void processUnprepared(@NonNull Unprepared errorMessage) {
                                     + "This usually happens when you run a 'USE...' query after "
                                     + "the statement was prepared.",
                                 Bytes.toHexString(idToReprepare), Bytes.toHexString(repreparedId)));
-                    trackNodeError(node, illegalStateException);
+                    trackNodeError(node, illegalStateException, null);
                     fatalError = illegalStateException;
                   } else {
                     LOG.trace(
@@ -1080,18 +1083,18 @@ private void processUnprepared(@NonNull Unprepared errorMessage) {
                           || prepareError instanceof FunctionFailureException
                           || prepareError instanceof ProtocolError) {
                         LOG.trace("[{}] Unrecoverable error on re-prepare, rethrowing", logPrefix);
-                        trackNodeError(node, prepareError);
+                        trackNodeError(node, prepareError, null);
                         fatalError = prepareError;
                       }
                     }
                   } else if (exception instanceof RequestThrottlingException) {
-                    trackNodeError(node, exception);
+                    trackNodeError(node, exception, null);
                     fatalError = exception;
                   }
                   if (fatalError == null) {
                     LOG.trace("[{}] Re-prepare failed, trying next node", logPrefix);
                     recordError(node, exception);
-                    trackNodeError(node, exception);
+                    trackNodeError(node, exception, null);
                     sendRequest(statement, null, executionIndex, retryCount, false);
                   }
                 }
@@ -1119,18 +1122,18 @@ private void processRetryVerdict(@NonNull RetryVerdict verdict, @NonNull Throwab
       switch (verdict.getRetryDecision()) {
         case RETRY_SAME:
           recordError(node, error);
-          trackNodeError(node, error);
+          trackNodeError(node, error, null);
           sendRequest(
               verdict.getRetryRequest(statement), node, executionIndex, retryCount + 1, false);
           break;
         case RETRY_NEXT:
           recordError(node, error);
-          trackNodeError(node, error);
+          trackNodeError(node, error, null);
           sendRequest(
               verdict.getRetryRequest(statement), null, executionIndex, retryCount + 1, false);
           break;
         case RETHROW:
-          trackNodeError(node, error);
+          trackNodeError(node, error, null);
           abort(error, true);
           break;
         case IGNORE:
@@ -1443,12 +1446,20 @@ private void reenableAutoReadIfNeeded() {
 
     // ERROR HANDLING
 
-    private void trackNodeError(@NonNull Node node, @NonNull Throwable error) {
+    private void trackNodeError(
+        @NonNull Node node, @NonNull Throwable error, @Nullable Frame frame) {
       if (nodeErrorReported.compareAndSet(false, true)) {
         long latencyNanos = System.nanoTime() - this.messageStartTimeNanos;
         context
             .getRequestTracker()
-            .onNodeError(this.statement, error, latencyNanos, executionProfile, node, logPrefix);
+            .onNodeError(
+                this.statement,
+                error,
+                latencyNanos,
+                executionProfile,
+                node,
+                createExecutionInfo(frame),
+                logPrefix);
       }
     }
 
@@ -1562,21 +1573,32 @@ private void completeResultSetFuture(
       if (resultSetClass.isInstance(pageOrError)) {
         if (future.complete(resultSetClass.cast(pageOrError))) {
           throttler.signalSuccess(ContinuousRequestHandlerBase.this);
+
+          ExecutionInfo executionInfo = null;
+          if (pageOrError instanceof AsyncPagingIterable) {
+            executionInfo = ((AsyncPagingIterable) pageOrError).getExecutionInfo();
+          } else if (pageOrError instanceof AsyncGraphResultSet) {
+            executionInfo = ((AsyncGraphResultSet) pageOrError).getRequestExecutionInfo();
+          }
+
           if (nodeSuccessReported.compareAndSet(false, true)) {
             context
                 .getRequestTracker()
-                .onNodeSuccess(statement, nodeLatencyNanos, executionProfile, node, logPrefix);
+                .onNodeSuccess(
+                    statement, nodeLatencyNanos, executionProfile, node, executionInfo, logPrefix);
           }
           context
               .getRequestTracker()
-              .onSuccess(statement, totalLatencyNanos, executionProfile, node, logPrefix);
+              .onSuccess(
+                  statement, totalLatencyNanos, executionProfile, node, executionInfo, logPrefix);
         }
       } else {
         Throwable error = (Throwable) pageOrError;
         if (future.completeExceptionally(error)) {
           context
               .getRequestTracker()
-              .onError(statement, error, totalLatencyNanos, executionProfile, node, logPrefix);
+              .onError(
+                  statement, error, totalLatencyNanos, executionProfile, node, null, logPrefix);
           if (error instanceof DriverTimeoutException) {
             throttler.signalTimeout(ContinuousRequestHandlerBase.this);
             session
@@ -1607,6 +1629,22 @@ private ExecutionInfo createExecutionInfo(@NonNull Result result, @Nullable Fram
           executionProfile);
     }
 
+    @NonNull
+    private ExecutionInfo createExecutionInfo(@Nullable Frame response) {
+      return new DefaultExecutionInfo(
+          statement,
+          node,
+          startedSpeculativeExecutionsCount.get(),
+          executionIndex,
+          errors,
+          null,
+          response,
+          true,
+          session,
+          context,
+          executionProfile);
+    }
+
     private void logTimeoutSchedulingError(IllegalStateException timeoutError) {
       // If we're racing with session shutdown, the timer might be stopped already. We don't want
       // to schedule more executions anyway, so swallow the error.

core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java

@@ -360,9 +360,19 @@ private void setFinalResult(
           totalLatencyNanos = completionTimeNanos - startTimeNanos;
           long nodeLatencyNanos = completionTimeNanos - callback.nodeStartTimeNanos;
           requestTracker.onNodeSuccess(
-              callback.statement, nodeLatencyNanos, executionProfile, callback.node, logPrefix);
+              callback.statement,
+              nodeLatencyNanos,
+              executionProfile,
+              callback.node,
+              executionInfo,
+              logPrefix);
           requestTracker.onSuccess(
-              callback.statement, totalLatencyNanos, executionProfile, callback.node, logPrefix);
+              callback.statement,
+              totalLatencyNanos,
+              executionProfile,
+              callback.node,
+              executionInfo,
+              logPrefix);
         }
         if (sessionMetricUpdater.isEnabled(
             DseSessionMetric.GRAPH_REQUESTS, executionProfile.getName())) {
@@ -446,27 +456,28 @@ private void setFinalError(
       GraphStatement<?> statement, Throwable error, Node node, int execution) {
     DriverExecutionProfile executionProfile =
         Conversions.resolveExecutionProfile(statement, context);
+    ExecutionInfo executionInfo =
+        new DefaultExecutionInfo(
+            statement,
+            node,
+            startedSpeculativeExecutionsCount.get(),
+            execution,
+            errors,
+            null,
+            null,
+            true,
+            session,
+            context,
+            executionProfile);
     if (error instanceof DriverException) {
-      ((DriverException) error)
-          .setExecutionInfo(
-              new DefaultExecutionInfo(
-                  statement,
-                  node,
-                  startedSpeculativeExecutionsCount.get(),
-                  execution,
-                  errors,
-                  null,
-                  null,
-                  true,
-                  session,
-                  context,
-                  executionProfile));
+      ((DriverException) error).setExecutionInfo(executionInfo);
     }
     if (result.completeExceptionally(error)) {
       cancelScheduledTasks();
       if (!(requestTracker instanceof NoopRequestTracker)) {
         long latencyNanos = System.nanoTime() - startTimeNanos;
-        requestTracker.onError(statement, error, latencyNanos, executionProfile, node, logPrefix);
+        requestTracker.onError(
+            statement, error, latencyNanos, executionProfile, node, executionInfo, logPrefix);
       }
       if (error instanceof DriverTimeoutException) {
         throttler.signalTimeout(this);
@@ -859,7 +870,8 @@ private void trackNodeError(Node node, Throwable error, long nodeResponseTimeNan
         nodeResponseTimeNanos = System.nanoTime();
       }
       long latencyNanos = nodeResponseTimeNanos - this.nodeStartTimeNanos;
-      requestTracker.onNodeError(statement, error, latencyNanos, executionProfile, node, logPrefix);
+      requestTracker.onNodeError(
+          statement, error, latencyNanos, executionProfile, node, null, logPrefix);
     }
 
     @Override