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/connection/ServerDescription.java

@@ -65,7 +65,7 @@ public class ServerDescription {
      * The maximum supported driver wire version
      * @since 3.8
      */
-    public static final int MAX_DRIVER_WIRE_VERSION = 25;
+    public static final int MAX_DRIVER_WIRE_VERSION = 29;
 
     private static final int DEFAULT_MAX_DOCUMENT_SIZE = 0x1000000;  // 16MB
 

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,6 +35,7 @@
 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.io.BsonOutput;
@@ -64,6 +66,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 +83,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 +109,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,36 +172,42 @@ 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 (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);
-                        int sequenceStart = byteBuf.position();
-                        int sequenceSizeInBytes = byteBuf.getInt();
-                        int sectionEnd = sequenceStart + sequenceSizeInBytes;
-
-                        String fieldName = getSequenceIdentifier(byteBuf);
-                        // If this assertion fires, it means that the driver has started using document sequences for nested fields.  If
-                        // so, this method will need to change in order to append the value to the correct nested document.
-                        assertFalse(fieldName.contains("."));
-
-                        ByteBuf documentsByteBufSlice = byteBuf.duplicate().limit(sectionEnd);
-                        try {
-                            commandBsonDocument.append(fieldName, new BsonArray(createList(documentsByteBufSlice)));
-                        } finally {
-                            documentsByteBufSlice.release();
-                        }
-                        byteBuf.position(sectionEnd);
+                // There may be 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.
+                // The command document is materialized lazily, only when a document sequence has to be folded into
+                // it; otherwise (no sections, or only a telemetry section) the cheap lazy document is returned.
+                BsonDocument commandBsonDocument = null;
+
+                // Each loop iteration processes one Document Sequence
+                // When there are no more bytes remaining, there are no more Document Sequences
+                while (byteBuf.hasRemaining()) {
+                    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;
+                    }
+                    if (commandBsonDocument == null) {
+                        commandBsonDocument = byteBufBsonDocument.toBaseBsonDocument();
+                    }
+                    int sequenceStart = byteBuf.position();
+                    int sequenceSizeInBytes = byteBuf.getInt();
+                    int sectionEnd = sequenceStart + sequenceSizeInBytes;
+
+                    String fieldName = getSequenceIdentifier(byteBuf);
+                    // If this assertion fires, it means that the driver has started using document sequences for nested fields.  If
+                    // so, this method will need to change in order to append the value to the correct nested document.
+                    assertFalse(fieldName.contains("."));
+
+                    ByteBuf documentsByteBufSlice = byteBuf.duplicate().limit(sectionEnd);
+                    try {
+                        commandBsonDocument.append(fieldName, new BsonArray(createList(documentsByteBufSlice)));
+                    } finally {
+                        documentsByteBufSlice.release();
                     }
-                    return commandBsonDocument;
-                } else {
-                    return byteBufBsonDocument;
+                    byteBuf.position(sectionEnd);
                 }
+                return commandBsonDocument != null ? commandBsonDocument : byteBufBsonDocument;
             } finally {
                 byteBuf.release();
             }
@@ -194,6 +216,24 @@ BsonDocument getCommandDocument(final ByteBufferBsonOutput bsonOutput) {
         }
     }
 
+    /**
+     * The command name (the 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 +275,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 +332,35 @@ 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);
+        // {otel: {traceparent: <value>}}
+        BsonBinaryWriter writer = new BsonBinaryWriter(bsonOutput);
+        writer.writeStartDocument();
+        writer.writeStartDocument("otel");
+        writer.writeString("traceparent", traceParent);
+        writer.writeEndDocument();
+        writer.writeEndDocument();
+    }
+
     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 {