Skip to content

Commit 5bc1a33

Browse files
vandonrvandonr
authored
fix bug on proto schema extraction (#8403)
* fix bug on proto schema extraction * formatting in test * use triple quoted strings * fix tests
1 parent 0dafd10 commit 5bc1a33

File tree

5 files changed

+237
-28
lines changed

5 files changed

+237
-28
lines changed

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

Lines changed: 12 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ public SchemaExtractor(Descriptor descriptor) {
4444
this.descriptor = descriptor;
4545
}
4646

47+
/** @return false if no more properties should be extracted */
4748
public static boolean extractProperty(
4849
FieldDescriptor field,
4950
String schemaName,
@@ -107,9 +108,7 @@ public static boolean extractProperty(
107108
case TYPE_MESSAGE:
108109
ref = "#/components/schemas/" + field.getMessageType().getFullName();
109110
// Recursively add nested message schemas
110-
if (!extractSchema(field.getMessageType(), builder, depth)) {
111-
return false;
112-
}
111+
extractSchema(field.getMessageType(), builder, depth);
113112
builder.addToHash(field.getMessageType().getFullName());
114113
break;
115114
case TYPE_BYTES:
@@ -154,22 +153,20 @@ public static boolean extractProperty(
154153
schemaName, fieldName, array, type, description, ref, format, enumValues, extensions);
155154
}
156155

157-
public static boolean extractSchema(Descriptor descriptor, SchemaBuilder builder, int depth) {
156+
public static void extractSchema(Descriptor descriptor, SchemaBuilder builder, int depth) {
158157
depth++;
159158
String schemaName = descriptor.getFullName();
160-
if (!builder.shouldExtractSchema(schemaName, depth)) {
161-
return false;
162-
}
163-
// iterate fields in number order to ensure hash stability
164-
for (FieldDescriptor field :
165-
descriptor.getFields().stream()
166-
.sorted(Comparator.comparingInt(FieldDescriptor::getNumber))
167-
.collect(Collectors.toList())) {
168-
if (!extractProperty(field, schemaName, field.getName(), builder, depth)) {
169-
return false;
159+
if (builder.shouldExtractSchema(schemaName, depth)) {
160+
for (FieldDescriptor field :
161+
descriptor.getFields().stream()
162+
// iterate fields in number order to ensure hash stability
163+
.sorted(Comparator.comparingInt(FieldDescriptor::getNumber))
164+
.collect(Collectors.toList())) {
165+
if (!extractProperty(field, schemaName, field.getName(), builder, depth)) {
166+
break; // we have reached the max nb of properties to extract
167+
}
170168
}
171169
}
172-
return true;
173170
}
174171

175172
public static Schema extractSchemas(Descriptor descriptor) {

‎dd-java-agent/instrumentation/protobuf/src/test/groovy/com/datadog/instrumentation/protobuf/AbstractMessageInstrumentationTest.groovy‎

Lines changed: 157 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package com.datadog.instrumentation.protobuf
22

33
import com.datadog.instrumentation.protobuf.generated.Message.MyMessage
44
import com.datadog.instrumentation.protobuf.generated.Message.OtherMessage
5+
import com.datadog.instrumentation.protobuf.generated.Message.RecursiveMessage
56
import com.google.protobuf.InvalidProtocolBufferException
67
import datadog.trace.agent.test.AgentTestRunner
78
import datadog.trace.api.DDTags
@@ -16,17 +17,74 @@ class AbstractMessageInstrumentationTest extends AgentTestRunner {
1617
return true
1718
}
1819

19-
String expectedSchema = "{\"components\":{\"schemas\":{\"com.datadog.instrumentation.protobuf.generated.MyMessage\":{\"properties\":{\"id\":{\"extensions\":{\"x-protobuf-number\":\"1\"},\"type\":\"string\"},\"value\":{\"extensions\":{\"x-protobuf-number\":\"2\"},\"type\":\"string\"},\"other_message\":{\"extensions\":{\"x-protobuf-number\":\"3\"},\"items\":{\"\$ref\":\"#/components/schemas/com.datadog.instrumentation.protobuf.generated.OtherMessage\"},\"type\":\"array\"}},\"type\":\"object\"},\"com.datadog.instrumentation.protobuf.generated.OtherMessage\":{\"properties\":{\"name\":{\"extensions\":{\"x-protobuf-number\":\"1\"},\"type\":\"string\"},\"age\":{\"extensions\":{\"x-protobuf-number\":\"2\"},\"format\":\"int32\",\"type\":\"integer\"}},\"type\":\"object\"}}},\"openapi\":\"3.0.0\"}"
20-
String expectedSchemaID = "4690647329509494987"
20+
void "test extract protobuf schema on serialize & deserialize"() {
2121

22+
String expectedSchema = """{
23+
"components":{
24+
"schemas":{
25+
"com.datadog.instrumentation.protobuf.generated.MyMessage":{
26+
"properties":{
27+
"id":{
28+
"extensions":{
29+
"x-protobuf-number":"1"
30+
},
31+
"type":"string"
32+
},
33+
"value":{
34+
"extensions":{
35+
"x-protobuf-number":"2"
36+
},
37+
"type":"string"
38+
},
39+
"other_message":{
40+
"extensions":{
41+
"x-protobuf-number":"3"
42+
},
43+
"items":{
44+
"\$ref":"#/components/schemas/com.datadog.instrumentation.protobuf.generated.OtherMessage"
45+
},
46+
"type":"array"
47+
},
48+
"nested":{
49+
"\$ref":"#/components/schemas/com.datadog.instrumentation.protobuf.generated.OtherMessage",
50+
"extensions":{
51+
"x-protobuf-number":"4"
52+
}
53+
}
54+
},
55+
"type":"object"
56+
},
57+
"com.datadog.instrumentation.protobuf.generated.OtherMessage":{
58+
"properties":{
59+
"name":{
60+
"extensions":{
61+
"x-protobuf-number":"1"
62+
},
63+
"type":"string"
64+
},
65+
"age":{
66+
"extensions":{
67+
"x-protobuf-number":"2"
68+
},
69+
"format":"int32",
70+
"type":"integer"
71+
}
72+
},
73+
"type":"object"
74+
}
75+
}
76+
},
77+
"openapi":"3.0.0"
78+
}"""
79+
expectedSchema = expectedSchema.replaceAll("[ \n]", "") // the spaces are just here to make it readable
80+
String expectedSchemaID = "2792908287829424040"
2281

23-
void 'test extract protobuf schema on serialize & deserialize'() {
2482
setup:
2583
MyMessage message = MyMessage.newBuilder()
26-
.setId("1")
27-
.setValue("Hello from Protobuf!")
28-
.setNested(OtherMessage.newBuilder().setName("hello").setAge(10).build())
29-
.build()
84+
.setId("1")
85+
.setValue("Hello from Protobuf!")
86+
.setNested(OtherMessage.newBuilder().setName("hello").setAge(10).build())
87+
.build()
3088
when:
3189
var bytes
3290
runUnderTrace("parent_serialize") {
@@ -81,12 +139,100 @@ class AbstractMessageInstrumentationTest extends AgentTestRunner {
81139
}
82140
}
83141

84-
void 'test error when de-serializing'() {
142+
void "test extract protobuf schema with recursive message"() {
143+
String expectedSchema = """{
144+
"components":{
145+
"schemas":{
146+
"com.datadog.instrumentation.protobuf.generated.RecursiveMessage":{
147+
"properties":{
148+
"value":{
149+
"extensions":{
150+
"x-protobuf-number":"1"
151+
},
152+
"format":"int32",
153+
"type":"integer"
154+
},
155+
"next":{
156+
"\$ref":"#/components/schemas/com.datadog.instrumentation.protobuf.generated.RecursiveMessage",
157+
"extensions":{
158+
"x-protobuf-number":"2"
159+
}
160+
}
161+
},
162+
"type":"object"
163+
}
164+
}
165+
},
166+
"openapi":"3.0.0"
167+
}"""
168+
expectedSchema = expectedSchema.replaceAll("[ \n]", "") // the spaces are just here to make it readable
169+
String expectedSchemaID = "8377547842972884891"
170+
171+
setup:
172+
getTEST_DATA_STREAMS_MONITORING()
173+
RecursiveMessage message = RecursiveMessage.newBuilder()
174+
.setValue(12)
175+
.build()
176+
when:
177+
byte[] bytes
178+
runUnderTrace("parent_serialize") {
179+
AgentSpan span = activeSpan()
180+
span.setTag(DDTags.MANUAL_KEEP, true)
181+
bytes = message.toByteArray()
182+
}
183+
runUnderTrace("parent_deserialize") {
184+
AgentSpan span = activeSpan()
185+
span.setTag(DDTags.MANUAL_KEEP, true)
186+
RecursiveMessage.parseFrom(bytes)
187+
}
188+
TEST_WRITER.waitForTraces(2)
189+
then:
190+
assertTraces(2, SORT_TRACES_BY_ID) {
191+
trace(1) {
192+
span {
193+
hasServiceName()
194+
operationName "parent_serialize"
195+
resourceName "parent_serialize"
196+
errored false
197+
measured false
198+
tags {
199+
"$DDTags.SCHEMA_DEFINITION" expectedSchema
200+
"$DDTags.SCHEMA_WEIGHT" 1
201+
"$DDTags.SCHEMA_TYPE" "protobuf"
202+
"$DDTags.SCHEMA_NAME" "com.datadog.instrumentation.protobuf.generated.RecursiveMessage"
203+
"$DDTags.SCHEMA_OPERATION" "serialization"
204+
"$DDTags.SCHEMA_ID" expectedSchemaID
205+
defaultTags(false)
206+
}
207+
}
208+
}
209+
trace(1) {
210+
span {
211+
hasServiceName()
212+
operationName "parent_deserialize"
213+
resourceName "parent_deserialize"
214+
errored false
215+
measured false
216+
tags {
217+
"$DDTags.SCHEMA_DEFINITION" expectedSchema
218+
"$DDTags.SCHEMA_WEIGHT" 1
219+
"$DDTags.SCHEMA_TYPE" "protobuf"
220+
"$DDTags.SCHEMA_NAME" "com.datadog.instrumentation.protobuf.generated.RecursiveMessage"
221+
"$DDTags.SCHEMA_OPERATION" "deserialization"
222+
"$DDTags.SCHEMA_ID" expectedSchemaID
223+
defaultTags(false)
224+
}
225+
}
226+
}
227+
}
228+
}
229+
230+
void "test error when de-serializing"() {
85231
setup:
86232
MyMessage message = MyMessage.newBuilder()
87-
.setId("1")
88-
.setValue("Hello from Protobuf!")
89-
.build()
233+
.setId("1")
234+
.setValue("Hello from Protobuf!")
235+
.build()
90236
when:
91237
runUnderTrace("parent_deserialize") {
92238
AgentSpan span = activeSpan()

‎dd-java-agent/instrumentation/protobuf/src/test/groovy/com/datadog/instrumentation/protobuf/DynamicMessageInstrumentationTest.groovy‎

Lines changed: 60 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,66 @@ class DynamicMessageInstrumentationTest extends AgentTestRunner {
2222
.setValue("Hello from Protobuf!")
2323
.build()
2424
when:
25-
String expectedSchema = "{\"components\":{\"schemas\":{\"com.datadog.instrumentation.protobuf.generated.MyMessage\":{\"properties\":{\"id\":{\"extensions\":{\"x-protobuf-number\":\"1\"},\"type\":\"string\"},\"value\":{\"extensions\":{\"x-protobuf-number\":\"2\"},\"type\":\"string\"},\"other_message\":{\"extensions\":{\"x-protobuf-number\":\"3\"},\"items\":{\"\$ref\":\"#/components/schemas/com.datadog.instrumentation.protobuf.generated.OtherMessage\"},\"type\":\"array\"}},\"type\":\"object\"},\"com.datadog.instrumentation.protobuf.generated.OtherMessage\":{\"properties\":{\"name\":{\"extensions\":{\"x-protobuf-number\":\"1\"},\"type\":\"string\"},\"age\":{\"extensions\":{\"x-protobuf-number\":\"2\"},\"format\":\"int32\",\"type\":\"integer\"}},\"type\":\"object\"}}},\"openapi\":\"3.0.0\"}"
26-
String expectedSchemaID = "4690647329509494987"
25+
String expectedSchema = """{
26+
"components":{
27+
"schemas":{
28+
"com.datadog.instrumentation.protobuf.generated.MyMessage":{
29+
"properties":{
30+
"id":{
31+
"extensions":{
32+
"x-protobuf-number":"1"
33+
},
34+
"type":"string"
35+
},
36+
"value":{
37+
"extensions":{
38+
"x-protobuf-number":"2"
39+
},
40+
"type":"string"
41+
},
42+
"other_message":{
43+
"extensions":{
44+
"x-protobuf-number":"3"
45+
},
46+
"items":{
47+
"\$ref":"#/components/schemas/com.datadog.instrumentation.protobuf.generated.OtherMessage"
48+
},
49+
"type":"array"
50+
},
51+
"nested":{
52+
"\$ref":"#/components/schemas/com.datadog.instrumentation.protobuf.generated.OtherMessage",
53+
"extensions":{
54+
"x-protobuf-number":"4"
55+
}
56+
}
57+
},
58+
"type":"object"
59+
},
60+
"com.datadog.instrumentation.protobuf.generated.OtherMessage":{
61+
"properties":{
62+
"name":{
63+
"extensions":{
64+
"x-protobuf-number":"1"
65+
},
66+
"type":"string"
67+
},
68+
"age":{
69+
"extensions":{
70+
"x-protobuf-number":"2"
71+
},
72+
"format":"int32",
73+
"type":"integer"
74+
}
75+
},
76+
"type":"object"
77+
}
78+
}
79+
},
80+
"openapi":"3.0.0"
81+
}"""
82+
expectedSchema = expectedSchema.replaceAll("[ \n]", "") // the spaces are just here to make it readable
83+
String expectedSchemaID = "2792908287829424040"
84+
2785
var bytes
2886
runUnderTrace("parent_serialize") {
2987
AgentSpan span = activeSpan()

‎dd-java-agent/instrumentation/protobuf/src/test/proto/message.proto‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
// run gradle target generateTestProto when you modify this
2+
13
syntax = "proto3";
24

35
package com.datadog.instrumentation.protobuf.generated;
@@ -13,3 +15,8 @@ message MyMessage {
1315
repeated OtherMessage other_message = 3;
1416
OtherMessage nested = 4;
1517
}
18+
19+
message RecursiveMessage {
20+
int32 value = 1;
21+
RecursiveMessage next = 2;
22+
}

‎dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoring.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -432,6 +432,7 @@ private void flush(long timestampNanos) {
432432
@Override
433433
public void clear() {
434434
timeToBucket.clear();
435+
schemaSamplers.clear();
435436
}
436437

437438
void report() {

0 commit comments

Comments
 (0)