Add tracing support using Micrometer

driver-core/build.gradle.kts

@@ -54,6 +54,7 @@ dependencies {
 
     optionalImplementation(libs.snappy.java)
     optionalImplementation(libs.zstd.jni)
+    optionalImplementation(libs.micrometer)
 
     testImplementation(project(path = ":bson", configuration = "testArtifacts"))
     testImplementation(libs.reflections)

driver-core/src/main/com/mongodb/MongoClientSettings.java

@@ -30,9 +30,11 @@
 import com.mongodb.connection.SslSettings;
 import com.mongodb.connection.TransportSettings;
 import com.mongodb.event.CommandListener;
+import com.mongodb.internal.tracing.TracingManager;
 import com.mongodb.lang.Nullable;
 import com.mongodb.spi.dns.DnsClient;
 import com.mongodb.spi.dns.InetAddressResolver;
+import com.mongodb.internal.tracing.Tracer;
 import org.bson.UuidRepresentation;
 import org.bson.codecs.BsonCodecProvider;
 import org.bson.codecs.BsonValueCodecProvider;
@@ -118,6 +120,7 @@ public final class MongoClientSettings {
     private final InetAddressResolver inetAddressResolver;
     @Nullable
     private final Long timeoutMS;
+    private final TracingManager tracingManager;
 
     /**
      * Gets the default codec registry.  It includes the following providers:
@@ -238,6 +241,7 @@ public static final class Builder {
         private ContextProvider contextProvider;
         private DnsClient dnsClient;
         private InetAddressResolver inetAddressResolver;
+        private TracingManager tracingManager;
 
         private Builder() {
         }
@@ -275,6 +279,7 @@ private Builder(final MongoClientSettings settings) {
             if (settings.heartbeatSocketTimeoutSetExplicitly) {
                 heartbeatSocketTimeoutMS = settings.heartbeatSocketSettings.getReadTimeout(MILLISECONDS);
             }
+            tracingManager = settings.tracingManager;
         }
 
         /**
@@ -723,6 +728,20 @@ Builder heartbeatSocketTimeoutMS(final int heartbeatSocketTimeoutMS) {
             return this;
         }
 
+        /**
+         * Sets the tracer to use for creating Spans for operations and commands.
+         *
+         * @param tracer the tracer
+         * @see com.mongodb.tracing.MicrometerTracer
+         * @return this
+         * @since 5.5
+         */
+        @Alpha(Reason.CLIENT)
+        public Builder tracer(final Tracer tracer) {
+            this.tracingManager = new TracingManager(tracer);
+            return this;
+        }
+
         /**
          * Build an instance of {@code MongoClientSettings}.
          *
@@ -1040,6 +1059,17 @@ public ContextProvider getContextProvider() {
         return contextProvider;
     }
 
+    /**
+     * Get the tracer to create Spans for operations and commands.
+     *
+     * @return this
+     * @since 5.5
+     */
+    @Alpha(Reason.CLIENT)
+    public TracingManager getTracingManager() {
+        return tracingManager;
+    }
+
     @Override
     public boolean equals(final Object o) {
         if (this == o) {
@@ -1156,5 +1186,6 @@ private MongoClientSettings(final Builder builder) {
         heartbeatConnectTimeoutSetExplicitly = builder.heartbeatConnectTimeoutMS != 0;
         contextProvider = builder.contextProvider;
         timeoutMS = builder.timeoutMS;
+        tracingManager = (builder.tracingManager == null) ? TracingManager.NO_OP : builder.tracingManager;
     }
 }

driver-core/src/main/com/mongodb/MongoNamespace.java

@@ -91,7 +91,7 @@ public static void checkCollectionNameValidity(final String collectionName) {
     public MongoNamespace(final String fullName) {
         notNull("fullName", fullName);
         this.fullName = fullName;
-        this.databaseName = getDatatabaseNameFromFullName(fullName);
+        this.databaseName = getDatabaseNameFromFullName(fullName);
         this.collectionName = getCollectionNameFullName(fullName);
         checkDatabaseNameValidity(databaseName);
         checkCollectionNameValidity(collectionName);
@@ -190,7 +190,7 @@ private static String getCollectionNameFullName(final String namespace) {
         return namespace.substring(firstDot + 1);
     }
 
-    private static String getDatatabaseNameFromFullName(final String namespace) {
+    private static String getDatabaseNameFromFullName(final String namespace) {
         int firstDot = namespace.indexOf('.');
         if (firstDot == -1) {
             return "";

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

@@ -186,6 +186,10 @@ BsonDocument getCommandDocument(final ByteBufferBsonOutput bsonOutput) {
         }
     }
 
+    BsonDocument getCommand() {
+        return command;
+    }
+
     /**
      * 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`.

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

@@ -51,6 +51,9 @@
 import com.mongodb.internal.logging.StructuredLogger;
 import com.mongodb.internal.session.SessionContext;
 import com.mongodb.internal.time.Timeout;
+import com.mongodb.internal.tracing.Span;
+import com.mongodb.internal.tracing.TraceContext;
+import com.mongodb.internal.tracing.TracingManager;
 import com.mongodb.lang.Nullable;
 import org.bson.BsonBinaryReader;
 import org.bson.BsonDocument;
@@ -94,6 +97,18 @@
 import static com.mongodb.internal.connection.ProtocolHelper.isCommandOk;
 import static com.mongodb.internal.logging.LogMessage.Level.DEBUG;
 import static com.mongodb.internal.thread.InterruptionUtil.translateInterruptedException;
+import static com.mongodb.internal.tracing.Tags.CLIENT_CONNECTION_ID;
+import static com.mongodb.internal.tracing.Tags.CURSOR_ID;
+import static com.mongodb.internal.tracing.Tags.NAMESPACE;
+import static com.mongodb.internal.tracing.Tags.QUERY_SUMMARY;
+import static com.mongodb.internal.tracing.Tags.QUERY_TEXT;
+import static com.mongodb.internal.tracing.Tags.SERVER_ADDRESS;
+import static com.mongodb.internal.tracing.Tags.SERVER_CONNECTION_ID;
+import static com.mongodb.internal.tracing.Tags.SERVER_PORT;
+import static com.mongodb.internal.tracing.Tags.SERVER_TYPE;
+import static com.mongodb.internal.tracing.Tags.SESSION_ID;
+import static com.mongodb.internal.tracing.Tags.SYSTEM;
+import static com.mongodb.internal.tracing.Tags.TRANSACTION_NUMBER;
 import static java.util.Arrays.asList;
 
 /**
@@ -376,7 +391,7 @@ public <T> T sendAndReceive(final CommandMessage message, final Decoder<T> decod
                 message, decoder, operationContext);
         try {
             return sendAndReceiveInternal.get();
-        } catch (MongoCommandException e) {
+        } catch (Throwable e) {
             if (reauthenticationIsTriggered(e)) {
                 return reauthenticateAndRetry(sendAndReceiveInternal, operationContext);
             }
@@ -391,9 +406,8 @@ public <T> void sendAndReceiveAsync(final CommandMessage message, final Decoder<
 
         AsyncSupplier<T> sendAndReceiveAsyncInternal = c -> sendAndReceiveAsyncInternal(
                 message, decoder, operationContext, c);
-        beginAsync().<T>thenSupply(c -> {
-            sendAndReceiveAsyncInternal.getAsync(c);
-        }).onErrorIf(e -> reauthenticationIsTriggered(e), (t, c) -> {
+
+        beginAsync().thenSupply(sendAndReceiveAsyncInternal::getAsync).onErrorIf(this::reauthenticationIsTriggered, (t, c) -> {
             reauthenticateAndRetryAsync(sendAndReceiveAsyncInternal, operationContext, c);
         }).finish(callback);
     }
@@ -432,15 +446,31 @@ public boolean reauthenticationIsTriggered(@Nullable final Throwable t) {
     private <T> T sendAndReceiveInternal(final CommandMessage message, final Decoder<T> decoder,
             final OperationContext operationContext) {
         CommandEventSender commandEventSender;
+        Span tracingSpan = createTracingSpan(message, operationContext);
+
         try (ByteBufferBsonOutput bsonOutput = new ByteBufferBsonOutput(this)) {
             message.encode(bsonOutput, operationContext);
-            commandEventSender = createCommandEventSender(message, bsonOutput, operationContext);
+            BsonDocument commandDocument = message.getCommandDocument(bsonOutput);
+
+            commandEventSender = createCommandEventSender(message, commandDocument, operationContext);
             commandEventSender.sendStartedEvent();
+
+            if (tracingSpan != null && operationContext.getTracingManager().isCommandPayloadEnabled()) {
+                tracingSpan.tag(QUERY_TEXT, commandDocument.toJson());
+            }
+
             try {
                 sendCommandMessage(message, bsonOutput, operationContext);
             } catch (Exception e) {
+                if (tracingSpan != null) {
+                    tracingSpan.error(e);
+                }
                 commandEventSender.sendFailedEvent(e);
                 throw e;
+            } finally {
+                if (tracingSpan != null) {
+                    tracingSpan.end();
+                }
             }
         }
 
@@ -553,7 +583,8 @@ private <T> void sendAndReceiveAsyncInternal(final CommandMessage message, final
 
         try {
             message.encode(bsonOutput, operationContext);
-            CommandEventSender commandEventSender = createCommandEventSender(message, bsonOutput, operationContext);
+            BsonDocument commandDocument = message.getCommandDocument(bsonOutput);
+            CommandEventSender commandEventSender = createCommandEventSender(message, commandDocument, operationContext);
             commandEventSender.sendStartedEvent();
             Compressor localSendCompressor = sendCompressor;
             if (localSendCompressor == null || SECURITY_SENSITIVE_COMMANDS.contains(message.getCommandDocument(bsonOutput).getFirstKey())) {
@@ -952,19 +983,87 @@ public void onResult(@Nullable final ByteBuf result, @Nullable final Throwable t
 
     private static final StructuredLogger COMMAND_PROTOCOL_LOGGER = new StructuredLogger("protocol.command");
 
-    private CommandEventSender createCommandEventSender(final CommandMessage message, final ByteBufferBsonOutput bsonOutput,
+    private CommandEventSender createCommandEventSender(final CommandMessage message, final BsonDocument commandDocument,
                                                         final OperationContext operationContext) {
         boolean listensOrLogs = commandListener != null || COMMAND_PROTOCOL_LOGGER.isRequired(DEBUG, getClusterId());
         if (!recordEverything && (isMonitoringConnection || !opened() || !authenticated.get() || !listensOrLogs)) {
             return new NoOpCommandEventSender();
         }
         return new LoggingCommandEventSender(
                 SECURITY_SENSITIVE_COMMANDS, SECURITY_SENSITIVE_HELLO_COMMANDS, description, commandListener,
-                operationContext, message, bsonOutput,
+                operationContext, message, commandDocument,
                 COMMAND_PROTOCOL_LOGGER, loggerSettings);
     }
 
     private ClusterId getClusterId() {
         return description.getConnectionId().getServerId().getClusterId();
     }
+
+    /**
+     * Creates a tracing span for the given command message.
+     * <p>
+     * The span is only created if tracing is enabled and the command is not security-sensitive.
+     * It attaches various tags to the span, such as database system, namespace, query summary, opcode,
+     * server address, port, server type, client and server connection IDs, and, if applicable,
+     * transaction number and session ID. For cursor fetching commands, the parent context is retrieved using the cursor ID.
+     * If command payload tracing is enabled, the command document is also attached as a tag.
+     *
+     * @param message          the command message to trace
+     * @param operationContext the operation context containing tracing and session information
+     * @return the created {@link Span}, or {@code null} if tracing is not enabled or the command is security-sensitive
+     */
+    @Nullable
+    private Span createTracingSpan(final CommandMessage message, final OperationContext operationContext) {
+        TracingManager tracingManager = operationContext.getTracingManager();
+        BsonDocument command = message.getCommand();
+        String commandName = command.getFirstKey();
+        if (!tracingManager.isEnabled()
+                || SECURITY_SENSITIVE_COMMANDS.contains(commandName)
+                || SECURITY_SENSITIVE_HELLO_COMMANDS.contains(commandName)) {
+            return null;
+        }
+
+        // Retrieving the appropriate parent context for the span.
+        TraceContext parentContext;
+        long cursorId = -1;
+        if (command.containsKey("getMore")) {
+            cursorId = command.getInt64("getMore").longValue();
+            parentContext = tracingManager.getCursorParentContext(cursorId);
+        } else {
+            parentContext = tracingManager.getParentContext(operationContext.getId());
+        }
+
+        Span span = tracingManager
+                .addSpan("Command " + commandName, parentContext)
+                .tag(SYSTEM, "mongodb")
+                .tag(NAMESPACE, message.getNamespace().getDatabaseName())
+                .tag(QUERY_SUMMARY, command.toString());
+
+        if (cursorId != -1) {
+            span.tag(CURSOR_ID, cursorId);
+        }
+
+        tagServerAndConnectionInfo(span, message);
+        tagSessionAndTransactionInfo(span, operationContext);
+
+        return span;
+    }
+
+    private void tagServerAndConnectionInfo(final Span span, final CommandMessage message) {
+        span.tag(SERVER_ADDRESS, serverId.getAddress().getHost())
+                .tag(SERVER_PORT, String.valueOf(serverId.getAddress().getPort()))
+                .tag(SERVER_TYPE, message.getSettings().getServerType().name())
+                .tag(CLIENT_CONNECTION_ID, this.description.getConnectionId().toString())
+                .tag(SERVER_CONNECTION_ID, String.valueOf(this.description.getConnectionId().getServerValue()));
+    }
+
+    private void tagSessionAndTransactionInfo(final Span span, final OperationContext operationContext) {
+        SessionContext sessionContext = operationContext.getSessionContext();
+        if (sessionContext.hasSession() && !sessionContext.isImplicitSession()) {
+            span.tag(TRANSACTION_NUMBER, String.valueOf(sessionContext.getTransactionNumber()))
+                    .tag(SESSION_ID, String.valueOf(sessionContext.getSessionId()
+                            .get(sessionContext.getSessionId().getFirstKey())
+                            .asBinary().asUuid()));
+        }
+    }
 }

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

@@ -78,7 +78,7 @@ class LoggingCommandEventSender implements CommandEventSender {
             @Nullable final CommandListener commandListener,
             final OperationContext operationContext,
             final CommandMessage message,
-            final ByteBufferBsonOutput bsonOutput,
+            final BsonDocument commandDocument,
             final StructuredLogger logger,
             final LoggerSettings loggerSettings) {
         this.description = description;
@@ -88,7 +88,7 @@ class LoggingCommandEventSender implements CommandEventSender {
         this.loggerSettings = loggerSettings;
         this.startTimeNanos = System.nanoTime();
         this.message = message;
-        this.commandDocument = message.getCommandDocument(bsonOutput);
+        this.commandDocument = commandDocument;
         this.commandName = commandDocument.getFirstKey();
         this.redactionRequired = securitySensitiveCommands.contains(commandName)
                 || (securitySensitiveHelloCommands.contains(commandName) && commandDocument.containsKey("speculativeAuthenticate"));