JAVA-6222 Support trace context propagation to the server

driver-core/build.gradle.kts

@@ -57,9 +57,15 @@ dependencies {
 
     optionalImplementation(platform(libs.micrometer.observation.bom))
     optionalImplementation(libs.micrometer.observation)
+    optionalImplementation(libs.micrometer.tracing)
 
     testImplementation(project(path = ":bson", configuration = "testArtifacts"))
     testImplementation(libs.reflections)
+
+    // Tracing testing
+    testImplementation(platform(libs.micrometer.tracing.integration.test.bom))
+    testImplementation(libs.micrometer.tracing.integration.test) { exclude(group = "org.junit.jupiter") }
+
     testImplementation(libs.netty.tcnative.boringssl.static)
     listOf("linux-x86_64", "linux-aarch_64", "osx-x86_64", "osx-aarch_64", "windows-x86_64").forEach { arch ->
         testImplementation("${libs.netty.tcnative.boringssl.static.get()}::$arch")

driver-core/src/main/com/mongodb/internal/connection/CommandMessage.java

@@ -25,6 +25,7 @@
 import com.mongodb.internal.MongoNamespaceHelper;
 import com.mongodb.internal.TimeoutContext;
 import com.mongodb.internal.connection.MessageSequences.EmptyMessageSequences;
+import com.mongodb.internal.observability.micrometer.Span;
 import com.mongodb.internal.session.SessionContext;
 import com.mongodb.lang.Nullable;
 import org.bson.BsonArray;
@@ -34,8 +35,11 @@
 import org.bson.BsonElement;
 import org.bson.BsonInt64;
 import org.bson.BsonString;
+import org.bson.BsonValue;
 import org.bson.ByteBuf;
 import org.bson.FieldNameValidator;
+import org.bson.codecs.BsonDocumentCodec;
+import org.bson.codecs.EncoderContext;
 import org.bson.io.BsonOutput;
 
 import java.io.ByteArrayOutputStream;
@@ -64,6 +68,7 @@
 import static com.mongodb.internal.connection.ByteBufBsonDocument.createList;
 import static com.mongodb.internal.connection.ByteBufBsonDocument.createOne;
 import static com.mongodb.internal.connection.ReadConcernHelper.getReadConcernDocument;
+import static com.mongodb.internal.operation.ServerVersionHelper.NINE_DOT_ZERO_WIRE_VERSION;
 import static com.mongodb.internal.operation.ServerVersionHelper.UNKNOWN_WIRE_VERSION;
 
 /**
@@ -80,6 +85,11 @@ public final class CommandMessage extends RequestMessage {
      * Specifies that the `OP_MSG` section payload is a sequence of BSON documents.
      */
     private static final byte PAYLOAD_TYPE_1_DOCUMENT_SEQUENCE = 1;
+    /**
+     * Specifies that the `OP_MSG` section payload is a BSON telemetry document
+     * ({@code {otel: {traceparent: <string>}}}).
+     */
+    private static final byte PAYLOAD_TYPE_3_TELEMETRY = 3;
 
     private static final int UNINITIALIZED_POSITION = -1;
 
@@ -101,6 +111,14 @@ public final class CommandMessage extends RequestMessage {
     private final ClusterConnectionMode clusterConnectionMode;
     private final ServerApi serverApi;
 
+    /**
+     * The command span to attach to the OP_MSG telemetry section during {@linkplain
+     * #encode(ByteBufferBsonOutput, OperationContext, Span) encoding}. Set for the duration of that overload's
+     * call and cleared afterward; {@code null} otherwise.
+     */
+    @Nullable
+    private Span commandSpanForEncoding;
+
     CommandMessage(final String database, final BsonDocument command, final FieldNameValidator commandFieldNameValidator,
                    final ReadPreference readPreference, final MessageSettings settings, final ClusterConnectionMode clusterConnectionMode,
                    @Nullable final ServerApi serverApi) {
@@ -156,15 +174,20 @@ BsonDocument getCommandDocument(final ByteBufferBsonOutput bsonOutput) {
                 byteBuf.position(firstDocumentPosition);
                 ByteBufBsonDocument byteBufBsonDocument = createOne(byteBuf);
 
-                // If true, it means there is at least one `PAYLOAD_TYPE_1_DOCUMENT_SEQUENCE` section in the OP_MSG
+                // If true, there are more sections after the `PAYLOAD_TYPE_0_DOCUMENT` section: either one or more
+                // `PAYLOAD_TYPE_1_DOCUMENT_SEQUENCE` sections, and/or a trailing `PAYLOAD_TYPE_3_TELEMETRY` section.
                 if (byteBuf.hasRemaining()) {
                     BsonDocument commandBsonDocument = byteBufBsonDocument.toBaseBsonDocument();
 
                     // Each loop iteration processes one Document Sequence
                     // When there are no more bytes remaining, there are no more Document Sequences
                     while (byteBuf.hasRemaining()) {
-                        // skip reading the payload type, we know it is `PAYLOAD_TYPE_1`
-                        byteBuf.position(byteBuf.position() + 1);
+                        byte payloadType = byteBuf.get();
+                        // Document-sequence sections always precede any trailing non-sequence section (e.g.
+                        // PAYLOAD_TYPE_3_TELEMETRY), and such sections carry no command-document fields, so stop here.
+                        if (payloadType != PAYLOAD_TYPE_1_DOCUMENT_SEQUENCE) {
+                            break;
+                        }
                         int sequenceStart = byteBuf.position();
                         int sequenceSizeInBytes = byteBuf.getInt();
                         int sectionEnd = sequenceStart + sequenceSizeInBytes;
@@ -194,6 +217,24 @@ BsonDocument getCommandDocument(final ByteBufferBsonOutput bsonOutput) {
         }
     }
 
+    /**
+     * The command name (first key of the command document. Available before {@code encode()} runs, unlike
+     * {@link #getCommandDocument(ByteBufferBsonOutput)}; identical to the encoded document's first
+     * key, since encoding only appends fields.
+     */
+    public String getCommandName() {
+        return command.getFirstKey();
+    }
+
+    /**
+     * The cursor id if this is a {@code getMore} command, or {@code null} otherwise.
+     */
+    @Nullable
+    public BsonInt64 getGetMoreCursorId() {
+        BsonValue value = command.get("getMore");
+        return value instanceof BsonInt64 ? (BsonInt64) value : null;
+    }
+
     /**
      * Get the field name from a buffer positioned at the start of the document sequence identifier of an OP_MSG Section of type
      * `PAYLOAD_TYPE_1_DOCUMENT_SEQUENCE`.
@@ -235,6 +276,21 @@ protected void encodeMessageBody(final ByteBufferBsonOutput bsonOutput, final Op
         this.firstDocumentPosition = useOpMsg() ? writeOpMsg(bsonOutput, operationContext) : writeOpQuery(bsonOutput);
     }
 
+    /**
+     * Encodes the message, attaching the OP_MSG telemetry section with {@code commandSpan}'s trace context when
+     * the gating conditions hold (see {@link #writeTelemetryContextSection}). The 2-arg {@link
+     * #encode(ByteBufferBsonOutput, OperationContext)} never attaches the section.
+     */
+    public void encode(final ByteBufferBsonOutput bsonOutput, final OperationContext operationContext,
+            @Nullable final Span commandSpan) {
+        this.commandSpanForEncoding = commandSpan;
+        try {
+            encode(bsonOutput, operationContext);
+        } finally {
+            this.commandSpanForEncoding = null;
+        }
+    }
+
     @SuppressWarnings("try")
     private int writeOpMsg(final ByteBufferBsonOutput bsonOutput, final OperationContext operationContext) {
         int messageStartPosition = bsonOutput.getPosition() - MESSAGE_PROLOGUE_LENGTH;
@@ -277,11 +333,32 @@ private int writeOpMsg(final ByteBufferBsonOutput bsonOutput, final OperationCon
             fail(sequences.toString());
         }
 
+        writeTelemetryContextSection(bsonOutput);
+
         // Write the flag bits
         bsonOutput.writeInt32(flagPosition, getOpMsgFlagBits());
         return commandStartPosition;
     }
 
+    private void writeTelemetryContextSection(final ByteBufferBsonOutput bsonOutput) {
+        if (getSettings().getMaxWireVersion() < NINE_DOT_ZERO_WIRE_VERSION) {
+            return;
+        }
+        Span commandSpan = this.commandSpanForEncoding;
+        if (commandSpan == null) {
+            return;
+        }
+        String traceParent = commandSpan.context().traceParent();
+        if (traceParent == null) {
+            return;
+        }
+        bsonOutput.writeByte(PAYLOAD_TYPE_3_TELEMETRY);
+        BsonDocument telemetry = new BsonDocument("otel",
+                new BsonDocument("traceparent", new BsonString(traceParent)));
+        new BsonDocumentCodec().encode(new BsonBinaryWriter(bsonOutput), telemetry,
+                EncoderContext.builder().build());
+    }
+
     private int writeOpQuery(final ByteBufferBsonOutput bsonOutput) {
         bsonOutput.writeInt32(0);
         bsonOutput.writeCString(new MongoNamespace(getDatabase(), MongoNamespaceHelper.COMMAND_COLLECTION_NAME).getFullName());

driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnection.java

@@ -439,20 +439,19 @@ public boolean reauthenticationIsTriggered(@Nullable final Throwable t) {
     @Nullable
     private <T> T sendAndReceiveInternal(final CommandMessage message, final Decoder<T> decoder,
             final OperationContext operationContext) {
-        CommandEventSender commandEventSender;
-        Span tracingSpan;
+        CommandEventSender commandEventSender = new NoOpCommandEventSender();
+        Span tracingSpan = null;
         try (ByteBufferBsonOutput bsonOutput = new ByteBufferBsonOutput(this)) {
-            message.encode(bsonOutput, operationContext);
             tracingSpan = operationContext
                     .getTracingManager()
                     .createTracingSpan(message,
                             operationContext,
-                            () -> message.getCommandDocument(bsonOutput),
                             cmdName -> SECURITY_SENSITIVE_COMMANDS.contains(cmdName)
                                     || SECURITY_SENSITIVE_HELLO_COMMANDS.contains(cmdName),
                             () -> getDescription().getServerAddress(),
                             () -> getDescription().getConnectionId()
                     );
+            message.encode(bsonOutput, operationContext, tracingSpan);
             boolean isLoggingCommandNeeded = isLoggingCommandNeeded();
             boolean isTracingCommandPayloadNeeded = tracingSpan != null && operationContext.getTracingManager().isCommandPayloadEnabled();
 
@@ -467,27 +466,22 @@ private <T> T sendAndReceiveInternal(final CommandMessage message, final Decoder
                         operationContext, message, commandDocument,
                         COMMAND_PROTOCOL_LOGGER, loggerSettings);
                 commandEventSender.sendStartedEvent();
-            } else {
-                commandEventSender = new NoOpCommandEventSender();
             }
             if (isTracingCommandPayloadNeeded) {
                 tracingSpan.setQueryText(commandDocument);
             }
             if (tracingSpan != null) {
                 tracingSpan.openScope();
             }
-
-            try {
-                sendCommandMessage(message, bsonOutput, operationContext);
-            } catch (Exception e) {
-                if (tracingSpan != null) {
-                    tracingSpan.error(e);
-                    tracingSpan.closeScope();
-                    tracingSpan.end();
-                }
-                commandEventSender.sendFailedEvent(e);
-                throw e;
+            sendCommandMessage(message, bsonOutput, operationContext);
+        } catch (Throwable t) {
+            if (tracingSpan != null) {
+                tracingSpan.error(t);
+                tracingSpan.closeScope();
+                tracingSpan.end();
             }
+            commandEventSender.sendFailedEvent(t);
+            throw t;
         }
 
         if (message.isResponseExpected()) {
@@ -528,7 +522,7 @@ private void sendCommandMessage(final CommandMessage message, final ByteBufferBs
             final OperationContext operationContext) {
 
         Compressor localSendCompressor = sendCompressor;
-        if (localSendCompressor == null || SECURITY_SENSITIVE_COMMANDS.contains(message.getCommandDocument(bsonOutput).getFirstKey())) {
+        if (localSendCompressor == null || SECURITY_SENSITIVE_COMMANDS.contains(message.getCommandName())) {
             trySendMessage(message, bsonOutput, operationContext);
         } else {
             ByteBufferBsonOutput compressedBsonOutput;
@@ -617,19 +611,18 @@ private <T> void sendAndReceiveAsyncInternal(final CommandMessage message, final
 
         Span tracingSpan = null;
         try {
-            message.encode(bsonOutput, operationContext);
-
             tracingSpan = operationContext
                     .getTracingManager()
                     .createTracingSpan(message,
                             operationContext,
-                            () -> message.getCommandDocument(bsonOutput),
                             cmdName -> SECURITY_SENSITIVE_COMMANDS.contains(cmdName)
                                     || SECURITY_SENSITIVE_HELLO_COMMANDS.contains(cmdName),
                             () -> getDescription().getServerAddress(),
                             () -> getDescription().getConnectionId()
                     );
 
+            message.encode(bsonOutput, operationContext, tracingSpan);
+
             CommandEventSender commandEventSender;
             boolean isLoggingCommandNeeded = isLoggingCommandNeeded();
             boolean isTracingCommandPayloadNeeded = tracingSpan != null && operationContext.getTracingManager().isCommandPayloadEnabled();
@@ -670,7 +663,7 @@ private <T> void sendAndReceiveAsyncInternal(final CommandMessage message, final
 
             commandEventSender.sendStartedEvent();
             Compressor localSendCompressor = sendCompressor;
-            if (localSendCompressor == null || SECURITY_SENSITIVE_COMMANDS.contains(message.getCommandDocument(bsonOutput).getFirstKey())) {
+            if (localSendCompressor == null || SECURITY_SENSITIVE_COMMANDS.contains(message.getCommandName())) {
                 sendCommandMessageAsync(message.getId(), decoder, operationContext, tracingCallback, bsonOutput, commandEventSender,
                         message.isResponseExpected());
             } else {

driver-core/src/main/com/mongodb/internal/observability/micrometer/MicrometerTracer.java

@@ -47,11 +47,28 @@
  * @since 5.7
  */
 public class MicrometerTracer implements Tracer {
+    /**
+     * {@code micrometer-tracing} is an optional dependency: consumers may have only {@code micrometer-observation}
+     * on the classpath (e.g. metrics/logging-only observability setups). Guard any use of {@code io.micrometer.tracing.*}
+     * types with this flag so that such setups never trigger a {@link NoClassDefFoundError} or {@link LinkageError}.
+     */
+    static final boolean MICROMETER_TRACING_ON_CLASSPATH = isMicrometerTracingOnClasspath();
+
     private final ObservationRegistry observationRegistry;
     private final boolean allowCommandPayload;
     private final int textMaxLength;
     private final ObservationConvention<MongodbObservationContext> convention;
 
+    private static boolean isMicrometerTracingOnClasspath() {
+        try {
+            Class.forName("io.micrometer.tracing.handler.TracingObservationHandler",
+                    false, MicrometerTracer.class.getClassLoader());
+            return true;
+        } catch (ClassNotFoundException | LinkageError e) {
+            return false;
+        }
+    }
+
     /**
      * Constructs a new {@link MicrometerTracer} instance.
      *
@@ -116,6 +133,68 @@ private static class MicrometerTraceContext implements TraceContext {
         MicrometerTraceContext(@Nullable final Observation observation) {
             this.observation = observation;
         }
+
+        @Override
+        @Nullable
+        public String traceParent() {
+            if (observation == null) {
+                return null;
+            }
+            if (!MICROMETER_TRACING_ON_CLASSPATH) {
+                return null;
+            }
+            return MicrometerTracingSupport.traceParent(observation);
+        }
+
+        /**
+         * The server ({@code validateW3CTraceparent}) rejects ids that are not exactly the expected
+         * length of lowercase hex, or that are all zeroes. Never emit a traceparent it would reject.
+         */
+        @SuppressWarnings("BooleanMethodIsAlwaysInverted")
+        static boolean isValidNonZeroLowercaseHex(@Nullable final String value, final int expectedLength) {
+            if (value == null || value.length() != expectedLength) {
+                return false;
+            }
+            boolean nonZero = false;
+            for (int i = 0; i < value.length(); i++) {
+                char c = value.charAt(i);
+                if (!((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f'))) {
+                    return false;
+                }
+                if (c != '0') {
+                    nonZero = true;
+                }
+            }
+            return nonZero;
+        }
+    }
+
+    /**
+     * Isolates all hard references to {@code io.micrometer.tracing.*} types. This class must only be loaded/invoked
+     * after {@link #MICROMETER_TRACING_ON_CLASSPATH} has been confirmed {@code true}, so that JVMs which resolve
+     * classes referenced by bytecode eagerly (e.g. during verification) do not trigger a {@link NoClassDefFoundError}
+     * when {@code micrometer-tracing} is absent from the classpath.
+     */
+    private static final class MicrometerTracingSupport {
+        @Nullable
+        static String traceParent(final Observation observation) {
+            io.micrometer.tracing.handler.TracingObservationHandler.TracingContext tracingContext =
+                    observation.getContextView().get(io.micrometer.tracing.handler.TracingObservationHandler.TracingContext.class);
+            if (tracingContext == null || tracingContext.getSpan() == null) {
+                return null;
+            }
+            // Span.context() is non-null by contract (the package is @NullMarked); a no-op span yields a
+            // NOOP context whose empty ids fail the hex validation below.
+            io.micrometer.tracing.TraceContext ctx = tracingContext.getSpan().context();
+            String traceId = ctx.traceId();
+            String spanId = ctx.spanId();
+            if (!MicrometerTraceContext.isValidNonZeroLowercaseHex(traceId, 32)
+                    || !MicrometerTraceContext.isValidNonZeroLowercaseHex(spanId, 16)) {
+                return null;
+            }
+            Boolean sampled = ctx.sampled();
+            return "00-" + traceId + "-" + spanId + (sampled != null && sampled ? "-01" : "-00");
+        }
     }
 
     /**

driver-core/src/main/com/mongodb/internal/observability/micrometer/Span.java

@@ -90,12 +90,14 @@ public MongodbObservationContext getMongodbObservationContext() {
 
     /**
      * Opens a scope for this span, making it the current observation on the thread.
-     * Must be paired with {@link #closeScope()} in a try-finally block.
+     * A successful {@code openScope()} must be followed by {@link #closeScope()}.
      */
     void openScope();
 
     /**
      * Closes the scope previously opened by {@link #openScope()}, restoring the previous observation.
+     * Calling this without a prior {@code openScope()}, or more than once, is a no-op — error-handling
+     * paths may therefore call it unconditionally.
      */
     void closeScope();