Skip to content

Commit

Permalink
Guard command completion listener against unsupported observation con…
Browse files Browse the repository at this point in the history
…text.

We now no longer attempt to complete the Observation if the context is not a MongoDB one. For commands that target the admin database and run within a parent observation, we still might have an Observation but that one points to the parent invocation and not the MongoDB one as we do not record commands for the admin database.

Closes #4481
  • Loading branch information
mp911de committed Aug 24, 2023
1 parent 31da8b1 commit 0be1e52
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 15 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,10 @@
*/
package org.springframework.data.mongodb.observability;

import io.micrometer.observation.Observation;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.lang.Nullable;
Expand All @@ -27,10 +31,6 @@
import com.mongodb.event.CommandStartedEvent;
import com.mongodb.event.CommandSucceededEvent;

import io.micrometer.observation.Observation;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;

/**
* Implement MongoDB's {@link CommandListener} using Micrometer's {@link Observation} API.
*
Expand Down Expand Up @@ -133,11 +133,10 @@ public void commandSucceeded(CommandSucceededEvent event) {
}

Observation observation = requestContext.getOrDefault(ObservationThreadLocalAccessor.KEY, null);
if (observation == null) {
if (observation == null || !(observation.getContext()instanceof MongoHandlerContext context)) {
return;
}

MongoHandlerContext context = (MongoHandlerContext) observation.getContext();
context.setCommandSucceededEvent(event);

if (log.isDebugEnabled()) {
Expand All @@ -157,11 +156,10 @@ public void commandFailed(CommandFailedEvent event) {
}

Observation observation = requestContext.getOrDefault(ObservationThreadLocalAccessor.KEY, null);
if (observation == null) {
if (observation == null || !(observation.getContext()instanceof MongoHandlerContext context)) {
return;
}

MongoHandlerContext context = (MongoHandlerContext) observation.getContext();
context.setCommandFailedEvent(event);

if (log.isDebugEnabled()) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,15 @@
package org.springframework.data.mongodb.observability;

import static io.micrometer.core.tck.MeterRegistryAssert.*;
import static org.mockito.Mockito.*;

import io.micrometer.common.KeyValues;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.observation.DefaultMeterObservationHandler;
import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
import io.micrometer.observation.Observation;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;

import org.bson.BsonDocument;
import org.bson.BsonString;
Expand All @@ -33,18 +42,12 @@
import com.mongodb.event.CommandStartedEvent;
import com.mongodb.event.CommandSucceededEvent;

import io.micrometer.common.KeyValues;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.observation.DefaultMeterObservationHandler;
import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
import io.micrometer.observation.Observation;
import io.micrometer.observation.ObservationRegistry;

/**
* Series of test cases exercising {@link MongoObservationCommandListener}.
*
* @author Marcin Grzejszczak
* @author Greg Turnquist
* @author Mark Paluch
*/
class MongoObservationCommandListenerTests {

Expand Down Expand Up @@ -176,6 +179,38 @@ void commandWithErrorShouldCreateTimerWhenParentSampleInRequestContext() {
assertThatTimerRegisteredWithTags();
}

@Test // GH-4481
void completionShouldIgnoreIncompatibleObservationContext() {

// given
RequestContext traceRequestContext = getContext();

Observation observation = mock(Observation.class);
traceRequestContext.put(ObservationThreadLocalAccessor.KEY, observation);

// when
listener.commandSucceeded(new CommandSucceededEvent(traceRequestContext, 0, null, "insert", null, 0));

verify(observation).getContext();
verifyNoMoreInteractions(observation);
}

@Test // GH-4481
void failureShouldIgnoreIncompatibleObservationContext() {

// given
RequestContext traceRequestContext = getContext();

Observation observation = mock(Observation.class);
traceRequestContext.put(ObservationThreadLocalAccessor.KEY, observation);

// when
listener.commandFailed(new CommandFailedEvent(traceRequestContext, 0, null, "insert", 0, null));

verify(observation).getContext();
verifyNoMoreInteractions(observation);
}

private RequestContext getContext() {
return ((SynchronousContextProvider) ContextProviderFactory.create(observationRegistry)).getContext();
}
Expand Down

0 comments on commit 0be1e52

Please sign in to comment.