Skip to content

Commit 2a8627d

Browse files
piochelepiotrpiochelepiotr
authored
Instrument Google protobuf (#6865)
Add Protobuf instrumentation for tracking schemas
1 parent e763318 commit 2a8627d

File tree

33 files changed

+3259
-60
lines changed

33 files changed

+3259
-60
lines changed

‎dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,7 @@
210210
2 com.google.logging.*
211211
2 com.google.longrunning.*
212212
2 com.google.protobuf.*
213+
0 com.google.protobuf.DynamicMessage
213214
2 com.google.rpc.*
214215
2 com.google.thirdparty.*
215216
2 com.google.type.*

‎dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/AvroSchemaExtractor.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ public static void tryExtractProducer(ProducerRecord record, AgentSpan span) {
1414
// don't extract schema if span is not sampled
1515
return;
1616
}
17-
int weight = AgentTracer.get().getDataStreamsMonitoring().shouldSampleSchema(record.topic());
17+
int weight = AgentTracer.get().getDataStreamsMonitoring().trySampleSchema(record.topic());
1818
if (weight == 0) {
1919
return;
2020
}

‎dd-java-agent/instrumentation/protobuf/build.gradle‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
muzzle {
2+
pass {
3+
group = "com.google.protobuf"
4+
module = "protobuf-java"
5+
versions = "[,)"
6+
assertInverse = true
7+
}
8+
}
9+
10+
apply from: "$rootDir/gradle/java.gradle"
11+
12+
dependencies {
13+
compileOnly group: 'com.google.protobuf', name: 'protobuf-java', version: '3.0.0'
14+
testImplementation group: 'com.google.protobuf', name: 'protobuf-java', version: '3.0.0'
15+
}

‎dd-java-agent/instrumentation/protobuf/src/main/java/datadog/trace/instrumentation/protobuf_java/AbstractMessageInstrumentation.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
package datadog.trace.instrumentation.protobuf_java;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.declaresMethod;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
5+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
6+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
7+
import static net.bytebuddy.matcher.ElementMatchers.*;
8+
9+
import com.google.auto.service.AutoService;
10+
import com.google.protobuf.AbstractMessage;
11+
import datadog.trace.agent.tooling.Instrumenter;
12+
import datadog.trace.agent.tooling.InstrumenterModule;
13+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
14+
import java.util.concurrent.ExecutionException;
15+
import net.bytebuddy.asm.Advice;
16+
import net.bytebuddy.description.type.TypeDescription;
17+
import net.bytebuddy.matcher.ElementMatcher;
18+
19+
@AutoService(InstrumenterModule.class)
20+
public final class AbstractMessageInstrumentation extends InstrumenterModule.Tracing
21+
implements Instrumenter.ForTypeHierarchy {
22+
23+
static final String instrumentationName = "protobuf";
24+
static final String TARGET_TYPE = "com.google.protobuf.AbstractMessage";
25+
26+
public AbstractMessageInstrumentation() {
27+
super(instrumentationName);
28+
}
29+
30+
@Override
31+
public String hierarchyMarkerType() {
32+
return TARGET_TYPE;
33+
}
34+
35+
@Override
36+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
37+
return declaresMethod(named("writeTo"))
38+
.and(extendsClass(named(hierarchyMarkerType())))
39+
.and(not(nameStartsWith("com.google.protobuf")));
40+
}
41+
42+
@Override
43+
public String[] helperClassNames() {
44+
return new String[] {
45+
packageName + ".SchemaExtractor",
46+
};
47+
}
48+
49+
@Override
50+
public void methodAdvice(MethodTransformer transformer) {
51+
transformer.applyAdvice(
52+
isMethod().and(named("writeTo")),
53+
AbstractMessageInstrumentation.class.getName() + "$WriteToAdvice");
54+
}
55+
56+
public static class WriteToAdvice {
57+
@Advice.OnMethodEnter(suppress = Throwable.class)
58+
public static void onEnter(@Advice.This AbstractMessage message) {
59+
SchemaExtractor.attachSchemaOnSpan(message, activeSpan(), SchemaExtractor.serialization);
60+
}
61+
62+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
63+
public static void stopSpan(@Advice.Thrown final Throwable throwable) {
64+
AgentSpan span = activeSpan();
65+
if (throwable != null) {
66+
span.addThrowable(
67+
throwable instanceof ExecutionException ? throwable.getCause() : throwable);
68+
}
69+
}
70+
}
71+
}

‎dd-java-agent/instrumentation/protobuf/src/main/java/datadog/trace/instrumentation/protobuf_java/AbstractParserInstrumentation.java‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
package datadog.trace.instrumentation.protobuf_java;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
6+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
7+
import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
8+
import static net.bytebuddy.matcher.ElementMatchers.not;
9+
10+
import com.google.auto.service.AutoService;
11+
import com.google.protobuf.AbstractMessage;
12+
import com.google.protobuf.MessageLite;
13+
import datadog.trace.agent.tooling.Instrumenter;
14+
import datadog.trace.agent.tooling.InstrumenterModule;
15+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
16+
import java.util.concurrent.ExecutionException;
17+
import net.bytebuddy.asm.Advice;
18+
import net.bytebuddy.description.type.TypeDescription;
19+
import net.bytebuddy.matcher.ElementMatcher;
20+
21+
@AutoService(InstrumenterModule.class)
22+
public final class AbstractParserInstrumentation extends InstrumenterModule.Tracing
23+
implements Instrumenter.ForTypeHierarchy {
24+
25+
static final String instrumentationName = "protobuf";
26+
static final String TARGET_TYPE = "com.google.protobuf.AbstractParser";
27+
28+
public AbstractParserInstrumentation() {
29+
super(instrumentationName);
30+
}
31+
32+
@Override
33+
public String hierarchyMarkerType() {
34+
return TARGET_TYPE;
35+
}
36+
37+
@Override
38+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
39+
return extendsClass(named(hierarchyMarkerType()))
40+
.and(not(nameStartsWith("com.google.protobuf")));
41+
}
42+
43+
@Override
44+
public String[] helperClassNames() {
45+
return new String[] {
46+
packageName + ".SchemaExtractor",
47+
};
48+
}
49+
50+
@Override
51+
public void methodAdvice(MethodTransformer transformer) {
52+
transformer.applyAdvice(
53+
isMethod().and(named("parsePartialFrom")),
54+
AbstractParserInstrumentation.class.getName() + "$ParseFromAdvice");
55+
}
56+
57+
public static class ParseFromAdvice {
58+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
59+
public static void stopSpan(
60+
@Advice.Thrown final Throwable throwable, @Advice.Return MessageLite message) {
61+
AgentSpan span = activeSpan();
62+
if (span == null) {
63+
return;
64+
}
65+
if (throwable != null) {
66+
span.addThrowable(
67+
throwable instanceof ExecutionException ? throwable.getCause() : throwable);
68+
}
69+
if (message instanceof AbstractMessage) {
70+
SchemaExtractor.attachSchemaOnSpan(
71+
(AbstractMessage) message, span, SchemaExtractor.deserialization);
72+
}
73+
}
74+
}
75+
}

‎dd-java-agent/instrumentation/protobuf/src/main/java/datadog/trace/instrumentation/protobuf_java/DynamicMessageInstrumentation.java‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
package datadog.trace.instrumentation.protobuf_java;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
5+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
6+
import static net.bytebuddy.matcher.ElementMatchers.isStatic;
7+
import static net.bytebuddy.matcher.ElementMatchers.returns;
8+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
9+
10+
import com.google.auto.service.AutoService;
11+
import com.google.protobuf.Descriptors.Descriptor;
12+
import datadog.trace.agent.tooling.Instrumenter;
13+
import datadog.trace.agent.tooling.InstrumenterModule;
14+
import net.bytebuddy.asm.Advice;
15+
16+
@AutoService(InstrumenterModule.class)
17+
public final class DynamicMessageInstrumentation extends InstrumenterModule.Tracing
18+
implements Instrumenter.ForSingleType {
19+
20+
static final String instrumentationName = "protobuf";
21+
static final String TARGET_TYPE = "com.google.protobuf.DynamicMessage";
22+
23+
public DynamicMessageInstrumentation() {
24+
super(instrumentationName);
25+
}
26+
27+
@Override
28+
public String instrumentedType() {
29+
return TARGET_TYPE;
30+
}
31+
32+
@Override
33+
public String[] helperClassNames() {
34+
return new String[] {
35+
packageName + ".SchemaExtractor",
36+
};
37+
}
38+
39+
@Override
40+
public void methodAdvice(MethodTransformer transformer) {
41+
transformer.applyAdvice(
42+
isMethod()
43+
.and(named("parseFrom"))
44+
.and(isStatic())
45+
.and(takesArgument(0, named("com.google.protobuf.Descriptors$Descriptor")))
46+
.and(returns(named(TARGET_TYPE))),
47+
DynamicMessageInstrumentation.class.getName() + "$ParseFromAdvice");
48+
}
49+
50+
public static class ParseFromAdvice {
51+
@Advice.OnMethodEnter(suppress = Throwable.class)
52+
public static void onEnter(@Advice.Argument(0) final Descriptor descriptor) {
53+
SchemaExtractor.attachSchemaOnSpan(descriptor, activeSpan(), SchemaExtractor.deserialization);
54+
}
55+
}
56+
}

0 commit comments

Comments
 (0)