Instrumenting RequestHandler with tracing: added empty spans.

Every request is covered by one "request" span, and every speculative
execution has its own "speculative_execution.n" span, where n is number
of that (possibly speculative) execution. Retries are performed in the
same execution and therefore do not increase that number, whereas
speculative executions do. Each retry is covered by "query" span,
effectively yielding a tree of form:
"request" -1--many-> "speculative_execution.n" -1--many-> "query"

Author: wprzytula

driver-core/src/main/java/com/datastax/driver/core/RequestHandler.java

@@ -40,6 +40,7 @@
 import com.datastax.driver.core.policies.RetryPolicy;
 import com.datastax.driver.core.policies.RetryPolicy.RetryDecision.Type;
 import com.datastax.driver.core.policies.SpeculativeExecutionPolicy.SpeculativeExecutionPlan;
+import com.datastax.driver.core.tracing.TracingInfo;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterators;
 import com.google.common.collect.Sets;
@@ -95,6 +96,8 @@ class RequestHandler {
   private final AtomicBoolean isDone = new AtomicBoolean();
   private final AtomicInteger executionIndex = new AtomicInteger();
 
+  private final TracingInfo tracingInfo;
+
   private Iterator<Host> getReplicas(
       String loggedKeyspace, Statement statement, Iterator<Host> fallback) {
     ProtocolVersion protocolVersion = manager.cluster.manager.protocolVersion();
@@ -120,7 +123,8 @@ private Iterator<Host> getReplicas(
     return replicas.iterator();
   }
 
-  public RequestHandler(SessionManager manager, Callback callback, Statement statement) {
+  public RequestHandler(
+      SessionManager manager, Callback callback, Statement statement, TracingInfo tracingInfo) {
     this.id = Long.toString(System.identityHashCode(this));
     if (logger.isTraceEnabled()) logger.trace("[{}] {}", id, statement);
     this.manager = manager;
@@ -156,6 +160,9 @@ public RequestHandler(SessionManager manager, Callback callback, Statement state
 
     this.timerContext = metricsEnabled() ? metrics().getRequestsTimer().time() : null;
     this.startTime = System.nanoTime();
+
+    this.tracingInfo = tracingInfo;
+    this.tracingInfo.setNameAndStartTime("request");
   }
 
   void sendRequest() {
@@ -274,13 +281,16 @@ private void setFinalResult(
         logServerWarnings(response.warnings);
       }
       callback.onSet(connection, response, info, statement, System.nanoTime() - startTime);
+      tracingInfo.tracingFinished();
     } catch (Exception e) {
       callback.onException(
           connection,
           new DriverInternalError(
               "Unexpected exception while setting final result from " + response, e),
           System.nanoTime() - startTime, /*unused*/
           0);
+
+      tracingInfo.tracingFinished();
     }
   }
 
@@ -305,6 +315,8 @@ private void setFinalException(
 
     cancelPendingExecutions(execution);
 
+    tracingInfo.tracingFinished();
+
     try {
       if (timerContext != null) timerContext.stop();
     } finally {
@@ -315,6 +327,7 @@ private void setFinalException(
   // Triggered when an execution reaches the end of the query plan.
   // This is only a failure if there are no other running executions.
   private void reportNoMoreHosts(SpeculativeExecution execution) {
+    execution.parentTracingInfo.tracingFinished();
     runningExecutions.remove(execution);
     if (runningExecutions.isEmpty())
       setFinalException(
@@ -383,11 +396,17 @@ class SpeculativeExecution implements Connection.ResponseCallback {
 
     private volatile Connection.ResponseHandler connectionHandler;
 
+    private final TracingInfo parentTracingInfo;
+    private TracingInfo currentChildTracingInfo;
+
     SpeculativeExecution(Message.Request request, int position) {
       this.id = RequestHandler.this.id + "-" + position;
       this.request = request;
       this.position = position;
       this.queryStateRef = new AtomicReference<QueryState>(QueryState.INITIAL);
+      this.parentTracingInfo =
+          manager.getTracingInfoFactory().buildTracingInfo(RequestHandler.this.tracingInfo);
+      this.parentTracingInfo.setNameAndStartTime("speculative_execution." + position);
       if (logger.isTraceEnabled()) logger.trace("[{}] Starting", id);
     }
 
@@ -429,6 +448,9 @@ private boolean query(final Host host) {
 
       if (logger.isTraceEnabled()) logger.trace("[{}] Querying node {}", id, host);
 
+      currentChildTracingInfo = manager.getTracingInfoFactory().buildTracingInfo(parentTracingInfo);
+      currentChildTracingInfo.setNameAndStartTime("query");
+
       if (allowSpeculativeExecutions && nextExecutionScheduled.compareAndSet(false, true))
         scheduleExecution(speculativeExecutionPlan.nextExecution(host));
 
@@ -647,6 +669,7 @@ void cancel() {
                 CancelledSpeculativeExecutionException.INSTANCE,
                 System.nanoTime() - startTime);
           }
+          parentTracingInfo.tracingFinished();
           return;
         } else if (!previous.inProgress
             && queryStateRef.compareAndSet(previous, QueryState.CANCELLED_WHILE_COMPLETE)) {
@@ -659,6 +682,7 @@ void cancel() {
                 CancelledSpeculativeExecutionException.INSTANCE,
                 System.nanoTime() - startTime);
           }
+          parentTracingInfo.tracingFinished();
           return;
         }
       }
@@ -674,6 +698,8 @@ public Message.Request request() {
     @Override
     public void onSet(
         Connection connection, Message.Response response, long latency, int retryCount) {
+      currentChildTracingInfo.tracingFinished();
+
       QueryState queryState = queryStateRef.get();
       if (!queryState.isInProgressAt(retryCount)
           || !queryStateRef.compareAndSet(queryState, queryState.complete())) {
@@ -832,7 +858,10 @@ public void onSet(
                     toPrepare.getQueryKeyspace(),
                     connection.endPoint);
 
-                write(connection, prepareAndRetry(toPrepare.getQueryString()));
+                TracingInfo prepareTracingInfo =
+                    manager.getTracingInfoFactory().buildTracingInfo(parentTracingInfo);
+                prepareTracingInfo.setNameAndStartTime("prepare");
+                write(connection, prepareAndRetry(toPrepare.getQueryString(), prepareTracingInfo));
                 // we're done for now, the prepareAndRetry callback will handle the rest
                 return;
               case READ_FAILURE:
@@ -878,7 +907,8 @@ public void onSet(
       }
     }
 
-    private Connection.ResponseCallback prepareAndRetry(final String toPrepare) {
+    private Connection.ResponseCallback prepareAndRetry(
+        final String toPrepare, final TracingInfo prepareTracingInfo) {
       // do not bother inspecting retry policy at this step, no other decision
       // makes sense than retry on the same host if the query was prepared,
       // or on another host, if an error/timeout occurred.
@@ -902,6 +932,8 @@ public int retryCount() {
         @Override
         public void onSet(
             Connection connection, Message.Response response, long latency, int retryCount) {
+          prepareTracingInfo.tracingFinished();
+
           QueryState queryState = queryStateRef.get();
           if (!queryState.isInProgressAt(retryCount)
               || !queryStateRef.compareAndSet(queryState, queryState.complete())) {
@@ -944,11 +976,14 @@ public void onSet(
         @Override
         public void onException(
             Connection connection, Exception exception, long latency, int retryCount) {
+          prepareTracingInfo.tracingFinished();
           SpeculativeExecution.this.onException(connection, exception, latency, retryCount);
         }
 
         @Override
         public boolean onTimeout(Connection connection, long latency, int retryCount) {
+          prepareTracingInfo.tracingFinished();
+
           QueryState queryState = queryStateRef.get();
           if (!queryState.isInProgressAt(retryCount)
               || !queryStateRef.compareAndSet(queryState, queryState.complete())) {
@@ -973,6 +1008,8 @@ public boolean onTimeout(Connection connection, long latency, int retryCount) {
     @Override
     public void onException(
         Connection connection, Exception exception, long latency, int retryCount) {
+      currentChildTracingInfo.tracingFinished();
+
       QueryState queryState = queryStateRef.get();
       if (!queryState.isInProgressAt(retryCount)
           || !queryStateRef.compareAndSet(queryState, queryState.complete())) {
@@ -1010,6 +1047,8 @@ public void onException(
 
     @Override
     public boolean onTimeout(Connection connection, long latency, int retryCount) {
+      currentChildTracingInfo.tracingFinished();
+
       QueryState queryState = queryStateRef.get();
       if (!queryState.isInProgressAt(retryCount)
           || !queryStateRef.compareAndSet(queryState, queryState.complete())) {
@@ -1051,10 +1090,12 @@ public int retryCount() {
     }
 
     private void setFinalException(Connection connection, Exception exception) {
+      parentTracingInfo.tracingFinished();
       RequestHandler.this.setFinalException(this, connection, exception);
     }
 
     private void setFinalResult(Connection connection, Message.Response response) {
+      parentTracingInfo.tracingFinished();
       RequestHandler.this.setFinalResult(this, connection, response);
     }
   }

driver-core/src/main/java/com/datastax/driver/core/SessionManager.java

@@ -29,6 +29,7 @@
 import com.datastax.driver.core.policies.LoadBalancingPolicy;
 import com.datastax.driver.core.policies.ReconnectionPolicy;
 import com.datastax.driver.core.policies.SpeculativeExecutionPolicy;
+import com.datastax.driver.core.tracing.TracingInfo;
 import com.datastax.driver.core.tracing.TracingInfoFactory;
 import com.datastax.driver.core.utils.MoreFutures;
 import com.google.common.base.Functions;
@@ -161,6 +162,7 @@ public ResultSetFuture executeAsync(final Statement statement) {
       // Because of the way the future is built, we need another 'proxy' future that we can return
       // now.
       final ChainedResultSetFuture chainedFuture = new ChainedResultSetFuture();
+      final TracingInfo tracingInfo = getTracingInfoFactory().buildTracingInfo();
       this.initAsync()
           .addListener(
               new Runnable() {
@@ -171,7 +173,7 @@ public void run() {
                           SessionManager.this,
                           cluster.manager.protocolVersion(),
                           makeRequestMessage(statement, null));
-                  execute(actualFuture, statement);
+                  execute(actualFuture, statement, tracingInfo);
                   chainedFuture.setSource(actualFuture);
                 }
               },
@@ -712,25 +714,34 @@ else if (fetchSize != Integer.MAX_VALUE)
    * <p>This method will find a suitable node to connect to using the {@link LoadBalancingPolicy}
    * and handle host failover.
    */
-  void execute(final RequestHandler.Callback callback, final Statement statement) {
+  void execute(
+      final RequestHandler.Callback callback,
+      final Statement statement,
+      final TracingInfo tracingInfo) {
     if (this.isClosed()) {
       callback.onException(
           null, new IllegalStateException("Could not send request, session is closed"), 0, 0);
       return;
     }
-    if (isInit) new RequestHandler(this, callback, statement).sendRequest();
+    if (isInit) new RequestHandler(this, callback, statement, tracingInfo).sendRequest();
     else
       this.initAsync()
           .addListener(
               new Runnable() {
                 @Override
                 public void run() {
-                  new RequestHandler(SessionManager.this, callback, statement).sendRequest();
+                  new RequestHandler(SessionManager.this, callback, statement, tracingInfo)
+                      .sendRequest();
                 }
               },
               executor());
   }
 
+  void execute(final RequestHandler.Callback callback, final Statement statement) {
+    final TracingInfo tracingInfo = getTracingInfoFactory().buildTracingInfo();
+    execute(callback, statement, tracingInfo);
+  }
+
   private ListenableFuture<PreparedStatement> prepare(
       final PreparedStatement statement, EndPoint toExclude) {
     final String query = statement.getQueryString();