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>
 
@@ -41,7 +41,7 @@
 				<executions>
 					<execution>
 						<id>generate-metrics-metadata</id>
-						<phase>prepare-package</phase>
+						<phase>generate-resources</phase>
 						<goals>
 							<goal>java</goal>
 						</goals>
@@ -52,7 +52,7 @@
 					</execution>
 					<execution>
 						<id>generate-tracing-metadata</id>
-						<phase>prepare-package</phase>
+						<phase>generate-resources</phase>
 						<goals>
 							<goal>java</goal>
 						</goals>

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/ContextProviderFactory.java

@@ -0,0 +1,139 @@
+/*
+ * Copyright 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.Map;
+import java.util.Map.Entry;
+import java.util.stream.Collectors;
+
+import org.reactivestreams.Subscriber;
+import org.springframework.data.repository.util.ReactiveWrappers;
+import org.springframework.data.repository.util.ReactiveWrappers.ReactiveLibrary;
+import org.springframework.util.ClassUtils;
+
+import com.mongodb.ContextProvider;
+import com.mongodb.RequestContext;
+import com.mongodb.client.SynchronousContextProvider;
+import com.mongodb.reactivestreams.client.ReactiveContextProvider;
+
+import io.micrometer.observation.Observation;
+import io.micrometer.observation.ObservationRegistry;
+import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
+import reactor.core.CoreSubscriber;
+
+/**
+ * Factory to create a {@link ContextProvider} to propagate the request context across tasks. Requires either
+ * {@link SynchronousContextProvider} or {@link ReactiveContextProvider} to be present.
+ *
+ * @author Mark Paluch
+ * @since 3.0
+ */
+public class ContextProviderFactory {
+
+	private static final boolean SYNCHRONOUS_PRESENT = ClassUtils
+			.isPresent("com.mongodb.client.SynchronousContextProvider", ContextProviderFactory.class.getClassLoader());
+
+	private static final boolean REACTIVE_PRESENT = ClassUtils.isPresent(
+			"com.mongodb.reactivestreams.client.ReactiveContextProvider", ContextProviderFactory.class.getClassLoader())
+			&& ReactiveWrappers.isAvailable(ReactiveLibrary.PROJECT_REACTOR);
+
+	/**
+	 * Create a {@link ContextProvider} given {@link ObservationRegistry}. The factory method attempts to create a
+	 * {@link ContextProvider} that is capable to propagate request contexts across imperative or reactive usage,
+	 * depending on their class path presence.
+	 *
+	 * @param observationRegistry must not be {@literal null}.
+	 * @return
+	 */
+	public static ContextProvider create(ObservationRegistry observationRegistry) {
+
+		if (SYNCHRONOUS_PRESENT && REACTIVE_PRESENT) {
+			return new CompositeContextProvider(observationRegistry);
+		}
+
+		if (SYNCHRONOUS_PRESENT) {
+			return new DefaultSynchronousContextProvider(observationRegistry);
+		}
+
+		if (REACTIVE_PRESENT) {
+			return DefaultReactiveContextProvider.INSTANCE;
+		}
+
+		throw new IllegalStateException(
+				"Cannot create ContextProvider. Neither SynchronousContextProvider nor ReactiveContextProvider is on the class path.");
+	}
+
+	record DefaultSynchronousContextProvider(
+			ObservationRegistry observationRegistry) implements SynchronousContextProvider {
+
+		@Override
+		public RequestContext getContext() {
+
+			MapRequestContext requestContext = new MapRequestContext();
+
+			Observation currentObservation = observationRegistry.getCurrentObservation();
+			if (currentObservation != null) {
+				requestContext.put(Observation.class, currentObservation);
+			}
+
+			return requestContext;
+		}
+
+	}
+
+	enum DefaultReactiveContextProvider implements ReactiveContextProvider {
+
+		INSTANCE;
+
+		@Override
+		public RequestContext getContext(Subscriber<?> subscriber) {
+
+			if (subscriber instanceof CoreSubscriber<?> cs) {
+
+				Map<Object, Object> map = cs.currentContext().stream()
+						.collect(Collectors.toConcurrentMap(Entry::getKey, Entry::getValue));
+				if (map.containsKey(ObservationThreadLocalAccessor.KEY)) {
+					map.put(Observation.class, map.get(ObservationThreadLocalAccessor.KEY));
+				}
+
+				return new MapRequestContext(map);
+			}
+
+			return new MapRequestContext();
+		}
+	}
+
+	record CompositeContextProvider(DefaultSynchronousContextProvider synchronousContextProvider)
+			implements
+				SynchronousContextProvider,
+				ReactiveContextProvider {
+
+		CompositeContextProvider(ObservationRegistry observationRegistry) {
+			this(new DefaultSynchronousContextProvider(observationRegistry));
+		}
+
+		@Override
+		public RequestContext getContext() {
+			return synchronousContextProvider.getContext();
+		}
+
+		@Override
+		public RequestContext getContext(Subscriber<?> subscriber) {
+			return DefaultReactiveContextProvider.INSTANCE.getContext(subscriber);
+		}
+	}
+
+}

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

@@ -15,34 +15,75 @@
  */
 package org.springframework.data.mongodb.observability;
 
-import io.micrometer.common.KeyValue;
-import io.micrometer.common.KeyValues;
+import java.net.InetSocketAddress;
 
 import org.springframework.data.mongodb.observability.MongoObservation.HighCardinalityCommandKeyNames;
 import org.springframework.data.mongodb.observability.MongoObservation.LowCardinalityCommandKeyNames;
+import org.springframework.lang.Nullable;
+import org.springframework.util.ObjectUtils;
 
+import com.mongodb.ConnectionString;
+import com.mongodb.ServerAddress;
 import com.mongodb.connection.ConnectionDescription;
 import com.mongodb.connection.ConnectionId;
 import com.mongodb.event.CommandStartedEvent;
 
+import io.micrometer.common.KeyValue;
+import io.micrometer.common.KeyValues;
+
 /**
  * Default {@link MongoHandlerObservationConvention} implementation.
  *
  * @author Greg Turnquist
- * @since 4.0.0
+ * @since 4
  */
-public class DefaultMongoHandlerObservationConvention implements MongoHandlerObservationConvention {
+class DefaultMongoHandlerObservationConvention implements MongoHandlerObservationConvention {
 
 	@Override
 	public KeyValues getLowCardinalityKeyValues(MongoHandlerContext context) {
 
-		KeyValues keyValues = KeyValues.empty();
+		KeyValues keyValues = KeyValues.of(LowCardinalityCommandKeyNames.DB_SYSTEM.withValue("mongodb"));
+
+		ConnectionString connectionString = context.getConnectionString();
+		if (connectionString != null) {
 
-		if (context.getCollectionName() != null) {
+			keyValues = keyValues
+					.and(LowCardinalityCommandKeyNames.DB_CONNECTION_STRING.withValue(connectionString.getConnectionString()));
+
+			String user = connectionString.getUsername();
+
+			if (!ObjectUtils.isEmpty(user)) {
+				keyValues = keyValues.and(LowCardinalityCommandKeyNames.DB_USER.withValue(user));
+			}
+
+		}
+
+		if (!ObjectUtils.isEmpty(context.getDatabaseName())) {
+			keyValues = keyValues.and(LowCardinalityCommandKeyNames.DB_NAME.withValue(context.getDatabaseName()));
+		}
+
+		if (!ObjectUtils.isEmpty(context.getCollectionName())) {
 			keyValues = keyValues
 					.and(LowCardinalityCommandKeyNames.MONGODB_COLLECTION.withValue(context.getCollectionName()));
 		}
 
+		ServerAddress serverAddress = context.getCommandStartedEvent().getConnectionDescription().getServerAddress();
+
+		if (serverAddress != null) {
+
+			keyValues = keyValues.and(LowCardinalityCommandKeyNames.NET_TRANSPORT.withValue("IP.TCP"),
+					LowCardinalityCommandKeyNames.NET_PEER_ADDR.withValue(serverAddress.getHost()));
+
+			InetSocketAddress socketAddress = serverAddress.getSocketAddress();
+
+			if (socketAddress != null) {
+
+				keyValues = keyValues.and(LowCardinalityCommandKeyNames.NET_PEER_NAME.withValue(socketAddress.getHostName()),
+						LowCardinalityCommandKeyNames.NET_PEER_PORT.withValue("" + socketAddress.getPort()));
+			}
+
+		}
+
 		KeyValue connectionTag = connectionTag(context.getCommandStartedEvent());
 		if (connectionTag != null) {
 			keyValues = keyValues.and(connectionTag);
@@ -54,8 +95,12 @@ public KeyValues getLowCardinalityKeyValues(MongoHandlerContext context) {
 	@Override
 	public KeyValues getHighCardinalityKeyValues(MongoHandlerContext context) {
 
-		return KeyValues.of(
-				HighCardinalityCommandKeyNames.MONGODB_COMMAND.withValue(context.getCommandStartedEvent().getCommandName()));
+		return KeyValues.of(HighCardinalityCommandKeyNames.MONGODB_COMMAND.withValue(context.getCommandName()));
+	}
+
+	@Override
+	public String getContextualName(MongoHandlerContext context) {
+		return context.getContextualName();
 	}
 
 	/**
@@ -64,6 +109,7 @@ public KeyValues getHighCardinalityKeyValues(MongoHandlerContext context) {
 	 * @param event
 	 * @return
 	 */
+	@Nullable
 	private static KeyValue connectionTag(CommandStartedEvent event) {
 
 		ConnectionDescription connectionDescription = event.getConnectionDescription();

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/observability/TestRequestContext.java → spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MapRequestContext.java

@@ -15,24 +15,30 @@
  */
 package org.springframework.data.mongodb.observability;
 
-import io.micrometer.observation.Observation;
-
 import java.util.HashMap;
 import java.util.Map;
 import java.util.stream.Stream;
 
 import com.mongodb.RequestContext;
 
 /**
- * A {@link Map}-based {@link RequestContext}. (For test purposes only).
+ * A {@link Map}-based {@link RequestContext}.
  *
  * @author Marcin Grzejszczak
  * @author Greg Turnquist
  * @since 4.0.0
  */
-class TestRequestContext implements RequestContext {
+class MapRequestContext implements RequestContext {
+
+	private final Map<Object, Object> map;
+
+	public MapRequestContext() {
+		this(new HashMap<>());
+	}
 
-	private final Map<Object, Object> map = new HashMap<>();
+	public MapRequestContext(Map<Object, Object> context) {
+		this.map = context;
+	}
 
 	@Override
 	public <T> T get(Object key) {
@@ -68,11 +74,4 @@ public int size() {
 	public Stream<Map.Entry<Object, Object>> stream() {
 		return map.entrySet().stream();
 	}
-
-	static TestRequestContext withObservation(Observation value) {
-
-		TestRequestContext testRequestContext = new TestRequestContext();
-		testRequestContext.put(Observation.class, value);
-		return testRequestContext;
-	}
 }