Skip to content

Commit cce064e

Browse files
Yoonena-ji
Yoone
and
na-ji
authored
Add support for DBM comment injection with MongoDB (#9589)
* Added support for DBM comment injection with MongoDB * Moved SharedDBCommenter from jdbc instrumentation to dd-trace-core * Fixed SharedDBCommenterTest import * Replaced hardcoded length with constant * Attempt at fixing existing tests * Updated INJECT_COMMENT logic * Fixed containsTraceComment for test cases * Switched MongoDBMInjectionTest to MongoBaseTest * Revamped Mongo comment tests * Attempt to fix SQL injection tests * Revert changes to gradle.lockfile * PR comments: build traceparent incl. sampling flag, exception logging * Add missing helperClassNames for mongo drivers 3.1 and 3.4 * Using unit testing for trace parent changes * Prevented cloning immutable Mongo objects * Fixed smoke test (revert bug) * add new mongo instrumentation to inject APM/DBM comment This is a draft that doesn't work yet. Tests are also missing * add support for mongo driver version 4.3 I assumed that since the code was working on 4.0 and 5.6, then it worked on all the intermediary versions. It seems that the API changed in 4.3, but then was restored later in 4.4 * rename MongoCommentInjector's getComment to buildComment * only create a mutable copy of the BsonDocument when needed * move shared sql commenter to bootstrap package * document why we check for existing span on command listener * reformat trace parent util * fix server connection instrumentation advice not working on write queries Write queries on the MongoDB driver >= 4.0 and <= 4.7 weren't correctly instrumented. I didn't notice it before because the tests weren't exhaustive, and my test app was using driver v5.6. * improve check for overload instrumentation * add missing helper class for mongo 3.1 instrumentation * move mongo < 4.0 to standalone package driver-3.6 In the end, the instrumentation only works on driver > 3.6, so we decided to make a standalone package for that. This version being 9 years old seems to be enough for DBM support * add mock method to force the DefaultServerConnection40Instrumentation to run on driver > 4.0 only * Move DefaultServerInstrumentation for mongo driver 4.0 to driver 3.8 We now have two instrumentations: one for 3.6 to 3.7, and another one for 3.8 to 5.6 --------- Co-authored-by: Naji Astier <[email protected]>
1 parent 45bc239 commit cce064e

File tree

30 files changed

+2107
-126
lines changed

30 files changed

+2107
-126
lines changed

‎.github/CODEOWNERS‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,9 @@ dd-java-agent/agent-llmobs/ @DataDog/ml-observability
154154
dd-trace-core/src/main/java/datadog/trace/llmobs/ @DataDog/ml-observability
155155
dd-trace-core/src/test/groovy/datadog/trace/llmobs/ @DataDog/ml-observability
156156

157+
# @DataDog/database-monitoring
158+
datadog/trace/bootstrap/instrumentation/dbm @DataDog/database-monitoring @DataDog/apm-idm-java
159+
157160
# @DataDog/rum
158161
/internal-api/src/main/java/datadog/trace/api/rum/ @DataDog/rum
159162
/internal-api/src/test/groovy/datadog/trace/api/rum/ @DataDog/rum

‎.gitignore‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,10 @@ out/
5050
##########
5151
.cursor
5252

53+
# Vim #
54+
#######
55+
*.sw[nop]
56+
5357
# Others #
5458
##########
5559
/logs/*

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/dbm/SharedDBCommenter.java‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
package datadog.trace.bootstrap.instrumentation.dbm;
2+
3+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
4+
5+
import datadog.trace.api.BaseHash;
6+
import datadog.trace.api.Config;
7+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
8+
import datadog.trace.bootstrap.instrumentation.api.Tags;
9+
import java.io.UnsupportedEncodingException;
10+
import java.net.URLEncoder;
11+
import java.nio.charset.StandardCharsets;
12+
import org.slf4j.Logger;
13+
import org.slf4j.LoggerFactory;
14+
15+
/** Shared database comment builder for generating trace context comments for SQL DBs and MongoDB */
16+
public class SharedDBCommenter {
17+
private static final Logger log = LoggerFactory.getLogger(SharedDBCommenter.class);
18+
private static final String UTF8 = StandardCharsets.UTF_8.toString();
19+
20+
private static final char EQUALS = '=';
21+
private static final char COMMA = ',';
22+
private static final char QUOTE = '\'';
23+
24+
// Injected fields. When adding a new one, be sure to update this and the methods below.
25+
private static final String PARENT_SERVICE = encode("ddps");
26+
private static final String DATABASE_SERVICE = encode("dddbs");
27+
private static final String DD_HOSTNAME = encode("ddh");
28+
private static final String DD_DB_NAME = encode("dddb");
29+
private static final String DD_PEER_SERVICE = "ddprs";
30+
private static final String DD_ENV = encode("dde");
31+
private static final String DD_VERSION = encode("ddpv");
32+
private static final String TRACEPARENT = encode("traceparent");
33+
private static final String DD_SERVICE_HASH = encode("ddsh");
34+
35+
// Used by SQLCommenter and MongoCommentInjector to avoid duplicate comment injection
36+
public static boolean containsTraceComment(String commentContent) {
37+
return commentContent.contains(PARENT_SERVICE + "=")
38+
|| commentContent.contains(DATABASE_SERVICE + "=")
39+
|| commentContent.contains(DD_HOSTNAME + "=")
40+
|| commentContent.contains(DD_DB_NAME + "=")
41+
|| commentContent.contains(DD_PEER_SERVICE + "=")
42+
|| commentContent.contains(DD_ENV + "=")
43+
|| commentContent.contains(DD_VERSION + "=")
44+
|| commentContent.contains(TRACEPARENT + "=")
45+
|| commentContent.contains(DD_SERVICE_HASH + "=");
46+
}
47+
48+
// Build database comment content without comment delimiters such as /* */
49+
public static String buildComment(
50+
String dbService, String dbType, String hostname, String dbName, String traceParent) {
51+
52+
Config config = Config.get();
53+
StringBuilder sb = new StringBuilder();
54+
55+
int initSize = 0; // No initial content for pure comment
56+
append(sb, PARENT_SERVICE, config.getServiceName(), initSize);
57+
append(sb, DATABASE_SERVICE, dbService, initSize);
58+
append(sb, DD_HOSTNAME, hostname, initSize);
59+
append(sb, DD_DB_NAME, dbName, initSize);
60+
append(sb, DD_PEER_SERVICE, getPeerService(), initSize);
61+
append(sb, DD_ENV, config.getEnv(), initSize);
62+
append(sb, DD_VERSION, config.getVersion(), initSize);
63+
append(sb, TRACEPARENT, traceParent, initSize);
64+
65+
if (config.isDbmInjectSqlBaseHash() && config.isExperimentalPropagateProcessTagsEnabled()) {
66+
append(sb, DD_SERVICE_HASH, BaseHash.getBaseHashStr(), initSize);
67+
}
68+
69+
return sb.length() > 0 ? sb.toString() : null;
70+
}
71+
72+
private static String getPeerService() {
73+
AgentSpan span = activeSpan();
74+
Object peerService = null;
75+
if (span != null) {
76+
peerService = span.getTag(Tags.PEER_SERVICE);
77+
}
78+
return peerService != null ? peerService.toString() : null;
79+
}
80+
81+
private static String encode(String val) {
82+
try {
83+
return URLEncoder.encode(val, UTF8);
84+
} catch (UnsupportedEncodingException exe) {
85+
if (log.isDebugEnabled()) {
86+
log.debug("exception thrown while encoding comment key {}", val, exe);
87+
}
88+
}
89+
return val;
90+
}
91+
92+
private static void append(StringBuilder sb, String key, String value, int initSize) {
93+
if (null == value || value.isEmpty()) {
94+
return;
95+
}
96+
String encodedValue;
97+
try {
98+
encodedValue = URLEncoder.encode(value, UTF8);
99+
} catch (UnsupportedEncodingException e) {
100+
encodedValue = value;
101+
}
102+
if (sb.length() > initSize) {
103+
sb.append(COMMA);
104+
}
105+
sb.append(key).append(EQUALS).append(QUOTE).append(encodedValue).append(QUOTE);
106+
}
107+
}

‎dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/DBMCompatibleConnectionInstrumentation.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
77
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
88
import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DECORATE;
9+
import static datadog.trace.instrumentation.jdbc.JDBCDecorator.INJECT_COMMENT;
910
import static datadog.trace.instrumentation.jdbc.JDBCDecorator.logQueryInfoInjection;
1011
import static net.bytebuddy.matcher.ElementMatchers.returns;
1112
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -77,7 +78,9 @@ public DBMCompatibleConnectionInstrumentation() {
7778
@Override
7879
public String[] helperClassNames() {
7980
return new String[] {
80-
packageName + ".JDBCDecorator", packageName + ".SQLCommenter",
81+
packageName + ".JDBCDecorator",
82+
packageName + ".SQLCommenter",
83+
"datadog.trace.bootstrap.instrumentation.dbm.SharedDBCommenter",
8184
};
8285
}
8386

@@ -110,8 +113,7 @@ public static class ConnectionAdvice {
110113
public static String onEnter(
111114
@Advice.This Connection connection,
112115
@Advice.Argument(value = 0, readOnly = false) String sql) {
113-
// Using INJECT_COMMENT fails to update when a test calls injectSysConfig
114-
if (!DECORATE.shouldInjectSQLComment()) {
116+
if (!INJECT_COMMENT) {
115117
return sql;
116118
}
117119
if (CallDepthThreadLocalMap.incrementCallDepth(Connection.class) > 0) {

‎dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package datadog.trace.instrumentation.jdbc;
22

3+
import static datadog.trace.api.Config.DBM_PROPAGATION_MODE_FULL;
4+
import static datadog.trace.api.Config.DBM_PROPAGATION_MODE_STATIC;
35
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
46
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DBM_TRACE_INJECTED;
57
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.INSTRUMENTATION_TIME_MS;
@@ -48,8 +50,6 @@ public class JDBCDecorator extends DatabaseClientDecorator<DBInfo> {
4850
UTF8BytesString.create("java-jdbc-prepared_statement");
4951
private static final String DEFAULT_SERVICE_NAME =
5052
SpanNaming.instance().namingSchema().database().service("jdbc");
51-
public static final String DBM_PROPAGATION_MODE_STATIC = "service";
52-
public static final String DBM_PROPAGATION_MODE_FULL = "full";
5353

5454
public static final String DD_INSTRUMENTATION_PREFIX = "_DD_";
5555

@@ -417,11 +417,6 @@ public boolean shouldInjectTraceContext(DBInfo dbInfo) {
417417
return INJECT_TRACE_CONTEXT;
418418
}
419419

420-
public boolean shouldInjectSQLComment() {
421-
return Config.get().getDbmPropagationMode().equals(DBM_PROPAGATION_MODE_FULL)
422-
|| Config.get().getDbmPropagationMode().equals(DBM_PROPAGATION_MODE_STATIC);
423-
}
424-
425420
private void logInjectionErrorOnce(String vessel, Throwable t) {
426421
if (!loggedInjectionError) {
427422
loggedInjectionError = true;

‎dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java‎

Lines changed: 31 additions & 109 deletions
Original file line numberDiff line numberDiff line change
@@ -1,47 +1,19 @@
11
package datadog.trace.instrumentation.jdbc;
22

3-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
4-
5-
import datadog.trace.api.BaseHash;
6-
import datadog.trace.api.Config;
7-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
8-
import datadog.trace.bootstrap.instrumentation.api.Tags;
9-
import java.io.UnsupportedEncodingException;
10-
import java.net.URLEncoder;
11-
import java.nio.charset.StandardCharsets;
12-
import org.slf4j.Logger;
13-
import org.slf4j.LoggerFactory;
3+
import datadog.trace.bootstrap.instrumentation.dbm.SharedDBCommenter;
144

155
public class SQLCommenter {
16-
private static final Logger log = LoggerFactory.getLogger(SQLCommenter.class);
17-
private static final String UTF8 = StandardCharsets.UTF_8.toString();
18-
19-
private static final char EQUALS = '=';
20-
private static final char COMMA = ',';
21-
private static final char QUOTE = '\'';
6+
// SQL-specific constants, rest defined in SharedDBCommenter
227
private static final char SPACE = ' ';
238
private static final String OPEN_COMMENT = "/*";
249
private static final int OPEN_COMMENT_LEN = OPEN_COMMENT.length();
2510
private static final String CLOSE_COMMENT = "*/";
2611

27-
// Injected fields. When adding a new one, be sure to update this and the methods below.
28-
private static final int NUMBER_OF_FIELDS = 9;
29-
private static final String PARENT_SERVICE = encode("ddps");
30-
private static final String DATABASE_SERVICE = encode("dddbs");
31-
private static final String DD_HOSTNAME = encode("ddh");
32-
private static final String DD_DB_NAME = encode("dddb");
33-
private static final String DD_PEER_SERVICE = "ddprs";
34-
private static final String DD_ENV = encode("dde");
35-
private static final String DD_VERSION = encode("ddpv");
36-
private static final String TRACEPARENT = encode("traceparent");
37-
private static final String DD_SERVICE_HASH = encode("ddsh");
38-
39-
private static final int KEY_AND_SEPARATORS_ESTIMATED_SIZE = 10;
40-
private static final int VALUE_ESTIMATED_SIZE = 10;
41-
private static final int TRACE_PARENT_EXTRA_ESTIMATED_SIZE = 50;
42-
private static final int INJECTED_COMMENT_ESTIMATED_SIZE =
43-
NUMBER_OF_FIELDS * (KEY_AND_SEPARATORS_ESTIMATED_SIZE + VALUE_ESTIMATED_SIZE)
44-
+ TRACE_PARENT_EXTRA_ESTIMATED_SIZE;
12+
// Size estimation for StringBuilder pre-allocation
13+
private static final int SPACE_CHARS = 2; // Leading and trailing spaces
14+
private static final int COMMENT_DELIMITERS = 4; // "/*" + "*/"
15+
private static final int BUFFER_EXTRA = 4;
16+
private static final int SQL_COMMENT_OVERHEAD = SPACE_CHARS + COMMENT_DELIMITERS + BUFFER_EXTRA;
4517

4618
protected static String getFirstWord(String sql) {
4719
int beginIndex = 0;
@@ -99,9 +71,16 @@ public static String inject(
9971
return sql;
10072
}
10173

102-
Config config = Config.get();
74+
String commentContent =
75+
SharedDBCommenter.buildComment(dbService, dbType, hostname, dbName, traceParent);
76+
77+
if (commentContent == null) {
78+
return sql;
79+
}
10380

104-
StringBuilder sb = new StringBuilder(sql.length() + INJECTED_COMMENT_ESTIMATED_SIZE);
81+
// SQL-specific wrapping with /* */
82+
StringBuilder sb =
83+
new StringBuilder(sql.length() + commentContent.length() + SQL_COMMENT_OVERHEAD);
10584
int closingSemicolonIndex = indexOfClosingSemicolon(sql);
10685
if (appendComment) {
10786
if (closingSemicolonIndex > -1) {
@@ -113,22 +92,7 @@ public static String inject(
11392
}
11493

11594
sb.append(OPEN_COMMENT);
116-
int initSize = sb.length();
117-
append(sb, PARENT_SERVICE, config.getServiceName(), initSize);
118-
append(sb, DATABASE_SERVICE, dbService, initSize);
119-
append(sb, DD_HOSTNAME, hostname, initSize);
120-
append(sb, DD_DB_NAME, dbName, initSize);
121-
append(sb, DD_PEER_SERVICE, getPeerService(), initSize);
122-
append(sb, DD_ENV, config.getEnv(), initSize);
123-
append(sb, DD_VERSION, config.getVersion(), initSize);
124-
append(sb, TRACEPARENT, traceParent, initSize);
125-
if (config.isDbmInjectSqlBaseHash() && config.isExperimentalPropagateProcessTagsEnabled()) {
126-
append(sb, DD_SERVICE_HASH, BaseHash.getBaseHashStr(), initSize);
127-
}
128-
if (initSize == sb.length()) {
129-
// no comment was added
130-
return sql;
131-
}
95+
sb.append(commentContent);
13296
sb.append(CLOSE_COMMENT);
13397
if (!appendComment) {
13498
sb.append(SPACE);
@@ -142,72 +106,30 @@ public static String inject(
142106
return sb.toString();
143107
}
144108

145-
private static String getPeerService() {
146-
AgentSpan span = activeSpan();
147-
Object peerService = null;
148-
if (span != null) {
149-
peerService = span.getTag(Tags.PEER_SERVICE);
150-
}
151-
return peerService != null ? peerService.toString() : null;
152-
}
153-
154109
private static boolean hasDDComment(String sql, boolean appendComment) {
155110
if ((!sql.endsWith(CLOSE_COMMENT) && appendComment)
156111
|| ((!sql.startsWith(OPEN_COMMENT)) && !appendComment)) {
157112
return false;
158113
}
159-
int startIdx = OPEN_COMMENT_LEN;
160-
if (appendComment) {
161-
startIdx += sql.lastIndexOf(OPEN_COMMENT);
162-
}
163-
return hasMatchingSubstring(sql, startIdx, PARENT_SERVICE)
164-
|| hasMatchingSubstring(sql, startIdx, DATABASE_SERVICE)
165-
|| hasMatchingSubstring(sql, startIdx, DD_HOSTNAME)
166-
|| hasMatchingSubstring(sql, startIdx, DD_DB_NAME)
167-
|| hasMatchingSubstring(sql, startIdx, DD_PEER_SERVICE)
168-
|| hasMatchingSubstring(sql, startIdx, DD_ENV)
169-
|| hasMatchingSubstring(sql, startIdx, DD_VERSION)
170-
|| hasMatchingSubstring(sql, startIdx, TRACEPARENT)
171-
|| hasMatchingSubstring(sql, startIdx, DD_SERVICE_HASH);
172-
}
173114

174-
private static boolean hasMatchingSubstring(String sql, int startIndex, String substring) {
175-
boolean tooLong = startIndex + substring.length() >= sql.length();
176-
if (tooLong || !(sql.charAt(startIndex + substring.length()) == EQUALS)) {
177-
return false;
178-
}
179-
return sql.startsWith(substring, startIndex);
115+
String commentContent = extractCommentContent(sql, appendComment);
116+
return SharedDBCommenter.containsTraceComment(commentContent);
180117
}
181118

182-
private static String encode(String val) {
183-
try {
184-
return URLEncoder.encode(val, UTF8);
185-
} catch (UnsupportedEncodingException exe) {
186-
if (log.isDebugEnabled()) {
187-
log.debug("exception thrown while encoding sql comment key %s", exe);
188-
}
189-
}
190-
return val;
191-
}
192-
193-
private static void append(StringBuilder sb, String key, String value, int initSize) {
194-
if (null == value || value.isEmpty()) {
195-
return;
196-
}
197-
String encodedValue;
198-
try {
199-
encodedValue = URLEncoder.encode(value, UTF8);
200-
} catch (UnsupportedEncodingException e) {
201-
if (log.isDebugEnabled()) {
202-
log.debug("exception thrown while encoding sql comment %s", e);
203-
}
204-
return;
119+
private static String extractCommentContent(String sql, boolean appendComment) {
120+
int startIdx;
121+
int endIdx;
122+
if (appendComment) {
123+
startIdx = sql.lastIndexOf(OPEN_COMMENT);
124+
endIdx = sql.lastIndexOf(CLOSE_COMMENT);
125+
} else {
126+
startIdx = sql.indexOf(OPEN_COMMENT);
127+
endIdx = sql.indexOf(CLOSE_COMMENT);
205128
}
206-
207-
if (sb.length() > initSize) {
208-
sb.append(COMMA);
129+
if (startIdx != -1 && endIdx != -1 && endIdx > startIdx) {
130+
return sql.substring(startIdx + OPEN_COMMENT_LEN, endIdx);
209131
}
210-
sb.append(key).append(EQUALS).append(QUOTE).append(encodedValue).append(QUOTE);
132+
return "";
211133
}
212134

213135
/**

‎dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,11 @@ public Map<String, String> contextStore() {
5858

5959
@Override
6060
public String[] helperClassNames() {
61-
return new String[] {packageName + ".JDBCDecorator", packageName + ".SQLCommenter"};
61+
return new String[] {
62+
packageName + ".JDBCDecorator",
63+
packageName + ".SQLCommenter",
64+
"datadog.trace.bootstrap.instrumentation.dbm.SharedDBCommenter",
65+
};
6266
}
6367

6468
@Override

‎dd-java-agent/instrumentation/mongo/common/build.gradle‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,9 @@ apply from: "$rootDir/gradle/java.gradle"
22

33
dependencies {
44
compileOnly group: 'org.mongodb', name: 'mongo-java-driver', version: '3.1.0'
5+
6+
implementation project(':dd-trace-core')
7+
58
testImplementation project(':dd-java-agent:instrumentation:mongo').sourceSets.test.output
9+
testImplementation group: 'org.mongodb', name: 'mongo-java-driver', version: '3.1.0'
610
}

0 commit comments

Comments
 (0)