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

patch by Jane He; reviewed by Abe Ratnofsky and Bret McGuire for CASSJAVA-97

Author: SiyaoIsHiding

core/revapi.json

@@ -7407,6 +7407,11 @@
       {
         "code": "java.method.varargOverloadsOnlyDifferInVarargParameter",
         "justification": "CASSJAVA-102: Migrate revapi config into dedicated config files, ported from pom.xml"
+      },
+      {
+        "code": "java.method.addedToInterface",
+        "new": "method java.util.Optional<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

@@ -995,6 +995,12 @@ public enum DefaultDriverOption implements DriverOption {
    * <p>Value-type: boolean
    */
   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"),
   /**
    * An address to always translate all node addresses to that same proxy hostname no matter what IP
    * address a node has, but still using its native transport port.

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;
@@ -47,6 +48,7 @@
 import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
 import com.datastax.oss.driver.internal.core.metadata.DefaultEndPoint;
 import com.datastax.oss.driver.internal.core.session.DefaultSession;
+import com.datastax.oss.driver.internal.core.tracker.W3CContextRequestIdGenerator;
 import com.datastax.oss.driver.internal.core.util.concurrent.BlockingOperation;
 import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
 import edu.umd.cs.findbugs.annotations.NonNull;
@@ -83,6 +85,8 @@
 @NotThreadSafe
 public abstract class SessionBuilder<SelfT extends SessionBuilder, SessionT> {
 
+  public static final String ASTRA_PAYLOAD_KEY = "traceparent";
+
   private static final Logger LOG = LoggerFactory.getLogger(SessionBuilder.class);
 
   @SuppressWarnings("unchecked")
@@ -318,6 +322,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.
    *
@@ -861,6 +876,13 @@ protected final CompletionStage<CqlSession> buildDefaultSessionAsync() {
       List<String> configContactPoints =
           defaultConfig.getStringList(DefaultDriverOption.CONTACT_POINTS, Collections.emptyList());
       if (cloudConfigInputStream != null) {
+        // override request id generator, unless user has already set it
+        if (programmaticArguments.getRequestIdGenerator() == null) {
+          programmaticArgumentsBuilder.withRequestIdGenerator(
+              new W3CContextRequestIdGenerator(ASTRA_PAYLOAD_KEY));
+          LOG.debug(
+              "A secure connect bundle is provided, using W3CContextRequestIdGenerator as request ID generator.");
+        }
         if (!programmaticContactPoints.isEmpty() || !configContactPoints.isEmpty()) {
           LOG.info(
               "Both a secure connect bundle and contact points were provided. These are mutually exclusive. The contact points from the secure bundle will have priority.");

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

@@ -0,0 +1,77 @@
+/*
+ * 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;
+
+/**
+ * Interface responsible for generating request IDs.
+ *
+ * <p>Note that all request IDs have a parent/child relationship. A "parent ID" can loosely be
+ * thought of as encompassing a sequence of a request + any attendant retries, speculative
+ * executions etc. It's scope is identical to that of a {@link
+ * com.datastax.oss.driver.internal.core.cql.CqlRequestHandler}. A "request ID" represents a single
+ * request within this larger scope. Note that a request corresponding to a request ID may be
+ * retried; in that case the retry count will be appended to the corresponding identifier in the
+ * logs.
+ */
+public interface RequestIdGenerator {
+
+  String DEFAULT_PAYLOAD_KEY = "request-id";
+
+  /**
+   * 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.
+   *
+   * @return a unique identifier for the session request
+   */
+  String getSessionRequestId();
+
+  /**
+   * 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 parentId the session request identifier
+   * @return a unique identifier for the node request
+   */
+  String getNodeRequestId(@NonNull Request statement, @NonNull String parentId);
+
+  default String getCustomPayloadKey() {
+    return DEFAULT_PAYLOAD_KEY;
+  }
+
+  default Statement<?> getDecoratedStatement(
+      @NonNull Statement<?> statement, @NonNull String requestId) {
+    Map<String, ByteBuffer> customPayload =
+        NullAllowingImmutableMap.<String, ByteBuffer>builder()
+            .putAll(statement.getCustomPayload())
+            .put(getCustomPayloadKey(), ByteBuffer.wrap(requestId.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);
   }