From d4daa305a88520b6cc8b47568d73c2898122b2a0 Mon Sep 17 00:00:00 2001 From: Marcin Grzejszczak Date: Fri, 21 Oct 2022 17:52:49 +0200 Subject: [PATCH] Align the context propagation entries with the rest of the portfolio. Closes #4218 --- .../observability/ContextProviderFactory.java | 3 ++- .../MongoObservationCommandListener.java | 20 +++++++++++-------- .../ReactiveIntegrationTests.java | 2 +- 3 files changed, 15 insertions(+), 10 deletions(-) diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/ContextProviderFactory.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/ContextProviderFactory.java index 87ee50483..b878a8e65 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/ContextProviderFactory.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/ContextProviderFactory.java @@ -85,7 +85,8 @@ public class ContextProviderFactory { Observation currentObservation = observationRegistry.getCurrentObservation(); if (currentObservation != null) { - requestContext.put(Observation.class, currentObservation); + // Aligned with ObservationThreadLocalAccessor.KEY + requestContext.put("micrometer.observation", currentObservation); } return requestContext; diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MongoObservationCommandListener.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MongoObservationCommandListener.java index 0a344570b..de0c13bce 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MongoObservationCommandListener.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MongoObservationCommandListener.java @@ -42,6 +42,11 @@ public class MongoObservationCommandListener implements CommandListener { private static final Log log = LogFactory.getLog(MongoObservationCommandListener.class); + /** + * Aligns with ObservationThreadLocalAccessor.KEY. + */ + private static final String MICROMETER_OBSERVATION_KEY = "micrometer.observation"; + private final ObservationRegistry observationRegistry; private final @Nullable ConnectionString connectionString; @@ -114,8 +119,7 @@ public class MongoObservationCommandListener implements CommandListener { observation.start(); - requestContext.put(Observation.class, observation); - requestContext.put(MongoHandlerContext.class, observationContext); + requestContext.put(MICROMETER_OBSERVATION_KEY, observation); if (log.isDebugEnabled()) { log.debug( @@ -132,12 +136,12 @@ public class MongoObservationCommandListener implements CommandListener { return; } - Observation observation = requestContext.getOrDefault(Observation.class, null); + Observation observation = requestContext.getOrDefault(MICROMETER_OBSERVATION_KEY, null); if (observation == null) { return; } - MongoHandlerContext context = requestContext.get(MongoHandlerContext.class); + MongoHandlerContext context = (MongoHandlerContext) observation.getContext(); context.setCommandSucceededEvent(event); if (log.isDebugEnabled()) { @@ -156,12 +160,12 @@ public class MongoObservationCommandListener implements CommandListener { return; } - Observation observation = requestContext.getOrDefault(Observation.class, null); + Observation observation = requestContext.getOrDefault(MICROMETER_OBSERVATION_KEY, null); if (observation == null) { return; } - MongoHandlerContext context = requestContext.get(MongoHandlerContext.class); + MongoHandlerContext context = (MongoHandlerContext) observation.getContext(); context.setCommandFailedEvent(event); if (log.isDebugEnabled()) { @@ -181,7 +185,7 @@ public class MongoObservationCommandListener implements CommandListener { @Nullable private static Observation observationFromContext(RequestContext context) { - Observation observation = context.getOrDefault(Observation.class, null); + Observation observation = context.getOrDefault(MICROMETER_OBSERVATION_KEY, null); if (observation != null) { @@ -192,7 +196,7 @@ public class MongoObservationCommandListener implements CommandListener { } if (log.isDebugEnabled()) { - log.debug("No observation was found - will not create any child spans"); + log.debug("No observation was found - will not create any child observations"); } return null; diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/observability/ReactiveIntegrationTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/observability/ReactiveIntegrationTests.java index 980c97554..c00c7d8ba 100644 --- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/observability/ReactiveIntegrationTests.java +++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/observability/ReactiveIntegrationTests.java @@ -73,7 +73,7 @@ public class ReactiveIntegrationTests extends SampleTestRunner { .verifyComplete(); repository.findByLastname("Matthews") // - .contextWrite(Context.of(Observation.class, intermediate)) // + .contextWrite(Context.of("micrometer.observation", intermediate)) // .as(StepVerifier::create).assertNext(actual -> { assertThat(actual).extracting("firstname", "lastname").containsExactly("Dave", "Matthews");