Skip to content

Commit bb5ec51

Browse files
mccullsmcculls
authored
Merge branch 'master' into mcculls/use-classloader-index
2 parents 7b225fa + c714e59 commit bb5ec51

File tree

26 files changed

+384
-95
lines changed

26 files changed

+384
-95
lines changed

‎buildSrc/src/main/kotlin/datadog/gradle/plugin/config/ConfigInversionLinter.kt‎

Lines changed: 54 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,40 @@ class ConfigInversionLinter : Plugin<Project> {
2626
}
2727
}
2828

29+
// Data class for fields from generated class
30+
private data class LoadedConfigFields(
31+
val supported: Set<String>,
32+
val aliasMapping: Map<String, String> = emptyMap()
33+
)
34+
35+
// Cache for fields from generated class
36+
private var cachedConfigFields: LoadedConfigFields? = null
37+
38+
// Helper function to load fields from the generated class
39+
private fun loadConfigFields(
40+
mainSourceSetOutput: org.gradle.api.file.FileCollection,
41+
generatedClassName: String
42+
): LoadedConfigFields {
43+
return cachedConfigFields ?: run {
44+
val urls = mainSourceSetOutput.files.map { it.toURI().toURL() }.toTypedArray()
45+
URLClassLoader(urls, LoadedConfigFields::class.java.classLoader).use { cl ->
46+
val clazz = Class.forName(generatedClassName, true, cl)
47+
48+
val supportedField = clazz.getField("SUPPORTED").get(null)
49+
@Suppress("UNCHECKED_CAST")
50+
val supportedSet = when (supportedField) {
51+
is Set<*> -> supportedField as Set<String>
52+
is Map<*, *> -> supportedField.keys as Set<String>
53+
else -> throw IllegalStateException("SUPPORTED field must be either Set<String> or Map<String, Any>, but was ${supportedField?.javaClass}")
54+
}
55+
56+
@Suppress("UNCHECKED_CAST")
57+
val aliasMappingMap = clazz.getField("ALIAS_MAPPING").get(null) as Map<String, String>
58+
LoadedConfigFields(supportedSet, aliasMappingMap)
59+
}.also { cachedConfigFields = it }
60+
}
61+
}
62+
2963
/** Registers `logEnvVarUsages` (scan for DD_/OTEL_ tokens and fail if unsupported). */
3064
private fun registerLogEnvVarUsages(target: Project, extension: SupportedTracerConfigurations) {
3165
val ownerPath = extension.configOwnerPath
@@ -52,16 +86,11 @@ private fun registerLogEnvVarUsages(target: Project, extension: SupportedTracerC
5286
inputs.files(javaFiles)
5387
outputs.upToDateWhen { true }
5488
doLast {
55-
// 1) Build classloader from the owner project’s runtime classpath
56-
val urls = mainSourceSetOutput.get().get().files.map { it.toURI().toURL() }.toTypedArray()
57-
val supported: Set<String> = URLClassLoader(urls, javaClass.classLoader).use { cl ->
58-
// 2) Load the generated class + read static field
59-
val clazz = Class.forName(generatedFile.get(), true, cl)
60-
@Suppress("UNCHECKED_CAST")
61-
clazz.getField("SUPPORTED").get(null) as Set<String>
62-
}
89+
// 1) Load configuration fields from the generated class
90+
val configFields = loadConfigFields(mainSourceSetOutput.get().get(), generatedFile.get())
91+
val supported = configFields.supported
6392

64-
// 3) Scan our sources and compare
93+
// 2) Scan our sources and compare
6594
val repoRoot = target.projectDir.toPath()
6695
val tokenRegex = Regex("\"(?:DD_|OTEL_)[A-Za-z0-9_]+\"")
6796

@@ -79,7 +108,7 @@ private fun registerLogEnvVarUsages(target: Project, extension: SupportedTracerC
79108
}
80109
tokenRegex.findAll(raw).forEach { m ->
81110
val token = m.value.trim('"')
82-
if (token !in supported) add("$rel:${i + 1} -> Unsupported token'$token'")
111+
if (token !in supported) add("$rel:${i + 1} -> Unsupported token '$token'")
83112
}
84113
}
85114
}
@@ -167,15 +196,9 @@ private fun registerCheckConfigStringsTask(project: Project, extension: Supporte
167196
throw GradleException("Config directory not found: ${configDir.absolutePath}")
168197
}
169198

170-
val urls = mainSourceSetOutput.get().get().files.map { it.toURI().toURL() }.toTypedArray()
171-
val (supported, aliasMapping) = URLClassLoader(urls, javaClass.classLoader).use { cl ->
172-
val clazz = Class.forName(generatedFile.get(), true, cl)
173-
@Suppress("UNCHECKED_CAST")
174-
val supportedSet = clazz.getField("SUPPORTED").get(null) as Set<String>
175-
@Suppress("UNCHECKED_CAST")
176-
val aliasMappingMap = clazz.getField("ALIAS_MAPPING").get(null) as Map<String, String>
177-
Pair(supportedSet, aliasMappingMap)
178-
}
199+
val configFields = loadConfigFields(mainSourceSetOutput.get().get(), generatedFile.get())
200+
val supported = configFields.supported
201+
val aliasMapping = configFields.aliasMapping
179202

180203
var parserConfig = ParserConfiguration()
181204
parserConfig.setLanguageLevel(ParserConfiguration.LanguageLevel.JAVA_8)
@@ -192,23 +215,23 @@ private fun registerCheckConfigStringsTask(project: Project, extension: Supporte
192215
.map { it as? FieldDeclaration }
193216
.ifPresent { field ->
194217
if (field.hasModifiers(Modifier.Keyword.PUBLIC, Modifier.Keyword.STATIC, Modifier.Keyword.FINAL) &&
195-
varDecl.typeAsString == "String") {
218+
varDecl.typeAsString == "String") {
196219

197-
val fieldName = varDecl.nameAsString
198-
if (fieldName.endsWith("_DEFAULT")) return@ifPresent
199-
val init = varDecl.initializer.orElse(null) ?: return@ifPresent
220+
val fieldName = varDecl.nameAsString
221+
if (fieldName.endsWith("_DEFAULT")) return@ifPresent
222+
val init = varDecl.initializer.orElse(null) ?: return@ifPresent
200223

201-
if (init !is StringLiteralExpr) return@ifPresent
202-
val rawValue = init.value
224+
if (init !is StringLiteralExpr) return@ifPresent
225+
val rawValue = init.value
203226

204-
val normalized = normalize(rawValue)
205-
if (normalized !in supported && normalized !in aliasMapping) {
206-
val line = varDecl.range.map { it.begin.line }.orElse(1)
207-
add("$fileName:$line -> Config '$rawValue' normalizes to '$normalized' " +
208-
"which is missing from '${extension.jsonFile.get()}'")
227+
val normalized = normalize(rawValue)
228+
if (normalized !in supported && normalized !in aliasMapping) {
229+
val line = varDecl.range.map { it.begin.line }.orElse(1)
230+
add("$fileName:$line -> Config '$rawValue' normalizes to '$normalized' " +
231+
"which is missing from '${extension.jsonFile.get()}'")
232+
}
209233
}
210234
}
211-
}
212235
}
213236
}
214237
}

‎dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/config/ConfigurationApiImpl.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -341,7 +341,8 @@ public Map<TestSetting, Map<String, Collection<TestFQN>>> getTestManagementTests
341341
tracerEnvironment.getRepositoryUrl(),
342342
commitMessage,
343343
tracerEnvironment.getConfigurations().getTestBundle(),
344-
commitSha)));
344+
commitSha,
345+
tracerEnvironment.getBranch())));
345346
String json = testManagementRequestAdapter.toJson(request);
346347
RequestBody requestBody = RequestBody.create(JSON, json);
347348
TestManagementTestsDto testManagementTestsDto =
@@ -509,13 +510,15 @@ private static final class TestManagementDto {
509510

510511
private final String module;
511512
private final String sha;
513+
private final String branch;
512514

513515
private TestManagementDto(
514-
String repositoryUrl, String commitMessage, String module, String sha) {
516+
String repositoryUrl, String commitMessage, String module, String sha, String branch) {
515517
this.repositoryUrl = repositoryUrl;
516518
this.commitMessage = commitMessage;
517519
this.module = module;
518520
this.sha = sha;
521+
this.branch = branch;
519522
}
520523
}
521524

‎dd-java-agent/agent-ci-visibility/src/test/resources/datadog/trace/civisibility/config/test-management-tests-request.ftl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@
55
"attributes": {
66
"repository_url" : "${tracerEnvironment.repositoryUrl}",
77
"commit_message" : "${tracerEnvironment.commitMessage}",
8-
"sha" : "${tracerEnvironment.sha}"
9-
}
8+
"sha" : "${tracerEnvironment.sha}",
9+
"branch" : "${tracerEnvironment.branch}"
1010
}
1111
}
1212
}

‎dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import static datadog.trace.instrumentation.aws.v2.eventbridge.TextMapInjectAdapter.SETTER;
77

88
import datadog.context.Context;
9+
import datadog.trace.api.Config;
910
import datadog.trace.api.datastreams.DataStreamsContext;
1011
import datadog.trace.api.datastreams.DataStreamsTags;
1112
import datadog.trace.api.datastreams.PathwayContext;
@@ -34,7 +35,8 @@ public class EventBridgeInterceptor implements ExecutionInterceptor {
3435

3536
@Override
3637
public SdkRequest modifyRequest(ModifyRequest context, ExecutionAttributes executionAttributes) {
37-
if (!(context.request() instanceof PutEventsRequest)) {
38+
if (!(context.request() instanceof PutEventsRequest)
39+
|| !Config.get().isEventbridgeInjectDatadogAttributeEnabled()) {
3840
return context.request();
3941
}
4042

‎dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/test/groovy/EventBridgeClientTest.groovy‎

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@ import datadog.trace.agent.test.InstrumentationSpecification
22
import datadog.trace.api.DDSpanTypes
33
import datadog.trace.api.config.GeneralConfig
44
import groovy.json.JsonSlurper
5+
import java.time.Duration
6+
import java.util.concurrent.CompletableFuture
57
import org.testcontainers.containers.GenericContainer
68
import org.testcontainers.utility.DockerImageName
79
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials
@@ -17,9 +19,6 @@ import software.amazon.awssdk.services.sqs.SqsClient
1719
import software.amazon.awssdk.services.sqs.model.QueueAttributeName
1820
import spock.lang.Shared
1921

20-
import java.time.Duration
21-
import java.util.concurrent.CompletableFuture
22-
2322
class EventBridgeClientTest extends InstrumentationSpecification {
2423
static final LOCALSTACK = new GenericContainer(DockerImageName.parse("localstack/localstack:4.2.0"))
2524
.withExposedPorts(4566)
@@ -495,4 +494,31 @@ class EventBridgeClientTest extends InstrumentationSpecification {
495494
'tracestate'
496495
]
497496
}
497+
498+
def "datadog context is not injected when eventbridgeInjectDatadogAttribute is disabled"() {
499+
setup:
500+
injectSysConfig("eventbridge.inject.datadog.attribute.enabled", "false")
501+
502+
when:
503+
TEST_WRITER.clear()
504+
eventBridgeClient.putEvents { req ->
505+
req.entries(
506+
PutEventsRequestEntry.builder()
507+
.source("com.example")
508+
.detailType("test-no-inject")
509+
.detail('{"message":"no-inject"}')
510+
.eventBusName(testBusARN)
511+
.build()
512+
)
513+
}
514+
515+
def message = sqsClient.receiveMessage { it.queueUrl(testQueueURL).waitTimeSeconds(3) }.messages().get(0)
516+
def messageBody = new JsonSlurper().parseText(message.body())
517+
518+
then:
519+
def detail = messageBody["detail"]
520+
assert detail instanceof Map
521+
assert detail["message"] == "no-inject"
522+
assert detail["_datadog"] == null
523+
}
498524
}

‎dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -88,15 +88,15 @@ static AgentScope enter(
8888
if (CallDepthThreadLocalMap.incrementCallDepth(RequestHandler.class) > 0) {
8989
return null;
9090
}
91-
92-
AgentSpanContext lambdaContext = AgentTracer.get().notifyExtensionStart(in);
91+
String lambdaRequestId = awsContext.getAwsRequestId();
92+
AgentSpanContext lambdaContext = AgentTracer.get().notifyExtensionStart(in, lambdaRequestId);
9393
final AgentSpan span;
9494
if (null == lambdaContext) {
9595
span = startSpan(INVOCATION_SPAN_NAME);
9696
} else {
9797
span = startSpan(INVOCATION_SPAN_NAME, lambdaContext);
9898
}
99-
span.setTag("request_id", awsContext.getAwsRequestId());
99+
span.setTag("request_id", lambdaRequestId);
100100

101101
final AgentScope scope = activateSpan(span);
102102
return scope;
@@ -107,6 +107,7 @@ static void exit(
107107
@Origin String method,
108108
@Enter final AgentScope scope,
109109
@Advice.Argument(1) final Object result,
110+
@Advice.Argument(2) final Context awsContext,
110111
@Advice.Thrown final Throwable throwable) {
111112

112113
if (scope == null) {
@@ -120,8 +121,10 @@ static void exit(
120121
if (throwable != null) {
121122
span.addThrowable(throwable);
122123
}
124+
String lambdaRequestId = awsContext.getAwsRequestId();
125+
123126
span.finish();
124-
AgentTracer.get().notifyExtensionEnd(span, result, null != throwable);
127+
AgentTracer.get().notifyExtensionEnd(span, result, null != throwable, lambdaRequestId);
125128
} finally {
126129
scope.close();
127130
}

‎dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
44

55
import datadog.context.Context;
6+
import datadog.trace.api.Config;
67
import datadog.trace.bootstrap.InstanceStore;
78
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
89
import software.amazon.awssdk.core.SdkRequest;
@@ -23,10 +24,14 @@ public SfnInterceptor() {}
2324

2425
@Override
2526
public SdkRequest modifyRequest(ModifyRequest context, ExecutionAttributes executionAttributes) {
27+
SdkRequest request = context.request();
28+
if (!Config.get().isSfnInjectDatadogAttributeEnabled()) {
29+
return request;
30+
}
2631
try {
2732
return modifyRequestImpl(context, executionAttributes);
2833
} catch (Exception e) {
29-
return context.request();
34+
return request;
3035
}
3136
}
3237

‎dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/test/groovy/SfnClientTest.groovy‎

Lines changed: 28 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,21 @@
1+
import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
2+
13
import datadog.trace.agent.test.naming.VersionedNamingTestBase
24
import datadog.trace.agent.test.utils.TraceUtils
35
import datadog.trace.api.DDSpanTypes
46
import datadog.trace.bootstrap.instrumentation.api.Tags
57
import groovy.json.JsonSlurper
8+
import java.time.Duration
69
import org.testcontainers.containers.GenericContainer
710
import org.testcontainers.utility.DockerImageName
11+
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials
12+
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider
13+
import software.amazon.awssdk.regions.Region
814
import software.amazon.awssdk.services.sfn.SfnClient
915
import software.amazon.awssdk.services.sfn.model.SfnException
1016
import software.amazon.awssdk.services.sfn.model.StartExecutionResponse
11-
import software.amazon.awssdk.regions.Region
12-
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider
13-
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials
1417
import spock.lang.Shared
1518

16-
import java.time.Duration
17-
18-
import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
19-
20-
2119
abstract class SfnClientTest extends VersionedNamingTestBase {
2220
@Shared GenericContainer localStack
2321
@Shared SfnClient sfnClient
@@ -117,6 +115,28 @@ abstract class SfnClientTest extends VersionedNamingTestBase {
117115
input["_datadog"]["x-datadog-tags"] != null
118116
}
119117

118+
def "datadog context is not injected when SfnInjectDatadogAttribute is disabled"() {
119+
setup:
120+
injectSysConfig("sfn.inject.datadog.attribute.enabled", "false")
121+
122+
when:
123+
StartExecutionResponse response = sfnClient.startExecution { builder ->
124+
builder.stateMachineArn(testStateMachineARN)
125+
.input("{\"key\": \"value\"}")
126+
.build()
127+
}
128+
129+
then:
130+
def execution = sfnClient.describeExecution { builder ->
131+
builder.executionArn(response.executionArn())
132+
.build()
133+
}
134+
135+
def input = new JsonSlurper().parseText(execution.input())
136+
assert input["key"] == "value"
137+
assert input["_datadog"] == null
138+
}
139+
120140
def "AWS rejects invalid JSON but instrumentation does not error"() {
121141
when:
122142
sfnClient.startExecution { b ->

‎dd-java-agent/instrumentation/aws-java/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import com.amazonaws.services.sns.model.PublishBatchRequestEntry;
1414
import com.amazonaws.services.sns.model.PublishRequest;
1515
import datadog.context.Context;
16+
import datadog.trace.api.Config;
1617
import datadog.trace.api.datastreams.DataStreamsContext;
1718
import datadog.trace.api.datastreams.DataStreamsTags;
1819
import datadog.trace.bootstrap.ContextStore;
@@ -49,6 +50,9 @@ private ByteBuffer getMessageAttributeValueToInject(
4950

5051
@Override
5152
public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request) {
53+
if (!Config.get().isSnsInjectDatadogAttributeEnabled()) {
54+
return request;
55+
}
5256
// Injecting the trace context into SNS messageAttributes.
5357
if (request instanceof PublishRequest) {
5458
PublishRequest pRequest = (PublishRequest) request;

0 commit comments

Comments
 (0)