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

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> 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

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
+	 * Stops the {@link Observation} after applying {@code action} given {@link MongoHandlerContext} if there is an
-	 * Observation. Exceptions thrown by the action are relayed to the caller.
+	 * 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<Observation, MongoHandlerContext> 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;

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();
 	}