Skip to content

Commit 3f8198c

Browse files
na-jina-ji
committed
add new mongo instrumentation to inject APM/DBM comment
This is a draft that doesn't work yet. Tests are also missing
1 parent 92a6054 commit 3f8198c

File tree

10 files changed

+444
-80
lines changed

10 files changed

+444
-80
lines changed

‎dd-java-agent/instrumentation/mongo/common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommandListener.java‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33
import static datadog.trace.api.Functions.UTF8_ENCODE;
44
import static datadog.trace.api.cache.RadixTreeCache.PORTS;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
6+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
7+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closeActive;
68
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
79

810
import com.mongodb.ServerAddress;
@@ -122,7 +124,13 @@ public void commandStarted(final CommandStartedEvent event) {
122124
if (listenerAccessor != null) {
123125
listenerAccessor.putIfAbsent(event.getConnectionDescription(), this);
124126
}
125-
final AgentSpan span = startSpan(MongoDecorator.OPERATION_NAME);
127+
AgentSpan span = activeSpan();
128+
boolean shouldForceCloseSpanScope = true;
129+
if (span == null || span.getSpanName() != MongoDecorator.OPERATION_NAME) {
130+
span = startSpan(MongoDecorator.OPERATION_NAME);
131+
shouldForceCloseSpanScope = false;
132+
}
133+
126134
try (final AgentScope scope = activateSpan(span)) {
127135
decorator.afterStart(span);
128136
decorator.onConnection(span, event);
@@ -145,31 +153,26 @@ public void commandStarted(final CommandStartedEvent event) {
145153
Tags.DB_OPERATION,
146154
COMMAND_NAMES.computeIfAbsent(event.getCommandName(), UTF8_ENCODE));
147155
}
156+
decorator.onStatement(span, event.getCommand(), byteBufAccessor);
157+
spanMap.put(event.getRequestId(), new SpanEntry(span));
148158

149-
BsonDocument commandToExecute = event.getCommand();
150-
151-
// Comment injection
152-
String dbmComment = MongoCommentInjector.getComment(span, event);
153-
if (dbmComment != null) {
154-
commandToExecute = MongoCommentInjector.injectComment(dbmComment, event);
159+
if (shouldForceCloseSpanScope) {
160+
closeActive();
155161
}
156-
157-
decorator.onStatement(span, commandToExecute, byteBufAccessor);
158-
spanMap.put(event.getRequestId(), new SpanEntry(span));
159162
}
160163
}
161164

162165
@Override
163166
public void commandSucceeded(final CommandSucceededEvent event) {
164-
finishSpah(event.getRequestId(), null);
167+
finishSpan(event.getRequestId(), null);
165168
}
166169

167170
@Override
168171
public void commandFailed(final CommandFailedEvent event) {
169-
finishSpah(event.getRequestId(), event.getThrowable());
172+
finishSpan(event.getRequestId(), event.getThrowable());
170173
}
171174

172-
private void finishSpah(int requestId, Throwable t) {
175+
private void finishSpan(int requestId, Throwable t) {
173176
final SpanEntry entry = spanMap.remove(requestId);
174177
final AgentSpan span = entry != null ? entry.span : null;
175178
if (span != null) {

‎dd-java-agent/instrumentation/mongo/common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommentInjector.java‎

Lines changed: 22 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,6 @@
33
import static datadog.trace.api.Config.DBM_PROPAGATION_MODE_FULL;
44
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DBM_TRACE_INJECTED;
55

6-
import com.mongodb.connection.ConnectionDescription;
7-
import com.mongodb.event.CommandStartedEvent;
86
import datadog.trace.api.Config;
97
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
108
import datadog.trace.core.database.SharedDBCommenter;
@@ -22,16 +20,18 @@
2220
public class MongoCommentInjector {
2321
private static final Logger log = LoggerFactory.getLogger(MongoCommentInjector.class);
2422

23+
public static final boolean INJECT_COMMENT = Config.get().isDbmCommentInjectionEnabled();
24+
2525
/** Main entry point for MongoDB command comment injection */
26-
public static BsonDocument injectComment(String dbmComment, CommandStartedEvent event) {
27-
if (!Config.get().isDbmCommentInjectionEnabled() || dbmComment == null || event == null) {
28-
return event != null ? event.getCommand() : null;
26+
public static BsonDocument injectComment(String dbmComment, BsonDocument originalBsonDocument) {
27+
if (!INJECT_COMMENT) {
28+
return originalBsonDocument;
2929
}
3030

3131
// Create a mutable copy by constructing a new BsonDocument and copying all entries
3232
// This handles both regular BsonDocument and immutable RawBsonDocument/ByteBufBsonDocument
3333
BsonDocument command = new BsonDocument();
34-
command.putAll(event.getCommand());
34+
command.putAll(originalBsonDocument);
3535

3636
try {
3737
for (String commentKey : new String[] {"comment", "$comment"}) {
@@ -49,14 +49,18 @@ public static BsonDocument injectComment(String dbmComment, CommandStartedEvent
4949
log.warn(
5050
"Linking Database Monitoring profiles to spans is not supported for the following query type: {}. "
5151
+ "To disable this feature please set the following environment variable: DD_DBM_PROPAGATION_MODE=disabled",
52-
event.getCommand().getClass().getSimpleName());
53-
return event.getCommand();
52+
originalBsonDocument.getClass().getSimpleName());
53+
return originalBsonDocument;
5454
}
5555
}
5656

5757
/** Build comment content using SharedDBCommenter */
58-
public static String getComment(AgentSpan dbSpan, CommandStartedEvent event) {
59-
if (!Config.get().isDbmCommentInjectionEnabled()) {
58+
public static String getComment(AgentSpan dbSpan, String hostname, String dbName) {
59+
if (!INJECT_COMMENT) {
60+
return null;
61+
}
62+
63+
if (dbSpan.forceSamplingDecision() == null) {
6064
return null;
6165
}
6266

@@ -65,8 +69,6 @@ public static String getComment(AgentSpan dbSpan, CommandStartedEvent event) {
6569

6670
// Extract connection details
6771
String dbService = dbSpan.getServiceName();
68-
String hostname = getHostnameFromEvent(event);
69-
String dbName = event.getDatabaseName();
7072
String traceParent =
7173
Config.get().getDbmPropagationMode().equals(DBM_PROPAGATION_MODE_FULL)
7274
? buildTraceParent(dbSpan)
@@ -107,22 +109,15 @@ private static BsonValue mergeComment(BsonValue existingComment, String dbmComme
107109
return existingComment;
108110
}
109111

110-
private static String getHostnameFromEvent(CommandStartedEvent event) {
111-
ConnectionDescription connectionDescription = event.getConnectionDescription();
112-
if (connectionDescription != null && connectionDescription.getServerAddress() != null) {
113-
return connectionDescription.getServerAddress().getHost();
114-
}
115-
return null;
116-
}
117-
118112
static String buildTraceParent(AgentSpan span) {
119113
// W3C traceparent format: version-traceId-spanId-flags
120-
StringBuilder sb = new StringBuilder(2 + 1 + 32 + 1 + 16 + 1 + 2);
121-
sb.append("00-"); // version
122-
sb.append(span.getTraceId().toHexStringPadded(32)); // traceId
123-
sb.append("-");
124-
sb.append(String.format("%016x", span.getSpanId())); // spanId
125-
sb.append(span.context().getSamplingPriority() > 0 ? "-01" : "-00");
126-
return sb.toString();
114+
return "00-"
115+
+ // version
116+
span.getTraceId().toHexStringPadded(32)
117+
+ // traceId
118+
'-'
119+
+ String.format("%016x", span.getSpanId())
120+
+ // spanId
121+
(span.context().getSamplingPriority() > 0 ? "-01" : "-00");
127122
}
128123
}

‎dd-java-agent/instrumentation/mongo/common/src/test/groovy/MongoCommentInjectorTest.groovy‎

Lines changed: 21 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,22 @@
1-
package datadog.trace.instrumentation.mongo
2-
3-
import com.mongodb.event.CommandStartedEvent
41
import datadog.trace.agent.test.InstrumentationSpecification
52
import datadog.trace.api.config.TraceInstrumentationConfig
63
import datadog.trace.api.sampling.PrioritySampling
4+
import datadog.trace.instrumentation.mongo.MongoCommentInjector
75
import org.bson.BsonDocument
86
import org.bson.BsonString
97
import org.bson.RawBsonDocument
108

11-
class MongoCommentInjectorTest extends InstrumentationSpecification {
9+
abstract class BaseMongoCommentInjectorTest extends InstrumentationSpecification {
1210
@Override
1311
void configurePreAgent() {
1412
super.configurePreAgent()
1513
injectSysConfig("service.name", "test-mongo-service")
1614
injectSysConfig("dd.env", "test")
1715
injectSysConfig("dd.version", "1.0.0")
1816
}
17+
}
1918

20-
def "getComment returns null when INJECT_COMMENT is false"() {
21-
setup:
22-
injectSysConfig(TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE, "disabled")
23-
24-
when:
25-
String comment = MongoCommentInjector.getComment(null, null)
26-
27-
then:
28-
comment == null
29-
}
30-
19+
class MongoCommentInjectorTest extends BaseMongoCommentInjectorTest {
3120
def "buildTraceParent with sampled flag (SAMPLER_KEEP)"() {
3221
setup:
3322
def span = TEST_TRACER.buildSpan("test-op").start()
@@ -101,7 +90,22 @@ class MongoCommentInjectorTest extends InstrumentationSpecification {
10190
comment.contains("ddh='localhost'")
10291
comment.contains("dddb='testdb'")
10392
}
93+
}
94+
95+
class MongoCommentInjectorDisabledModeForkedTest extends BaseMongoCommentInjectorTest {
96+
def "getComment returns null when INJECT_COMMENT is false"() {
97+
setup:
98+
injectSysConfig(TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE, "disabled")
99+
100+
when:
101+
String comment = MongoCommentInjector.getComment(null, null, null)
102+
103+
then:
104+
comment == null
105+
}
106+
}
104107
108+
class MongoCommentInjectorServiceModeForkedTest extends BaseMongoCommentInjectorTest {
105109
def "injectComment handles immutable RawBsonDocument"() {
106110
setup:
107111
// Enable DBM propagation for this test
@@ -120,17 +124,8 @@ class MongoCommentInjectorTest extends InstrumentationSpecification {
120124
thrown(UnsupportedOperationException)
121125
122126
when:
123-
// Create CommandStartedEvent with RawBsonDocument
124-
def event = new CommandStartedEvent(
125-
1, // requestId
126-
null, // connectionDescription
127-
"testdb", // databaseName
128-
"find", // commandName
129-
rawDoc // command (immutable)
130-
)
131-
132127
// This should NOT throw UnsupportedOperationException with the fix
133-
BsonDocument result = MongoCommentInjector.injectComment(dbmComment, event)
128+
BsonDocument result = MongoCommentInjector.injectComment(dbmComment, rawDoc)
134129
135130
then:
136131
// Should successfully inject comment
@@ -149,16 +144,8 @@ class MongoCommentInjectorTest extends InstrumentationSpecification {
149144
def originalCommand = new BsonDocument("find", new BsonString("collection"))
150145
def dbmComment = "dddbs='test-service',dde='test'"
151146
152-
def event = new CommandStartedEvent(
153-
1, // requestId
154-
null, // connectionDescription
155-
"testdb", // databaseName
156-
"find", // commandName
157-
originalCommand // command (mutable)
158-
)
159-
160147
when:
161-
BsonDocument result = MongoCommentInjector.injectComment(dbmComment, event)
148+
BsonDocument result = MongoCommentInjector.injectComment(dbmComment, originalCommand)
162149
163150
then:
164151
result != null

‎dd-java-agent/instrumentation/mongo/driver-3.1/src/main/java/datadog/trace/instrumentation/mongo/DefaultServerConnection31Instrumentation.java‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
package datadog.trace.instrumentation.mongo;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpanWithoutScope;
5+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
6+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
7+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
8+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
9+
10+
import com.google.auto.service.AutoService;
11+
import com.mongodb.connection.Connection;
12+
import com.mongodb.connection.ConnectionDescription;
13+
import datadog.trace.agent.tooling.Instrumenter;
14+
import datadog.trace.agent.tooling.InstrumenterModule;
15+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
16+
import datadog.trace.core.database.SharedDBCommenter;
17+
import net.bytebuddy.asm.Advice;
18+
import org.bson.BsonDocument;
19+
20+
@AutoService(InstrumenterModule.class)
21+
public class DefaultServerConnection31Instrumentation extends InstrumenterModule.Tracing
22+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
23+
24+
public DefaultServerConnection31Instrumentation() {
25+
super("mongo", "mongo-3.1");
26+
}
27+
28+
@Override
29+
public String instrumentedType() {
30+
return "com.mongodb.connection.DefaultServerConnection";
31+
}
32+
33+
@Override
34+
public String[] helperClassNames() {
35+
return new String[] {
36+
SharedDBCommenter.class.getName(),
37+
MongoCommentInjector.class.getName(),
38+
MongoDecorator.class.getName(),
39+
};
40+
}
41+
42+
@Override
43+
public void methodAdvice(MethodTransformer transformer) {
44+
transformer.applyAdvice(
45+
isMethod()
46+
.and(named("command"))
47+
.and(takesArgument(0, String.class))
48+
.and(takesArgument(1, named("org.bson.BsonDocument")))
49+
.and(takesArguments(5)),
50+
DefaultServerConnection31Instrumentation.class.getName() + "$CommandAdvice");
51+
52+
transformer.applyAdvice(
53+
isMethod()
54+
.and(named("commandAsync"))
55+
.and(takesArgument(0, String.class))
56+
.and(takesArgument(1, named("org.bson.BsonDocument")))
57+
.and(takesArguments(6)),
58+
DefaultServerConnection31Instrumentation.class.getName() + "$CommandAdvice");
59+
}
60+
61+
public static class CommandAdvice {
62+
@Advice.OnMethodEnter(suppress = Throwable.class)
63+
public static void onEnter(
64+
@Advice.This Connection connection,
65+
@Advice.Argument(value = 0) String dbName,
66+
@Advice.Argument(value = 1, readOnly = false) BsonDocument originalBsonDocument) {
67+
if (!MongoCommentInjector.INJECT_COMMENT) {
68+
return;
69+
}
70+
71+
AgentSpan span = startSpan(MongoDecorator.OPERATION_NAME);
72+
// scope is going to be closed by the MongoCommandListener
73+
activateSpanWithoutScope(span);
74+
75+
String hostname = null;
76+
ConnectionDescription connectionDescription = connection.getDescription();
77+
if (connectionDescription != null && connectionDescription.getServerAddress() != null) {
78+
hostname = connectionDescription.getServerAddress().getHost();
79+
}
80+
81+
String dbmComment = MongoCommentInjector.getComment(span, hostname, dbName);
82+
if (dbmComment != null) {
83+
originalBsonDocument = MongoCommentInjector.injectComment(dbmComment, originalBsonDocument);
84+
}
85+
}
86+
}
87+
}

‎dd-java-agent/instrumentation/mongo/driver-3.1/src/main/java/datadog/trace/instrumentation/mongo/MongoClient31Instrumentation.java‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -57,8 +57,6 @@ public String[] helperClassNames() {
5757
packageName + ".MongoDecorator",
5858
packageName + ".MongoDecorator31",
5959
packageName + ".Context",
60-
packageName + ".MongoCommentInjector",
61-
"datadog.trace.core.database.SharedDBCommenter",
6260
packageName + ".MongoCommandListener",
6361
packageName + ".MongoCommandListener$SpanEntry"
6462
};

‎dd-java-agent/instrumentation/mongo/driver-3.4/src/main/java/datadog/trace/instrumentation/mongo/MongoClient34Instrumentation.java‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,6 @@ public String[] helperClassNames() {
6969
packageName + ".MongoDecorator",
7070
packageName + ".MongoDecorator34",
7171
packageName + ".Context",
72-
packageName + ".MongoCommentInjector",
73-
"datadog.trace.core.database.SharedDBCommenter",
7472
packageName + ".MongoCommandListener",
7573
packageName + ".MongoCommandListener$SpanEntry"
7674
};

‎dd-java-agent/instrumentation/mongo/driver-4.0/build.gradle‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ dependencies {
2727
compileOnly group: 'org.mongodb', name: 'mongodb-driver-sync', version: '4.0.0'
2828
compileOnly group: 'org.mongodb', name: 'mongodb-driver-reactivestreams', version: '4.0.0'
2929

30-
testImplementation(project(':dd-java-agent:instrumentation:mongo:common')) {
30+
implementation(project(':dd-java-agent:instrumentation:mongo:common')) {
3131
transitive = false
3232
}
3333
testImplementation(project(':dd-java-agent:instrumentation:mongo:driver-3.1')) {

0 commit comments

Comments
 (0)