Skip to content

Commit 86a76ec

Browse files
author
Andrew Kent
committed
Convert Mongo instrumentation to ByteBuddy
1 parent b4ba6d8 commit 86a76ec

File tree

8 files changed

+52
-80
lines changed

8 files changed

+52
-80
lines changed

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,3 +51,4 @@ Thumbs.db
5151
*/out
5252
dd-java-agent/integrations/*/out
5353
dd-trace-examples/*/out
54+
derby.log

‎dd-java-agent/dd-java-agent.gradle‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,9 @@ dependencies {
3030
compile(project(':dd-java-agent:integrations:jms-2')) {
3131
transitive = false
3232
}
33+
compile(project(':dd-java-agent:integrations:mongo')) {
34+
transitive = false
35+
}
3336
compile(project(':dd-java-agent:integrations:servlet-2')) {
3437
transitive = false
3538
}
@@ -51,7 +54,6 @@ dependencies {
5154

5255
testCompile deps.testLogging
5356
testCompile deps.opentracingMock
54-
testCompile group: 'org.mongodb', name: 'mongo-java-driver', version: '3.4.2'
5557

5658
testCompile(project(path: ':dd-java-agent:integrations:helpers')) {
5759
transitive = false

‎dd-java-agent/integrations/mongo/mongo.gradle‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,3 +9,15 @@ versionScan {
99
"com.mongodb.client.model.MapReduceAction": null,
1010
]
1111
}
12+
13+
apply from: "${rootDir}/gradle/java.gradle"
14+
15+
dependencies {
16+
compile group: 'org.mongodb', name: 'mongo-java-driver', version: '3.4.2'
17+
18+
compile project(':dd-trace')
19+
compile project(':dd-java-agent:tooling')
20+
21+
compile deps.bytebuddy
22+
compile deps.opentracing
23+
}

dd-java-agent/integrations/helpers/src/main/java/com/datadoghq/agent/integration/MongoHelper.java renamed to dd-java-agent/integrations/mongo/src/main/java/dd/inst/mongo/MongoClientInstrumentation.java

Lines changed: 31 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,20 @@
1-
package com.datadoghq.agent.integration;
1+
package dd.inst.mongo;
2+
3+
import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
4+
import static net.bytebuddy.matcher.ElementMatchers.*;
25

36
import com.datadoghq.trace.DDTags;
7+
import com.google.auto.service.AutoService;
48
import com.mongodb.MongoClientOptions;
59
import com.mongodb.event.CommandFailedEvent;
610
import com.mongodb.event.CommandListener;
711
import com.mongodb.event.CommandStartedEvent;
812
import com.mongodb.event.CommandSucceededEvent;
13+
import dd.trace.Instrumenter;
914
import io.opentracing.Span;
1015
import io.opentracing.Tracer;
1116
import io.opentracing.tag.Tags;
17+
import io.opentracing.util.GlobalTracer;
1218
import java.net.Inet4Address;
1319
import java.net.InetAddress;
1420
import java.nio.ByteBuffer;
@@ -18,39 +24,40 @@
1824
import java.util.Map;
1925
import java.util.concurrent.ConcurrentHashMap;
2026
import lombok.extern.slf4j.Slf4j;
27+
import net.bytebuddy.agent.builder.AgentBuilder;
28+
import net.bytebuddy.asm.Advice;
2129
import org.bson.BsonArray;
2230
import org.bson.BsonDocument;
2331
import org.bson.BsonString;
2432
import org.bson.BsonValue;
25-
import org.jboss.byteman.rule.Rule;
2633

27-
/** Patch the Mongo builder before constructing the final client */
2834
@Slf4j
29-
public class MongoHelper extends DDAgentTracingHelper<MongoClientOptions.Builder> {
30-
31-
public MongoHelper(final Rule rule) {
32-
super(rule);
33-
}
35+
@AutoService(Instrumenter.class)
36+
public final class MongoClientInstrumentation implements Instrumenter {
3437

35-
/**
36-
* Strategy: Just before com.mongodb.MongoClientOptions$Builder.build() method is called, we add a
37-
* new command listener in charge of the tracing.
38-
*
39-
* @param builder The builder instance
40-
* @return The same builder instance with a new tracing command listener that will be use for the
41-
* client construction
42-
* @throws Exception
43-
*/
4438
@Override
45-
protected MongoClientOptions.Builder doPatch(final MongoClientOptions.Builder builder)
46-
throws Exception {
47-
48-
final DDTracingCommandListener listener = new DDTracingCommandListener(tracer);
49-
builder.addCommandListener(listener);
39+
public AgentBuilder instrument(AgentBuilder agentBuilder) {
40+
return agentBuilder
41+
.type(named("com.mongodb.MongoClientOptions$Builder"))
42+
.transform(
43+
new AgentBuilder.Transformer.ForAdvice()
44+
.advice(
45+
isMethod().and(isPublic()).and(named("build")).and(takesArguments(0)),
46+
MongoClientAdvice.class.getName())
47+
.withExceptionHandler(defaultExceptionHandler()))
48+
.asDecorator();
49+
}
5050

51-
setState(builder, 1);
51+
public static class MongoClientAdvice {
5252

53-
return builder;
53+
@Advice.OnMethodEnter(suppress = Throwable.class)
54+
public static void injectTraceListener(@Advice.This final Object dis) {
55+
// referencing "this" in the method args causes the class to load under a transformer.
56+
// This bypasses the Builder instrumentation. Casting as a workaround.
57+
MongoClientOptions.Builder builder = (MongoClientOptions.Builder) dis;
58+
final DDTracingCommandListener listener = new DDTracingCommandListener(GlobalTracer.get());
59+
builder.addCommandListener(listener);
60+
}
5461
}
5562

5663
public static class DDTracingCommandListener implements CommandListener {

dd-java-agent/src/test/java/com/datadoghq/agent/integration/MongoHelperTest.java renamed to dd-java-agent/integrations/mongo/src/test/java/dd/inst/mongo/MongoClientInstrumentationTest.java

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package com.datadoghq.agent.integration;
1+
package dd.inst.mongo;
22

33
import static org.assertj.core.api.Java6Assertions.assertThat;
44

@@ -9,6 +9,7 @@
99
import com.mongodb.connection.ConnectionDescription;
1010
import com.mongodb.connection.ServerId;
1111
import com.mongodb.event.CommandStartedEvent;
12+
import dd.inst.mongo.MongoClientInstrumentation.DDTracingCommandListener;
1213
import io.opentracing.tag.Tags;
1314
import java.util.Arrays;
1415
import java.util.List;
@@ -17,7 +18,7 @@
1718
import org.bson.BsonString;
1819
import org.junit.Test;
1920

20-
public class MongoHelperTest {
21+
public class MongoClientInstrumentationTest {
2122

2223
private static ConnectionDescription makeConnection() {
2324
return new ConnectionDescription(new ServerId(new ClusterId(), new ServerAddress()));
@@ -29,7 +30,7 @@ public void mongoSpan() {
2930
new CommandStartedEvent(1, makeConnection(), "databasename", "query", new BsonDocument());
3031

3132
final DDSpan span = new DDTracer().buildSpan("foo").startManual();
32-
MongoHelper.DDTracingCommandListener.decorate(span, cmd);
33+
DDTracingCommandListener.decorate(span, cmd);
3334

3435
assertThat(span.context().getSpanType()).isEqualTo("mongodb");
3536
assertThat(span.context().getResourceName())
@@ -51,7 +52,7 @@ public void queryScrubbing() {
5152
new CommandStartedEvent(1, makeConnection(), "databasename", "query", query);
5253

5354
final DDSpan span = new DDTracer().buildSpan("foo").startManual();
54-
MongoHelper.DDTracingCommandListener.decorate(span, cmd);
55+
DDTracingCommandListener.decorate(span, cmd);
5556

5657
assertThat(span.getTags().get(Tags.DB_STATEMENT.getKey()))
5758
.isEqualTo(query.toString().replaceAll("secret", "?"));

‎dd-java-agent/src/main/resources/dd-trace-supported-framework.yaml‎

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -47,32 +47,6 @@ opentracing-jms-2_consumer:
4747
javax.jms.JMSContext:
4848
javax.jms.CompletionListener:
4949

50-
opentracing-mongo-driver:
51-
- artifact: mongo-java-driver
52-
supported_version: 3\..*
53-
identifying_present_classes:
54-
com.mongodb.operation.AsyncReadOperation:
55-
com.mongodb.client.model.MapReduceAction:
56-
57-
- artifact: mongodb-driver-async
58-
supported_version: 3\..*
59-
identifying_present_classes:
60-
com.mongodb.operation.AsyncReadOperation:
61-
com.mongodb.client.model.MapReduceAction:
62-
63-
opentracing-mongo-driver-helper:
64-
- artifact: mongo-java-driver
65-
supported_version: 3\..*
66-
identifying_present_classes:
67-
com.mongodb.operation.AsyncReadOperation:
68-
com.mongodb.client.model.MapReduceAction:
69-
70-
- artifact: mongodb-driver-async
71-
supported_version: 3\..*
72-
identifying_present_classes:
73-
com.mongodb.operation.AsyncReadOperation:
74-
com.mongodb.client.model.MapReduceAction:
75-
7650
opentracing-okhttp3:
7751
- artifact: okhttp
7852
supported_version: 3\..*

‎dd-java-agent/src/main/resources/initializer-rules.btm‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -25,18 +25,6 @@ DO
2525
ENDRULE
2626

2727

28-
# Instrument Mongo client
29-
# ========================
30-
RULE MongoClientOptions$Builder-init
31-
CLASS com.mongodb.MongoClientOptions$Builder
32-
METHOD <init>
33-
AT EXIT
34-
IF TRUE
35-
DO
36-
com.datadoghq.agent.InstrumentationRulesManager.registerClassLoad($0);
37-
ENDRULE
38-
39-
4028
# Instrument OkHttp
4129
# ===========================
4230
RULE OkHttpClient$Builder-init

‎dd-java-agent/src/main/resources/integration-rules.btm‎

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -34,19 +34,6 @@ DO
3434
ENDRULE
3535

3636

37-
# Instrument Mongo client
38-
# ========================
39-
RULE opentracing-mongo-driver
40-
CLASS com.mongodb.MongoClientOptions$Builder
41-
METHOD build
42-
HELPER com.datadoghq.agent.integration.MongoHelper
43-
AT ENTRY
44-
IF getState($0) == 0
45-
DO
46-
patch($0);
47-
ENDRULE
48-
49-
5037
# Instrument OkHttp
5138
# ===========================
5239
RULE opentracing-okhttp3

0 commit comments

Comments
 (0)