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.RequestIdGenerator com.datastax.oss.driver.api.core.context.DriverContext::getRequestIdGenerator()",
+        "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,14 @@ 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 request IDs.
+   *
+   * <p>Value-type: {@link String}
+   */
+  REQUEST_ID_GENERATOR_CLASS("advanced.request-id.generator.class");
 
   private final String path;
 

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

@@ -281,6 +281,10 @@ public String toString() {
       new TypedDriverOption<>(
           DefaultDriverOption.REQUEST_TRACKER_CLASSES, GenericType.listOf(String.class));
 
+  /** The class of a session-wide component that generates request IDs. */
+  public static final TypedDriverOption<String> REQUEST_ID_GENERATOR_CLASS =
+      new TypedDriverOption<>(DefaultDriverOption.REQUEST_ID_GENERATOR_CLASS, 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.RequestIdGenerator;
 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 request ID generator; never {@code null}. */
+  @NonNull
+  Optional<RequestIdGenerator> getRequestIdGenerator();
+
   /** @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.RequestIdGenerator;
 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 RequestIdGenerator requestIdGenerator;
   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 RequestIdGenerator requestIdGenerator,
       @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.requestIdGenerator = requestIdGenerator;
     this.localDatacenters = localDatacenters;
     this.nodeFilters = nodeFilters;
     this.nodeDistanceEvaluators = nodeDistanceEvaluators;
@@ -128,6 +132,11 @@ public RequestTracker getRequestTracker() {
     return requestTracker;
   }
 
+  @Nullable
+  public RequestIdGenerator getRequestIdGenerator() {
+    return requestIdGenerator;
+  }
+
   @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 RequestIdGenerator requestIdGenerator;
     private ImmutableMap.Builder<String, String> localDatacentersBuilder = ImmutableMap.builder();
     private final ImmutableMap.Builder<String, Predicate<Node>> nodeFiltersBuilder =
         ImmutableMap.builder();
@@ -294,6 +304,12 @@ public Builder addRequestTracker(@NonNull RequestTracker requestTracker) {
       return this;
     }
 
+    @NonNull
+    public Builder withRequestIdGenerator(@Nullable RequestIdGenerator requestIdGenerator) {
+      this.requestIdGenerator = requestIdGenerator;
+      return this;
+    }
+
     @NonNull
     public Builder withLocalDatacenter(
         @NonNull String profileName, @NonNull String localDatacenter) {
@@ -417,6 +433,7 @@ public ProgrammaticArguments build() {
           nodeStateListener,
           schemaChangeListener,
           requestTracker,
+          requestIdGenerator,
           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.RequestIdGenerator;
 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 request ID generator. The driver will use the generated ID in the logs and
+   * optionally add to the custom payload so that users can correlate logs about the same request
+   * from the Cassandra side.
+   */
+  @NonNull
+  public SelfT withRequestIdGenerator(@NonNull RequestIdGenerator requestIdGenerator) {
+    this.programmaticArgumentsBuilder.withRequestIdGenerator(requestIdGenerator);
+    return self;
+  }
+
   /**
    * Registers an authentication provider to use with the session.
    *

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

@@ -0,0 +1,63 @@
+/*
+ * 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.cql.Statement;
+import com.datastax.oss.driver.api.core.session.Request;
+import com.datastax.oss.protocol.internal.util.collection.NullAllowingImmutableMap;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.Map;
+
+public interface RequestIdGenerator {
+  /**
+   * 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
+   * @return a unique identifier for the session request
+   */
+  String getSessionRequestId(@NonNull Request statement);
+
+  /**
+   * 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);
+
+  default Statement<?> getDecoratedStatement(
+      @NonNull Statement<?> statement, @NonNull String nodeRequestId) {
+    Map<String, ByteBuffer> customPayload =
+        NullAllowingImmutableMap.<String, ByteBuffer>builder()
+            .putAll(statement.getCustomPayload())
+            .put("request-id", ByteBuffer.wrap(nodeRequestId.getBytes(StandardCharsets.UTF_8)))
+            .build();
+    return statement.setCustomPayload(customPayload);
+  }
+}

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);
   }

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.RequestIdGenerator;
 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<Optional<RequestIdGenerator>> requestIdGeneratorRef;
   private final LazyReference<Optional<AuthProvider>> authProviderRef;
   private final LazyReference<List<LifecycleListener>> lifecycleListenersRef =
       new LazyReference<>("lifecycleListeners", this::buildLifecycleListeners, cycleDetector);
@@ -282,6 +284,11 @@ public DefaultDriverContext(
     this.requestTrackerRef =
         new LazyReference<>(
             "requestTracker", () -> buildRequestTracker(requestTrackerFromBuilder), cycleDetector);
+    this.requestIdGeneratorRef =
+        new LazyReference<>(
+            "requestIdGenerator",
+            () -> buildRequestIdGenerator(programmaticArguments.getRequestIdGenerator()),
+            cycleDetector);
     this.sslEngineFactoryRef =
         new LazyReference<>(
             "sslEngineFactory",
@@ -709,6 +716,17 @@ protected RequestTracker buildRequestTracker(RequestTracker requestTrackerFromBu
     }
   }
 
+  protected Optional<RequestIdGenerator> buildRequestIdGenerator(
+      RequestIdGenerator requestIdGenerator) {
+    return (requestIdGenerator != null)
+        ? Optional.of(requestIdGenerator)
+        : Reflection.buildFromConfig(
+            this,
+            DefaultDriverOption.REQUEST_ID_GENERATOR_CLASS,
+            RequestIdGenerator.class,
+            "com.datastax.oss.driver.internal.core.tracker");
+  }
+
   protected Optional<AuthProvider> buildAuthProvider(AuthProvider authProviderFromBuilder) {
     return (authProviderFromBuilder != null)
         ? Optional.of(authProviderFromBuilder)
@@ -973,6 +991,12 @@ public RequestTracker getRequestTracker() {
     return requestTrackerRef.get();
   }
 
+  @NonNull
+  @Override
+  public Optional<RequestIdGenerator> getRequestIdGenerator() {
+    return requestIdGeneratorRef.get();
+  }
+
   @Nullable
   @Override
   public String getLocalDatacenter(@NonNull String profileName) {