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,19 @@ 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.
+   *
+   * @param distributedTraceIdGenerator
+   * @return
+   */
+  @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,46 @@
+/*
+ * 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 identifier will be added to logs.
+   *
+   * @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 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/internal/core/context/DefaultDriverContext.java

@@ -44,6 +44,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 com.datastax.oss.driver.api.core.type.codec.TypeCodec;
 import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry;
@@ -221,6 +222,7 @@ public class DefaultDriverContext implements InternalDriverContext {
   private final LazyReference<NodeStateListener> nodeStateListenerRef;
   private final LazyReference<SchemaChangeListener> schemaChangeListenerRef;
   private final LazyReference<RequestTracker> requestTrackerRef;
+  private final LazyReference<DistributedTraceIdGenerator> distributedTraceIdGeneratorRef;
   private final LazyReference<Optional<AuthProvider>> authProviderRef;
   private final LazyReference<List<LifecycleListener>> lifecycleListenersRef =
       new LazyReference<>("lifecycleListeners", this::buildLifecycleListeners, cycleDetector);
@@ -282,6 +284,13 @@ public DefaultDriverContext(
     this.requestTrackerRef =
         new LazyReference<>(
             "requestTracker", () -> buildRequestTracker(requestTrackerFromBuilder), cycleDetector);
+    this.distributedTraceIdGeneratorRef =
+        new LazyReference<>(
+            "distributedTraceIdGenerator",
+            () ->
+                buildDistributedTraceIdGenerator(
+                    programmaticArguments.getDistributedTraceIdGenerator()),
+            cycleDetector);
     this.sslEngineFactoryRef =
         new LazyReference<>(
             "sslEngineFactory",
@@ -709,6 +718,23 @@ protected RequestTracker buildRequestTracker(RequestTracker requestTrackerFromBu
     }
   }
 
+  protected DistributedTraceIdGenerator buildDistributedTraceIdGenerator(
+      DistributedTraceIdGenerator distributedTraceIdGenerator) {
+    return (distributedTraceIdGenerator != null)
+        ? distributedTraceIdGenerator
+        : Reflection.buildFromConfig(
+                this,
+                DefaultDriverOption.DISTRIBUTED_TRACE_ID_GENERATOR_CLASS,
+                DistributedTraceIdGenerator.class,
+                "com.datastax.oss.driver.internal.core.tracker")
+            .orElseThrow(
+                () ->
+                    new IllegalArgumentException(
+                        String.format(
+                            "Missing distributed trace ID generator, check your configuration (%s)",
+                            DefaultDriverOption.DISTRIBUTED_TRACE_ID_GENERATOR_CLASS)));
+  }
+
   protected Optional<AuthProvider> buildAuthProvider(AuthProvider authProviderFromBuilder) {
     return (authProviderFromBuilder != null)
         ? Optional.of(authProviderFromBuilder)
@@ -973,6 +999,12 @@ public RequestTracker getRequestTracker() {
     return requestTrackerRef.get();
   }
 
+  @NonNull
+  @Override
+  public DistributedTraceIdGenerator getDistributedTraceIdGenerator() {
+    return distributedTraceIdGeneratorRef.get();
+  }
+
   @Nullable
   @Override
   public String getLocalDatacenter(@NonNull String profileName) {

core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java

@@ -44,6 +44,7 @@
 import com.datastax.oss.driver.api.core.servererrors.WriteTimeoutException;
 import com.datastax.oss.driver.api.core.session.throttling.RequestThrottler;
 import com.datastax.oss.driver.api.core.session.throttling.Throttled;
+import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator;
 import com.datastax.oss.driver.api.core.tracker.RequestTracker;
 import com.datastax.oss.driver.internal.core.adminrequest.ThrottledAdminRequestHandler;
 import com.datastax.oss.driver.internal.core.adminrequest.UnexpectedResponseException;
@@ -78,8 +79,10 @@
 import io.netty.util.concurrent.Future;
 import io.netty.util.concurrent.GenericFutureListener;
 import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.AbstractMap;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Queue;
@@ -125,12 +128,14 @@ public class CqlRequestHandler implements Throttled {
   private final List<NodeResponseCallback> inFlightCallbacks;
   private final RequestThrottler throttler;
   private final RequestTracker requestTracker;
+  private final DistributedTraceIdGenerator distributedTraceIdGenerator;
   private final SessionMetricUpdater sessionMetricUpdater;
   private final DriverExecutionProfile executionProfile;
 
   // The errors on the nodes that were already tried (lazily initialized on the first error).
   // We don't use a map because nodes can appear multiple times.
   private volatile List<Map.Entry<Node, Throwable>> errors;
+  private final String customPayloadKey;
 
   protected CqlRequestHandler(
       Statement<?> statement,
@@ -139,7 +144,10 @@ protected CqlRequestHandler(
       String sessionLogPrefix) {
 
     this.startTimeNanos = System.nanoTime();
-    this.logPrefix = sessionLogPrefix + "|" + this.hashCode();
+    this.distributedTraceIdGenerator = context.getDistributedTraceIdGenerator();
+    this.logPrefix =
+        this.distributedTraceIdGenerator.getSessionRequestId(
+            statement, sessionLogPrefix, this.hashCode());
     LOG.trace("[{}] Creating new handler for request {}", logPrefix, statement);
 
     this.initialStatement = statement;
@@ -170,6 +178,11 @@ protected CqlRequestHandler(
 
     this.timer = context.getNettyOptions().getTimer();
     this.executionProfile = Conversions.resolveExecutionProfile(initialStatement, context);
+
+    this.customPayloadKey =
+        this.executionProfile.getString(
+            DefaultDriverOption.DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY);
+
     Duration timeout = Conversions.resolveRequestTimeout(statement, executionProfile);
     this.scheduledTimeout = scheduleTimeout(timeout);
 
@@ -248,6 +261,18 @@ private void sendRequest(
     if (result.isDone()) {
       return;
     }
+    String nodeRequestId =
+        this.distributedTraceIdGenerator.getNodeRequestId(
+            statement, logPrefix, currentExecutionIndex);
+    if (!this.customPayloadKey.isEmpty()) {
+      // We cannot do statement.getCustomPayload().put() because the default empty map is abstract
+      // But this will create new Statement instance for every request. We might want to optimize
+      // this
+      Map<String, ByteBuffer> existingMap = new HashMap<>(statement.getCustomPayload());
+      existingMap.put(
+          this.customPayloadKey, ByteBuffer.wrap(nodeRequestId.getBytes(StandardCharsets.UTF_8)));
+      statement = statement.setCustomPayload(existingMap);
+    }
     Node node = retriedNode;
     DriverChannel channel = null;
     if (node == null || (channel = session.getChannel(node, logPrefix)) == null) {
@@ -276,7 +301,7 @@ private void sendRequest(
               currentExecutionIndex,
               retryCount,
               scheduleNextExecution,
-              logPrefix);
+              nodeRequestId);
       Message message = Conversions.toMessage(statement, executionProfile, context);
       channel
           .write(message, statement.isTracing(), statement.getCustomPayload(), nodeResponseCallback)
@@ -489,7 +514,7 @@ private NodeResponseCallback(
       this.execution = execution;
       this.retryCount = retryCount;
       this.scheduleNextExecution = scheduleNextExecution;
-      this.logPrefix = logPrefix + "|" + execution;
+      this.logPrefix = logPrefix;
     }
 
     // this gets invoked once the write completes.