diff --git a/dd-java-agent/instrumentation/protobuf/src/main/java/datadog/trace/instrumentation/protobuf_java/SchemaExtractor.java b/dd-java-agent/instrumentation/protobuf/src/main/java/datadog/trace/instrumentation/protobuf_java/SchemaExtractor.java index 03593389199..a4f9d56c64d 100644 --- a/dd-java-agent/instrumentation/protobuf/src/main/java/datadog/trace/instrumentation/protobuf_java/SchemaExtractor.java +++ b/dd-java-agent/instrumentation/protobuf/src/main/java/datadog/trace/instrumentation/protobuf_java/SchemaExtractor.java @@ -44,6 +44,7 @@ public SchemaExtractor(Descriptor descriptor) { this.descriptor = descriptor; } + /** @return false if no more properties should be extracted */ public static boolean extractProperty( FieldDescriptor field, String schemaName, @@ -107,9 +108,7 @@ public static boolean extractProperty( case TYPE_MESSAGE: ref = "#/components/schemas/" + field.getMessageType().getFullName(); // Recursively add nested message schemas - if (!extractSchema(field.getMessageType(), builder, depth)) { - return false; - } + extractSchema(field.getMessageType(), builder, depth); builder.addToHash(field.getMessageType().getFullName()); break; case TYPE_BYTES: @@ -154,22 +153,20 @@ public static boolean extractProperty( schemaName, fieldName, array, type, description, ref, format, enumValues, extensions); } - public static boolean extractSchema(Descriptor descriptor, SchemaBuilder builder, int depth) { + public static void extractSchema(Descriptor descriptor, SchemaBuilder builder, int depth) { depth++; String schemaName = descriptor.getFullName(); - if (!builder.shouldExtractSchema(schemaName, depth)) { - return false; - } - // iterate fields in number order to ensure hash stability - for (FieldDescriptor field : - descriptor.getFields().stream() - .sorted(Comparator.comparingInt(FieldDescriptor::getNumber)) - .collect(Collectors.toList())) { - if (!extractProperty(field, schemaName, field.getName(), builder, depth)) { - return false; + if (builder.shouldExtractSchema(schemaName, depth)) { + for (FieldDescriptor field : + descriptor.getFields().stream() + // iterate fields in number order to ensure hash stability + .sorted(Comparator.comparingInt(FieldDescriptor::getNumber)) + .collect(Collectors.toList())) { + if (!extractProperty(field, schemaName, field.getName(), builder, depth)) { + break; // we have reached the max nb of properties to extract + } } } - return true; } public static Schema extractSchemas(Descriptor descriptor) { diff --git a/dd-java-agent/instrumentation/protobuf/src/test/groovy/com/datadog/instrumentation/protobuf/AbstractMessageInstrumentationTest.groovy b/dd-java-agent/instrumentation/protobuf/src/test/groovy/com/datadog/instrumentation/protobuf/AbstractMessageInstrumentationTest.groovy index 381c24ff3d4..551a0329ad1 100644 --- a/dd-java-agent/instrumentation/protobuf/src/test/groovy/com/datadog/instrumentation/protobuf/AbstractMessageInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/protobuf/src/test/groovy/com/datadog/instrumentation/protobuf/AbstractMessageInstrumentationTest.groovy @@ -2,6 +2,7 @@ package com.datadog.instrumentation.protobuf import com.datadog.instrumentation.protobuf.generated.Message.MyMessage import com.datadog.instrumentation.protobuf.generated.Message.OtherMessage +import com.datadog.instrumentation.protobuf.generated.Message.RecursiveMessage import com.google.protobuf.InvalidProtocolBufferException import datadog.trace.agent.test.AgentTestRunner import datadog.trace.api.DDTags @@ -16,17 +17,74 @@ class AbstractMessageInstrumentationTest extends AgentTestRunner { return true } - 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\"}" - String expectedSchemaID = "4690647329509494987" + void "test extract protobuf schema on serialize & deserialize"() { + 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" + }, + "nested":{ + "\$ref":"#/components/schemas/com.datadog.instrumentation.protobuf.generated.OtherMessage", + "extensions":{ + "x-protobuf-number":"4" + } + } + }, + "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" + }""" + expectedSchema = expectedSchema.replaceAll("[ \n]", "") // the spaces are just here to make it readable + String expectedSchemaID = "2792908287829424040" - void 'test extract protobuf schema on serialize & deserialize'() { setup: MyMessage message = MyMessage.newBuilder() - .setId("1") - .setValue("Hello from Protobuf!") - .setNested(OtherMessage.newBuilder().setName("hello").setAge(10).build()) - .build() + .setId("1") + .setValue("Hello from Protobuf!") + .setNested(OtherMessage.newBuilder().setName("hello").setAge(10).build()) + .build() when: var bytes runUnderTrace("parent_serialize") { @@ -81,12 +139,100 @@ class AbstractMessageInstrumentationTest extends AgentTestRunner { } } - void 'test error when de-serializing'() { + void "test extract protobuf schema with recursive message"() { + String expectedSchema = """{ + "components":{ + "schemas":{ + "com.datadog.instrumentation.protobuf.generated.RecursiveMessage":{ + "properties":{ + "value":{ + "extensions":{ + "x-protobuf-number":"1" + }, + "format":"int32", + "type":"integer" + }, + "next":{ + "\$ref":"#/components/schemas/com.datadog.instrumentation.protobuf.generated.RecursiveMessage", + "extensions":{ + "x-protobuf-number":"2" + } + } + }, + "type":"object" + } + } + }, + "openapi":"3.0.0" + }""" + expectedSchema = expectedSchema.replaceAll("[ \n]", "") // the spaces are just here to make it readable + String expectedSchemaID = "8377547842972884891" + + setup: + getTEST_DATA_STREAMS_MONITORING() + RecursiveMessage message = RecursiveMessage.newBuilder() + .setValue(12) + .build() + when: + byte[] bytes + runUnderTrace("parent_serialize") { + AgentSpan span = activeSpan() + span.setTag(DDTags.MANUAL_KEEP, true) + bytes = message.toByteArray() + } + runUnderTrace("parent_deserialize") { + AgentSpan span = activeSpan() + span.setTag(DDTags.MANUAL_KEEP, true) + RecursiveMessage.parseFrom(bytes) + } + TEST_WRITER.waitForTraces(2) + then: + assertTraces(2, SORT_TRACES_BY_ID) { + trace(1) { + span { + hasServiceName() + operationName "parent_serialize" + resourceName "parent_serialize" + errored false + measured false + tags { + "$DDTags.SCHEMA_DEFINITION" expectedSchema + "$DDTags.SCHEMA_WEIGHT" 1 + "$DDTags.SCHEMA_TYPE" "protobuf" + "$DDTags.SCHEMA_NAME" "com.datadog.instrumentation.protobuf.generated.RecursiveMessage" + "$DDTags.SCHEMA_OPERATION" "serialization" + "$DDTags.SCHEMA_ID" expectedSchemaID + defaultTags(false) + } + } + } + trace(1) { + span { + hasServiceName() + operationName "parent_deserialize" + resourceName "parent_deserialize" + errored false + measured false + tags { + "$DDTags.SCHEMA_DEFINITION" expectedSchema + "$DDTags.SCHEMA_WEIGHT" 1 + "$DDTags.SCHEMA_TYPE" "protobuf" + "$DDTags.SCHEMA_NAME" "com.datadog.instrumentation.protobuf.generated.RecursiveMessage" + "$DDTags.SCHEMA_OPERATION" "deserialization" + "$DDTags.SCHEMA_ID" expectedSchemaID + defaultTags(false) + } + } + } + } + } + + void "test error when de-serializing"() { setup: MyMessage message = MyMessage.newBuilder() - .setId("1") - .setValue("Hello from Protobuf!") - .build() + .setId("1") + .setValue("Hello from Protobuf!") + .build() when: runUnderTrace("parent_deserialize") { AgentSpan span = activeSpan() diff --git a/dd-java-agent/instrumentation/protobuf/src/test/groovy/com/datadog/instrumentation/protobuf/DynamicMessageInstrumentationTest.groovy b/dd-java-agent/instrumentation/protobuf/src/test/groovy/com/datadog/instrumentation/protobuf/DynamicMessageInstrumentationTest.groovy index 7393dc369ce..d4b26ec0c34 100644 --- a/dd-java-agent/instrumentation/protobuf/src/test/groovy/com/datadog/instrumentation/protobuf/DynamicMessageInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/protobuf/src/test/groovy/com/datadog/instrumentation/protobuf/DynamicMessageInstrumentationTest.groovy @@ -22,8 +22,66 @@ class DynamicMessageInstrumentationTest extends AgentTestRunner { .setValue("Hello from Protobuf!") .build() when: - 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\"}" - String expectedSchemaID = "4690647329509494987" + 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" + }, + "nested":{ + "\$ref":"#/components/schemas/com.datadog.instrumentation.protobuf.generated.OtherMessage", + "extensions":{ + "x-protobuf-number":"4" + } + } + }, + "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" + }""" + expectedSchema = expectedSchema.replaceAll("[ \n]", "") // the spaces are just here to make it readable + String expectedSchemaID = "2792908287829424040" + var bytes runUnderTrace("parent_serialize") { AgentSpan span = activeSpan() diff --git a/dd-java-agent/instrumentation/protobuf/src/test/proto/message.proto b/dd-java-agent/instrumentation/protobuf/src/test/proto/message.proto index d78165d2da7..4dc248a4332 100644 --- a/dd-java-agent/instrumentation/protobuf/src/test/proto/message.proto +++ b/dd-java-agent/instrumentation/protobuf/src/test/proto/message.proto @@ -1,3 +1,5 @@ +// run gradle target generateTestProto when you modify this + syntax = "proto3"; package com.datadog.instrumentation.protobuf.generated; @@ -13,3 +15,8 @@ message MyMessage { repeated OtherMessage other_message = 3; OtherMessage nested = 4; } + +message RecursiveMessage { + int32 value = 1; + RecursiveMessage next = 2; +} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoring.java b/dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoring.java index 2c8a897e762..f37d84695b6 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoring.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoring.java @@ -432,6 +432,7 @@ private void flush(long timestampNanos) { @Override public void clear() { timeToBucket.clear(); + schemaSamplers.clear(); } void report() {