Skip to content

Clean up Observation from context after command completion #5067

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Clean up Observation from context after command completion #5067

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e1281a3
Prepare issue branch.
mp911de Oct 9, 2025
95735d0
Clean up Observation from context after command completion.
mp911de Oct 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

<groupId>org.springframework.data</groupId>
<artifactId>spring-data-mongodb-parent</artifactId>
<version>5.0.0-SNAPSHOT</version>
<version>5.0.0-GH-5064-SNAPSHOT</version>
<packaging>pom</packaging>

<name>Spring Data MongoDB</name>
Expand Down
2 changes: 1 addition & 1 deletion spring-data-mongodb-distribution/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
<parent>
<groupId>org.springframework.data</groupId>
<artifactId>spring-data-mongodb-parent</artifactId>
<version>5.0.0-SNAPSHOT</version>
<version>5.0.0-GH-5064-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>

Expand Down
2 changes: 1 addition & 1 deletion spring-data-mongodb/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
<parent>
<groupId>org.springframework.data</groupId>
<artifactId>spring-data-mongodb-parent</artifactId>
<version>5.0.0-SNAPSHOT</version>
<version>5.0.0-GH-5064-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>

Expand Down
14 changes: 3 additions & 11 deletions ...ngodb/src/main/java/org/springframework/data/mongodb/observability/MapRequestContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@
import java.util.stream.Stream;

import com.mongodb.RequestContext;
import org.jspecify.annotations.Nullable;

/**
* A {@link Map}-based {@link RequestContext}.
Expand All @@ -30,32 +29,25 @@
* @author Greg Turnquist
* @since 4.0.0
*/
class MapRequestContext implements RequestContext {

private final Map<Object, Object> map;
record MapRequestContext(Map<Object, Object> map) implements RequestContext {

public MapRequestContext() {
this(new HashMap<>());
}

public MapRequestContext(Map<Object, Object> context) {
this.map = context;
}

@Override
public <T> 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));
}

@Override
public boolean hasKey(Object key) {
return map.containsKey(key);
return map.get(key) != null;
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is hard to understand policy for null values.
For stream, size, isEmpty it is regular value.
For get, hasKey it is "hidden" value

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change aligns with getOrDefault(…), it gets more pronounced when you take a look at RequestContext.getOrEmpty(…). The interface assumes that a key is only present if its value is not null.

}

@Override
Expand Down
32 changes: 18 additions & 14 deletions .../java/org/springframework/data/mongodb/observability/MongoObservationCommandListener.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,10 +113,6 @@ public void commandStarted(CommandStartedEvent event) {

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");

Expand All @@ -141,22 +137,20 @@ public void commandStarted(CommandStartedEvent event) {
@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);

Expand All @@ -165,18 +159,17 @@ public void commandFailed(CommandFailedEvent event) {
}

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<Observation, MongoHandlerContext> action) {

if (requestContext == null) {
Expand All @@ -188,7 +181,18 @@ private void doInObservation(@Nullable RequestContext requestContext,
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());
}
}

/**
Expand All @@ -210,7 +214,7 @@ private void doInObservation(@Nullable RequestContext requestContext,
}

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;
Expand Down
41 changes: 41 additions & 0 deletions .../org/springframework/data/mongodb/observability/MongoObservationCommandListenerTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
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;
Expand Down Expand Up @@ -251,6 +252,46 @@ public String getName() {
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();
}
Expand Down