CASSJAVA-97: Let users inject an ID for each request and write to the custom payload

core/revapi.json

@@ -7386,6 +7386,11 @@
         "old": "method <T extends java.lang.Number> com.datastax.oss.driver.api.core.type.reflect.GenericType<com.datastax.oss.driver.api.core.data.CqlVector<T>> com.datastax.oss.driver.api.core.type.reflect.GenericType<T>::vectorOf(java.lang.Class<T>)",
         "new": "method <T> com.datastax.oss.driver.api.core.type.reflect.GenericType<com.datastax.oss.driver.api.core.data.CqlVector<T>> com.datastax.oss.driver.api.core.type.reflect.GenericType<T>::vectorOf(java.lang.Class<T>)",
         "justification": "JAVA-3143: Extend driver vector support to arbitrary subtypes and fix handling of variable length types (OSS C* 5.0)"
+      },
+      {
+        "code": "java.method.addedToInterface",
+        "new": "method com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator com.datastax.oss.driver.api.core.context.DriverContext::getDistributedTraceIdGenerator()",
+        "justification": "CASSJAVA-97: Let users inject an ID for each request and write to the custom payload"
       }
     ]
   }

core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java

@@ -994,7 +994,21 @@ public enum DefaultDriverOption implements DriverOption {
    *
    * <p>Value-type: boolean
    */
-  SSL_ALLOW_DNS_REVERSE_LOOKUP_SAN("advanced.ssl-engine-factory.allow-dns-reverse-lookup-san");
+  SSL_ALLOW_DNS_REVERSE_LOOKUP_SAN("advanced.ssl-engine-factory.allow-dns-reverse-lookup-san"),
+
+  /**
+   * The class of session-wide component that generates distributed trace IDs.
+   *
+   * <p>Value-type: {@link String}
+   */
+  DISTRIBUTED_TRACE_ID_GENERATOR_CLASS("advanced.distributed-tracing.id-generator.class"),
+
+  /**
+   * If not empty, the driver will write the distributed trace ID to this key in the custom payload
+   *
+   * <p>Value-type: {@link String}
+   */
+  DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY("advanced.distributed-tracing.custom-payload-with-key");
 
   private final String path;
 

core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java

@@ -291,6 +291,10 @@ protected static void fillWithDriverDefaults(OptionsMap map) {
     map.put(TypedDriverOption.REQUEST_TRACE_INTERVAL, Duration.ofMillis(3));
     map.put(TypedDriverOption.REQUEST_TRACE_CONSISTENCY, "ONE");
     map.put(TypedDriverOption.REQUEST_LOG_WARNINGS, true);
+    map.put(
+        TypedDriverOption.DISTRIBUTED_TRACE_ID_GENERATOR_CLASS,
+        "DefaultDistributedTraceIdGenerator");
+    map.put(TypedDriverOption.DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY, "");
     map.put(TypedDriverOption.GRAPH_PAGING_ENABLED, "AUTO");
     map.put(TypedDriverOption.GRAPH_CONTINUOUS_PAGING_PAGE_SIZE, requestPageSize);
     map.put(TypedDriverOption.GRAPH_CONTINUOUS_PAGING_MAX_PAGES, continuousMaxPages);

core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java

@@ -281,6 +281,18 @@ public String toString() {
       new TypedDriverOption<>(
           DefaultDriverOption.REQUEST_TRACKER_CLASSES, GenericType.listOf(String.class));
 
+  /** The class of a session-wide component that generates distributed trace IDs. */
+  public static final TypedDriverOption<String> DISTRIBUTED_TRACE_ID_GENERATOR_CLASS =
+      new TypedDriverOption<>(
+          DefaultDriverOption.DISTRIBUTED_TRACE_ID_GENERATOR_CLASS, GenericType.STRING);
+
+  /**
+   * If not empty, the driver will write the distributed trace ID to this key in the custom payload
+   */
+  public static final TypedDriverOption<String> DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY =
+      new TypedDriverOption<>(
+          DefaultDriverOption.DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY, GenericType.STRING);
+
   /** Whether to log successful requests. */
   public static final TypedDriverOption<Boolean> REQUEST_LOGGER_SUCCESS_ENABLED =
       new TypedDriverOption<>(

core/src/main/java/com/datastax/oss/driver/api/core/context/DriverContext.java

@@ -33,6 +33,7 @@
 import com.datastax.oss.driver.api.core.specex.SpeculativeExecutionPolicy;
 import com.datastax.oss.driver.api.core.ssl.SslEngineFactory;
 import com.datastax.oss.driver.api.core.time.TimestampGenerator;
+import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator;
 import com.datastax.oss.driver.api.core.tracker.RequestTracker;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import java.util.Map;
@@ -139,6 +140,10 @@ default SpeculativeExecutionPolicy getSpeculativeExecutionPolicy(@NonNull String
   @NonNull
   RequestTracker getRequestTracker();
 
+  /** @return The driver's distributed trace ID generator; never {@code null}. */
+  @NonNull
+  DistributedTraceIdGenerator getDistributedTraceIdGenerator();
+
   /** @return The driver's request throttler; never {@code null}. */
   @NonNull
   RequestThrottler getRequestThrottler();

core/src/main/java/com/datastax/oss/driver/api/core/session/ProgrammaticArguments.java

@@ -23,6 +23,7 @@
 import com.datastax.oss.driver.api.core.metadata.NodeStateListener;
 import com.datastax.oss.driver.api.core.metadata.schema.SchemaChangeListener;
 import com.datastax.oss.driver.api.core.ssl.SslEngineFactory;
+import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator;
 import com.datastax.oss.driver.api.core.tracker.RequestTracker;
 import com.datastax.oss.driver.api.core.type.codec.TypeCodec;
 import com.datastax.oss.driver.api.core.type.codec.registry.MutableCodecRegistry;
@@ -59,6 +60,7 @@ public static Builder builder() {
   private final NodeStateListener nodeStateListener;
   private final SchemaChangeListener schemaChangeListener;
   private final RequestTracker requestTracker;
+  private final DistributedTraceIdGenerator distributedTraceIdGenerator;
   private final Map<String, String> localDatacenters;
   private final Map<String, Predicate<Node>> nodeFilters;
   private final Map<String, NodeDistanceEvaluator> nodeDistanceEvaluators;
@@ -77,6 +79,7 @@ private ProgrammaticArguments(
       @Nullable NodeStateListener nodeStateListener,
       @Nullable SchemaChangeListener schemaChangeListener,
       @Nullable RequestTracker requestTracker,
+      @Nullable DistributedTraceIdGenerator distributedTraceIdGenerator,
       @NonNull Map<String, String> localDatacenters,
       @NonNull Map<String, Predicate<Node>> nodeFilters,
       @NonNull Map<String, NodeDistanceEvaluator> nodeDistanceEvaluators,
@@ -94,6 +97,7 @@ private ProgrammaticArguments(
     this.nodeStateListener = nodeStateListener;
     this.schemaChangeListener = schemaChangeListener;
     this.requestTracker = requestTracker;
+    this.distributedTraceIdGenerator = distributedTraceIdGenerator;
     this.localDatacenters = localDatacenters;
     this.nodeFilters = nodeFilters;
     this.nodeDistanceEvaluators = nodeDistanceEvaluators;
@@ -128,6 +132,11 @@ public RequestTracker getRequestTracker() {
     return requestTracker;
   }
 
+  @Nullable
+  public DistributedTraceIdGenerator getDistributedTraceIdGenerator() {
+    return distributedTraceIdGenerator;
+  }
+
   @NonNull
   public Map<String, String> getLocalDatacenters() {
     return localDatacenters;
@@ -196,6 +205,7 @@ public static class Builder {
     private NodeStateListener nodeStateListener;
     private SchemaChangeListener schemaChangeListener;
     private RequestTracker requestTracker;
+    private DistributedTraceIdGenerator distributedTraceIdGenerator;
     private ImmutableMap.Builder<String, String> localDatacentersBuilder = ImmutableMap.builder();
     private final ImmutableMap.Builder<String, Predicate<Node>> nodeFiltersBuilder =
         ImmutableMap.builder();
@@ -294,6 +304,13 @@ public Builder addRequestTracker(@NonNull RequestTracker requestTracker) {
       return this;
     }
 
+    @NonNull
+    public Builder withDistributedTraceIdGenerator(
+        @Nullable DistributedTraceIdGenerator distributedTraceIdGenerator) {
+      this.distributedTraceIdGenerator = distributedTraceIdGenerator;
+      return this;
+    }
+
     @NonNull
     public Builder withLocalDatacenter(
         @NonNull String profileName, @NonNull String localDatacenter) {
@@ -417,6 +434,7 @@ public ProgrammaticArguments build() {
           nodeStateListener,
           schemaChangeListener,
           requestTracker,
+          distributedTraceIdGenerator,
           localDatacentersBuilder.build(),
           nodeFiltersBuilder.build(),
           nodeDistanceEvaluatorsBuilder.build(),

core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java

@@ -35,6 +35,7 @@
 import com.datastax.oss.driver.api.core.metadata.schema.SchemaChangeListener;
 import com.datastax.oss.driver.api.core.ssl.ProgrammaticSslEngineFactory;
 import com.datastax.oss.driver.api.core.ssl.SslEngineFactory;
+import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator;
 import com.datastax.oss.driver.api.core.tracker.RequestTracker;
 import com.datastax.oss.driver.api.core.type.codec.TypeCodec;
 import com.datastax.oss.driver.api.core.type.codec.registry.MutableCodecRegistry;
@@ -318,6 +319,17 @@ public SelfT addRequestTracker(@NonNull RequestTracker requestTracker) {
     return self;
   }
 
+  /**
+   * Registers a distributed trace ID generator. The driver will use the distributed trace ID in the
+   * logs So that users can correlate logs about the same request from different loggers.
+   */
+  @NonNull
+  public SelfT withDistributedTraceIdGenerator(
+      @NonNull DistributedTraceIdGenerator distributedTraceIdGenerator) {
+    this.programmaticArgumentsBuilder.withDistributedTraceIdGenerator(distributedTraceIdGenerator);
+    return self;
+  }
+
   /**
    * Registers an authentication provider to use with the session.
    *

core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java

@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.datastax.oss.driver.api.core.tracker;
+
+import com.datastax.oss.driver.api.core.session.Request;
+import edu.umd.cs.findbugs.annotations.NonNull;
+
+public interface DistributedTraceIdGenerator {
+  /**
+   * Generates a unique identifier for the session request. This will be the identifier for the
+   * entire `session.execute()` call. This identifier will be added to logs, and propagated to
+   * request trackers.
+   *
+   * @param statement the statement to be executed
+   * @param sessionName the name of the session
+   * @param hashCode the hashcode of the CqlRequestHandler
+   * @return a unique identifier for the session request
+   */
+  String getSessionRequestId(@NonNull Request statement, @NonNull String sessionName, int hashCode);
+
+  /**
+   * Generates a unique identifier for the node request. This will be the identifier for the CQL
+   * request against a particular node. There can be one or more node requests for a single session
+   * request, due to retries or speculative executions. This identifier will be added to logs, and
+   * propagated to request trackers.
+   *
+   * @param statement the statement to be executed
+   * @param sessionRequestId the session request identifier
+   * @param executionCount the number of previous node requests for this session request, due to
+   *     retries or speculative executions
+   * @return a unique identifier for the node request
+   */
+  String getNodeRequestId(
+      @NonNull Request statement, @NonNull String sessionRequestId, int executionCount);
+}

core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestTracker.java

@@ -47,21 +47,22 @@ default void onSuccess(
       @NonNull Node node) {}
 
   /**
-   * Invoked each time a request succeeds.
+   * Invoked each time a session request succeeds. A session request is a `session.execute()` call
    *
    * @param latencyNanos the overall execution time (from the {@link Session#execute(Request,
    *     GenericType) session.execute} call until the result is made available to the client).
    * @param executionProfile the execution profile of this request.
    * @param node the node that returned the successful response.
-   * @param requestLogPrefix the dedicated log prefix for this request
+   * @param sessionRequestLogPrefix the dedicated log prefix for this request
    */
   default void onSuccess(
       @NonNull Request request,
       long latencyNanos,
       @NonNull DriverExecutionProfile executionProfile,
       @NonNull Node node,
-      @NonNull String requestLogPrefix) {
-    // If client doesn't override onSuccess with requestLogPrefix delegate call to the old method
+      @NonNull String sessionRequestLogPrefix) {
+    // If client doesn't override onSuccess with sessionRequestLogPrefix delegate call to the old
+    // method
     onSuccess(request, latencyNanos, executionProfile, node);
   }
 
@@ -78,22 +79,23 @@ default void onError(
       @Nullable Node node) {}
 
   /**
-   * Invoked each time a request fails.
+   * Invoked each time a session request fails. A session request is a `session.execute()` call
    *
    * @param latencyNanos the overall execution time (from the {@link Session#execute(Request,
    *     GenericType) session.execute} call until the error is propagated to the client).
    * @param executionProfile the execution profile of this request.
    * @param node the node that returned the error response, or {@code null} if the error occurred
-   * @param requestLogPrefix the dedicated log prefix for this request
+   * @param sessionRequestLogPrefix the dedicated log prefix for this request
    */
   default void onError(
       @NonNull Request request,
       @NonNull Throwable error,
       long latencyNanos,
       @NonNull DriverExecutionProfile executionProfile,
       @Nullable Node node,
-      @NonNull String requestLogPrefix) {
-    // If client doesn't override onError with requestLogPrefix delegate call to the old method
+      @NonNull String sessionRequestLogPrefix) {
+    // If client doesn't override onError with sessionRequestLogPrefix delegate call to the old
+    // method
     onError(request, error, latencyNanos, executionProfile, node);
   }
 
@@ -110,23 +112,25 @@ default void onNodeError(
       @NonNull Node node) {}
 
   /**
-   * Invoked each time a request fails at the node level. Similar to {@link #onError(Request,
-   * Throwable, long, DriverExecutionProfile, Node, String)} but at a per node level.
+   * Invoked each time a node request fails. A node request is a CQL request sent to a particular
+   * node. There can be one or more node requests for a single session request, due to retries or
+   * speculative executions.
    *
    * @param latencyNanos the overall execution time (from the {@link Session#execute(Request,
    *     GenericType) session.execute} call until the error is propagated to the client).
    * @param executionProfile the execution profile of this request.
    * @param node the node that returned the error response.
-   * @param requestLogPrefix the dedicated log prefix for this request
+   * @param nodeRequestLogPrefix the dedicated log prefix for this request
    */
   default void onNodeError(
       @NonNull Request request,
       @NonNull Throwable error,
       long latencyNanos,
       @NonNull DriverExecutionProfile executionProfile,
       @NonNull Node node,
-      @NonNull String requestLogPrefix) {
-    // If client doesn't override onNodeError with requestLogPrefix delegate call to the old method
+      @NonNull String nodeRequestLogPrefix) {
+    // If client doesn't override onNodeError with nodeRequestLogPrefix delegate call to the old
+    // method
     onNodeError(request, error, latencyNanos, executionProfile, node);
   }
 
@@ -142,22 +146,23 @@ default void onNodeSuccess(
       @NonNull Node node) {}
 
   /**
-   * Invoked each time a request succeeds at the node level. Similar to {@link #onSuccess(Request,
-   * long, DriverExecutionProfile, Node, String)} but at per node level.
+   * Invoked each time a node request succeeds. A node request is a CQL request sent to a particular
+   * node. There can be one or more node requests for a single session request, due to retries or
+   * speculative executions.
    *
    * @param latencyNanos the overall execution time (from the {@link Session#execute(Request,
    *     GenericType) session.execute} call until the result is made available to the client).
    * @param executionProfile the execution profile of this request.
    * @param node the node that returned the successful response.
-   * @param requestLogPrefix the dedicated log prefix for this request
+   * @param nodeRequestLogPrefix the dedicated log prefix for this request
    */
   default void onNodeSuccess(
       @NonNull Request request,
       long latencyNanos,
       @NonNull DriverExecutionProfile executionProfile,
       @NonNull Node node,
-      @NonNull String requestLogPrefix) {
-    // If client doesn't override onNodeSuccess with requestLogPrefix delegate call to the old
+      @NonNull String nodeRequestLogPrefix) {
+    // If client doesn't override onNodeSuccess with nodeRequestLogPrefix delegate call to the old
     // method
     onNodeSuccess(request, latencyNanos, executionProfile, node);
   }