From b774dffa1586ae8fc12eddaa77c70d269d1f4296 Mon Sep 17 00:00:00 2001 From: janehe Date: Tue, 15 Apr 2025 13:18:34 -0700 Subject: [PATCH 01/16] DistributedTraceIdGenerator --- core/revapi.json | 5 +++ .../api/core/config/DefaultDriverOption.java | 9 ++++- .../driver/api/core/config/OptionsMap.java | 2 + .../api/core/config/TypedDriverOption.java | 5 +++ .../api/core/context/DriverContext.java | 5 +++ .../core/session/ProgrammaticArguments.java | 18 +++++++++ .../api/core/session/SessionBuilder.java | 14 +++++++ .../tracker/DistributedTraceIdGenerator.java | 28 +++++++++++++ .../core/context/DefaultDriverContext.java | 32 +++++++++++++++ .../internal/core/cql/CqlRequestHandler.java | 8 +++- .../NoopDistributedTraceIdGenerator.java | 39 +++++++++++++++++++ .../UuidDistributedTraceIdGenerator.java | 38 ++++++++++++++++++ core/src/main/resources/reference.conf | 4 ++ 13 files changed, 204 insertions(+), 3 deletions(-) create mode 100644 core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java create mode 100644 core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopDistributedTraceIdGenerator.java create mode 100644 core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java diff --git a/core/revapi.json b/core/revapi.json index 5aa46a3ccad..1f5ec4bc2bb 100644 --- a/core/revapi.json +++ b/core/revapi.json @@ -7386,6 +7386,11 @@ "old": "method com.datastax.oss.driver.api.core.type.reflect.GenericType> com.datastax.oss.driver.api.core.type.reflect.GenericType::vectorOf(java.lang.Class)", "new": "method com.datastax.oss.driver.api.core.type.reflect.GenericType> com.datastax.oss.driver.api.core.type.reflect.GenericType::vectorOf(java.lang.Class)", "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": "DistributedRequestID" } ] } diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java index 6ffd51d86ef..f0403396c9f 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java @@ -994,7 +994,14 @@ public enum DefaultDriverOption implements DriverOption { * *

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. + * + *

Value-type: {@link String} + */ + DISTRIBUTED_TRACE_ID_GENERATOR_CLASS("advanced.distributed-tracing.id-generator.class"); private final String path; diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java index 98faf3e590c..636803a044c 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java @@ -291,6 +291,8 @@ 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, "NoopDistributedTraceIdGenerator"); 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); diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java index 93e2b468461..217d46bad3f 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java @@ -281,6 +281,11 @@ 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 DISTRIBUTED_TRACE_ID_GENERATOR_CLASS = + new TypedDriverOption<>( + DefaultDriverOption.DISTRIBUTED_TRACE_ID_GENERATOR_CLASS, GenericType.STRING); + /** Whether to log successful requests. */ public static final TypedDriverOption REQUEST_LOGGER_SUCCESS_ENABLED = new TypedDriverOption<>( diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/context/DriverContext.java b/core/src/main/java/com/datastax/oss/driver/api/core/context/DriverContext.java index 5b32389e362..a0d047025f4 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/context/DriverContext.java +++ b/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(); diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/session/ProgrammaticArguments.java b/core/src/main/java/com/datastax/oss/driver/api/core/session/ProgrammaticArguments.java index 4e08bd5434c..07fc1e65926 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/session/ProgrammaticArguments.java +++ b/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 localDatacenters; private final Map> nodeFilters; private final Map nodeDistanceEvaluators; @@ -77,6 +79,7 @@ private ProgrammaticArguments( @Nullable NodeStateListener nodeStateListener, @Nullable SchemaChangeListener schemaChangeListener, @Nullable RequestTracker requestTracker, + @Nullable DistributedTraceIdGenerator distributedTraceIdGenerator, @NonNull Map localDatacenters, @NonNull Map> nodeFilters, @NonNull Map 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 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 localDatacentersBuilder = ImmutableMap.builder(); private final ImmutableMap.Builder> 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(), diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java b/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java index cbf896a0873..d13c8ca7361 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java +++ b/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. * diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java new file mode 100644 index 00000000000..20fe3b6ad75 --- /dev/null +++ b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java @@ -0,0 +1,28 @@ +/* + * 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; +import edu.umd.cs.findbugs.annotations.Nullable; + +public interface DistributedTraceIdGenerator { + String getSessionRequestId(@NonNull Request statement); + + String getNodeRequestId(@NonNull Request statement, @Nullable String sessionRequestId); +} diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java b/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java index a24b632f640..eeb541975ea 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java +++ b/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 nodeStateListenerRef; private final LazyReference schemaChangeListenerRef; private final LazyReference requestTrackerRef; + private final LazyReference distributedTraceIdGeneratorRef; private final LazyReference> authProviderRef; private final LazyReference> 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 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) { diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java index 0808bdce63f..1aa7a04b945 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java +++ b/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; @@ -125,6 +126,7 @@ public class CqlRequestHandler implements Throttled { private final List inFlightCallbacks; private final RequestThrottler throttler; private final RequestTracker requestTracker; + private final DistributedTraceIdGenerator distributedTraceIdGenerator; private final SessionMetricUpdater sessionMetricUpdater; private final DriverExecutionProfile executionProfile; @@ -139,7 +141,8 @@ protected CqlRequestHandler( String sessionLogPrefix) { this.startTimeNanos = System.nanoTime(); - this.logPrefix = sessionLogPrefix + "|" + this.hashCode(); + this.distributedTraceIdGenerator = context.getDistributedTraceIdGenerator(); + this.logPrefix = this.distributedTraceIdGenerator.getSessionRequestId(statement); LOG.trace("[{}] Creating new handler for request {}", logPrefix, statement); this.initialStatement = statement; @@ -489,7 +492,8 @@ private NodeResponseCallback( this.execution = execution; this.retryCount = retryCount; this.scheduleNextExecution = scheduleNextExecution; - this.logPrefix = logPrefix + "|" + execution; + this.logPrefix = + CqlRequestHandler.this.distributedTraceIdGenerator.getNodeRequestId(statement, logPrefix); } // this gets invoked once the write completes. diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopDistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopDistributedTraceIdGenerator.java new file mode 100644 index 00000000000..db7106b8c38 --- /dev/null +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopDistributedTraceIdGenerator.java @@ -0,0 +1,39 @@ +/* + * 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.internal.core.tracker; + +import com.datastax.oss.driver.api.core.context.DriverContext; +import com.datastax.oss.driver.api.core.session.Request; +import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator; +import edu.umd.cs.findbugs.annotations.NonNull; +import edu.umd.cs.findbugs.annotations.Nullable; + +public class NoopDistributedTraceIdGenerator implements DistributedTraceIdGenerator { + + public NoopDistributedTraceIdGenerator(DriverContext context) {} + + @Override + public String getSessionRequestId(@NonNull Request statement) { + return ""; + } + + @Override + public String getNodeRequestId(@NonNull Request statement, @Nullable String sessionRequestId) { + return ""; + } +} diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java new file mode 100644 index 00000000000..62e5a2e44a6 --- /dev/null +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java @@ -0,0 +1,38 @@ +/* + * 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.internal.core.tracker; + +import com.datastax.oss.driver.api.core.context.DriverContext; +import com.datastax.oss.driver.api.core.session.Request; +import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator; +import com.datastax.oss.driver.api.core.uuid.Uuids; +import edu.umd.cs.findbugs.annotations.NonNull; +import edu.umd.cs.findbugs.annotations.Nullable; + +public class UuidDistributedTraceIdGenerator implements DistributedTraceIdGenerator { + public UuidDistributedTraceIdGenerator(DriverContext context) {} + @Override + public String getSessionRequestId(@NonNull Request statement) { + return Uuids.random().toString(); + } + + @Override + public String getNodeRequestId(@NonNull Request statement, @Nullable String sessionRequestId) { + return sessionRequestId + "-" + Uuids.random().toString(); + } +} diff --git a/core/src/main/resources/reference.conf b/core/src/main/resources/reference.conf index f09ffd18a10..8092dd878b2 100644 --- a/core/src/main/resources/reference.conf +++ b/core/src/main/resources/reference.conf @@ -918,6 +918,10 @@ datastax-java-driver { } } + # The component that generates unique identifiers for distributed tracing. + advanced.distributed-tracing.id-generator{ + class = NoopDistributedTraceIdGenerator + } # A session-wide component that controls the rate at which requests are executed. # # Implementations vary, but throttlers generally track a metric that represents the level of From 5b126951757d317870b6b913c717e05fa5973fdf Mon Sep 17 00:00:00 2001 From: janehe Date: Tue, 15 Apr 2025 17:09:34 -0700 Subject: [PATCH 02/16] CustomPayloadKey and W3CContextDistributedTraceIdGenerator --- .../api/core/config/DefaultDriverOption.java | 9 +++- .../driver/api/core/config/OptionsMap.java | 1 + .../api/core/config/TypedDriverOption.java | 7 +++ .../tracker/DistributedTraceIdGenerator.java | 3 +- .../internal/core/cql/CqlRequestHandler.java | 23 ++++++++-- .../NoopDistributedTraceIdGenerator.java | 3 +- .../UuidDistributedTraceIdGenerator.java | 20 ++++---- ...W3CContextDistributedTraceIdGenerator.java | 46 +++++++++++++++++++ core/src/main/resources/reference.conf | 11 +++-- 9 files changed, 102 insertions(+), 21 deletions(-) create mode 100644 core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextDistributedTraceIdGenerator.java diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java index f0403396c9f..20aa6745137 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java @@ -1001,7 +1001,14 @@ public enum DefaultDriverOption implements DriverOption { * *

Value-type: {@link String} */ - DISTRIBUTED_TRACE_ID_GENERATOR_CLASS("advanced.distributed-tracing.id-generator.class"); + 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 + * + *

Value-type: {@link String} + */ + DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY("advanced.distributed-tracing.custom-payload-with-key"); private final String path; diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java index 636803a044c..00144440ab1 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java @@ -293,6 +293,7 @@ protected static void fillWithDriverDefaults(OptionsMap map) { map.put(TypedDriverOption.REQUEST_LOG_WARNINGS, true); map.put( TypedDriverOption.DISTRIBUTED_TRACE_ID_GENERATOR_CLASS, "NoopDistributedTraceIdGenerator"); + 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); diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java index 217d46bad3f..51cd1e17939 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java @@ -286,6 +286,13 @@ public String toString() { 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 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 REQUEST_LOGGER_SUCCESS_ENABLED = new TypedDriverOption<>( diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java index 20fe3b6ad75..17af1a7459b 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java @@ -19,10 +19,9 @@ import com.datastax.oss.driver.api.core.session.Request; import edu.umd.cs.findbugs.annotations.NonNull; -import edu.umd.cs.findbugs.annotations.Nullable; public interface DistributedTraceIdGenerator { String getSessionRequestId(@NonNull Request statement); - String getNodeRequestId(@NonNull Request statement, @Nullable String sessionRequestId); + String getNodeRequestId(@NonNull Request statement, @NonNull String sessionRequestId); } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java index 1aa7a04b945..d28d35a05d0 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java @@ -79,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; @@ -133,6 +135,7 @@ public class CqlRequestHandler implements Throttled { // 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> errors; + private final String customPayloadKey; protected CqlRequestHandler( Statement statement, @@ -173,6 +176,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); @@ -251,6 +259,16 @@ private void sendRequest( if (result.isDone()) { return; } + String nodeRequestId = this.distributedTraceIdGenerator.getNodeRequestId(statement, logPrefix); + 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 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) { @@ -279,7 +297,7 @@ private void sendRequest( currentExecutionIndex, retryCount, scheduleNextExecution, - logPrefix); + nodeRequestId); Message message = Conversions.toMessage(statement, executionProfile, context); channel .write(message, statement.isTracing(), statement.getCustomPayload(), nodeResponseCallback) @@ -492,8 +510,7 @@ private NodeResponseCallback( this.execution = execution; this.retryCount = retryCount; this.scheduleNextExecution = scheduleNextExecution; - this.logPrefix = - CqlRequestHandler.this.distributedTraceIdGenerator.getNodeRequestId(statement, logPrefix); + this.logPrefix = logPrefix; } // this gets invoked once the write completes. diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopDistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopDistributedTraceIdGenerator.java index db7106b8c38..00cc31f6d9d 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopDistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopDistributedTraceIdGenerator.java @@ -21,7 +21,6 @@ import com.datastax.oss.driver.api.core.session.Request; import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator; import edu.umd.cs.findbugs.annotations.NonNull; -import edu.umd.cs.findbugs.annotations.Nullable; public class NoopDistributedTraceIdGenerator implements DistributedTraceIdGenerator { @@ -33,7 +32,7 @@ public String getSessionRequestId(@NonNull Request statement) { } @Override - public String getNodeRequestId(@NonNull Request statement, @Nullable String sessionRequestId) { + public String getNodeRequestId(@NonNull Request statement, @NonNull String sessionRequestId) { return ""; } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java index 62e5a2e44a6..d4e9c11f75f 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java @@ -22,17 +22,17 @@ import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator; import com.datastax.oss.driver.api.core.uuid.Uuids; import edu.umd.cs.findbugs.annotations.NonNull; -import edu.umd.cs.findbugs.annotations.Nullable; public class UuidDistributedTraceIdGenerator implements DistributedTraceIdGenerator { - public UuidDistributedTraceIdGenerator(DriverContext context) {} - @Override - public String getSessionRequestId(@NonNull Request statement) { - return Uuids.random().toString(); - } + public UuidDistributedTraceIdGenerator(DriverContext context) {} - @Override - public String getNodeRequestId(@NonNull Request statement, @Nullable String sessionRequestId) { - return sessionRequestId + "-" + Uuids.random().toString(); - } + @Override + public String getSessionRequestId(@NonNull Request statement) { + return Uuids.random().toString(); + } + + @Override + public String getNodeRequestId(@NonNull Request statement, @NonNull String sessionRequestId) { + return sessionRequestId + "-" + Uuids.random().toString(); + } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextDistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextDistributedTraceIdGenerator.java new file mode 100644 index 00000000000..26e733dcf80 --- /dev/null +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextDistributedTraceIdGenerator.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.internal.core.tracker; + +import com.datastax.oss.driver.api.core.context.DriverContext; +import com.datastax.oss.driver.api.core.session.Request; +import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator; +import com.datastax.oss.driver.shaded.guava.common.io.BaseEncoding; +import edu.umd.cs.findbugs.annotations.NonNull; +import java.util.Random; + +public class W3CContextDistributedTraceIdGenerator implements DistributedTraceIdGenerator { + Random random = new Random(); + BaseEncoding baseEncoding = BaseEncoding.base16().lowerCase(); + + public W3CContextDistributedTraceIdGenerator(DriverContext context) {} + + @Override + public String getSessionRequestId(@NonNull Request statement) { + byte[] bytes = new byte[16]; + random.nextBytes(bytes); + return baseEncoding.encode(bytes); + } + + @Override + public String getNodeRequestId(@NonNull Request statement, @NonNull String sessionRequestId) { + byte[] bytes = new byte[8]; + random.nextBytes(bytes); + return String.format("00-%s-%s-00", sessionRequestId, baseEncoding.encode(bytes)); + } +} diff --git a/core/src/main/resources/reference.conf b/core/src/main/resources/reference.conf index 8092dd878b2..0820a151ea5 100644 --- a/core/src/main/resources/reference.conf +++ b/core/src/main/resources/reference.conf @@ -918,9 +918,14 @@ datastax-java-driver { } } - # The component that generates unique identifiers for distributed tracing. - advanced.distributed-tracing.id-generator{ - class = NoopDistributedTraceIdGenerator + advanced.distributed-tracing{ + id-generator{ + # The component that generates unique identifiers for distributed tracing. + class = NoopDistributedTraceIdGenerator + } + # add the trace-id to the custom payload with the given key + # if empty, the trace-id will not be added to the custom payload + custom-payload-with-key = "" } # A session-wide component that controls the rate at which requests are executed. # From 6660c0416bdfa748360e03c33ece5d3a4b2af490 Mon Sep 17 00:00:00 2001 From: janehe Date: Wed, 16 Apr 2025 11:19:58 -0700 Subject: [PATCH 03/16] Change Noop to DefaultDistributedTraceIdGenerator, preserve the existing logprefix behavior --- core/revapi.json | 2 +- .../oss/driver/api/core/config/OptionsMap.java | 3 ++- .../core/tracker/DistributedTraceIdGenerator.java | 5 +++-- .../internal/core/cql/CqlRequestHandler.java | 8 ++++++-- ...ava => DefaultDistributedTraceIdGenerator.java} | 14 ++++++++------ .../tracker/UuidDistributedTraceIdGenerator.java | 6 ++++-- .../W3CContextDistributedTraceIdGenerator.java | 9 ++++++--- core/src/main/resources/reference.conf | 2 +- 8 files changed, 31 insertions(+), 18 deletions(-) rename core/src/main/java/com/datastax/oss/driver/internal/core/tracker/{NoopDistributedTraceIdGenerator.java => DefaultDistributedTraceIdGenerator.java} (69%) diff --git a/core/revapi.json b/core/revapi.json index 1f5ec4bc2bb..90774ba5c7c 100644 --- a/core/revapi.json +++ b/core/revapi.json @@ -7390,7 +7390,7 @@ { "code": "java.method.addedToInterface", "new": "method com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator com.datastax.oss.driver.api.core.context.DriverContext::getDistributedTraceIdGenerator()", - "justification": "DistributedRequestID" + "justification": "CASSJAVA-97: Let users inject an ID for each request and write to the custom payload" } ] } diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java index 00144440ab1..07eeb8b8e6e 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java @@ -292,7 +292,8 @@ protected static void fillWithDriverDefaults(OptionsMap map) { map.put(TypedDriverOption.REQUEST_TRACE_CONSISTENCY, "ONE"); map.put(TypedDriverOption.REQUEST_LOG_WARNINGS, true); map.put( - TypedDriverOption.DISTRIBUTED_TRACE_ID_GENERATOR_CLASS, "NoopDistributedTraceIdGenerator"); + 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); diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java index 17af1a7459b..4720b9b88c4 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java @@ -21,7 +21,8 @@ import edu.umd.cs.findbugs.annotations.NonNull; public interface DistributedTraceIdGenerator { - String getSessionRequestId(@NonNull Request statement); + String getSessionRequestId(@NonNull Request statement, @NonNull String sessionName, int hashCode); - String getNodeRequestId(@NonNull Request statement, @NonNull String sessionRequestId); + String getNodeRequestId( + @NonNull Request statement, @NonNull String sessionRequestId, int executionCount); } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java index d28d35a05d0..accf20e8ba6 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java @@ -145,7 +145,9 @@ protected CqlRequestHandler( this.startTimeNanos = System.nanoTime(); this.distributedTraceIdGenerator = context.getDistributedTraceIdGenerator(); - this.logPrefix = this.distributedTraceIdGenerator.getSessionRequestId(statement); + this.logPrefix = + this.distributedTraceIdGenerator.getSessionRequestId( + statement, sessionLogPrefix, this.hashCode()); LOG.trace("[{}] Creating new handler for request {}", logPrefix, statement); this.initialStatement = statement; @@ -259,7 +261,9 @@ private void sendRequest( if (result.isDone()) { return; } - String nodeRequestId = this.distributedTraceIdGenerator.getNodeRequestId(statement, logPrefix); + 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 diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopDistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/DefaultDistributedTraceIdGenerator.java similarity index 69% rename from core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopDistributedTraceIdGenerator.java rename to core/src/main/java/com/datastax/oss/driver/internal/core/tracker/DefaultDistributedTraceIdGenerator.java index 00cc31f6d9d..7142d0502aa 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopDistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/DefaultDistributedTraceIdGenerator.java @@ -22,17 +22,19 @@ import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator; import edu.umd.cs.findbugs.annotations.NonNull; -public class NoopDistributedTraceIdGenerator implements DistributedTraceIdGenerator { +public class DefaultDistributedTraceIdGenerator implements DistributedTraceIdGenerator { - public NoopDistributedTraceIdGenerator(DriverContext context) {} + public DefaultDistributedTraceIdGenerator(DriverContext context) {} @Override - public String getSessionRequestId(@NonNull Request statement) { - return ""; + public String getSessionRequestId( + @NonNull Request statement, @NonNull String sessionName, int hashCode) { + return sessionName + "|" + hashCode; } @Override - public String getNodeRequestId(@NonNull Request statement, @NonNull String sessionRequestId) { - return ""; + public String getNodeRequestId( + @NonNull Request statement, @NonNull String sessionRequestId, int executionCount) { + return sessionRequestId + "|" + executionCount; } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java index d4e9c11f75f..7948140aa2c 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java @@ -27,12 +27,14 @@ public class UuidDistributedTraceIdGenerator implements DistributedTraceIdGenera public UuidDistributedTraceIdGenerator(DriverContext context) {} @Override - public String getSessionRequestId(@NonNull Request statement) { + public String getSessionRequestId( + @NonNull Request statement, @NonNull String sessionName, int hashCode) { return Uuids.random().toString(); } @Override - public String getNodeRequestId(@NonNull Request statement, @NonNull String sessionRequestId) { + public String getNodeRequestId( + @NonNull Request statement, @NonNull String sessionRequestId, int executionCount) { return sessionRequestId + "-" + Uuids.random().toString(); } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextDistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextDistributedTraceIdGenerator.java index 26e733dcf80..c186577ade7 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextDistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextDistributedTraceIdGenerator.java @@ -22,23 +22,26 @@ import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator; import com.datastax.oss.driver.shaded.guava.common.io.BaseEncoding; import edu.umd.cs.findbugs.annotations.NonNull; +import java.security.SecureRandom; import java.util.Random; public class W3CContextDistributedTraceIdGenerator implements DistributedTraceIdGenerator { - Random random = new Random(); + Random random = new SecureRandom(); BaseEncoding baseEncoding = BaseEncoding.base16().lowerCase(); public W3CContextDistributedTraceIdGenerator(DriverContext context) {} @Override - public String getSessionRequestId(@NonNull Request statement) { + public String getSessionRequestId( + @NonNull Request statement, @NonNull String sessionName, int hashCode) { byte[] bytes = new byte[16]; random.nextBytes(bytes); return baseEncoding.encode(bytes); } @Override - public String getNodeRequestId(@NonNull Request statement, @NonNull String sessionRequestId) { + public String getNodeRequestId( + @NonNull Request statement, @NonNull String sessionRequestId, int executionCount) { byte[] bytes = new byte[8]; random.nextBytes(bytes); return String.format("00-%s-%s-00", sessionRequestId, baseEncoding.encode(bytes)); diff --git a/core/src/main/resources/reference.conf b/core/src/main/resources/reference.conf index 0820a151ea5..ad292d4d2eb 100644 --- a/core/src/main/resources/reference.conf +++ b/core/src/main/resources/reference.conf @@ -921,7 +921,7 @@ datastax-java-driver { advanced.distributed-tracing{ id-generator{ # The component that generates unique identifiers for distributed tracing. - class = NoopDistributedTraceIdGenerator + class = DefaultDistributedTraceIdGenerator } # add the trace-id to the custom payload with the given key # if empty, the trace-id will not be added to the custom payload From 15d03969bcad2e28fd6e67956567246303a501cb Mon Sep 17 00:00:00 2001 From: janehe Date: Wed, 16 Apr 2025 14:29:51 -0700 Subject: [PATCH 04/16] fix tests --- .../internal/core/cql/RequestHandlerTestHarness.java | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java index 9d86302aabf..bbee1e4e3d6 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java @@ -37,6 +37,7 @@ import com.datastax.oss.driver.api.core.session.Session; import com.datastax.oss.driver.api.core.specex.SpeculativeExecutionPolicy; 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.type.codec.registry.CodecRegistry; import com.datastax.oss.driver.internal.core.DefaultConsistencyLevelRegistry; import com.datastax.oss.driver.internal.core.ProtocolFeature; @@ -51,6 +52,7 @@ import com.datastax.oss.driver.internal.core.servererrors.DefaultWriteTypeRegistry; import com.datastax.oss.driver.internal.core.session.DefaultSession; import com.datastax.oss.driver.internal.core.session.throttling.PassThroughRequestThrottler; +import com.datastax.oss.driver.internal.core.tracker.DefaultDistributedTraceIdGenerator; import com.datastax.oss.driver.internal.core.tracker.NoopRequestTracker; import com.datastax.oss.driver.internal.core.util.concurrent.CapturingTimer; import com.datastax.oss.driver.internal.core.util.concurrent.CapturingTimer.CapturedTimeout; @@ -114,6 +116,8 @@ protected RequestHandlerTestHarness(Builder builder) { when(defaultProfile.getBoolean(DefaultDriverOption.REQUEST_DEFAULT_IDEMPOTENCE)) .thenReturn(builder.defaultIdempotence); when(defaultProfile.getBoolean(DefaultDriverOption.PREPARE_ON_ALL_NODES)).thenReturn(true); + when(defaultProfile.getString(DefaultDriverOption.DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY)) + .thenReturn(""); when(config.getDefaultProfile()).thenReturn(defaultProfile); when(context.getConfig()).thenReturn(config); @@ -125,6 +129,10 @@ protected RequestHandlerTestHarness(Builder builder) { when(context.getRetryPolicy(anyString())).thenReturn(retryPolicy); + DistributedTraceIdGenerator distributedTraceIdGenerator = + new DefaultDistributedTraceIdGenerator(context); + when(context.getDistributedTraceIdGenerator()).thenReturn(distributedTraceIdGenerator); + // Disable speculative executions by default when(speculativeExecutionPolicy.nextExecution( any(Node.class), any(CqlIdentifier.class), any(Request.class), anyInt())) From c075c0fe7b98bbb1e98360f48e61add474e39974 Mon Sep 17 00:00:00 2001 From: janehe Date: Thu, 17 Apr 2025 01:00:25 -0700 Subject: [PATCH 05/16] add tests, add doc --- .../tracker/DistributedTraceIdGenerator.java | 18 ++++ .../DistributedTraceIdGeneratorTest.java | 70 +++++++++++++ .../DistributedTraceIdGeneratorIT.java | 99 +++++++++++++++++++ manual/core/distributed_tracing/README.md | 64 ++++++++++++ 4 files changed, 251 insertions(+) create mode 100644 core/src/test/java/com/datastax/oss/driver/internal/core/tracker/DistributedTraceIdGeneratorTest.java create mode 100644 integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/DistributedTraceIdGeneratorIT.java create mode 100644 manual/core/distributed_tracing/README.md diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java index 4720b9b88c4..8f8504b1969 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java @@ -21,8 +21,26 @@ 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); } diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/DistributedTraceIdGeneratorTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/DistributedTraceIdGeneratorTest.java new file mode 100644 index 00000000000..bec1aa739a6 --- /dev/null +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/DistributedTraceIdGeneratorTest.java @@ -0,0 +1,70 @@ +/* + * 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.internal.core.tracker; + +import static org.assertj.core.api.Assertions.assertThat; + +import com.datastax.oss.driver.api.core.cql.Statement; +import com.datastax.oss.driver.internal.core.context.InternalDriverContext; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.junit.MockitoJUnitRunner; + +@RunWith(MockitoJUnitRunner.Strict.class) +public class DistributedTraceIdGeneratorTest { + @Mock private InternalDriverContext context; + @Mock private Statement statement; + + @Test + public void default_generator_should_generate() { + // given + DefaultDistributedTraceIdGenerator generator = new DefaultDistributedTraceIdGenerator(context); + // when + String sessionRequestId = generator.getSessionRequestId(statement, "sessionName", 123); + String nodeRequestId = generator.getNodeRequestId(statement, sessionRequestId, 1); + // then + assertThat(sessionRequestId).isEqualTo("sessionName|123"); + assertThat(nodeRequestId).isEqualTo("sessionName|123|1"); + } + + @Test + public void uuid_generator_should_generate() { + // given + UuidDistributedTraceIdGenerator generator = new UuidDistributedTraceIdGenerator(context); + // when + String sessionRequestId = generator.getSessionRequestId(statement, "sessionName", 123); + String nodeRequestId = generator.getNodeRequestId(statement, sessionRequestId, 1); + // then + assertThat(sessionRequestId.length()).isEqualTo(36); + assertThat(nodeRequestId.length()).isEqualTo(73); + } + + @Test + public void w3c_generator_should_generate() { + // given + W3CContextDistributedTraceIdGenerator generator = + new W3CContextDistributedTraceIdGenerator(context); + // when + String sessionRequestId = generator.getSessionRequestId(statement, "sessionName", 123); + String nodeRequestId = generator.getNodeRequestId(statement, sessionRequestId, 1); + // then + assertThat(sessionRequestId.length()).isEqualTo(32); + assertThat(nodeRequestId.length()).isEqualTo(55); + } +} diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/DistributedTraceIdGeneratorIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/DistributedTraceIdGeneratorIT.java new file mode 100644 index 00000000000..18f3c76d940 --- /dev/null +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/DistributedTraceIdGeneratorIT.java @@ -0,0 +1,99 @@ +/* + * 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.core.tracker; + +import static org.assertj.core.api.Assertions.assertThat; + +import com.datastax.oss.driver.api.core.CqlSession; +import com.datastax.oss.driver.api.core.config.DefaultDriverOption; +import com.datastax.oss.driver.api.core.config.DriverConfigLoader; +import com.datastax.oss.driver.api.core.cql.ResultSet; +import com.datastax.oss.driver.api.testinfra.ccm.CcmRule; +import com.datastax.oss.driver.api.testinfra.session.SessionUtils; +import com.datastax.oss.driver.categories.ParallelizableTests; +import java.nio.ByteBuffer; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.RuleChain; +import org.junit.rules.TestRule; + +@Category(ParallelizableTests.class) +public class DistributedTraceIdGeneratorIT { + private CcmRule ccmRule = CcmRule.getInstance(); + + @Rule public TestRule chain = RuleChain.outerRule(ccmRule); + + @Test + public void should_write_default_id_to_custom_payload_with_key() { + DriverConfigLoader loader = + SessionUtils.configLoaderBuilder() + .withString(DefaultDriverOption.DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY, "trace_key") + .build(); + try (CqlSession session = SessionUtils.newSession(ccmRule, loader)) { + String query = "SELECT * FROM system.local"; + ResultSet rs = session.execute(query); + assertThat(rs.getExecutionInfo().getRequest().getCustomPayload().get("trace_key")) + .isNotNull(); + } + } + + @Test + public void should_write_uuid_to_custom_payload_with_key() { + DriverConfigLoader loader = + SessionUtils.configLoaderBuilder() + .withString( + DefaultDriverOption.DISTRIBUTED_TRACE_ID_GENERATOR_CLASS, + "UuidDistributedTraceIdGenerator") + .withString(DefaultDriverOption.DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY, "trace_key") + .build(); + try (CqlSession session = SessionUtils.newSession(ccmRule, loader)) { + String query = "SELECT * FROM system.local"; + ResultSet rs = session.execute(query); + ByteBuffer id = rs.getExecutionInfo().getRequest().getCustomPayload().get("trace_key"); + assertThat(id.array().length).isEqualTo(73); + } + } + + @Test + public void should_write_w3c_context_to_custom_payload_with_key() { + DriverConfigLoader loader = + SessionUtils.configLoaderBuilder() + .withString( + DefaultDriverOption.DISTRIBUTED_TRACE_ID_GENERATOR_CLASS, + "W3CContextDistributedTraceIdGenerator") + .withString(DefaultDriverOption.DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY, "trace_key") + .build(); + try (CqlSession session = SessionUtils.newSession(ccmRule, loader)) { + String query = "SELECT * FROM system.local"; + ResultSet rs = session.execute(query); + ByteBuffer id = rs.getExecutionInfo().getRequest().getCustomPayload().get("trace_key"); + assertThat(id.array().length).isEqualTo(55); + } + } + + @Test + public void should_not_write_id_to_custom_payload_when_key_is_not_set() { + DriverConfigLoader loader = SessionUtils.configLoaderBuilder().build(); + try (CqlSession session = SessionUtils.newSession(ccmRule, loader)) { + String query = "SELECT * FROM system.local"; + ResultSet rs = session.execute(query); + assertThat(rs.getExecutionInfo().getRequest().getCustomPayload().get("trace_key")).isNull(); + } + } +} diff --git a/manual/core/distributed_tracing/README.md b/manual/core/distributed_tracing/README.md new file mode 100644 index 00000000000..e67bb538ea6 --- /dev/null +++ b/manual/core/distributed_tracing/README.md @@ -0,0 +1,64 @@ + + +## Distributed tracing + +### Quick overview + +Users can inject an identifier for each individual CQL request, and such ID can be written in to the custom payload to +correlate a request across the driver and the Apache Cassandra server. + +* Inject ID generator: set the desired `DistributedTraceIdGenerator` in `advanced.distributed-tracing.id-generator.class`. + The default implementation generates ID as `{session_name}|{hash_code}|{execution_count}`. +* Add ID to custom payload: disabled by default. Set the desired key in `advanced.distributed-tracing.custom-payload-with-key`, + then the driver will add the generated ID to the custom payload with the specified key. + +### Distributed Trace Id Generator Configuration + +Distributed trace ID generator can be declared in the [configuration](../configuration/) as follows: + +``` +datastax-java-driver.advanced.distributed-tracing.id-generator { + class = com.example.app.MyGenerator +} +``` + +To register your own trackers, specify the name of a class +that implements `DistributedTraceIdGenerator`. + +By default, the build-in implementation `DefaultDistributedTraceIdGenerator` is used. It generates the ID as +`{session_name}|{hash_code}|{execution_count}`. Note that this ID is not guaranteed to be unique. +Other built-in implementations include `UUIDDistributedTraceIdGenerator` and `W3CContextDistributedTraceIdGenerator`. + +The generated ID will be added to the log message of `CqlRequestHandler`, and propagated to the request trackers. + +### Custom Payload Configuration + +Users can opt in to add the generated ID to the custom payload to achieve request end-to-end tracing. +Custom payload is a map of string to `ByteBuffer` pairs, and the driver will add the generated ID to the custom payload with the specified key. + +``` +datastax-java-driver.advanced.distributed-tracing{ + custom-payload-with-key = my-trace-key +} +``` + +Users can then run Apache Cassandra with customized query handler to extract the trace ID from the custom payload, to achieve end-to-end tracing. + +When this key is set to an empty string (default), the driver will not add the ID to the custom payload. From ce2ae9a99461a91897996d324da93ecc1596371a Mon Sep 17 00:00:00 2001 From: janehe Date: Fri, 18 Apr 2025 14:15:48 -0700 Subject: [PATCH 06/16] Use ByteBuffer.remaining() --- .../driver/core/tracker/DistributedTraceIdGeneratorIT.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/DistributedTraceIdGeneratorIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/DistributedTraceIdGeneratorIT.java index 18f3c76d940..49ed63a32b2 100644 --- a/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/DistributedTraceIdGeneratorIT.java +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/DistributedTraceIdGeneratorIT.java @@ -66,7 +66,7 @@ public void should_write_uuid_to_custom_payload_with_key() { String query = "SELECT * FROM system.local"; ResultSet rs = session.execute(query); ByteBuffer id = rs.getExecutionInfo().getRequest().getCustomPayload().get("trace_key"); - assertThat(id.array().length).isEqualTo(73); + assertThat(id.remaining()).isEqualTo(73); } } @@ -83,7 +83,7 @@ public void should_write_w3c_context_to_custom_payload_with_key() { String query = "SELECT * FROM system.local"; ResultSet rs = session.execute(query); ByteBuffer id = rs.getExecutionInfo().getRequest().getCustomPayload().get("trace_key"); - assertThat(id.array().length).isEqualTo(55); + assertThat(id.remaining()).isEqualTo(55); } } From 5cfcbc2a4a7cb1374818864ee37fd6c61f40e795 Mon Sep 17 00:00:00 2001 From: janehe Date: Tue, 29 Apr 2025 13:05:45 -0700 Subject: [PATCH 07/16] copy the custom payload, add doc. Integration tests fail because execution info does not have actual custom payload --- .../api/core/session/SessionBuilder.java | 6 +-- .../tracker/DistributedTraceIdGenerator.java | 8 +++- .../api/core/tracker/RequestTracker.java | 41 +++++++++++-------- .../internal/core/cql/CqlRequestHandler.java | 19 +++++---- .../DefaultLoadBalancingPolicy.java | 4 +- .../tracker/MultiplexingRequestTracker.java | 31 ++++++++------ .../core/tracker/NoopRequestTracker.java | 8 ++-- .../internal/core/tracker/RequestLogger.java | 12 +++--- .../UuidDistributedTraceIdGenerator.java | 2 +- .../tracker/RequestNodeLoggerExample.java | 8 ++-- 10 files changed, 77 insertions(+), 62 deletions(-) diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java b/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java index d13c8ca7361..abdb884d838 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java @@ -320,11 +320,8 @@ public SelfT addRequestTracker(@NonNull RequestTracker requestTracker) { } /** - * Registers a distributed trace ID generator The driver will use the distributed trace ID in the + * 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( @@ -332,6 +329,7 @@ public SelfT withDistributedTraceIdGenerator( this.programmaticArgumentsBuilder.withDistributedTraceIdGenerator(distributedTraceIdGenerator); return self; } + /** * Registers an authentication provider to use with the session. * diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java index 8f8504b1969..338b85f1dc2 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java @@ -22,7 +22,9 @@ public interface DistributedTraceIdGenerator { /** - * Generates a unique identifier for the session request. This identifier will be added to logs. + * 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 @@ -32,7 +34,9 @@ public interface DistributedTraceIdGenerator { 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 + * 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 diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestTracker.java b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestTracker.java index d29ee48d352..065b41e496a 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestTracker.java +++ b/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,13 +79,13 @@ 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, @@ -92,8 +93,9 @@ default void onError( 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,14 +112,15 @@ 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, @@ -125,8 +128,9 @@ default void onNodeError( 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); } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java index accf20e8ba6..23047c4ac68 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java @@ -72,6 +72,7 @@ import com.datastax.oss.protocol.internal.response.result.SetKeyspace; import com.datastax.oss.protocol.internal.response.result.Void; import com.datastax.oss.protocol.internal.util.Bytes; +import com.datastax.oss.protocol.internal.util.collection.NullAllowingImmutableMap; import edu.umd.cs.findbugs.annotations.NonNull; import io.netty.handler.codec.EncoderException; import io.netty.util.Timeout; @@ -82,7 +83,6 @@ 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; @@ -264,14 +264,15 @@ private void sendRequest( String nodeRequestId = this.distributedTraceIdGenerator.getNodeRequestId( statement, logPrefix, currentExecutionIndex); + Map customPayload = statement.getCustomPayload(); 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 existingMap = new HashMap<>(statement.getCustomPayload()); - existingMap.put( - this.customPayloadKey, ByteBuffer.wrap(nodeRequestId.getBytes(StandardCharsets.UTF_8))); - statement = statement.setCustomPayload(existingMap); + customPayload = + NullAllowingImmutableMap.builder() + .putAll(customPayload) + .put( + this.customPayloadKey, + ByteBuffer.wrap(nodeRequestId.getBytes(StandardCharsets.UTF_8))) + .build(); } Node node = retriedNode; DriverChannel channel = null; @@ -304,7 +305,7 @@ private void sendRequest( nodeRequestId); Message message = Conversions.toMessage(statement, executionProfile, context); channel - .write(message, statement.isTracing(), statement.getCustomPayload(), nodeResponseCallback) + .write(message, statement.isTracing(), customPayload, nodeResponseCallback) .addListener(nodeResponseCallback); } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java b/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java index 0f03cbb3643..ea2342314e5 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java @@ -245,7 +245,7 @@ public void onNodeSuccess( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @NonNull Node node, - @NonNull String logPrefix) { + @NonNull String nodeRequestLogPrefix) { updateResponseTimes(node); } @@ -256,7 +256,7 @@ public void onNodeError( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @NonNull Node node, - @NonNull String logPrefix) { + @NonNull String nodeRequestLogPrefix) { updateResponseTimes(node); } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/MultiplexingRequestTracker.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/MultiplexingRequestTracker.java index d4d20f3eb78..6fe2ba059bd 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/MultiplexingRequestTracker.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/MultiplexingRequestTracker.java @@ -82,10 +82,12 @@ public void onSuccess( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @NonNull Node node, - @NonNull String logPrefix) { + @NonNull String sessionRequestLogPrefix) { invokeTrackers( - tracker -> tracker.onSuccess(request, latencyNanos, executionProfile, node, logPrefix), - logPrefix, + tracker -> + tracker.onSuccess( + request, latencyNanos, executionProfile, node, sessionRequestLogPrefix), + sessionRequestLogPrefix, "onSuccess"); } @@ -96,10 +98,12 @@ public void onError( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @Nullable Node node, - @NonNull String logPrefix) { + @NonNull String sessionRequestLogPrefix) { invokeTrackers( - tracker -> tracker.onError(request, error, latencyNanos, executionProfile, node, logPrefix), - logPrefix, + tracker -> + tracker.onError( + request, error, latencyNanos, executionProfile, node, sessionRequestLogPrefix), + sessionRequestLogPrefix, "onError"); } @@ -109,10 +113,12 @@ public void onNodeSuccess( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @NonNull Node node, - @NonNull String logPrefix) { + @NonNull String nodeRequestLogPrefix) { invokeTrackers( - tracker -> tracker.onNodeSuccess(request, latencyNanos, executionProfile, node, logPrefix), - logPrefix, + tracker -> + tracker.onNodeSuccess( + request, latencyNanos, executionProfile, node, nodeRequestLogPrefix), + nodeRequestLogPrefix, "onNodeSuccess"); } @@ -123,11 +129,12 @@ public void onNodeError( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @NonNull Node node, - @NonNull String logPrefix) { + @NonNull String nodeRequestLogPrefix) { invokeTrackers( tracker -> - tracker.onNodeError(request, error, latencyNanos, executionProfile, node, logPrefix), - logPrefix, + tracker.onNodeError( + request, error, latencyNanos, executionProfile, node, nodeRequestLogPrefix), + nodeRequestLogPrefix, "onNodeError"); } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopRequestTracker.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopRequestTracker.java index 09ac27e5e75..3821c6ace2d 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopRequestTracker.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/NoopRequestTracker.java @@ -42,7 +42,7 @@ public void onSuccess( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @NonNull Node node, - @NonNull String requestPrefix) { + @NonNull String sessionRequestLogPrefix) { // nothing to do } @@ -53,7 +53,7 @@ public void onError( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, Node node, - @NonNull String requestPrefix) { + @NonNull String sessionRequestLogPrefix) { // nothing to do } @@ -64,7 +64,7 @@ public void onNodeError( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @NonNull Node node, - @NonNull String requestPrefix) { + @NonNull String nodeRequestLogPrefix) { // nothing to do } @@ -74,7 +74,7 @@ public void onNodeSuccess( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @NonNull Node node, - @NonNull String requestPrefix) { + @NonNull String nodeRequestLogPrefix) { // nothing to do } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/RequestLogger.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/RequestLogger.java index 235ef051b40..f242ff89c54 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/RequestLogger.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/RequestLogger.java @@ -86,7 +86,7 @@ public void onSuccess( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @NonNull Node node, - @NonNull String logPrefix) { + @NonNull String sessionRequestLogPrefix) { boolean successEnabled = executionProfile.getBoolean(DefaultDriverOption.REQUEST_LOGGER_SUCCESS_ENABLED, false); @@ -129,7 +129,7 @@ public void onSuccess( showValues, maxValues, maxValueLength, - logPrefix); + sessionRequestLogPrefix); } @Override @@ -139,7 +139,7 @@ public void onError( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, Node node, - @NonNull String logPrefix) { + @NonNull String sessionRequestLogPrefix) { if (!executionProfile.getBoolean(DefaultDriverOption.REQUEST_LOGGER_ERROR_ENABLED, false)) { return; @@ -173,7 +173,7 @@ public void onError( maxValues, maxValueLength, showStackTraces, - logPrefix); + sessionRequestLogPrefix); } @Override @@ -183,7 +183,7 @@ public void onNodeError( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @NonNull Node node, - @NonNull String logPrefix) { + @NonNull String nodeRequestLogPrefix) { // Nothing to do } @@ -193,7 +193,7 @@ public void onNodeSuccess( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @NonNull Node node, - @NonNull String logPrefix) { + @NonNull String nodeRequestLogPrefix) { // Nothing to do } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java index 7948140aa2c..d05cf799b98 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java @@ -35,6 +35,6 @@ public String getSessionRequestId( @Override public String getNodeRequestId( @NonNull Request statement, @NonNull String sessionRequestId, int executionCount) { - return sessionRequestId + "-" + Uuids.random().toString(); + return sessionRequestId + "-" + Uuids.random(); } } diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestNodeLoggerExample.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestNodeLoggerExample.java index eae98339637..8eb2fb80a73 100644 --- a/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestNodeLoggerExample.java +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestNodeLoggerExample.java @@ -39,7 +39,7 @@ public void onNodeError( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @NonNull Node node, - @NonNull String logPrefix) { + @NonNull String nodeRequestLogPrefix) { if (!executionProfile.getBoolean(DefaultDriverOption.REQUEST_LOGGER_ERROR_ENABLED)) { return; } @@ -66,7 +66,7 @@ public void onNodeError( maxValues, maxValueLength, showStackTraces, - logPrefix); + nodeRequestLogPrefix); } @Override @@ -75,7 +75,7 @@ public void onNodeSuccess( long latencyNanos, @NonNull DriverExecutionProfile executionProfile, @NonNull Node node, - @NonNull String logPrefix) { + @NonNull String nodeRequestLogPrefix) { boolean successEnabled = executionProfile.getBoolean(DefaultDriverOption.REQUEST_LOGGER_SUCCESS_ENABLED); boolean slowEnabled = @@ -114,6 +114,6 @@ public void onNodeSuccess( showValues, maxValues, maxValueLength, - logPrefix); + nodeRequestLogPrefix); } } From 5f95b14385adcb2dc7e5f851c6dd15ef531f9937 Mon Sep 17 00:00:00 2001 From: janehe Date: Mon, 19 May 2025 16:43:08 -0700 Subject: [PATCH 08/16] rename to Request Id --- core/revapi.json | 2 +- .../api/core/config/DefaultDriverOption.java | 8 ++--- .../driver/api/core/config/OptionsMap.java | 6 ++-- .../api/core/config/TypedDriverOption.java | 15 ++++---- .../api/core/context/DriverContext.java | 6 ++-- .../core/session/ProgrammaticArguments.java | 21 ++++++----- .../api/core/session/SessionBuilder.java | 12 +++---- ...Generator.java => RequestIdGenerator.java} | 2 +- .../core/context/DefaultDriverContext.java | 31 ++++++++-------- .../internal/core/cql/CqlRequestHandler.java | 24 ++++++------- ...or.java => DefaultRequestIdGenerator.java} | 6 ++-- ...rator.java => UuidRequestIdGenerator.java} | 6 ++-- ...java => W3CContextRequestIdGenerator.java} | 6 ++-- core/src/main/resources/reference.conf | 12 +++---- .../core/cql/RequestHandlerTestHarness.java | 11 +++--- ...rTest.java => RequestIdGeneratorTest.java} | 9 +++-- ...ratorIT.java => RequestIdGeneratorIT.java} | 15 ++++---- .../README.md | 35 +++++++++++-------- 18 files changed, 108 insertions(+), 119 deletions(-) rename core/src/main/java/com/datastax/oss/driver/api/core/tracker/{DistributedTraceIdGenerator.java => RequestIdGenerator.java} (97%) rename core/src/main/java/com/datastax/oss/driver/internal/core/tracker/{DefaultDistributedTraceIdGenerator.java => DefaultRequestIdGenerator.java} (85%) rename core/src/main/java/com/datastax/oss/driver/internal/core/tracker/{UuidDistributedTraceIdGenerator.java => UuidRequestIdGenerator.java} (86%) rename core/src/main/java/com/datastax/oss/driver/internal/core/tracker/{W3CContextDistributedTraceIdGenerator.java => W3CContextRequestIdGenerator.java} (88%) rename core/src/test/java/com/datastax/oss/driver/internal/core/tracker/{DistributedTraceIdGeneratorTest.java => RequestIdGeneratorTest.java} (87%) rename integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/{DistributedTraceIdGeneratorIT.java => RequestIdGeneratorIT.java} (85%) rename manual/core/{distributed_tracing => request_id}/README.md (54%) diff --git a/core/revapi.json b/core/revapi.json index 90774ba5c7c..a2f4c16e651 100644 --- a/core/revapi.json +++ b/core/revapi.json @@ -7389,7 +7389,7 @@ }, { "code": "java.method.addedToInterface", - "new": "method com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator com.datastax.oss.driver.api.core.context.DriverContext::getDistributedTraceIdGenerator()", + "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" } ] diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java index 20aa6745137..5e2f1252a95 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java @@ -997,18 +997,18 @@ public enum DefaultDriverOption implements DriverOption { 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. + * The class of session-wide component that generates request IDs. * *

Value-type: {@link String} */ - DISTRIBUTED_TRACE_ID_GENERATOR_CLASS("advanced.distributed-tracing.id-generator.class"), + REQUEST_ID_GENERATOR_CLASS("advanced.request-id.generator.class"), /** - * If not empty, the driver will write the distributed trace ID to this key in the custom payload + * If not empty, the driver will write the node request ID to this key in the custom payload * *

Value-type: {@link String} */ - DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY("advanced.distributed-tracing.custom-payload-with-key"); + REQUEST_ID_CUSTOM_PAYLOAD_KEY("advanced.request-id.custom-payload-with-key"); private final String path; diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java index 07eeb8b8e6e..27df2d4f025 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java @@ -291,10 +291,8 @@ 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.REQUEST_ID_GENERATOR_CLASS, "DefaultRequestIdGenerator"); + map.put(TypedDriverOption.REQUEST_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); diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java index 51cd1e17939..76c4aea3864 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java @@ -281,17 +281,14 @@ 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 DISTRIBUTED_TRACE_ID_GENERATOR_CLASS = - new TypedDriverOption<>( - DefaultDriverOption.DISTRIBUTED_TRACE_ID_GENERATOR_CLASS, GenericType.STRING); + /** The class of a session-wide component that generates request IDs. */ + public static final TypedDriverOption REQUEST_ID_GENERATOR_CLASS = + new TypedDriverOption<>(DefaultDriverOption.REQUEST_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 DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY = + /** If not empty, the driver will write the node request ID to this key in the custom payload */ + public static final TypedDriverOption REQUEST_ID_CUSTOM_PAYLOAD_KEY = new TypedDriverOption<>( - DefaultDriverOption.DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY, GenericType.STRING); + DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY, GenericType.STRING); /** Whether to log successful requests. */ public static final TypedDriverOption REQUEST_LOGGER_SUCCESS_ENABLED = diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/context/DriverContext.java b/core/src/main/java/com/datastax/oss/driver/api/core/context/DriverContext.java index a0d047025f4..8d19ec6a149 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/context/DriverContext.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/context/DriverContext.java @@ -33,7 +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.RequestIdGenerator; import com.datastax.oss.driver.api.core.tracker.RequestTracker; import edu.umd.cs.findbugs.annotations.NonNull; import java.util.Map; @@ -140,9 +140,9 @@ default SpeculativeExecutionPolicy getSpeculativeExecutionPolicy(@NonNull String @NonNull RequestTracker getRequestTracker(); - /** @return The driver's distributed trace ID generator; never {@code null}. */ + /** @return The driver's request ID generator; never {@code null}. */ @NonNull - DistributedTraceIdGenerator getDistributedTraceIdGenerator(); + RequestIdGenerator getRequestIdGenerator(); /** @return The driver's request throttler; never {@code null}. */ @NonNull diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/session/ProgrammaticArguments.java b/core/src/main/java/com/datastax/oss/driver/api/core/session/ProgrammaticArguments.java index 07fc1e65926..5e10fb4d915 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/session/ProgrammaticArguments.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/session/ProgrammaticArguments.java @@ -23,7 +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.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; @@ -60,7 +60,7 @@ public static Builder builder() { private final NodeStateListener nodeStateListener; private final SchemaChangeListener schemaChangeListener; private final RequestTracker requestTracker; - private final DistributedTraceIdGenerator distributedTraceIdGenerator; + private final RequestIdGenerator requestIdGenerator; private final Map localDatacenters; private final Map> nodeFilters; private final Map nodeDistanceEvaluators; @@ -79,7 +79,7 @@ private ProgrammaticArguments( @Nullable NodeStateListener nodeStateListener, @Nullable SchemaChangeListener schemaChangeListener, @Nullable RequestTracker requestTracker, - @Nullable DistributedTraceIdGenerator distributedTraceIdGenerator, + @Nullable RequestIdGenerator requestIdGenerator, @NonNull Map localDatacenters, @NonNull Map> nodeFilters, @NonNull Map nodeDistanceEvaluators, @@ -97,7 +97,7 @@ private ProgrammaticArguments( this.nodeStateListener = nodeStateListener; this.schemaChangeListener = schemaChangeListener; this.requestTracker = requestTracker; - this.distributedTraceIdGenerator = distributedTraceIdGenerator; + this.requestIdGenerator = requestIdGenerator; this.localDatacenters = localDatacenters; this.nodeFilters = nodeFilters; this.nodeDistanceEvaluators = nodeDistanceEvaluators; @@ -133,8 +133,8 @@ public RequestTracker getRequestTracker() { } @Nullable - public DistributedTraceIdGenerator getDistributedTraceIdGenerator() { - return distributedTraceIdGenerator; + public RequestIdGenerator getRequestIdGenerator() { + return requestIdGenerator; } @NonNull @@ -205,7 +205,7 @@ public static class Builder { private NodeStateListener nodeStateListener; private SchemaChangeListener schemaChangeListener; private RequestTracker requestTracker; - private DistributedTraceIdGenerator distributedTraceIdGenerator; + private RequestIdGenerator requestIdGenerator; private ImmutableMap.Builder localDatacentersBuilder = ImmutableMap.builder(); private final ImmutableMap.Builder> nodeFiltersBuilder = ImmutableMap.builder(); @@ -305,9 +305,8 @@ public Builder addRequestTracker(@NonNull RequestTracker requestTracker) { } @NonNull - public Builder withDistributedTraceIdGenerator( - @Nullable DistributedTraceIdGenerator distributedTraceIdGenerator) { - this.distributedTraceIdGenerator = distributedTraceIdGenerator; + public Builder withRequestIdGenerator(@Nullable RequestIdGenerator requestIdGenerator) { + this.requestIdGenerator = requestIdGenerator; return this; } @@ -434,7 +433,7 @@ public ProgrammaticArguments build() { nodeStateListener, schemaChangeListener, requestTracker, - distributedTraceIdGenerator, + requestIdGenerator, localDatacentersBuilder.build(), nodeFiltersBuilder.build(), nodeDistanceEvaluatorsBuilder.build(), diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java b/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java index abdb884d838..3d829c1a17f 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java @@ -35,7 +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.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; @@ -320,13 +320,13 @@ public SelfT addRequestTracker(@NonNull RequestTracker requestTracker) { } /** - * 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. + * 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 withDistributedTraceIdGenerator( - @NonNull DistributedTraceIdGenerator distributedTraceIdGenerator) { - this.programmaticArgumentsBuilder.withDistributedTraceIdGenerator(distributedTraceIdGenerator); + public SelfT withRequestIdGenerator(@NonNull RequestIdGenerator requestIdGenerator) { + this.programmaticArgumentsBuilder.withRequestIdGenerator(requestIdGenerator); return self; } diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java similarity index 97% rename from core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java rename to core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java index 338b85f1dc2..e620cc2bbad 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/DistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java @@ -20,7 +20,7 @@ import com.datastax.oss.driver.api.core.session.Request; import edu.umd.cs.findbugs.annotations.NonNull; -public interface DistributedTraceIdGenerator { +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 diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java b/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java index eeb541975ea..dfbd9685bfd 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java @@ -44,7 +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.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; @@ -222,7 +222,7 @@ public class DefaultDriverContext implements InternalDriverContext { private final LazyReference nodeStateListenerRef; private final LazyReference schemaChangeListenerRef; private final LazyReference requestTrackerRef; - private final LazyReference distributedTraceIdGeneratorRef; + private final LazyReference requestIdGeneratorRef; private final LazyReference> authProviderRef; private final LazyReference> lifecycleListenersRef = new LazyReference<>("lifecycleListeners", this::buildLifecycleListeners, cycleDetector); @@ -284,12 +284,10 @@ public DefaultDriverContext( this.requestTrackerRef = new LazyReference<>( "requestTracker", () -> buildRequestTracker(requestTrackerFromBuilder), cycleDetector); - this.distributedTraceIdGeneratorRef = + this.requestIdGeneratorRef = new LazyReference<>( - "distributedTraceIdGenerator", - () -> - buildDistributedTraceIdGenerator( - programmaticArguments.getDistributedTraceIdGenerator()), + "requestIdGenerator", + () -> buildRequestIdGenerator(programmaticArguments.getRequestIdGenerator()), cycleDetector); this.sslEngineFactoryRef = new LazyReference<>( @@ -718,21 +716,20 @@ protected RequestTracker buildRequestTracker(RequestTracker requestTrackerFromBu } } - protected DistributedTraceIdGenerator buildDistributedTraceIdGenerator( - DistributedTraceIdGenerator distributedTraceIdGenerator) { - return (distributedTraceIdGenerator != null) - ? distributedTraceIdGenerator + protected RequestIdGenerator buildRequestIdGenerator(RequestIdGenerator requestIdGenerator) { + return (requestIdGenerator != null) + ? requestIdGenerator : Reflection.buildFromConfig( this, - DefaultDriverOption.DISTRIBUTED_TRACE_ID_GENERATOR_CLASS, - DistributedTraceIdGenerator.class, + DefaultDriverOption.REQUEST_ID_GENERATOR_CLASS, + RequestIdGenerator.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))); + "Missing request ID generator, check your configuration (%s)", + DefaultDriverOption.REQUEST_ID_GENERATOR_CLASS))); } protected Optional buildAuthProvider(AuthProvider authProviderFromBuilder) { @@ -1001,8 +998,8 @@ public RequestTracker getRequestTracker() { @NonNull @Override - public DistributedTraceIdGenerator getDistributedTraceIdGenerator() { - return distributedTraceIdGeneratorRef.get(); + public RequestIdGenerator getRequestIdGenerator() { + return requestIdGeneratorRef.get(); } @Nullable diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java index 23047c4ac68..90b0fc8d6f4 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java @@ -44,7 +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.RequestIdGenerator; 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; @@ -128,7 +128,7 @@ public class CqlRequestHandler implements Throttled { private final List inFlightCallbacks; private final RequestThrottler throttler; private final RequestTracker requestTracker; - private final DistributedTraceIdGenerator distributedTraceIdGenerator; + private final RequestIdGenerator requestIdGenerator; private final SessionMetricUpdater sessionMetricUpdater; private final DriverExecutionProfile executionProfile; @@ -144,10 +144,9 @@ protected CqlRequestHandler( String sessionLogPrefix) { this.startTimeNanos = System.nanoTime(); - this.distributedTraceIdGenerator = context.getDistributedTraceIdGenerator(); + this.requestIdGenerator = context.getRequestIdGenerator(); this.logPrefix = - this.distributedTraceIdGenerator.getSessionRequestId( - statement, sessionLogPrefix, this.hashCode()); + this.requestIdGenerator.getSessionRequestId(statement, sessionLogPrefix, this.hashCode()); LOG.trace("[{}] Creating new handler for request {}", logPrefix, statement); this.initialStatement = statement; @@ -180,8 +179,7 @@ protected CqlRequestHandler( this.executionProfile = Conversions.resolveExecutionProfile(initialStatement, context); this.customPayloadKey = - this.executionProfile.getString( - DefaultDriverOption.DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY); + this.executionProfile.getString(DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY); Duration timeout = Conversions.resolveRequestTimeout(statement, executionProfile); this.scheduledTimeout = scheduleTimeout(timeout); @@ -262,17 +260,17 @@ private void sendRequest( return; } String nodeRequestId = - this.distributedTraceIdGenerator.getNodeRequestId( - statement, logPrefix, currentExecutionIndex); - Map customPayload = statement.getCustomPayload(); + this.requestIdGenerator.getNodeRequestId(statement, logPrefix, currentExecutionIndex); if (!this.customPayloadKey.isEmpty()) { - customPayload = + Map customPayload = NullAllowingImmutableMap.builder() - .putAll(customPayload) + .putAll(statement.getCustomPayload()) .put( this.customPayloadKey, ByteBuffer.wrap(nodeRequestId.getBytes(StandardCharsets.UTF_8))) .build(); + // TODO: we are creating a new statement object for every request. We should optimize this. + statement = statement.setCustomPayload(customPayload); } Node node = retriedNode; DriverChannel channel = null; @@ -305,7 +303,7 @@ private void sendRequest( nodeRequestId); Message message = Conversions.toMessage(statement, executionProfile, context); channel - .write(message, statement.isTracing(), customPayload, nodeResponseCallback) + .write(message, statement.isTracing(), statement.getCustomPayload(), nodeResponseCallback) .addListener(nodeResponseCallback); } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/DefaultDistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/DefaultRequestIdGenerator.java similarity index 85% rename from core/src/main/java/com/datastax/oss/driver/internal/core/tracker/DefaultDistributedTraceIdGenerator.java rename to core/src/main/java/com/datastax/oss/driver/internal/core/tracker/DefaultRequestIdGenerator.java index 7142d0502aa..64e5eddacd9 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/DefaultDistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/DefaultRequestIdGenerator.java @@ -19,12 +19,12 @@ import com.datastax.oss.driver.api.core.context.DriverContext; import com.datastax.oss.driver.api.core.session.Request; -import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator; +import com.datastax.oss.driver.api.core.tracker.RequestIdGenerator; import edu.umd.cs.findbugs.annotations.NonNull; -public class DefaultDistributedTraceIdGenerator implements DistributedTraceIdGenerator { +public class DefaultRequestIdGenerator implements RequestIdGenerator { - public DefaultDistributedTraceIdGenerator(DriverContext context) {} + public DefaultRequestIdGenerator(DriverContext context) {} @Override public String getSessionRequestId( diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java similarity index 86% rename from core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java rename to core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java index d05cf799b98..0a8b8a7004b 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidDistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java @@ -19,12 +19,12 @@ import com.datastax.oss.driver.api.core.context.DriverContext; import com.datastax.oss.driver.api.core.session.Request; -import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator; +import com.datastax.oss.driver.api.core.tracker.RequestIdGenerator; import com.datastax.oss.driver.api.core.uuid.Uuids; import edu.umd.cs.findbugs.annotations.NonNull; -public class UuidDistributedTraceIdGenerator implements DistributedTraceIdGenerator { - public UuidDistributedTraceIdGenerator(DriverContext context) {} +public class UuidRequestIdGenerator implements RequestIdGenerator { + public UuidRequestIdGenerator(DriverContext context) {} @Override public String getSessionRequestId( diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextDistributedTraceIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java similarity index 88% rename from core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextDistributedTraceIdGenerator.java rename to core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java index c186577ade7..c16afdc398b 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextDistributedTraceIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java @@ -19,17 +19,17 @@ import com.datastax.oss.driver.api.core.context.DriverContext; import com.datastax.oss.driver.api.core.session.Request; -import com.datastax.oss.driver.api.core.tracker.DistributedTraceIdGenerator; +import com.datastax.oss.driver.api.core.tracker.RequestIdGenerator; import com.datastax.oss.driver.shaded.guava.common.io.BaseEncoding; import edu.umd.cs.findbugs.annotations.NonNull; import java.security.SecureRandom; import java.util.Random; -public class W3CContextDistributedTraceIdGenerator implements DistributedTraceIdGenerator { +public class W3CContextRequestIdGenerator implements RequestIdGenerator { Random random = new SecureRandom(); BaseEncoding baseEncoding = BaseEncoding.base16().lowerCase(); - public W3CContextDistributedTraceIdGenerator(DriverContext context) {} + public W3CContextRequestIdGenerator(DriverContext context) {} @Override public String getSessionRequestId( diff --git a/core/src/main/resources/reference.conf b/core/src/main/resources/reference.conf index ad292d4d2eb..be7d27e21ac 100644 --- a/core/src/main/resources/reference.conf +++ b/core/src/main/resources/reference.conf @@ -918,13 +918,13 @@ datastax-java-driver { } } - advanced.distributed-tracing{ - id-generator{ - # The component that generates unique identifiers for distributed tracing. - class = DefaultDistributedTraceIdGenerator + advanced.request-id{ + generator{ + # The component that generates a unique identifier for each CQL request. + class = DefaultRequestIdGenerator } - # add the trace-id to the custom payload with the given key - # if empty, the trace-id will not be added to the custom payload + # add the request id to the custom payload with the given key + # if empty, the request id will not be added to the custom payload custom-payload-with-key = "" } # A session-wide component that controls the rate at which requests are executed. diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java index bbee1e4e3d6..676c9ac56c5 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java @@ -37,7 +37,7 @@ import com.datastax.oss.driver.api.core.session.Session; import com.datastax.oss.driver.api.core.specex.SpeculativeExecutionPolicy; 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.RequestIdGenerator; import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry; import com.datastax.oss.driver.internal.core.DefaultConsistencyLevelRegistry; import com.datastax.oss.driver.internal.core.ProtocolFeature; @@ -52,7 +52,7 @@ import com.datastax.oss.driver.internal.core.servererrors.DefaultWriteTypeRegistry; import com.datastax.oss.driver.internal.core.session.DefaultSession; import com.datastax.oss.driver.internal.core.session.throttling.PassThroughRequestThrottler; -import com.datastax.oss.driver.internal.core.tracker.DefaultDistributedTraceIdGenerator; +import com.datastax.oss.driver.internal.core.tracker.DefaultRequestIdGenerator; import com.datastax.oss.driver.internal.core.tracker.NoopRequestTracker; import com.datastax.oss.driver.internal.core.util.concurrent.CapturingTimer; import com.datastax.oss.driver.internal.core.util.concurrent.CapturingTimer.CapturedTimeout; @@ -116,7 +116,7 @@ protected RequestHandlerTestHarness(Builder builder) { when(defaultProfile.getBoolean(DefaultDriverOption.REQUEST_DEFAULT_IDEMPOTENCE)) .thenReturn(builder.defaultIdempotence); when(defaultProfile.getBoolean(DefaultDriverOption.PREPARE_ON_ALL_NODES)).thenReturn(true); - when(defaultProfile.getString(DefaultDriverOption.DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY)) + when(defaultProfile.getString(DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY)) .thenReturn(""); when(config.getDefaultProfile()).thenReturn(defaultProfile); @@ -129,9 +129,8 @@ protected RequestHandlerTestHarness(Builder builder) { when(context.getRetryPolicy(anyString())).thenReturn(retryPolicy); - DistributedTraceIdGenerator distributedTraceIdGenerator = - new DefaultDistributedTraceIdGenerator(context); - when(context.getDistributedTraceIdGenerator()).thenReturn(distributedTraceIdGenerator); + RequestIdGenerator requestIdGenerator = new DefaultRequestIdGenerator(context); + when(context.getRequestIdGenerator()).thenReturn(requestIdGenerator); // Disable speculative executions by default when(speculativeExecutionPolicy.nextExecution( diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/DistributedTraceIdGeneratorTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java similarity index 87% rename from core/src/test/java/com/datastax/oss/driver/internal/core/tracker/DistributedTraceIdGeneratorTest.java rename to core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java index bec1aa739a6..ca761f025bd 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/DistributedTraceIdGeneratorTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java @@ -27,14 +27,14 @@ import org.mockito.junit.MockitoJUnitRunner; @RunWith(MockitoJUnitRunner.Strict.class) -public class DistributedTraceIdGeneratorTest { +public class RequestIdGeneratorTest { @Mock private InternalDriverContext context; @Mock private Statement statement; @Test public void default_generator_should_generate() { // given - DefaultDistributedTraceIdGenerator generator = new DefaultDistributedTraceIdGenerator(context); + DefaultRequestIdGenerator generator = new DefaultRequestIdGenerator(context); // when String sessionRequestId = generator.getSessionRequestId(statement, "sessionName", 123); String nodeRequestId = generator.getNodeRequestId(statement, sessionRequestId, 1); @@ -46,7 +46,7 @@ public void default_generator_should_generate() { @Test public void uuid_generator_should_generate() { // given - UuidDistributedTraceIdGenerator generator = new UuidDistributedTraceIdGenerator(context); + UuidRequestIdGenerator generator = new UuidRequestIdGenerator(context); // when String sessionRequestId = generator.getSessionRequestId(statement, "sessionName", 123); String nodeRequestId = generator.getNodeRequestId(statement, sessionRequestId, 1); @@ -58,8 +58,7 @@ public void uuid_generator_should_generate() { @Test public void w3c_generator_should_generate() { // given - W3CContextDistributedTraceIdGenerator generator = - new W3CContextDistributedTraceIdGenerator(context); + W3CContextRequestIdGenerator generator = new W3CContextRequestIdGenerator(context); // when String sessionRequestId = generator.getSessionRequestId(statement, "sessionName", 123); String nodeRequestId = generator.getNodeRequestId(statement, sessionRequestId, 1); diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/DistributedTraceIdGeneratorIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java similarity index 85% rename from integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/DistributedTraceIdGeneratorIT.java rename to integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java index 49ed63a32b2..3b4b0871335 100644 --- a/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/DistributedTraceIdGeneratorIT.java +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java @@ -34,7 +34,7 @@ import org.junit.rules.TestRule; @Category(ParallelizableTests.class) -public class DistributedTraceIdGeneratorIT { +public class RequestIdGeneratorIT { private CcmRule ccmRule = CcmRule.getInstance(); @Rule public TestRule chain = RuleChain.outerRule(ccmRule); @@ -43,7 +43,7 @@ public class DistributedTraceIdGeneratorIT { public void should_write_default_id_to_custom_payload_with_key() { DriverConfigLoader loader = SessionUtils.configLoaderBuilder() - .withString(DefaultDriverOption.DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY, "trace_key") + .withString(DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY, "trace_key") .build(); try (CqlSession session = SessionUtils.newSession(ccmRule, loader)) { String query = "SELECT * FROM system.local"; @@ -57,10 +57,8 @@ public void should_write_default_id_to_custom_payload_with_key() { public void should_write_uuid_to_custom_payload_with_key() { DriverConfigLoader loader = SessionUtils.configLoaderBuilder() - .withString( - DefaultDriverOption.DISTRIBUTED_TRACE_ID_GENERATOR_CLASS, - "UuidDistributedTraceIdGenerator") - .withString(DefaultDriverOption.DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY, "trace_key") + .withString(DefaultDriverOption.REQUEST_ID_GENERATOR_CLASS, "UuidRequestIdGenerator") + .withString(DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY, "trace_key") .build(); try (CqlSession session = SessionUtils.newSession(ccmRule, loader)) { String query = "SELECT * FROM system.local"; @@ -75,9 +73,8 @@ public void should_write_w3c_context_to_custom_payload_with_key() { DriverConfigLoader loader = SessionUtils.configLoaderBuilder() .withString( - DefaultDriverOption.DISTRIBUTED_TRACE_ID_GENERATOR_CLASS, - "W3CContextDistributedTraceIdGenerator") - .withString(DefaultDriverOption.DISTRIBUTED_TRACE_ID_CUSTOM_PAYLOAD_KEY, "trace_key") + DefaultDriverOption.REQUEST_ID_GENERATOR_CLASS, "W3CContextRequestIdGenerator") + .withString(DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY, "trace_key") .build(); try (CqlSession session = SessionUtils.newSession(ccmRule, loader)) { String query = "SELECT * FROM system.local"; diff --git a/manual/core/distributed_tracing/README.md b/manual/core/request_id/README.md similarity index 54% rename from manual/core/distributed_tracing/README.md rename to manual/core/request_id/README.md index e67bb538ea6..dcebf5fe0a1 100644 --- a/manual/core/distributed_tracing/README.md +++ b/manual/core/request_id/README.md @@ -17,48 +17,53 @@ specific language governing permissions and limitations under the License. --> -## Distributed tracing +## Request Id ### Quick overview Users can inject an identifier for each individual CQL request, and such ID can be written in to the custom payload to correlate a request across the driver and the Apache Cassandra server. -* Inject ID generator: set the desired `DistributedTraceIdGenerator` in `advanced.distributed-tracing.id-generator.class`. - The default implementation generates ID as `{session_name}|{hash_code}|{execution_count}`. -* Add ID to custom payload: disabled by default. Set the desired key in `advanced.distributed-tracing.custom-payload-with-key`, +A request ID generator needs to generate both: +- Session request ID: an identifier for an entire session.execute() call +- Node request ID: an identifier for the execution of a CQL statement against a particular node. There can be one or more node requests for a single session request, due to retries or speculative executions. + +Usage: +* Inject ID generator: set the desired `RequestIdGenerator` in `advanced.request-id.generator.class`. + The default implementation generates the session request ID as `{session_name}|{hash_code}`, and node request ID as `{session_name}|{hash_code}|{execution_count}`. +* Add ID to custom payload: disabled by default. Set the desired key in `advanced.request-id.custom-payload-with-key`, then the driver will add the generated ID to the custom payload with the specified key. -### Distributed Trace Id Generator Configuration +### Request Id Generator Configuration -Distributed trace ID generator can be declared in the [configuration](../configuration/) as follows: +Request ID generator can be declared in the [configuration](../configuration/) as follows: ``` -datastax-java-driver.advanced.distributed-tracing.id-generator { +datastax-java-driver.advanced.request-id.generator { class = com.example.app.MyGenerator } ``` -To register your own trackers, specify the name of a class -that implements `DistributedTraceIdGenerator`. +To register your own request ID generator, specify the name of the class +that implements `RequestIdGenerator`. -By default, the build-in implementation `DefaultDistributedTraceIdGenerator` is used. It generates the ID as +By default, the build-in implementation `DefaultRequestIdGenerator` is used. It generates the ID as `{session_name}|{hash_code}|{execution_count}`. Note that this ID is not guaranteed to be unique. -Other built-in implementations include `UUIDDistributedTraceIdGenerator` and `W3CContextDistributedTraceIdGenerator`. +Other built-in implementations include `UUIDRequestIdGenerator` and `W3CContextRequestIdGenerator`. The generated ID will be added to the log message of `CqlRequestHandler`, and propagated to the request trackers. ### Custom Payload Configuration -Users can opt in to add the generated ID to the custom payload to achieve request end-to-end tracing. +Users can opt in to add the generated node request ID to the custom payload to achieve request end-to-end tracing. Custom payload is a map of string to `ByteBuffer` pairs, and the driver will add the generated ID to the custom payload with the specified key. ``` -datastax-java-driver.advanced.distributed-tracing{ - custom-payload-with-key = my-trace-key +datastax-java-driver.advanced.request-id{ + custom-payload-with-key = my-request-id } ``` -Users can then run Apache Cassandra with customized query handler to extract the trace ID from the custom payload, to achieve end-to-end tracing. +Users can then run Apache Cassandra with customized query handler to extract the request ID from the custom payload, to achieve end-to-end tracing. When this key is set to an empty string (default), the driver will not add the ID to the custom payload. From 2f40b1eaed48f6f48d8f2bdc4f806feee7f06685 Mon Sep 17 00:00:00 2001 From: janehe Date: Mon, 19 May 2025 17:45:09 -0700 Subject: [PATCH 09/16] Address PR review from Yuqi DU --- .../core/tracker/UuidRequestIdGenerator.java | 5 +++ .../tracker/W3CContextRequestIdGenerator.java | 7 ++++ .../core/tracker/RequestIdGeneratorTest.java | 7 ++++ .../core/tracker/RequestIdGeneratorIT.java | 38 +++++++++++++++++++ 4 files changed, 57 insertions(+) diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java index 0a8b8a7004b..bbda42ab8da 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java @@ -26,12 +26,17 @@ public class UuidRequestIdGenerator implements RequestIdGenerator { public UuidRequestIdGenerator(DriverContext context) {} + /** Generates a random v4 UUID. */ @Override public String getSessionRequestId( @NonNull Request statement, @NonNull String sessionName, int hashCode) { return Uuids.random().toString(); } + /** + * {session-request-id}-{random-uuid} All node requests for a session request will have the same + * session request id + */ @Override public String getNodeRequestId( @NonNull Request statement, @NonNull String sessionRequestId, int executionCount) { diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java index c16afdc398b..bde9fb4af55 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java @@ -31,6 +31,7 @@ public class W3CContextRequestIdGenerator implements RequestIdGenerator { public W3CContextRequestIdGenerator(DriverContext context) {} + /** Random 16 bytes, e.g. "4bf92f3577b34da6a3ce929d0e0e4736" */ @Override public String getSessionRequestId( @NonNull Request statement, @NonNull String sessionName, int hashCode) { @@ -39,6 +40,12 @@ public String getSessionRequestId( return baseEncoding.encode(bytes); } + /** + * Following the format of W3C "traceparent" spec, + * https://www.w3.org/TR/trace-context/#traceparent-header-field-values e.g. + * "00-4bf92f3577b34da6a3ce929d0e0e4736-a3ce929d0e0e4736-01" All node requests in the same session + * request share the same "trace-id" field value + */ @Override public String getNodeRequestId( @NonNull Request statement, @NonNull String sessionRequestId, int executionCount) { diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java index ca761f025bd..3322d260788 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java @@ -51,7 +51,10 @@ public void uuid_generator_should_generate() { String sessionRequestId = generator.getSessionRequestId(statement, "sessionName", 123); String nodeRequestId = generator.getNodeRequestId(statement, sessionRequestId, 1); // then + // e.g. "550e8400-e29b-41d4-a716-446655440000", which is 36 characters long assertThat(sessionRequestId.length()).isEqualTo(36); + // e.g. "550e8400-e29b-41d4-a716-446655440000-550e8400-e29b-41d4-a716-446655440000", which is 73 + // characters long assertThat(nodeRequestId.length()).isEqualTo(73); } @@ -63,7 +66,11 @@ public void w3c_generator_should_generate() { String sessionRequestId = generator.getSessionRequestId(statement, "sessionName", 123); String nodeRequestId = generator.getNodeRequestId(statement, sessionRequestId, 1); // then + // e.g. "4bf92f3577b34da6a3ce929d0e0e4736", which is 32 characters long assertThat(sessionRequestId.length()).isEqualTo(32); + // According to W3C "traceparent" spec, + // https://www.w3.org/TR/trace-context/#traceparent-header-field-values + // e.g. "00-4bf92f3577b34da6a3ce929d0e0e4736-a3ce929d0e0e4736-01", which 55 characters long assertThat(nodeRequestId.length()).isEqualTo(55); } } diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java index 3b4b0871335..c9ecb807a7a 100644 --- a/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java @@ -23,10 +23,14 @@ import com.datastax.oss.driver.api.core.config.DefaultDriverOption; import com.datastax.oss.driver.api.core.config.DriverConfigLoader; import com.datastax.oss.driver.api.core.cql.ResultSet; +import com.datastax.oss.driver.api.core.session.Request; +import com.datastax.oss.driver.api.core.tracker.RequestIdGenerator; import com.datastax.oss.driver.api.testinfra.ccm.CcmRule; import com.datastax.oss.driver.api.testinfra.session.SessionUtils; import com.datastax.oss.driver.categories.ParallelizableTests; +import edu.umd.cs.findbugs.annotations.NonNull; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; import org.junit.Rule; import org.junit.Test; import org.junit.experimental.categories.Category; @@ -84,6 +88,40 @@ public void should_write_w3c_context_to_custom_payload_with_key() { } } + @Test + public void should_use_customized_request_id_generator() { + DriverConfigLoader loader = + SessionUtils.configLoaderBuilder() + .withString(DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY, "trace_key") + .build(); + RequestIdGenerator myRequestIdGenerator = + new RequestIdGenerator() { + @Override + public String getSessionRequestId( + @NonNull Request statement, @NonNull String sessionName, int hashCode) { + return "123"; + } + + @Override + public String getNodeRequestId( + @NonNull Request statement, @NonNull String sessionRequestId, int executionCount) { + return "456"; + } + }; + try (CqlSession session = + (CqlSession) + SessionUtils.baseBuilder() + .addContactEndPoints(ccmRule.getContactPoints()) + .withRequestIdGenerator(myRequestIdGenerator) + .withConfigLoader(loader) + .build()) { + String query = "SELECT * FROM system.local"; + ResultSet rs = session.execute(query); + ByteBuffer id = rs.getExecutionInfo().getRequest().getCustomPayload().get("trace_key"); + assertThat(id).isEqualTo(ByteBuffer.wrap("456".getBytes(StandardCharsets.UTF_8))); + } + } + @Test public void should_not_write_id_to_custom_payload_when_key_is_not_set() { DriverConfigLoader loader = SessionUtils.configLoaderBuilder().build(); From 457d5825c2d3de68d5650bc1dee6cf6245a291be Mon Sep 17 00:00:00 2001 From: janehe Date: Thu, 22 May 2025 00:43:47 -0700 Subject: [PATCH 10/16] add doc --- core/src/main/resources/reference.conf | 4 ++-- manual/core/request_id/README.md | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/core/src/main/resources/reference.conf b/core/src/main/resources/reference.conf index be7d27e21ac..31a97a6798e 100644 --- a/core/src/main/resources/reference.conf +++ b/core/src/main/resources/reference.conf @@ -918,8 +918,8 @@ datastax-java-driver { } } - advanced.request-id{ - generator{ + advanced.request-id { + generator { # The component that generates a unique identifier for each CQL request. class = DefaultRequestIdGenerator } diff --git a/manual/core/request_id/README.md b/manual/core/request_id/README.md index dcebf5fe0a1..fdee6a11563 100644 --- a/manual/core/request_id/README.md +++ b/manual/core/request_id/README.md @@ -21,7 +21,7 @@ under the License. ### Quick overview -Users can inject an identifier for each individual CQL request, and such ID can be written in to the custom payload to +Users can inject an identifier for each individual CQL request, and such ID can be written in to the [custom payload](https://github.com/apache/cassandra/blob/trunk/doc/native_protocol_v5.spec) to correlate a request across the driver and the Apache Cassandra server. A request ID generator needs to generate both: @@ -30,7 +30,9 @@ A request ID generator needs to generate both: Usage: * Inject ID generator: set the desired `RequestIdGenerator` in `advanced.request-id.generator.class`. - The default implementation generates the session request ID as `{session_name}|{hash_code}`, and node request ID as `{session_name}|{hash_code}|{execution_count}`. + The default implementation generates the session request ID as `{session_name}|{hash_code}`, and node request ID as `{session_name}|{hash_code}|{execution_count}`, + where "hash_code" is the hash code of the `CqlRequestHandler` object, and "execution_count" is the zero-based index of the node request in the session request. + For example, if there is a retry or speculative execution right after the first node request, the second node request will have the ID `{session_name}|{hash_code}|1`. * Add ID to custom payload: disabled by default. Set the desired key in `advanced.request-id.custom-payload-with-key`, then the driver will add the generated ID to the custom payload with the specified key. From 05159d8a2e16af0f380bb85a96d22cadebed359a Mon Sep 17 00:00:00 2001 From: janehe Date: Thu, 29 May 2025 00:23:43 -0700 Subject: [PATCH 11/16] remove default request id generator --- .../driver/api/core/config/OptionsMap.java | 1 - .../api/core/context/DriverContext.java | 2 +- .../api/core/tracker/RequestIdGenerator.java | 19 +++++++++ .../internal/core/cql/CqlRequestHandler.java | 30 ++------------ .../tracker/DefaultRequestIdGenerator.java | 40 ------------------- .../core/tracker/UuidRequestIdGenerator.java | 2 +- core/src/main/resources/reference.conf | 8 ++-- .../core/cql/RequestHandlerTestHarness.java | 5 --- .../core/tracker/RequestIdGeneratorTest.java | 12 ------ 9 files changed, 27 insertions(+), 92 deletions(-) delete mode 100644 core/src/main/java/com/datastax/oss/driver/internal/core/tracker/DefaultRequestIdGenerator.java diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java index 27df2d4f025..5135524ce99 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java @@ -291,7 +291,6 @@ 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.REQUEST_ID_GENERATOR_CLASS, "DefaultRequestIdGenerator"); map.put(TypedDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY, ""); map.put(TypedDriverOption.GRAPH_PAGING_ENABLED, "AUTO"); map.put(TypedDriverOption.GRAPH_CONTINUOUS_PAGING_PAGE_SIZE, requestPageSize); diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/context/DriverContext.java b/core/src/main/java/com/datastax/oss/driver/api/core/context/DriverContext.java index 8d19ec6a149..6f0afd3df8a 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/context/DriverContext.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/context/DriverContext.java @@ -142,7 +142,7 @@ default SpeculativeExecutionPolicy getSpeculativeExecutionPolicy(@NonNull String /** @return The driver's request ID generator; never {@code null}. */ @NonNull - RequestIdGenerator getRequestIdGenerator(); + Optional getRequestIdGenerator(); /** @return The driver's request throttler; never {@code null}. */ @NonNull diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java index e620cc2bbad..cb8a4196a62 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java @@ -17,9 +17,15 @@ */ 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 @@ -47,4 +53,17 @@ public interface RequestIdGenerator { */ String getNodeRequestId( @NonNull Request statement, @NonNull String sessionRequestId, int executionCount); + + + default Statement getDecoratedStatement( + @NonNull Statement statement, @NonNull String nodeRequestId) { + Map customPayload = + NullAllowingImmutableMap.builder() + .putAll(statement.getCustomPayload()) + .put( + "request-id", + ByteBuffer.wrap(nodeRequestId.getBytes(StandardCharsets.UTF_8))) + .build(); + return statement.setCustomPayload(customPayload); + } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java index 90b0fc8d6f4..0808bdce63f 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java @@ -44,7 +44,6 @@ 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.RequestIdGenerator; 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; @@ -72,7 +71,6 @@ import com.datastax.oss.protocol.internal.response.result.SetKeyspace; import com.datastax.oss.protocol.internal.response.result.Void; import com.datastax.oss.protocol.internal.util.Bytes; -import com.datastax.oss.protocol.internal.util.collection.NullAllowingImmutableMap; import edu.umd.cs.findbugs.annotations.NonNull; import io.netty.handler.codec.EncoderException; import io.netty.util.Timeout; @@ -80,7 +78,6 @@ 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.List; @@ -128,14 +125,12 @@ public class CqlRequestHandler implements Throttled { private final List inFlightCallbacks; private final RequestThrottler throttler; private final RequestTracker requestTracker; - private final RequestIdGenerator requestIdGenerator; 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> errors; - private final String customPayloadKey; protected CqlRequestHandler( Statement statement, @@ -144,9 +139,7 @@ protected CqlRequestHandler( String sessionLogPrefix) { this.startTimeNanos = System.nanoTime(); - this.requestIdGenerator = context.getRequestIdGenerator(); - this.logPrefix = - this.requestIdGenerator.getSessionRequestId(statement, sessionLogPrefix, this.hashCode()); + this.logPrefix = sessionLogPrefix + "|" + this.hashCode(); LOG.trace("[{}] Creating new handler for request {}", logPrefix, statement); this.initialStatement = statement; @@ -177,10 +170,6 @@ protected CqlRequestHandler( this.timer = context.getNettyOptions().getTimer(); this.executionProfile = Conversions.resolveExecutionProfile(initialStatement, context); - - this.customPayloadKey = - this.executionProfile.getString(DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY); - Duration timeout = Conversions.resolveRequestTimeout(statement, executionProfile); this.scheduledTimeout = scheduleTimeout(timeout); @@ -259,19 +248,6 @@ private void sendRequest( if (result.isDone()) { return; } - String nodeRequestId = - this.requestIdGenerator.getNodeRequestId(statement, logPrefix, currentExecutionIndex); - if (!this.customPayloadKey.isEmpty()) { - Map customPayload = - NullAllowingImmutableMap.builder() - .putAll(statement.getCustomPayload()) - .put( - this.customPayloadKey, - ByteBuffer.wrap(nodeRequestId.getBytes(StandardCharsets.UTF_8))) - .build(); - // TODO: we are creating a new statement object for every request. We should optimize this. - statement = statement.setCustomPayload(customPayload); - } Node node = retriedNode; DriverChannel channel = null; if (node == null || (channel = session.getChannel(node, logPrefix)) == null) { @@ -300,7 +276,7 @@ private void sendRequest( currentExecutionIndex, retryCount, scheduleNextExecution, - nodeRequestId); + logPrefix); Message message = Conversions.toMessage(statement, executionProfile, context); channel .write(message, statement.isTracing(), statement.getCustomPayload(), nodeResponseCallback) @@ -513,7 +489,7 @@ private NodeResponseCallback( this.execution = execution; this.retryCount = retryCount; this.scheduleNextExecution = scheduleNextExecution; - this.logPrefix = logPrefix; + this.logPrefix = logPrefix + "|" + execution; } // this gets invoked once the write completes. diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/DefaultRequestIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/DefaultRequestIdGenerator.java deleted file mode 100644 index 64e5eddacd9..00000000000 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/DefaultRequestIdGenerator.java +++ /dev/null @@ -1,40 +0,0 @@ -/* - * 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.internal.core.tracker; - -import com.datastax.oss.driver.api.core.context.DriverContext; -import com.datastax.oss.driver.api.core.session.Request; -import com.datastax.oss.driver.api.core.tracker.RequestIdGenerator; -import edu.umd.cs.findbugs.annotations.NonNull; - -public class DefaultRequestIdGenerator implements RequestIdGenerator { - - public DefaultRequestIdGenerator(DriverContext context) {} - - @Override - public String getSessionRequestId( - @NonNull Request statement, @NonNull String sessionName, int hashCode) { - return sessionName + "|" + hashCode; - } - - @Override - public String getNodeRequestId( - @NonNull Request statement, @NonNull String sessionRequestId, int executionCount) { - return sessionRequestId + "|" + executionCount; - } -} diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java index bbda42ab8da..eaf7f5b80d6 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java @@ -29,7 +29,7 @@ public UuidRequestIdGenerator(DriverContext context) {} /** Generates a random v4 UUID. */ @Override public String getSessionRequestId( - @NonNull Request statement, @NonNull String sessionName, int hashCode) { + @NonNull Request statement) { return Uuids.random().toString(); } diff --git a/core/src/main/resources/reference.conf b/core/src/main/resources/reference.conf index 31a97a6798e..24dff6a052a 100644 --- a/core/src/main/resources/reference.conf +++ b/core/src/main/resources/reference.conf @@ -920,13 +920,11 @@ datastax-java-driver { advanced.request-id { generator { - # The component that generates a unique identifier for each CQL request. - class = DefaultRequestIdGenerator + # The component that generates a unique identifier for each CQL request, and possibly write the id to the custom payload . + // class = W3CContextRequestIdGenerator } - # add the request id to the custom payload with the given key - # if empty, the request id will not be added to the custom payload - custom-payload-with-key = "" } + # A session-wide component that controls the rate at which requests are executed. # # Implementations vary, but throttlers generally track a metric that represents the level of diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java index 676c9ac56c5..a6fe45458ff 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java @@ -37,7 +37,6 @@ import com.datastax.oss.driver.api.core.session.Session; import com.datastax.oss.driver.api.core.specex.SpeculativeExecutionPolicy; 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.type.codec.registry.CodecRegistry; import com.datastax.oss.driver.internal.core.DefaultConsistencyLevelRegistry; import com.datastax.oss.driver.internal.core.ProtocolFeature; @@ -52,7 +51,6 @@ import com.datastax.oss.driver.internal.core.servererrors.DefaultWriteTypeRegistry; import com.datastax.oss.driver.internal.core.session.DefaultSession; import com.datastax.oss.driver.internal.core.session.throttling.PassThroughRequestThrottler; -import com.datastax.oss.driver.internal.core.tracker.DefaultRequestIdGenerator; import com.datastax.oss.driver.internal.core.tracker.NoopRequestTracker; import com.datastax.oss.driver.internal.core.util.concurrent.CapturingTimer; import com.datastax.oss.driver.internal.core.util.concurrent.CapturingTimer.CapturedTimeout; @@ -129,9 +127,6 @@ protected RequestHandlerTestHarness(Builder builder) { when(context.getRetryPolicy(anyString())).thenReturn(retryPolicy); - RequestIdGenerator requestIdGenerator = new DefaultRequestIdGenerator(context); - when(context.getRequestIdGenerator()).thenReturn(requestIdGenerator); - // Disable speculative executions by default when(speculativeExecutionPolicy.nextExecution( any(Node.class), any(CqlIdentifier.class), any(Request.class), anyInt())) diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java index 3322d260788..1fad60ac154 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java @@ -31,18 +31,6 @@ public class RequestIdGeneratorTest { @Mock private InternalDriverContext context; @Mock private Statement statement; - @Test - public void default_generator_should_generate() { - // given - DefaultRequestIdGenerator generator = new DefaultRequestIdGenerator(context); - // when - String sessionRequestId = generator.getSessionRequestId(statement, "sessionName", 123); - String nodeRequestId = generator.getNodeRequestId(statement, sessionRequestId, 1); - // then - assertThat(sessionRequestId).isEqualTo("sessionName|123"); - assertThat(nodeRequestId).isEqualTo("sessionName|123|1"); - } - @Test public void uuid_generator_should_generate() { // given From 3505d7a41eaf282ad032e6b74066bc3ca3643553 Mon Sep 17 00:00:00 2001 From: janehe Date: Thu, 29 May 2025 01:14:42 -0700 Subject: [PATCH 12/16] Optional of request id generator --- .../api/core/config/DefaultDriverOption.java | 9 +--- .../driver/api/core/config/OptionsMap.java | 1 - .../api/core/config/TypedDriverOption.java | 5 --- .../api/core/tracker/RequestIdGenerator.java | 20 +++------ .../core/context/DefaultDriverContext.java | 23 ++++------ .../internal/core/cql/CqlRequestHandler.java | 23 ++++++++-- .../core/tracker/UuidRequestIdGenerator.java | 3 +- .../tracker/W3CContextRequestIdGenerator.java | 3 +- .../core/cql/RequestHandlerTestHarness.java | 2 - .../core/tracker/RequestIdGeneratorTest.java | 4 +- .../core/tracker/RequestIdGeneratorIT.java | 44 ++++++++----------- 11 files changed, 60 insertions(+), 77 deletions(-) diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java index 5e2f1252a95..06cfe286819 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java @@ -1001,14 +1001,7 @@ public enum DefaultDriverOption implements DriverOption { * *

Value-type: {@link String} */ - REQUEST_ID_GENERATOR_CLASS("advanced.request-id.generator.class"), - - /** - * If not empty, the driver will write the node request ID to this key in the custom payload - * - *

Value-type: {@link String} - */ - REQUEST_ID_CUSTOM_PAYLOAD_KEY("advanced.request-id.custom-payload-with-key"); + REQUEST_ID_GENERATOR_CLASS("advanced.request-id.generator.class"); private final String path; diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java index 5135524ce99..98faf3e590c 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java @@ -291,7 +291,6 @@ 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.REQUEST_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); diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java index 76c4aea3864..997f0a09272 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java @@ -285,11 +285,6 @@ public String toString() { public static final TypedDriverOption REQUEST_ID_GENERATOR_CLASS = new TypedDriverOption<>(DefaultDriverOption.REQUEST_ID_GENERATOR_CLASS, GenericType.STRING); - /** If not empty, the driver will write the node request ID to this key in the custom payload */ - public static final TypedDriverOption REQUEST_ID_CUSTOM_PAYLOAD_KEY = - new TypedDriverOption<>( - DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY, GenericType.STRING); - /** Whether to log successful requests. */ public static final TypedDriverOption REQUEST_LOGGER_SUCCESS_ENABLED = new TypedDriverOption<>( diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java index cb8a4196a62..58e567f4246 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java @@ -21,7 +21,6 @@ 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; @@ -33,11 +32,9 @@ public interface RequestIdGenerator { * 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); + String getSessionRequestId(@NonNull Request statement); /** * Generates a unique identifier for the node request. This will be the identifier for the CQL @@ -54,16 +51,13 @@ public interface RequestIdGenerator { String getNodeRequestId( @NonNull Request statement, @NonNull String sessionRequestId, int executionCount); - default Statement getDecoratedStatement( @NonNull Statement statement, @NonNull String nodeRequestId) { - Map customPayload = - NullAllowingImmutableMap.builder() - .putAll(statement.getCustomPayload()) - .put( - "request-id", - ByteBuffer.wrap(nodeRequestId.getBytes(StandardCharsets.UTF_8))) - .build(); - return statement.setCustomPayload(customPayload); + Map customPayload = + NullAllowingImmutableMap.builder() + .putAll(statement.getCustomPayload()) + .put("request-id", ByteBuffer.wrap(nodeRequestId.getBytes(StandardCharsets.UTF_8))) + .build(); + return statement.setCustomPayload(customPayload); } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java b/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java index dfbd9685bfd..56d9e0f4eac 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java @@ -222,7 +222,7 @@ public class DefaultDriverContext implements InternalDriverContext { private final LazyReference nodeStateListenerRef; private final LazyReference schemaChangeListenerRef; private final LazyReference requestTrackerRef; - private final LazyReference requestIdGeneratorRef; + private final LazyReference> requestIdGeneratorRef; private final LazyReference> authProviderRef; private final LazyReference> lifecycleListenersRef = new LazyReference<>("lifecycleListeners", this::buildLifecycleListeners, cycleDetector); @@ -716,20 +716,15 @@ protected RequestTracker buildRequestTracker(RequestTracker requestTrackerFromBu } } - protected RequestIdGenerator buildRequestIdGenerator(RequestIdGenerator requestIdGenerator) { + protected Optional buildRequestIdGenerator( + RequestIdGenerator requestIdGenerator) { return (requestIdGenerator != null) - ? requestIdGenerator + ? Optional.of(requestIdGenerator) : Reflection.buildFromConfig( - this, - DefaultDriverOption.REQUEST_ID_GENERATOR_CLASS, - RequestIdGenerator.class, - "com.datastax.oss.driver.internal.core.tracker") - .orElseThrow( - () -> - new IllegalArgumentException( - String.format( - "Missing request ID generator, check your configuration (%s)", - DefaultDriverOption.REQUEST_ID_GENERATOR_CLASS))); + this, + DefaultDriverOption.REQUEST_ID_GENERATOR_CLASS, + RequestIdGenerator.class, + "com.datastax.oss.driver.internal.core.tracker"); } protected Optional buildAuthProvider(AuthProvider authProviderFromBuilder) { @@ -998,7 +993,7 @@ public RequestTracker getRequestTracker() { @NonNull @Override - public RequestIdGenerator getRequestIdGenerator() { + public Optional getRequestIdGenerator() { return requestIdGeneratorRef.get(); } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java index 0808bdce63f..5cb0ebacd18 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/CqlRequestHandler.java +++ b/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.RequestIdGenerator; 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; @@ -82,6 +83,7 @@ import java.util.AbstractMap; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Queue; import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; @@ -125,6 +127,7 @@ public class CqlRequestHandler implements Throttled { private final List inFlightCallbacks; private final RequestThrottler throttler; private final RequestTracker requestTracker; + private final Optional requestIdGenerator; private final SessionMetricUpdater sessionMetricUpdater; private final DriverExecutionProfile executionProfile; @@ -139,7 +142,11 @@ protected CqlRequestHandler( String sessionLogPrefix) { this.startTimeNanos = System.nanoTime(); - this.logPrefix = sessionLogPrefix + "|" + this.hashCode(); + this.requestIdGenerator = context.getRequestIdGenerator(); + this.logPrefix = + this.requestIdGenerator.isPresent() + ? this.requestIdGenerator.get().getSessionRequestId(statement) + : sessionLogPrefix + "|" + this.hashCode(); LOG.trace("[{}] Creating new handler for request {}", logPrefix, statement); this.initialStatement = statement; @@ -267,6 +274,16 @@ private void sendRequest( setFinalError(statement, AllNodesFailedException.fromErrors(this.errors), null, -1); } } else { + String nodeLogPrefix; + if (this.requestIdGenerator.isPresent()) { + nodeLogPrefix = + this.requestIdGenerator + .get() + .getNodeRequestId(statement, logPrefix, currentExecutionIndex); + statement = this.requestIdGenerator.get().getDecoratedStatement(statement, nodeLogPrefix); + } else { + nodeLogPrefix = logPrefix + "|" + currentExecutionIndex; + } NodeResponseCallback nodeResponseCallback = new NodeResponseCallback( statement, @@ -276,7 +293,7 @@ private void sendRequest( currentExecutionIndex, retryCount, scheduleNextExecution, - logPrefix); + nodeLogPrefix); Message message = Conversions.toMessage(statement, executionProfile, context); channel .write(message, statement.isTracing(), statement.getCustomPayload(), nodeResponseCallback) @@ -489,7 +506,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. diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java index eaf7f5b80d6..cafe278cf85 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/UuidRequestIdGenerator.java @@ -28,8 +28,7 @@ public UuidRequestIdGenerator(DriverContext context) {} /** Generates a random v4 UUID. */ @Override - public String getSessionRequestId( - @NonNull Request statement) { + public String getSessionRequestId(@NonNull Request statement) { return Uuids.random().toString(); } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java index bde9fb4af55..71a895e9c4e 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java @@ -33,8 +33,7 @@ public W3CContextRequestIdGenerator(DriverContext context) {} /** Random 16 bytes, e.g. "4bf92f3577b34da6a3ce929d0e0e4736" */ @Override - public String getSessionRequestId( - @NonNull Request statement, @NonNull String sessionName, int hashCode) { + public String getSessionRequestId(@NonNull Request statement) { byte[] bytes = new byte[16]; random.nextBytes(bytes); return baseEncoding.encode(bytes); diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java index a6fe45458ff..9d86302aabf 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/RequestHandlerTestHarness.java @@ -114,8 +114,6 @@ protected RequestHandlerTestHarness(Builder builder) { when(defaultProfile.getBoolean(DefaultDriverOption.REQUEST_DEFAULT_IDEMPOTENCE)) .thenReturn(builder.defaultIdempotence); when(defaultProfile.getBoolean(DefaultDriverOption.PREPARE_ON_ALL_NODES)).thenReturn(true); - when(defaultProfile.getString(DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY)) - .thenReturn(""); when(config.getDefaultProfile()).thenReturn(defaultProfile); when(context.getConfig()).thenReturn(config); diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java index 1fad60ac154..448aaaa45d2 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/tracker/RequestIdGeneratorTest.java @@ -36,7 +36,7 @@ public void uuid_generator_should_generate() { // given UuidRequestIdGenerator generator = new UuidRequestIdGenerator(context); // when - String sessionRequestId = generator.getSessionRequestId(statement, "sessionName", 123); + String sessionRequestId = generator.getSessionRequestId(statement); String nodeRequestId = generator.getNodeRequestId(statement, sessionRequestId, 1); // then // e.g. "550e8400-e29b-41d4-a716-446655440000", which is 36 characters long @@ -51,7 +51,7 @@ public void w3c_generator_should_generate() { // given W3CContextRequestIdGenerator generator = new W3CContextRequestIdGenerator(context); // when - String sessionRequestId = generator.getSessionRequestId(statement, "sessionName", 123); + String sessionRequestId = generator.getSessionRequestId(statement); String nodeRequestId = generator.getNodeRequestId(statement, sessionRequestId, 1); // then // e.g. "4bf92f3577b34da6a3ce929d0e0e4736", which is 32 characters long diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java index c9ecb807a7a..d5bb80c5b46 100644 --- a/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java @@ -23,14 +23,17 @@ import com.datastax.oss.driver.api.core.config.DefaultDriverOption; import com.datastax.oss.driver.api.core.config.DriverConfigLoader; import com.datastax.oss.driver.api.core.cql.ResultSet; +import com.datastax.oss.driver.api.core.cql.Statement; import com.datastax.oss.driver.api.core.session.Request; import com.datastax.oss.driver.api.core.tracker.RequestIdGenerator; import com.datastax.oss.driver.api.testinfra.ccm.CcmRule; import com.datastax.oss.driver.api.testinfra.session.SessionUtils; import com.datastax.oss.driver.categories.ParallelizableTests; +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; import org.junit.Rule; import org.junit.Test; import org.junit.experimental.categories.Category; @@ -43,31 +46,16 @@ public class RequestIdGeneratorIT { @Rule public TestRule chain = RuleChain.outerRule(ccmRule); - @Test - public void should_write_default_id_to_custom_payload_with_key() { - DriverConfigLoader loader = - SessionUtils.configLoaderBuilder() - .withString(DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY, "trace_key") - .build(); - try (CqlSession session = SessionUtils.newSession(ccmRule, loader)) { - String query = "SELECT * FROM system.local"; - ResultSet rs = session.execute(query); - assertThat(rs.getExecutionInfo().getRequest().getCustomPayload().get("trace_key")) - .isNotNull(); - } - } - @Test public void should_write_uuid_to_custom_payload_with_key() { DriverConfigLoader loader = SessionUtils.configLoaderBuilder() .withString(DefaultDriverOption.REQUEST_ID_GENERATOR_CLASS, "UuidRequestIdGenerator") - .withString(DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY, "trace_key") .build(); try (CqlSession session = SessionUtils.newSession(ccmRule, loader)) { String query = "SELECT * FROM system.local"; ResultSet rs = session.execute(query); - ByteBuffer id = rs.getExecutionInfo().getRequest().getCustomPayload().get("trace_key"); + ByteBuffer id = rs.getExecutionInfo().getRequest().getCustomPayload().get("request-id"); assertThat(id.remaining()).isEqualTo(73); } } @@ -78,27 +66,21 @@ public void should_write_w3c_context_to_custom_payload_with_key() { SessionUtils.configLoaderBuilder() .withString( DefaultDriverOption.REQUEST_ID_GENERATOR_CLASS, "W3CContextRequestIdGenerator") - .withString(DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY, "trace_key") .build(); try (CqlSession session = SessionUtils.newSession(ccmRule, loader)) { String query = "SELECT * FROM system.local"; ResultSet rs = session.execute(query); - ByteBuffer id = rs.getExecutionInfo().getRequest().getCustomPayload().get("trace_key"); + ByteBuffer id = rs.getExecutionInfo().getRequest().getCustomPayload().get("request-id"); assertThat(id.remaining()).isEqualTo(55); } } @Test public void should_use_customized_request_id_generator() { - DriverConfigLoader loader = - SessionUtils.configLoaderBuilder() - .withString(DefaultDriverOption.REQUEST_ID_CUSTOM_PAYLOAD_KEY, "trace_key") - .build(); RequestIdGenerator myRequestIdGenerator = new RequestIdGenerator() { @Override - public String getSessionRequestId( - @NonNull Request statement, @NonNull String sessionName, int hashCode) { + public String getSessionRequestId(@NonNull Request statement) { return "123"; } @@ -107,13 +89,25 @@ public String getNodeRequestId( @NonNull Request statement, @NonNull String sessionRequestId, int executionCount) { return "456"; } + + @Override + public Statement getDecoratedStatement( + @NonNull Statement statement, @NonNull String nodeRequestId) { + Map customPayload = + NullAllowingImmutableMap.builder() + .putAll(statement.getCustomPayload()) + .put( + "trace_key", + ByteBuffer.wrap(nodeRequestId.getBytes(StandardCharsets.UTF_8))) + .build(); + return statement.setCustomPayload(customPayload); + } }; try (CqlSession session = (CqlSession) SessionUtils.baseBuilder() .addContactEndPoints(ccmRule.getContactPoints()) .withRequestIdGenerator(myRequestIdGenerator) - .withConfigLoader(loader) .build()) { String query = "SELECT * FROM system.local"; ResultSet rs = session.execute(query); From 4fa88a0ce30be6560c35ae2348f758fe55ec364a Mon Sep 17 00:00:00 2001 From: janehe Date: Thu, 29 May 2025 18:35:40 -0700 Subject: [PATCH 13/16] update manual --- manual/core/request_id/README.md | 29 +++-------------------------- 1 file changed, 3 insertions(+), 26 deletions(-) diff --git a/manual/core/request_id/README.md b/manual/core/request_id/README.md index fdee6a11563..a766a4419af 100644 --- a/manual/core/request_id/README.md +++ b/manual/core/request_id/README.md @@ -29,12 +29,8 @@ A request ID generator needs to generate both: - Node request ID: an identifier for the execution of a CQL statement against a particular node. There can be one or more node requests for a single session request, due to retries or speculative executions. Usage: -* Inject ID generator: set the desired `RequestIdGenerator` in `advanced.request-id.generator.class`. - The default implementation generates the session request ID as `{session_name}|{hash_code}`, and node request ID as `{session_name}|{hash_code}|{execution_count}`, - where "hash_code" is the hash code of the `CqlRequestHandler` object, and "execution_count" is the zero-based index of the node request in the session request. - For example, if there is a retry or speculative execution right after the first node request, the second node request will have the ID `{session_name}|{hash_code}|1`. -* Add ID to custom payload: disabled by default. Set the desired key in `advanced.request-id.custom-payload-with-key`, - then the driver will add the generated ID to the custom payload with the specified key. +* Inject ID generator: set the desired `RequestIdGenerator` in `advanced.request-id.generator.class`. +* Add ID to custom payload: the default behavior of a `RequestIdGenerator` is to add the request ID into the custom payload with the key `request-id`. Override `RequestIdGenerator.getDecoratedStatement` to customize the behavior. ### Request Id Generator Configuration @@ -49,23 +45,4 @@ datastax-java-driver.advanced.request-id.generator { To register your own request ID generator, specify the name of the class that implements `RequestIdGenerator`. -By default, the build-in implementation `DefaultRequestIdGenerator` is used. It generates the ID as -`{session_name}|{hash_code}|{execution_count}`. Note that this ID is not guaranteed to be unique. -Other built-in implementations include `UUIDRequestIdGenerator` and `W3CContextRequestIdGenerator`. - -The generated ID will be added to the log message of `CqlRequestHandler`, and propagated to the request trackers. - -### Custom Payload Configuration - -Users can opt in to add the generated node request ID to the custom payload to achieve request end-to-end tracing. -Custom payload is a map of string to `ByteBuffer` pairs, and the driver will add the generated ID to the custom payload with the specified key. - -``` -datastax-java-driver.advanced.request-id{ - custom-payload-with-key = my-request-id -} -``` - -Users can then run Apache Cassandra with customized query handler to extract the request ID from the custom payload, to achieve end-to-end tracing. - -When this key is set to an empty string (default), the driver will not add the ID to the custom payload. +The generated ID will be added to the log message of `CqlRequestHandler`, and propagated to other classes, e.g. the request trackers. \ No newline at end of file From e2a2482900e3b7c3a7f9c1f8b958f5bd87da84c9 Mon Sep 17 00:00:00 2001 From: janehe Date: Thu, 12 Jun 2025 19:26:40 -0700 Subject: [PATCH 14/16] use W3C generator when Astra --- .../oss/driver/api/core/session/SessionBuilder.java | 7 +++++++ .../oss/driver/api/core/tracker/RequestIdGenerator.java | 8 +++++++- .../core/tracker/W3CContextRequestIdGenerator.java | 2 ++ .../oss/driver/core/tracker/RequestIdGeneratorIT.java | 2 +- 4 files changed, 17 insertions(+), 2 deletions(-) diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java b/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java index 3d829c1a17f..15e23fc70b5 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/session/SessionBuilder.java @@ -48,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; @@ -873,6 +874,12 @@ protected final CompletionStage buildDefaultSessionAsync() { List 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()); + 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."); diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java index 58e567f4246..7184c6d30e9 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/tracker/RequestIdGenerator.java @@ -51,12 +51,18 @@ public interface RequestIdGenerator { String getNodeRequestId( @NonNull Request statement, @NonNull String sessionRequestId, int executionCount); + default String getCustomPayloadKey() { + return "request-id"; + } + default Statement getDecoratedStatement( @NonNull Statement statement, @NonNull String nodeRequestId) { Map customPayload = NullAllowingImmutableMap.builder() .putAll(statement.getCustomPayload()) - .put("request-id", ByteBuffer.wrap(nodeRequestId.getBytes(StandardCharsets.UTF_8))) + .put( + getCustomPayloadKey(), + ByteBuffer.wrap(nodeRequestId.getBytes(StandardCharsets.UTF_8))) .build(); return statement.setCustomPayload(customPayload); } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java index 71a895e9c4e..67f9ea43aa4 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/tracker/W3CContextRequestIdGenerator.java @@ -31,6 +31,8 @@ public class W3CContextRequestIdGenerator implements RequestIdGenerator { public W3CContextRequestIdGenerator(DriverContext context) {} + public W3CContextRequestIdGenerator() {} + /** Random 16 bytes, e.g. "4bf92f3577b34da6a3ce929d0e0e4736" */ @Override public String getSessionRequestId(@NonNull Request statement) { diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java index d5bb80c5b46..4f50114b8a9 100644 --- a/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/tracker/RequestIdGeneratorIT.java @@ -61,7 +61,7 @@ public void should_write_uuid_to_custom_payload_with_key() { } @Test - public void should_write_w3c_context_to_custom_payload_with_key() { + public void should_write_default_request_id_to_custom_payload_with_key() { DriverConfigLoader loader = SessionUtils.configLoaderBuilder() .withString( From 6368097487969346a7feee52fa144a9408576d68 Mon Sep 17 00:00:00 2001 From: janehe Date: Wed, 18 Jun 2025 15:35:47 -0700 Subject: [PATCH 15/16] revapi --- core/revapi.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/revapi.json b/core/revapi.json index a2f4c16e651..f333d22e661 100644 --- a/core/revapi.json +++ b/core/revapi.json @@ -7389,7 +7389,7 @@ }, { "code": "java.method.addedToInterface", - "new": "method com.datastax.oss.driver.api.core.tracker.RequestIdGenerator com.datastax.oss.driver.api.core.context.DriverContext::getRequestIdGenerator()", + "new": "method java.util.Optional 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" } ] From 69d06660bdf442fda0803393b8d1abec20e729bf Mon Sep 17 00:00:00 2001 From: janehe Date: Tue, 1 Jul 2025 15:55:58 -0700 Subject: [PATCH 16/16] empty