Polishing.

See: spring-projects#4216

Author: mp911de

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2022 the original author or authors.
+ * 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.
@@ -15,30 +15,32 @@
  */
 package org.springframework.data.mongodb.observability;
 
-import io.micrometer.common.KeyValue;
-import io.micrometer.common.KeyValues;
-
 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.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();
 
-		if (context.getCollectionName() != null) {
+		if (!ObjectUtils.isEmpty(context.getCollectionName())) {
 			keyValues = keyValues
 					.and(LowCardinalityCommandKeyNames.MONGODB_COLLECTION.withValue(context.getCollectionName()));
 		}
@@ -58,12 +60,18 @@ public KeyValues getHighCardinalityKeyValues(MongoHandlerContext context) {
 				HighCardinalityCommandKeyNames.MONGODB_COMMAND.withValue(context.getCommandStartedEvent().getCommandName()));
 	}
 
+	@Override
+	public String getContextualName(MongoHandlerContext context) {
+		return context.getContextualName();
+	}
+
 	/**
 	 * Extract connection details for a MongoDB connection into a {@link KeyValue}.
 	 *
 	 * @param event
 	 * @return
 	 */
+	@Nullable
 	private static KeyValue connectionTag(CommandStartedEvent event) {
 
 		ConnectionDescription connectionDescription = event.getConnectionDescription();

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2022 the original author or authors.
+ * 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.
@@ -28,15 +28,15 @@
  * @author Greg Turnquist
  * @since 4.0.0
  */
-class TraceRequestContext implements RequestContext {
+class MapRequestContext implements RequestContext {
 
 	private final Map<Object, Object> map;
 
-	public TraceRequestContext() {
+	public MapRequestContext() {
 		this(new HashMap<>());
 	}
 
-	public TraceRequestContext(Map<Object, Object> context) {
+	public MapRequestContext(Map<Object, Object> context) {
 		this.map = context;
 	}
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2022 the original author or authors.
+ * 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.
@@ -15,8 +15,6 @@
  */
 package org.springframework.data.mongodb.observability;
 
-import io.micrometer.observation.Observation;
-
 import java.util.Arrays;
 import java.util.LinkedHashSet;
 import java.util.Set;
@@ -30,17 +28,23 @@
 import com.mongodb.event.CommandStartedEvent;
 import com.mongodb.event.CommandSucceededEvent;
 
+import io.micrometer.observation.Observation;
+import io.micrometer.observation.transport.Kind;
+import io.micrometer.observation.transport.SenderContext;
+
 /**
  * A {@link Observation.Context} that contains MongoDB events.
  *
  * @author Marcin Grzejszczak
  * @author Greg Turnquist
  * @since 4.0.0
  */
-public class MongoHandlerContext extends Observation.Context {
+public class MongoHandlerContext extends SenderContext<Object> {
 
 	/**
-	 * @see https://docs.mongodb.com/manual/reference/command for the command reference
+	 * @see <a href=
+	 *      "https://docs.mongodb.com/manual/reference/command">https://docs.mongodb.com/manual/reference/command</a> for
+	 *      the command reference
 	 */
 	private static final Set<String> COMMANDS_WITH_COLLECTION_NAME = new LinkedHashSet<>(
 			Arrays.asList("aggregate", "count", "distinct", "mapReduce", "geoSearch", "delete", "find", "findAndModify",
@@ -55,7 +59,7 @@ public class MongoHandlerContext extends Observation.Context {
 	private CommandFailedEvent commandFailedEvent;
 
 	public MongoHandlerContext(CommandStartedEvent commandStartedEvent, RequestContext requestContext) {
-
+		super((carrier, key, value) -> {}, Kind.CLIENT);
 		this.commandStartedEvent = commandStartedEvent;
 		this.requestContext = requestContext;
 		this.collectionName = getCollectionName(commandStartedEvent);

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2022 the original author or authors.
+ * 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.
@@ -22,7 +22,7 @@
  * {@link ObservationConvention} for {@link MongoHandlerContext}.
  *
  * @author Greg Turnquist
- * @since 4.0.0
+ * @since 4
  */
 public interface MongoHandlerObservationConvention extends ObservationConvention<MongoHandlerContext> {
 

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

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