Add tracing support using Micrometer

buildSrc/src/main/kotlin/project/Companion.kt

@@ -23,4 +23,4 @@ import org.gradle.kotlin.dsl.getByType
 internal val Project.libs: LibrariesForLibs
     get() = extensions.getByType()
 
-internal const val DEFAULT_JAVA_VERSION = 17
+const val DEFAULT_JAVA_VERSION = 17

config/checkstyle/suppressions.xml

@@ -60,6 +60,7 @@
 
     <!-- Allow printStackTrace in this file -->
     <suppress checks="Regexp" files="CallbackResultHolder"/>
+    <suppress checks="Regexp" files="MicrometerTracer"/>
 
     <!--Do not check documentation tests classes -->
     <suppress checks="Javadoc*" files=".*documentation.*"/>

driver-core/build.gradle.kts

@@ -55,6 +55,9 @@ dependencies {
     optionalImplementation(libs.snappy.java)
     optionalImplementation(libs.zstd.jni)
 
+    optionalImplementation(platform(libs.micrometer.observation.bom))
+    optionalImplementation(libs.micrometer.observation)
+
     testImplementation(project(path = ":bson", configuration = "testArtifacts"))
     testImplementation(libs.reflections)
     testImplementation(libs.netty.tcnative.boringssl.static)

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.VisibleForTesting;
 import com.mongodb.lang.Nullable;
 import com.mongodb.spi.dns.DnsClient;
 import com.mongodb.spi.dns.InetAddressResolver;
+import io.micrometer.observation.ObservationRegistry;
 import org.bson.UuidRepresentation;
 import org.bson.codecs.BsonCodecProvider;
 import org.bson.codecs.BsonValueCodecProvider;
@@ -57,6 +59,7 @@
 import static com.mongodb.assertions.Assertions.isTrueArgument;
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.internal.TimeoutSettings.convertAndValidateTimeout;
+import static com.mongodb.internal.VisibleForTesting.AccessModifier.PRIVATE;
 import static java.util.Arrays.asList;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static org.bson.codecs.configuration.CodecRegistries.fromProviders;
@@ -119,6 +122,15 @@ public final class MongoClientSettings {
     @Nullable
     private final Long timeoutMS;
 
+    @VisibleForTesting(otherwise = PRIVATE)
+    // If set, this will enable/disable tracing even when an observationRegistry has been passed.
+    public static final String ENV_OTEL_ENABLED = "OTEL_JAVA_INSTRUMENTATION_MONGODB_ENABLED";
+    @VisibleForTesting(otherwise = PRIVATE)
+    // If set, this will truncate the command payload captured in the tracing span to the specified length.
+    public static final String ENV_OTEL_QUERY_TEXT_MAX_LENGTH = "OTEL_JAVA_INSTRUMENTATION_MONGODB_QUERY_TEXT_MAX_LENGTH";
+    private final ObservationRegistry observationRegistry;
+    private final boolean enableCommandPayloadTracing;
+
     /**
      * Gets the default codec registry.  It includes the following providers:
      *
@@ -238,6 +250,8 @@ public static final class Builder {
         private ContextProvider contextProvider;
         private DnsClient dnsClient;
         private InetAddressResolver inetAddressResolver;
+        private ObservationRegistry observationRegistry;
+        private boolean enableCommandPayloadTracing;
 
         private Builder() {
         }
@@ -275,6 +289,8 @@ private Builder(final MongoClientSettings settings) {
             if (settings.heartbeatSocketTimeoutSetExplicitly) {
                 heartbeatSocketTimeoutMS = settings.heartbeatSocketSettings.getReadTimeout(MILLISECONDS);
             }
+            observationRegistry = settings.observationRegistry;
+            enableCommandPayloadTracing = settings.enableCommandPayloadTracing;
         }
 
         /**
@@ -723,6 +739,47 @@ Builder heartbeatSocketTimeoutMS(final int heartbeatSocketTimeoutMS) {
             return this;
         }
 
+        /**
+         * Sets the observation registry to use for creating tracing Spans for operations, commands and transactions.
+         *
+         * <p> If set the environment variable {@value ENV_OTEL_ENABLED} is used to enable or disable the creation of tracing spans.
+         *
+         * <p> If set the environment variable {@value ENV_OTEL_QUERY_TEXT_MAX_LENGTH} is used to determine the maximum length
+         * of command payloads captured in tracing spans. If the environment variable is not set, the entire command payloads is
+         * captured.
+         *
+         * @param observationRegistry the observation registry
+         * @return this
+         * @since 5.7
+         */
+        @Alpha(Reason.CLIENT)
+        public Builder observationRegistry(final ObservationRegistry observationRegistry) {
+            this.observationRegistry = observationRegistry;
+            return this;
+        }
+
+        /**
+         * Sets the observation registry to use for creating tracing Spans for operations, commands and transactions.
+         *
+         * <p> If set the environment variable {@value ENV_OTEL_ENABLED} is used to enable or disable the creation of tracing spans.
+         *
+         * <p> If set the environment variable {@value ENV_OTEL_QUERY_TEXT_MAX_LENGTH} is used to determine the maximum length
+         * of command payloads captured in tracing spans. If the environment variable is not set, the entire command payloads is
+         * captured.
+         *
+         * @param observationRegistry the observation registry
+         * @param enableCommandPayload whether command payloads should be captured in tracing spans. This may have performance
+         *                             implications so should be used with care.
+         * @return this
+         * @since 5.7
+         */
+        @Alpha(Reason.CLIENT)
+        public Builder observationRegistry(final ObservationRegistry observationRegistry, final boolean enableCommandPayload) {
+            this.observationRegistry = observationRegistry;
+            this.enableCommandPayloadTracing = enableCommandPayload;
+            return this;
+        }
+
         /**
          * Build an instance of {@code MongoClientSettings}.
          *
@@ -1040,6 +1097,27 @@ public ContextProvider getContextProvider() {
         return contextProvider;
     }
 
+    /** Get the ObservationRegistry to create tracing Spans for operations, commands and transactions.
+     *
+     * @return the configured ObservationRegistry
+     * @since 5.7
+     */
+    @Nullable
+    @Alpha(Reason.CLIENT)
+    public ObservationRegistry getObservationRegistry() {
+        return observationRegistry;
+    }
+
+    /** Returns true if command payload tracing is enabled.
+     *
+     * @return the enableCommandPayloadTracing value
+     * @since 5.7
+     */
+    @Alpha(Reason.CLIENT)
+    public boolean isCommandPayloadTracingEnabled() {
+        return enableCommandPayloadTracing;
+    }
+
     @Override
     public boolean equals(final Object o) {
         if (this == o) {
@@ -1156,5 +1234,8 @@ private MongoClientSettings(final Builder builder) {
         heartbeatConnectTimeoutSetExplicitly = builder.heartbeatConnectTimeoutMS != 0;
         contextProvider = builder.contextProvider;
         timeoutMS = builder.timeoutMS;
+
+        observationRegistry = builder.observationRegistry;
+        enableCommandPayloadTracing = builder.enableCommandPayloadTracing;
     }
 }

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

@@ -51,6 +51,7 @@
 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.lang.Nullable;
 import org.bson.BsonBinaryReader;
 import org.bson.BsonDocument;
@@ -94,6 +95,8 @@
 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.MongodbObservation.HighCardinalityKeyNames.QUERY_TEXT;
+import static com.mongodb.internal.tracing.MongodbObservation.LowCardinalityKeyNames.RESPONSE_STATUS_CODE;
 import static java.util.Arrays.asList;
 
 /**
@@ -432,22 +435,59 @@ 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;
         try (ByteBufferBsonOutput bsonOutput = new ByteBufferBsonOutput(this)) {
             message.encode(bsonOutput, operationContext);
-            commandEventSender = createCommandEventSender(message, bsonOutput, operationContext);
-            commandEventSender.sendStartedEvent();
+            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()
+                    );
+
+            boolean isLoggingCommandNeeded = isLoggingCommandNeeded();
+            boolean isTracingCommandPayloadNeeded = tracingSpan != null && operationContext.getTracingManager().isCommandPayloadEnabled();
+
+            // Only hydrate the command document if necessary
+            BsonDocument commandDocument = null;
+            if (isLoggingCommandNeeded || isTracingCommandPayloadNeeded) {
+                commandDocument = message.getCommandDocument(bsonOutput);
+            }
+            if (isLoggingCommandNeeded) {
+                commandEventSender = new LoggingCommandEventSender(
+                        SECURITY_SENSITIVE_COMMANDS, SECURITY_SENSITIVE_HELLO_COMMANDS, description, commandListener,
+                        operationContext, message, commandDocument,
+                        COMMAND_PROTOCOL_LOGGER, loggerSettings);
+                commandEventSender.sendStartedEvent();
+            } else {
+                commandEventSender = new NoOpCommandEventSender();
+            }
+            if (isTracingCommandPayloadNeeded) {
+                tracingSpan.tagHighCardinality(QUERY_TEXT.asString(), commandDocument);
+            }
+
             try {
                 sendCommandMessage(message, bsonOutput, operationContext);
             } catch (Exception e) {
+                if (tracingSpan != null) {
+                    tracingSpan.error(e);
+                }
                 commandEventSender.sendFailedEvent(e);
                 throw e;
             }
         }
 
         if (message.isResponseExpected()) {
-            return receiveCommandMessageResponse(decoder, commandEventSender, operationContext);
+            return receiveCommandMessageResponse(decoder, commandEventSender, operationContext, tracingSpan);
         } else {
             commandEventSender.sendSucceededEventForOneWayCommand();
+            if (tracingSpan != null) {
+                tracingSpan.end();
+            }
             return null;
         }
     }
@@ -466,7 +506,7 @@ public <T> void send(final CommandMessage message, final Decoder<T> decoder, fin
     @Override
     public <T> T receive(final Decoder<T> decoder, final OperationContext operationContext) {
         isTrue("Response is expected", hasMoreToCome);
-        return receiveCommandMessageResponse(decoder, new NoOpCommandEventSender(), operationContext);
+        return receiveCommandMessageResponse(decoder, new NoOpCommandEventSender(), operationContext, null);
     }
 
     @Override
@@ -512,7 +552,7 @@ private void trySendMessage(final CommandMessage message, final ByteBufferBsonOu
     }
 
     private <T> T receiveCommandMessageResponse(final Decoder<T> decoder, final CommandEventSender commandEventSender,
-            final OperationContext operationContext) {
+            final OperationContext operationContext, @Nullable final Span tracingSpan) {
         boolean commandSuccessful = false;
         try (ResponseBuffers responseBuffers = receiveResponseBuffers(operationContext)) {
             updateSessionContext(operationContext.getSessionContext(), responseBuffers);
@@ -537,7 +577,17 @@ private <T> T receiveCommandMessageResponse(final Decoder<T> decoder, final Comm
             if (!commandSuccessful) {
                 commandEventSender.sendFailedEvent(e);
             }
+            if (tracingSpan != null) {
+                if (e instanceof MongoCommandException) {
+                    tracingSpan.tagLowCardinality(RESPONSE_STATUS_CODE.withValue(String.valueOf(((MongoCommandException) e).getErrorCode())));
+                }
+                tracingSpan.error(e);
+            }
             throw e;
+        } finally {
+            if (tracingSpan != null) {
+                tracingSpan.end();
+            }
         }
     }
 
@@ -553,7 +603,18 @@ private <T> void sendAndReceiveAsyncInternal(final CommandMessage message, final
 
         try {
             message.encode(bsonOutput, operationContext);
-            CommandEventSender commandEventSender = createCommandEventSender(message, bsonOutput, operationContext);
+
+            CommandEventSender commandEventSender;
+            if (isLoggingCommandNeeded()) {
+                BsonDocument commandDocument = message.getCommandDocument(bsonOutput);
+                commandEventSender = new LoggingCommandEventSender(
+                        SECURITY_SENSITIVE_COMMANDS, SECURITY_SENSITIVE_HELLO_COMMANDS, description, commandListener,
+                        operationContext, message, commandDocument,
+                        COMMAND_PROTOCOL_LOGGER, loggerSettings);
+            } else {
+                commandEventSender = new NoOpCommandEventSender();
+            }
+
             commandEventSender.sendStartedEvent();
             Compressor localSendCompressor = sendCompressor;
             if (localSendCompressor == null || SECURITY_SENSITIVE_COMMANDS.contains(message.getCommandDocument(bsonOutput).getFirstKey())) {
@@ -952,19 +1013,13 @@ 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,
-                                                        final OperationContext operationContext) {
+    private boolean isLoggingCommandNeeded() {
         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,
-                COMMAND_PROTOCOL_LOGGER, loggerSettings);
+        return recordEverything || (!isMonitoringConnection && opened() && authenticated.get() && listensOrLogs);
     }
 
     private ClusterId getClusterId() {
         return description.getConnectionId().getServerId().getClusterId();
     }
+
 }

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"));