From 070f6132960f5788323f0d14bc015a236c891e0d Mon Sep 17 00:00:00 2001 From: Mark Paluch Date: Thu, 9 Oct 2025 10:32:36 +0200 Subject: [PATCH] Clean up Observation from context after command completion. We now restore the previous Observation parent after completing the Observation to avoid unintentional nesting of observations. Previously, a command became the parent of an earlier command on the same thread regardless of whether the earlier command has been finished or whether it was active. Observation nesting now considers only commands that are still active restoring their parent after completion to avoid invalid nesting and with that, to not allocate observation graphs that allocate infinite memory. Closes: #5064 Original Pull Request: #5067 --- .../observability/MapRequestContext.java | 15 ++----- .../MongoObservationCommandListener.java | 32 ++++++++------- .../MongoObservationCommandListenerTests.java | 41 +++++++++++++++++++ 3 files changed, 63 insertions(+), 25 deletions(-) diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MapRequestContext.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MapRequestContext.java index 6185c95db..11c229350 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MapRequestContext.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/observability/MapRequestContext.java @@ -21,7 +21,6 @@ import java.util.NoSuchElementException; import java.util.stream.Stream; import com.mongodb.RequestContext; -import org.jspecify.annotations.Nullable; /** * A {@link Map}-based {@link RequestContext}. @@ -30,24 +29,18 @@ import org.jspecify.annotations.Nullable; * @author Greg Turnquist * @since 4.0.0 */ -class MapRequestContext implements RequestContext { - - private final Map map; +record MapRequestContext(Map map) implements RequestContext { public MapRequestContext() { this(new HashMap<>()); } - public MapRequestContext(Map context) { - this.map = context; - } - @Override + @SuppressWarnings("unchecked") public T get(Object key) { - T value = (T) map.get(key); - if(value != null) { + if (value != null) { return value; } throw new NoSuchElementException("%s is missing".formatted(key)); @@ -55,7 +48,7 @@ class MapRequestContext implements RequestContext { @Override public boolean hasKey(Object key) { - return map.containsKey(key); + return map.get(key) != null; } @Override 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 914396ab9..fcd477804 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 @@ -113,10 +113,6 @@ public class MongoObservationCommandListener implements CommandListener { Observation parent = observationFromContext(requestContext); - if (log.isDebugEnabled()) { - log.debug("Found the following observation passed from the mongo context [" + parent + "]"); - } - MongoHandlerContext observationContext = new MongoHandlerContext(connectionString, event, requestContext); observationContext.setRemoteServiceName("mongo"); @@ -141,22 +137,20 @@ public class MongoObservationCommandListener implements CommandListener { @Override public void commandSucceeded(CommandSucceededEvent event) { - doInObservation(event.getRequestContext(), (observation, context) -> { + stopObservation(event.getRequestContext(), (observation, context) -> { context.setCommandSucceededEvent(event); if (log.isDebugEnabled()) { log.debug("Command succeeded - will stop observation [" + observation + "]"); } - - observation.stop(); }); } @Override public void commandFailed(CommandFailedEvent event) { - doInObservation(event.getRequestContext(), (observation, context) -> { + stopObservation(event.getRequestContext(), (observation, context) -> { context.setCommandFailedEvent(event); @@ -165,18 +159,17 @@ public class MongoObservationCommandListener implements CommandListener { } observation.error(event.getThrowable()); - observation.stop(); }); } /** - * Performs the given action for the {@link Observation} and {@link MongoHandlerContext} if there is an ongoing Mongo - * Observation. Exceptions thrown by the action are relayed to the caller. + * Stops the {@link Observation} after applying {@code action} given {@link MongoHandlerContext} if there is an + * ongoing Mongo Observation. Exceptions thrown by the action are relayed to the caller. * * @param requestContext the context to extract the Observation from. * @param action the action to invoke. */ - private void doInObservation(@Nullable RequestContext requestContext, + private void stopObservation(@Nullable RequestContext requestContext, BiConsumer action) { if (requestContext == null) { @@ -188,7 +181,18 @@ public class MongoObservationCommandListener implements CommandListener { return; } - action.accept(observation, context); + try { + action.accept(observation, context); + } finally { + + observation.stop(); + + if (log.isDebugEnabled()) { + log.debug( + "Restoring parent observation [" + observation + "] for Mongo instrumentation and put it in Mongo context"); + } + requestContext.put(ObservationThreadLocalAccessor.KEY, observation.getContext().getParentObservation()); + } } /** @@ -210,7 +214,7 @@ public class MongoObservationCommandListener implements CommandListener { } if (log.isDebugEnabled()) { - log.debug("No observation was found - will not create any child observations"); + log.debug("No observation was found: Creating a new root observation"); } return null; diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/observability/MongoObservationCommandListenerTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/observability/MongoObservationCommandListenerTests.java index fe74a03bd..dadb98ce2 100644 --- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/observability/MongoObservationCommandListenerTests.java +++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/observability/MongoObservationCommandListenerTests.java @@ -27,6 +27,7 @@ import io.micrometer.observation.Observation; import io.micrometer.observation.ObservationRegistry; import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor; +import org.assertj.core.api.Assertions; import org.bson.BsonDocument; import org.bson.BsonString; import org.junit.jupiter.api.BeforeEach; @@ -251,6 +252,46 @@ class MongoObservationCommandListenerTests { assertThat(meterRegistry).hasMeterWithName("custom.name.active"); } + @Test // GH-5064 + void completionRestoresParentObservation() { + + // given + Observation parent = Observation.start("name", observationRegistry); + observationRegistry.setCurrentObservationScope(parent.openScope()); + RequestContext traceRequestContext = getContext(); + + // when + listener.commandStarted(new CommandStartedEvent(traceRequestContext, 0, 0, null, "database", "insert", + new BsonDocument("collection", new BsonString("user")))); + + Assertions.assertThat((Observation) traceRequestContext.get(ObservationThreadLocalAccessor.KEY)).isNotNull() + .isNotEqualTo(parent); + + listener.commandSucceeded(new CommandSucceededEvent(traceRequestContext, 0, 0, null, "insert", null, null, 0)); + + Assertions.assertThat((Observation) traceRequestContext.get(ObservationThreadLocalAccessor.KEY)).isEqualTo(parent); + } + + @Test // GH-5064 + void failureRestoresParentObservation() { + + // given + Observation parent = Observation.start("name", observationRegistry); + observationRegistry.setCurrentObservationScope(parent.openScope()); + RequestContext traceRequestContext = getContext(); + + // when + listener.commandStarted(new CommandStartedEvent(traceRequestContext, 0, 0, null, "database", "insert", + new BsonDocument("collection", new BsonString("user")))); + + Assertions.assertThat((Observation) traceRequestContext.get(ObservationThreadLocalAccessor.KEY)).isNotNull() + .isNotEqualTo(parent); + + listener.commandFailed(new CommandFailedEvent(traceRequestContext, 0, 0, null, "insert", null, 0, null)); + + Assertions.assertThat((Observation) traceRequestContext.get(ObservationThreadLocalAccessor.KEY)).isEqualTo(parent); + } + private RequestContext getContext() { return ((SynchronousContextProvider) ContextProviderFactory.create(observationRegistry)).getContext(); }