Tracing: added empty spans covering request

Every request (in RequestHandler) is covered by one "request" span,
and every (speculative) execution has its own "speculative_execution"
span. Each retry attempt is covered by "attempt" span, effectively
yielding a tree of form:
"request" -1--many-> "speculative_execution" -1--many-> "attempt"

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");
       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("attempt");
+
       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();