Make sure span is opened and closed with scope in Mongo

This commit is contained in:
Nikolay Martynov 2019-04-04 11:36:24 -04:00
parent f3e1eb4f5d
commit bddee3d6e2
5 changed files with 19 additions and 16 deletions

View File

@ -48,7 +48,7 @@ public final class MongoClientInstrumentation extends Instrumenter.Default {
"datadog.trace.agent.decorator.ClientDecorator",
"datadog.trace.agent.decorator.DatabaseClientDecorator",
packageName + ".MongoClientDecorator",
packageName + ".DDTracingCommandListener"
packageName + ".TracingCommandListener"
};
}
@ -66,7 +66,7 @@ public final class MongoClientInstrumentation extends Instrumenter.Default {
// referencing "this" in the method args causes the class to load under a transformer.
// This bypasses the Builder instrumentation. Casting as a workaround.
final MongoClientOptions.Builder builder = (MongoClientOptions.Builder) dis;
final DDTracingCommandListener listener = new DDTracingCommandListener();
final TracingCommandListener listener = new TracingCommandListener();
builder.addCommandListener(listener);
}
}

View File

@ -6,6 +6,7 @@ import com.mongodb.event.CommandFailedEvent;
import com.mongodb.event.CommandListener;
import com.mongodb.event.CommandStartedEvent;
import com.mongodb.event.CommandSucceededEvent;
import io.opentracing.Scope;
import io.opentracing.Span;
import io.opentracing.util.GlobalTracer;
import java.util.Map;
@ -13,23 +14,25 @@ import java.util.concurrent.ConcurrentHashMap;
import lombok.extern.slf4j.Slf4j;
@Slf4j
public class DDTracingCommandListener implements CommandListener {
public class TracingCommandListener implements CommandListener {
private final Map<Integer, Span> spanMap = new ConcurrentHashMap<>();
@Override
public void commandStarted(final CommandStartedEvent event) {
final Span span = GlobalTracer.get().buildSpan("mongo.query").start();
DECORATE.afterStart(span);
DECORATE.onConnection(span, event);
if (event.getConnectionDescription() != null
&& event.getConnectionDescription() != null
&& event.getConnectionDescription().getServerAddress() != null) {
DECORATE.onPeerConnection(
span, event.getConnectionDescription().getServerAddress().getSocketAddress());
try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) {
DECORATE.afterStart(span);
DECORATE.onConnection(span, event);
if (event.getConnectionDescription() != null
&& event.getConnectionDescription() != null
&& event.getConnectionDescription().getServerAddress() != null) {
DECORATE.onPeerConnection(
span, event.getConnectionDescription().getServerAddress().getSocketAddress());
}
DECORATE.onStatement(span, event.getCommand());
spanMap.put(event.getRequestId(), span);
}
DECORATE.onStatement(span, event.getCommand());
spanMap.put(event.getRequestId(), span);
}
@Override

View File

@ -48,7 +48,7 @@ public final class MongoAsyncClientInstrumentation extends Instrumenter.Default
"datadog.trace.agent.decorator.ClientDecorator",
"datadog.trace.agent.decorator.DatabaseClientDecorator",
packageName + ".MongoClientDecorator",
packageName + ".DDTracingCommandListener"
packageName + ".TracingCommandListener"
};
}
@ -66,7 +66,7 @@ public final class MongoAsyncClientInstrumentation extends Instrumenter.Default
// referencing "this" in the method args causes the class to load under a transformer.
// This bypasses the Builder instrumentation. Casting as a workaround.
final MongoClientSettings.Builder builder = (MongoClientSettings.Builder) dis;
final DDTracingCommandListener listener = new DDTracingCommandListener();
final TracingCommandListener listener = new TracingCommandListener();
builder.addCommandListener(listener);
}
}

View File

@ -49,7 +49,7 @@ public class MongoAsyncClientInstrumentationTest {
public void asyncClientHasListener() {
Assert.assertEquals(1, client.getSettings().getCommandListeners().size());
Assert.assertEquals(
"DDTracingCommandListener",
"TracingCommandListener",
client.getSettings().getCommandListeners().get(0).getClass().getSimpleName());
}

View File

@ -81,7 +81,7 @@ public class MongoClientInstrumentationTest {
public void syncClientHasListener() {
Assert.assertEquals(1, client.getMongoClientOptions().getCommandListeners().size());
Assert.assertEquals(
"DDTracingCommandListener",
"TracingCommandListener",
client.getMongoClientOptions().getCommandListeners().get(0).getClass().getSimpleName());
}