Tracing: added various request driver-side tags

All data that constitute added tags are collected offline,
i.e. not fetched from the cluster, but known instantly by the driver.

Author: wprzytula

clirr-ignores.xml

@@ -16,6 +16,12 @@
     Modified by ScyllaDB
 -->
 <differences>
+    <difference>
+        <differenceType>7012</differenceType> <!-- method added to interface -->
+        <className>com/datastax/driver/core/PreparedStatement</className>
+        <method>java.lang.String getOperationType()</method>
+        <justification>New method to get the type of operation performed by this PreparedStatement</justification>
+    </difference>
     <difference>
         <differenceType>7012</differenceType> <!-- method added to interface -->
         <className>com/datastax/driver/core/Session</className>

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

@@ -74,6 +74,8 @@ public class BoundStatement extends Statement
 
   private ByteBuffer routingKey;
 
+  private final String operationType;
+
   /**
    * Creates a new {@code BoundStatement} from the provided prepared statement.
    *
@@ -92,6 +94,7 @@ public BoundStatement(PreparedStatement statement) {
       this.setSerialConsistencyLevel(statement.getSerialConsistencyLevel());
     if (statement.isTracing()) this.enableTracing();
     if (statement.getRetryPolicy() != null) this.setRetryPolicy(statement.getRetryPolicy());
+    this.operationType = statement.getOperationType();
     if (statement.getOutgoingPayload() != null)
       this.setOutgoingPayload(statement.getOutgoingPayload());
     else
@@ -104,6 +107,10 @@ public BoundStatement(PreparedStatement statement) {
     }
   }
 
+  public String getOperationType() {
+    return operationType;
+  }
+
   @Override
   public boolean isLWT() {
     return statement.isLWT();

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

@@ -41,6 +41,7 @@ public class DefaultPreparedStatement implements PreparedStatement {
   final Cluster cluster;
   final boolean isLWT;
   final Token.Factory partitioner;
+  final String operationType;
 
   volatile ByteBuffer routingKey;
 
@@ -66,6 +67,7 @@ private DefaultPreparedStatement(
     this.cluster = cluster;
     this.isLWT = isLWT;
     this.partitioner = partitioner;
+    this.operationType = null;
   }
 
   static DefaultPreparedStatement fromMessage(
@@ -315,4 +317,10 @@ public Boolean isIdempotent() {
   public boolean isLWT() {
     return isLWT;
   }
+
+  /** {@inheritDoc} */
+  @Override
+  public String getOperationType() {
+    return operationType;
+  }
 }

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

@@ -361,4 +361,7 @@ public interface PreparedStatement {
 
   /** Whether a prepared statement is LWT statement */
   public boolean isLWT();
+
+  /** Type of prepared operation (e.g. SELECT) */
+  public String getOperationType();
 }

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

@@ -48,6 +48,7 @@
 import com.google.common.util.concurrent.ListenableFuture;
 import io.netty.util.Timeout;
 import io.netty.util.TimerTask;
+import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
 import java.util.Collections;
 import java.util.Iterator;
@@ -74,6 +75,10 @@ class RequestHandler {
   private static final QueryLogger QUERY_LOGGER = QueryLogger.builder().build();
   static final String DISABLE_QUERY_WARNING_LOGS = "com.datastax.driver.DISABLE_QUERY_WARNING_LOGS";
 
+  private static final int STATEMENT_MAX_LENGTH = 1000;
+  private static final int PARTITION_KEY_MAX_LENGTH = 1000;
+  private static final int REPLICAS_MAX_LENGTH = 1000;
+
   final String id;
 
   private final SessionManager manager;
@@ -161,8 +166,71 @@ public RequestHandler(
     this.timerContext = metricsEnabled() ? metrics().getRequestsTimer().time() : null;
     this.startTime = System.nanoTime();
 
+    ConsistencyLevel consistency = statement.getConsistencyLevel();
+    if (consistency == null) consistency = Statement.DEFAULT.getConsistencyLevel();
+
+    String statementType = null;
+    String statementText = null;
+    int batchSize = 1;
+
+    String keyspace = null;
+    String partitionKey = null;
+    String table = null;
+    String operationType = null;
+
+    if (statement instanceof BatchStatement) {
+      statementType = "batch";
+      batchSize = ((BatchStatement) statement).size();
+      StringBuilder statementTextBuilder = new StringBuilder(STATEMENT_MAX_LENGTH);
+      for (Statement subStatement : ((BatchStatement) statement).getStatements()) {
+        if (subStatement instanceof BoundStatement)
+          statementTextBuilder.append(((BoundStatement) subStatement).statement.getQueryString());
+        else statementTextBuilder.append(subStatement.toString());
+      }
+      statementText = statementTextBuilder.toString();
+    } else if (statement instanceof BoundStatement) {
+      statementType = "prepared";
+      statementText = ((BoundStatement) statement).statement.getQueryString();
+      keyspace = ((BoundStatement) statement).getKeyspace();
+      operationType = ((BoundStatement) statement).getOperationType();
+
+      ColumnDefinitions boundColumns =
+          ((BoundStatement) statement).statement.getPreparedId().boundValuesMetadata.variables;
+
+      StringBuilder partitionKeyBuilder = new StringBuilder(PARTITION_KEY_MAX_LENGTH);
+      int[] rkIndexes = ((BoundStatement) statement).statement.getPreparedId().routingKeyIndexes;
+      if (rkIndexes != null) {
+        for (int i : rkIndexes) {
+          Object value = ((BoundStatement) statement).getObject(i);
+          String valueString = (value == null) ? "NULL" : value.toString();
+          if (partitionKeyBuilder.length() > 0) partitionKeyBuilder.append(", ");
+          String columnName = boundColumns.getName(i);
+          partitionKeyBuilder.append(columnName);
+          partitionKeyBuilder.append('=');
+          partitionKeyBuilder.append(valueString);
+        }
+      }
+      partitionKey = partitionKeyBuilder.toString();
+
+      if (boundColumns.size() > 0) table = boundColumns.getTable(0);
+    } else if (statement instanceof RegularStatement) {
+      statementType = "regular";
+      statementText = ((RegularStatement) statement).toString();
+    }
+
     this.tracingInfo = tracingInfo;
     this.tracingInfo.setNameAndStartTime("request");
+    this.tracingInfo.setConsistencyLevel(consistency);
+    this.tracingInfo.setRetryPolicy(retryPolicy());
+    this.tracingInfo.setBatchSize(batchSize);
+    this.tracingInfo.setLoadBalancingPolicy(manager.loadBalancingPolicy());
+    this.tracingInfo.setSpeculativeExecutionPolicy(manager.speculativeExecutionPolicy());
+    if (statementType != null) this.tracingInfo.setStatementType(statementType);
+    if (statementText != null) this.tracingInfo.setStatement(statementText, STATEMENT_MAX_LENGTH);
+    if (keyspace != null) this.tracingInfo.setKeyspace(keyspace);
+    if (partitionKey != null) this.tracingInfo.setPartitionKey(partitionKey);
+    if (table != null) this.tracingInfo.setTable(table);
+    if (operationType != null) this.tracingInfo.setOperationType(operationType);
   }
 
   void sendRequest() {
@@ -282,6 +350,14 @@ private void setFinalResult(
       }
       callback.onSet(connection, response, info, statement, System.nanoTime() - startTime);
 
+      if (response.type == Message.Response.Type.RESULT) {
+        Responses.Result rm = (Responses.Result) response;
+        if (rm.kind == Responses.Result.Kind.ROWS) {
+          Responses.Result.Rows r = (Responses.Result.Rows) rm;
+          tracingInfo.setRowsCount(r.data.size());
+        }
+      }
+      tracingInfo.setQueryPaged(info.getPagingState() != null);
       tracingInfo.setStatus(
           response.type == Message.Response.Type.ERROR
               ? TracingInfo.StatusCode.ERROR
@@ -336,6 +412,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.setAttemptCount(execution.retryCount() + 1);
     execution.parentTracingInfo.setStatus(TracingInfo.StatusCode.ERROR);
     execution.parentTracingInfo.tracingFinished();
     runningExecutions.remove(execution);
@@ -460,6 +537,10 @@ private boolean query(final Host host) {
 
       currentChildTracingInfo = manager.getTracingInfoFactory().buildTracingInfo(parentTracingInfo);
       currentChildTracingInfo.setNameAndStartTime("attempt");
+      InetSocketAddress hostAddress = host.getEndPoint().resolve();
+      currentChildTracingInfo.setPeerName(hostAddress.getHostName());
+      currentChildTracingInfo.setPeerIP(hostAddress.getAddress());
+      currentChildTracingInfo.setPeerPort(hostAddress.getPort());
 
       if (allowSpeculativeExecutions && nextExecutionScheduled.compareAndSet(false, true))
         scheduleExecution(speculativeExecutionPlan.nextExecution(host));
@@ -679,6 +760,7 @@ void cancel() {
                 CancelledSpeculativeExecutionException.INSTANCE,
                 System.nanoTime() - startTime);
           }
+          parentTracingInfo.setAttemptCount(previous.retryCount + 1);
           parentTracingInfo.setStatus(TracingInfo.StatusCode.OK);
           parentTracingInfo.tracingFinished();
           return;
@@ -693,6 +775,7 @@ void cancel() {
                 CancelledSpeculativeExecutionException.INSTANCE,
                 System.nanoTime() - startTime);
           }
+          parentTracingInfo.setAttemptCount(previous.retryCount + 1);
           parentTracingInfo.setStatus(TracingInfo.StatusCode.OK);
           parentTracingInfo.tracingFinished();
           return;
@@ -710,6 +793,7 @@ public Message.Request request() {
     @Override
     public void onSet(
         Connection connection, Message.Response response, long latency, int retryCount) {
+      currentChildTracingInfo.setShardID(connection.shardId());
       currentChildTracingInfo.setStatus(TracingInfo.StatusCode.OK);
       currentChildTracingInfo.tracingFinished();
 
@@ -1021,6 +1105,7 @@ public boolean onTimeout(Connection connection, long latency, int retryCount) {
     @Override
     public void onException(
         Connection connection, Exception exception, long latency, int retryCount) {
+      currentChildTracingInfo.setShardID(connection.shardId());
       currentChildTracingInfo.recordException(exception);
       currentChildTracingInfo.setStatus(TracingInfo.StatusCode.ERROR);
       currentChildTracingInfo.tracingFinished();
@@ -1062,6 +1147,7 @@ public void onException(
 
     @Override
     public boolean onTimeout(Connection connection, long latency, int retryCount) {
+      currentChildTracingInfo.setShardID(connection.shardId());
       currentChildTracingInfo.setStatus(TracingInfo.StatusCode.ERROR, "timeout");
       currentChildTracingInfo.tracingFinished();
 
@@ -1106,12 +1192,14 @@ public int retryCount() {
     }
 
     private void setFinalException(Connection connection, Exception exception) {
+      parentTracingInfo.setAttemptCount(retryCount() + 1);
       parentTracingInfo.setStatus(TracingInfo.StatusCode.ERROR);
       parentTracingInfo.tracingFinished();
       RequestHandler.this.setFinalException(this, connection, exception);
     }
 
     private void setFinalResult(Connection connection, Message.Response response) {
+      parentTracingInfo.setAttemptCount(retryCount() + 1);
       parentTracingInfo.setStatus(TracingInfo.StatusCode.OK);
       parentTracingInfo.tracingFinished();
       RequestHandler.this.setFinalResult(this, connection, response);

driver-core/src/main/java/com/datastax/driver/core/tracing/NoopTracingInfoFactory.java

@@ -16,12 +16,76 @@
 
 package com.datastax.driver.core.tracing;
 
+import com.datastax.driver.core.ConsistencyLevel;
+import com.datastax.driver.core.policies.LoadBalancingPolicy;
+import com.datastax.driver.core.policies.RetryPolicy;
+import com.datastax.driver.core.policies.SpeculativeExecutionPolicy;
+import java.net.InetAddress;
+
 public class NoopTracingInfoFactory implements TracingInfoFactory {
 
   private static class NoopTracingInfo implements TracingInfo {
     @Override
     public void setNameAndStartTime(String name) {}
 
+    @Override
+    public void setConsistencyLevel(ConsistencyLevel consistency) {}
+
+    @Override
+    public void setStatementType(String statementType) {}
+
+    @Override
+    public void setRetryPolicy(RetryPolicy retryPolicy) {}
+
+    @Override
+    public void setLoadBalancingPolicy(LoadBalancingPolicy loadBalancingPolicy) {};
+
+    @Override
+    public void setSpeculativeExecutionPolicy(
+        SpeculativeExecutionPolicy speculativeExecutionPolicy) {}
+
+    @Override
+    public void setBatchSize(int batchSize) {}
+
+    @Override
+    public void setAttemptCount(int attemptCount) {}
+
+    @Override
+    public void setShardID(int shardID) {}
+
+    @Override
+    public void setPeerName(String peerName) {}
+
+    @Override
+    public void setPeerIP(InetAddress peerIP) {}
+
+    @Override
+    public void setPeerPort(int peerPort) {}
+
+    @Override
+    public void setQueryPaged(Boolean queryPaged) {}
+
+    @Override
+    public void setRowsCount(int rowsCount) {}
+
+    @Override
+    public void setStatement(String statement, int limit) {}
+
+    @Override
+    public void setKeyspace(String keyspace) {}
+
+    @Override
+    public void setPartitionKey(String partitionKey) {}
+
+    @Override
+    public void setTable(String table) {}
+
+    @Override
+    public void setOperationType(String operationType) {}
+
+    @Override
+    public void setReplicas(String replicas) {}
+
     @Override
     public void recordException(Exception exception) {}