Open and close observation scopes for sync driver context propagation

The driver called observation.start()/stop() but never openScope()/
closeScope(). Without scopes, registry.getCurrentObservation() returned
null during MongoDB operations, breaking context propagation for any
downstream code (Spring interceptors, user observations, MDC logging).

For example, in withTransaction, a user observation created inside the
callback would attach to the Spring HTTP parent instead of the MongoDB
transaction span, because the transaction observation was never made "current" on the thread.
Added openScope()/closeScope() to the Span interface with scope lifecycle management in MongoClusterImpl (operation spans), InternalStreamConnection (command spans), and TransactionSpan.

Author: nhachicha

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

@@ -454,6 +454,9 @@ private <T> T sendAndReceiveInternal(final CommandMessage message, final Decoder
                             () -> getDescription().getServerAddress(),
                             () -> getDescription().getConnectionId()
                     );
+            if (tracingSpan != null) {
+                tracingSpan.openScope();
+            }
 
             boolean isLoggingCommandNeeded = isLoggingCommandNeeded();
             boolean isTracingCommandPayloadNeeded = tracingSpan != null && operationContext.getTracingManager().isCommandPayloadEnabled();
@@ -481,6 +484,8 @@ private <T> T sendAndReceiveInternal(final CommandMessage message, final Decoder
             } catch (Exception e) {
                 if (tracingSpan != null) {
                     tracingSpan.error(e);
+                    tracingSpan.closeScope();
+                    tracingSpan.end();
                 }
                 commandEventSender.sendFailedEvent(e);
                 throw e;
@@ -492,6 +497,7 @@ private <T> T sendAndReceiveInternal(final CommandMessage message, final Decoder
         } else {
             commandEventSender.sendSucceededEventForOneWayCommand();
             if (tracingSpan != null) {
+                tracingSpan.closeScope();
                 tracingSpan.end();
             }
             return null;
@@ -595,6 +601,7 @@ private <T> T receiveCommandMessageResponse(final Decoder<T> decoder, final Comm
             throw e;
         } finally {
             if (tracingSpan != null) {
+                tracingSpan.closeScope();
                 tracingSpan.end();
             }
         }

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

@@ -122,6 +122,8 @@ private static class MicrometerSpan implements Span {
         @Nullable
         private final MongoNamespace namespace;
         private final int queryTextLength;
+        @Nullable
+        private Observation.Scope scope;
 
         /**
          * Constructs a new {@link MicrometerSpan} instance with an associated Observation and MongoDB namespace.
@@ -136,6 +138,19 @@ private static class MicrometerSpan implements Span {
             this.queryTextLength = queryTextLength;
         }
 
+        @Override
+        public void openScope() {
+            this.scope = observation.openScope();
+        }
+
+        @Override
+        public void closeScope() {
+            if (scope != null) {
+                scope.close();
+                scope = null;
+            }
+        }
+
         @Override
         public void setQueryText(final BsonDocument commandDocument) {
             MongodbContext ctx = getMongodbContext();

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

@@ -45,6 +45,14 @@ public interface Span {
      * </p>
      */
     Span EMPTY = new Span() {
+        @Override
+        public void openScope() {
+        }
+
+        @Override
+        public void closeScope() {
+        }
+
         @Override
         public void setQueryText(final BsonDocument commandDocument) {
         }
@@ -79,6 +87,17 @@ public MongodbContext getMongodbContext() {
         }
     };
 
+    /**
+     * 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.
+     */
+    void openScope();
+
+    /**
+     * Closes the scope previously opened by {@link #openScope()}, restoring the previous observation.
+     */
+    void closeScope();
+
     /**
      * Sets the query text on the observation context from the given command document.
      * The document is converted to a JSON string and may be truncated based on configuration.

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

@@ -31,6 +31,7 @@ public class TransactionSpan {
 
     public TransactionSpan(final TracingManager tracingManager) {
         this.span = tracingManager.addTransactionSpan();
+        this.span.openScope();
     }
 
     /**
@@ -54,6 +55,7 @@ public void handleTransactionSpanError(final Throwable e) {
         }
 
         if (!isConvenientTransaction) {
+            span.closeScope();
             span.end();
         }
     }
@@ -67,6 +69,7 @@ public void finalizeTransactionSpan(final String status) {
         span.event(status);
         // clear previous commit error if any
         if (!isConvenientTransaction) {
+            span.closeScope();
             span.end();
         }
         reportedError = null; // clear previous commit error if any
@@ -82,6 +85,7 @@ public void spanFinalizing(final boolean cleanupTransactionContext) {
         if (reportedError != null) {
             span.error(reportedError);
         }
+        span.closeScope();
         span.end();
         reportedError = null;
         // Don't clean up transaction context if we're still retrying (we want the retries to fold under the original transaction span)

driver-sync/src/main/com/mongodb/client/internal/MongoClusterImpl.java

@@ -426,9 +426,11 @@ public <T> T execute(final ReadOperation<T, ?> operation, final ReadPreference r
                     .withSessionContext(new ClientSessionBinding.SyncClientSessionContext(actualClientSession, readConcern, implicitSession));
             Span span = operationContext.getTracingManager().createOperationSpan(
                     actualClientSession.getTransactionSpan(), operationContext, operation.getCommandName(), operation.getNamespace());
+            if (span != null) {
+                span.openScope();
+            }
             ReadBinding binding = getReadBinding(readPreference, actualClientSession, implicitSession);
 
-
             try {
                 if (actualClientSession.hasActiveTransaction() && !binding.getReadPreference().equals(primary())) {
                     throw new MongoClientException("Read preference in a transaction must be primary");
@@ -445,6 +447,7 @@ public <T> T execute(final ReadOperation<T, ?> operation, final ReadPreference r
             } finally {
                 binding.release();
                 if (span != null) {
+                    span.closeScope();
                     span.end();
                 }
             }
@@ -462,6 +465,9 @@ public <T> T execute(final WriteOperation<T> operation, final ReadConcern readCo
                     .withSessionContext(new ClientSessionBinding.SyncClientSessionContext(actualClientSession, readConcern, isImplicitSession(session)));
             Span span = operationContext.getTracingManager().createOperationSpan(
                     actualClientSession.getTransactionSpan(), operationContext, operation.getCommandName(), operation.getNamespace());
+            if (span != null) {
+                span.openScope();
+            }
             WriteBinding binding = getWriteBinding(actualClientSession, isImplicitSession(session));
 
             try {
@@ -477,6 +483,7 @@ public <T> T execute(final WriteOperation<T> operation, final ReadConcern readCo
             } finally {
                 binding.release();
                 if (span != null) {
+                    span.closeScope();
                     span.end();
                 }
             }

driver-sync/src/test/functional/com/mongodb/client/unified/Entities.java

@@ -49,6 +49,7 @@
 import com.mongodb.lang.Nullable;
 import com.mongodb.logging.TestLoggingInterceptor;
 import com.mongodb.observability.ObservabilitySettings;
+import io.micrometer.observation.Observation;
 import io.micrometer.observation.ObservationRegistry;
 import io.micrometer.tracing.test.reporter.inmemory.InMemoryOtelSetup;
 import org.bson.BsonArray;
@@ -113,6 +114,7 @@ public final class Entities {
     private final Map<String, TestLoggingInterceptor> clientLoggingInterceptors = new HashMap<>();
     private final Map<String, InMemoryOtelSetup.Builder.OtelBuildingBlocks> clientTracing = new HashMap<>();
     private final Set<InMemoryOtelSetup> inMemoryOTelInstances = new HashSet<>();
+    private final Set<Observation.Scope> observationScopes = new HashSet<>();
     private final Map<String, TestConnectionPoolListener> clientConnectionPoolListeners = new HashMap<>();
     private final Map<String, TestServerListener> clientServerListeners = new HashMap<>();
     private final Map<String, TestClusterListener> clientClusterListeners = new HashMap<>();
@@ -593,6 +595,12 @@ private void initClient(final BsonDocument entity, final String id,
                     .observabilitySettings(ObservabilitySettings.micrometerBuilder()
                     .observationRegistry(observationRegistry)
                     .enableCommandPayloadTracing(enableCommandPayload).build());
+
+            // Simulate what Spring Boot's observation does
+            // open a parent observation's scope before running the MongoDB operation
+            Observation parentObservation = Observation.createNotStarted("http.request", observationRegistry)
+                    .start();
+            observationScopes.add(parentObservation.openScope());
         }
 
         MongoClientSettings clientSettings = clientSettingsBuilder.build();
@@ -816,5 +824,6 @@ public void close() {
         clientLoggingInterceptors.values().forEach(TestLoggingInterceptor::close);
         threads.values().forEach(ExecutorService::shutdownNow);
         inMemoryOTelInstances.forEach(InMemoryOtelSetup::close);
+        observationScopes.forEach(Observation.Scope::close);
     }
 }