Improve configuration support for Observability integration

pom.xml

@@ -5,7 +5,7 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-mongodb-parent</artifactId>
-	<version>4.0.0-SNAPSHOT</version>
+	<version>4.0.0-mongo-micrometer-SNAPSHOT</version>
 	<packaging>pom</packaging>
 
 	<name>Spring Data MongoDB</name>

spring-data-mongodb-benchmarks/pom.xml

@@ -7,7 +7,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-mongodb-parent</artifactId>
-		<version>4.0.0-SNAPSHOT</version>
+		<version>4.0.0-mongo-micrometer-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-mongodb-distribution/pom.xml

@@ -15,7 +15,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-mongodb-parent</artifactId>
-		<version>4.0.0-SNAPSHOT</version>
+		<version>4.0.0-mongo-micrometer-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-mongodb/pom.xml

@@ -13,7 +13,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-mongodb-parent</artifactId>
-		<version>4.0.0-SNAPSHOT</version>
+		<version>4.0.0-mongo-micrometer-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/EnableMongoObservability.java

@@ -0,0 +1,16 @@
+package org.springframework.data.mongodb.observability;
+
+import java.lang.annotation.*;
+
+import org.springframework.context.annotation.Import;
+
+/**
+ * Annotation to active Spring Data MongoDB's usage of Micrometer's Observation API.
+ */
+@Inherited
+@Documented
+@Target(ElementType.TYPE)
+@Retention(RetentionPolicy.RUNTIME)
+@Import(MongoMetricsConfiguration.class)
+public @interface EnableMongoObservability {
+}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MongoMetricsConfiguration.java

@@ -0,0 +1,21 @@
+package org.springframework.data.mongodb.observability;
+
+import io.micrometer.observation.ObservationRegistry;
+import io.micrometer.tracing.Tracer;
+import org.springframework.context.annotation.Bean;
+
+/**
+ * Class to configure needed beans for MongoDB + Micrometer.
+ */
+public class MongoMetricsConfiguration {
+
+    @Bean
+    MongoObservationCommandListener mongoObservationCommandListener(ObservationRegistry registry) {
+        return new MongoObservationCommandListener(registry);
+    }
+
+    @Bean
+    MongoTracingObservationHandler mongoTracingObservationHandler(Tracer tracer) {
+        return new MongoTracingObservationHandler(tracer);
+    }
+}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MongoMetricsConfigurationHelper.java

@@ -0,0 +1,21 @@
+package org.springframework.data.mongodb.observability;
+
+import io.micrometer.observation.Observation;
+import io.micrometer.observation.ObservationRegistry;
+import io.micrometer.tracing.Tracer;
+
+import com.mongodb.client.SynchronousContextProvider;
+
+/**
+ * Helper functions to ease registration of Spring Data MongoDB's observability.
+ */
+public class MongoMetricsConfigurationHelper {
+
+	public static SynchronousContextProvider synchronousContextProvider(Tracer tracer, ObservationRegistry registry) {
+		return () -> new SynchronousTraceRequestContext(tracer).withObservation(Observation.start("name", registry));
+	}
+
+	public static void addObservationHandler(ObservationRegistry registry, Tracer tracer) {
+		registry.observationConfig().observationHandler(new MongoTracingObservationHandler(tracer));
+	}
+}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MongoMetricsReactiveConfigurationHelper.java

@@ -0,0 +1,24 @@
+package org.springframework.data.mongodb.observability;
+
+import io.micrometer.observation.Observation;
+import io.micrometer.observation.ObservationRegistry;
+import reactor.core.CoreSubscriber;
+import reactor.util.context.Context;
+
+import com.mongodb.reactivestreams.client.ReactiveContextProvider;
+
+/**
+ * Helper functions to ease registration of Spring Data MongoDB's observability.
+ */
+public class MongoMetricsReactiveConfigurationHelper {
+
+	public static ReactiveContextProvider reactiveContextProvider(ObservationRegistry registry) {
+		return subscriber -> {
+			if (subscriber instanceof CoreSubscriber<?> coreSubscriber) {
+				return new ReactiveTraceRequestContext(coreSubscriber.currentContext())
+						.withObservation(Observation.start("name", registry));
+			}
+			return new ReactiveTraceRequestContext(Context.empty()).withObservation(Observation.start("name", registry));
+		};
+	}
+}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MongoTracingObservationHandler.java

@@ -16,6 +16,7 @@
 package org.springframework.data.mongodb.observability;
 
 import io.micrometer.observation.Observation;
+import io.micrometer.observation.ObservationRegistry;
 import io.micrometer.tracing.Span;
 import io.micrometer.tracing.Tracer;
 import io.micrometer.tracing.handler.TracingObservationHandler;
@@ -49,6 +50,10 @@ public MongoTracingObservationHandler(Tracer tracer) {
 		this.tracer = tracer;
 	}
 
+	public void register(ObservationRegistry observationRegistry) {
+		observationRegistry.observationConfig().observationHandler(this);
+	}
+
 	@Override
 	public Tracer getTracer() {
 		return this.tracer;

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/ReactiveTraceRequestContext.java

@@ -0,0 +1,20 @@
+package org.springframework.data.mongodb.observability;
+
+import io.micrometer.observation.Observation;
+import reactor.util.context.ContextView;
+
+import java.util.Map;
+import java.util.stream.Collectors;
+
+class ReactiveTraceRequestContext extends TraceRequestContext {
+
+	ReactiveTraceRequestContext withObservation(Observation value) {
+
+		put(Observation.class, value);
+		return this;
+	}
+
+	ReactiveTraceRequestContext(ContextView context) {
+		super(context.stream().collect(Collectors.toConcurrentMap(Map.Entry::getKey, Map.Entry::getValue)));
+	}
+}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/SynchronousTraceRequestContext.java

@@ -0,0 +1,38 @@
+package org.springframework.data.mongodb.observability;
+
+import io.micrometer.observation.Observation;
+import io.micrometer.tracing.Span;
+import io.micrometer.tracing.TraceContext;
+import io.micrometer.tracing.Tracer;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+class SynchronousTraceRequestContext extends TraceRequestContext {
+
+	SynchronousTraceRequestContext(Tracer tracer) {
+		super(context(tracer));
+	}
+
+	SynchronousTraceRequestContext withObservation(Observation value) {
+
+		put(Observation.class, value);
+		return this;
+	}
+
+	private static Map<Object, Object> context(Tracer tracer) {
+
+		Map<Object, Object> map = new ConcurrentHashMap<>();
+
+		Span currentSpan = tracer.currentSpan();
+
+		if (currentSpan == null) {
+			return map;
+		}
+
+		map.put(Span.class, currentSpan);
+		map.put(TraceContext.class, currentSpan.context());
+
+		return map;
+	}
+}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/TraceRequestContext.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2013-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.observability;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.stream.Stream;
+
+import com.mongodb.RequestContext;
+
+/**
+ * A {@link Map}-based {@link RequestContext}.
+ *
+ * @author Marcin Grzejszczak
+ * @author Greg Turnquist
+ * @since 4.0.0
+ */
+class TraceRequestContext implements RequestContext {
+
+	private final Map<Object, Object> map;
+
+	public TraceRequestContext() {
+		this(new HashMap<>());
+	}
+
+	public TraceRequestContext(Map<Object, Object> context) {
+		this.map = context;
+	}
+
+	@Override
+	public <T> T get(Object key) {
+		return (T) map.get(key);
+	}
+
+	@Override
+	public boolean hasKey(Object key) {
+		return map.containsKey(key);
+	}
+
+	@Override
+	public boolean isEmpty() {
+		return map.isEmpty();
+	}
+
+	@Override
+	public void put(Object key, Object value) {
+		map.put(key, value);
+	}
+
+	@Override
+	public void delete(Object key) {
+		map.remove(key);
+	}
+
+	@Override
+	public int size() {
+		return map.size();
+	}
+
+	@Override
+	public Stream<Map.Entry<Object, Object>> stream() {
+		return map.entrySet().stream();
+	}
+}

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/observability/MongoObservationCommandListenerForTracingTests.java

@@ -77,10 +77,10 @@ void setup() {
 	void successfullyCompletedCommandShouldCreateSpanWhenParentSampleInRequestContext() {
 
 		// given
-		TestRequestContext testRequestContext = createTestRequestContextWithParentObservationAndStartIt();
+		TraceRequestContext traceRequestContext = createTestRequestContextWithParentObservationAndStartIt();
 
 		// when
-		commandStartedAndSucceeded(testRequestContext);
+		commandStartedAndSucceeded(traceRequestContext);
 
 		// then
 		assertThatMongoSpanIsClientWithTags().hasIpThatIsBlank().hasPortThatIsNotSet();
@@ -91,10 +91,10 @@ void successfullyCompletedCommandShouldCreateSpanWithAddressInfoWhenParentSample
 
 		// given
 		handler.setSetRemoteIpAndPortEnabled(true);
-		TestRequestContext testRequestContext = createTestRequestContextWithParentObservationAndStartIt();
+		TraceRequestContext traceRequestContext = createTestRequestContextWithParentObservationAndStartIt();
 
 		// when
-		commandStartedAndSucceeded(testRequestContext);
+		commandStartedAndSucceeded(traceRequestContext);
 
 		// then
 		assertThatMongoSpanIsClientWithTags().hasIpThatIsNotBlank().hasPortThatIsSet();
@@ -104,28 +104,28 @@ void successfullyCompletedCommandShouldCreateSpanWithAddressInfoWhenParentSample
 	void commandWithErrorShouldCreateTimerWhenParentSampleInRequestContext() {
 
 		// given
-		TestRequestContext testRequestContext = createTestRequestContextWithParentObservationAndStartIt();
+		TraceRequestContext traceRequestContext = createTestRequestContextWithParentObservationAndStartIt();
 
 		// when
-		listener.commandStarted(new CommandStartedEvent(testRequestContext, 0, //
+		listener.commandStarted(new CommandStartedEvent(traceRequestContext, 0, //
 				new ConnectionDescription( //
 						new ServerId( //
 								new ClusterId("description"), //
 								new ServerAddress("localhost", 1234))), //
 				"database", "insert", //
 				new BsonDocument("collection", new BsonString("user"))));
 		listener.commandFailed( //
-				new CommandFailedEvent(testRequestContext, 0, null, "insert", 0, new IllegalAccessException()));
+				new CommandFailedEvent(traceRequestContext, 0, null, "insert", 0, new IllegalAccessException()));
 
 		// then
 		assertThatMongoSpanIsClientWithTags().assertThatThrowable().isInstanceOf(IllegalAccessException.class);
 	}
 
 	/**
-	 * Create a parent {@link Observation} then wrap it inside a {@link TestRequestContext}.
+	 * Create a parent {@link Observation} then wrap it inside a {@link TraceRequestContext}.
 	 */
 	@NotNull
-	private TestRequestContext createTestRequestContextWithParentObservationAndStartIt() {
+	private TraceRequestContext createTestRequestContextWithParentObservationAndStartIt() {
 
 		Observation parent = Observation.start("name", observationRegistry);
 		return TestRequestContext.withObservation(parent);
@@ -134,21 +134,21 @@ private TestRequestContext createTestRequestContextWithParentObservationAndStart
 	/**
 	 * Execute MongoDB's {@link com.mongodb.event.CommandListener#commandStarted(CommandStartedEvent)} and
 	 * {@link com.mongodb.event.CommandListener#commandSucceeded(CommandSucceededEvent)} operations against the
-	 * {@link TestRequestContext} in order to inject some test data.
+	 * {@link TraceRequestContext} in order to inject some test data.
 	 *
-	 * @param testRequestContext
+	 * @param traceRequestContext
 	 */
-	private void commandStartedAndSucceeded(TestRequestContext testRequestContext) {
+	private void commandStartedAndSucceeded(TraceRequestContext traceRequestContext) {
 
-		listener.commandStarted(new CommandStartedEvent(testRequestContext, 0, //
+		listener.commandStarted(new CommandStartedEvent(traceRequestContext, 0, //
 				new ConnectionDescription( //
 						new ServerId( //
 								new ClusterId("description"), //
 								new ServerAddress("localhost", 1234))), //
 				"database", "insert", //
 				new BsonDocument("collection", new BsonString("user"))));
 
-		listener.commandSucceeded(new CommandSucceededEvent(testRequestContext, 0, null, "insert", null, 0));
+		listener.commandSucceeded(new CommandSucceededEvent(traceRequestContext, 0, null, "insert", null, 0));
 	}
 
 	/**