From 0bc50ecce1c70855abcf4f33622317bc6b3c81ac 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 | 3 +- .../MongoObservationCommandListener.java | 32 ++++++++------- .../MongoObservationCommandListenerTests.java | 41 +++++++++++++++++++ 3 files changed, 61 insertions(+), 15 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 854e1481f..bd9208273 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 @@ -41,13 +41,14 @@ class MapRequestContext implements RequestContext { } @Override + @SuppressWarnings("unchecked") public T get(Object key) { return (T) map.get(key); } @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 9360a95de..45a2ba701 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()); + } } /** @@ -211,7 +215,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(); }