Skip to content

Incorporating TagMap into the tracer #8589

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jul 15, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ public AbstractTestSession(
AgentSpanContext traceContext =
new TagContext(
CIConstants.CIAPP_TEST_ORIGIN,
Collections.emptyMap(),
Copy link
Contributor Author

@dougqh dougqh Mar 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have mixed feelings about this particular change. In effect, the constructor previously required the user to pass a mutable map. However if the provided Map was empty, the class would lazily construction a mutable Map to take place of the empty Map.

Because TagMap does not have an O(1) isEmpty, I didn't want to stick with this pattern.

What could be done instead is to pass TagMap.EMPTY and then check via a reference equality check. If others prefer that, I can adjust accordingly.

null,
null,
null,
PrioritySampling.UNSET,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,6 @@
import datadog.trace.civisibility.test.ExecutionResults;
import java.lang.reflect.Method;
import java.util.Collection;
import java.util.Collections;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import javax.annotation.Nonnull;
Expand Down Expand Up @@ -105,8 +104,7 @@ public TestImpl(

this.context = new TestContextImpl(coverageStore);

AgentSpanContext traceContext =
new TagContext(CIConstants.CIAPP_TEST_ORIGIN, Collections.emptyMap());
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same mutable empty Map issue

AgentSpanContext traceContext = new TagContext(CIConstants.CIAPP_TEST_ORIGIN, null);
AgentTracer.SpanBuilder spanBuilder =
AgentTracer.get()
.buildSpan(CI_VISIBILITY_INSTRUMENTATION_NAME, testDecorator.component() + ".test")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import com.datadog.debugger.util.ExceptionHelper;
import com.datadog.debugger.util.TestSnapshotListener;
import datadog.trace.api.Config;
import datadog.trace.api.TagMap;
import datadog.trace.bootstrap.debugger.CapturedContext;
import datadog.trace.bootstrap.debugger.CapturedStackFrame;
import datadog.trace.bootstrap.debugger.MethodLocation;
Expand All @@ -41,7 +42,6 @@
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Deque;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
Expand All @@ -57,7 +57,7 @@ public class DefaultExceptionDebuggerTest {
private ConfigurationUpdater configurationUpdater;
private DefaultExceptionDebugger exceptionDebugger;
private TestSnapshotListener listener;
private Map<String, Object> spanTags = new HashMap<>();
Copy link
Contributor Author

@dougqh dougqh Mar 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the primary type of change that I've made throughout -- replacing HashMaps with TagMap-s.
To get the benefit of TagMap's quick Map-to-Map copying ability, both the source and destination Map need to be TagMap-s.

private TagMap spanTags = TagMap.create();

@BeforeEach
public void setUp() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import com.datadog.iast.model.Vulnerability
import com.datadog.iast.model.VulnerabilityType
import com.datadog.iast.overhead.Operation
import com.datadog.iast.overhead.OverheadController
import datadog.trace.api.TagMap
import datadog.trace.api.gateway.Flow
import datadog.trace.api.iast.InstrumentationBridge
import datadog.trace.api.internal.TraceSegment
Expand Down Expand Up @@ -45,10 +46,10 @@ class HstsMissingHeaderModuleTest extends IastModuleImplTestBase {
final handler = new RequestEndedHandler(dependencies)
ctx.xForwardedProto = 'https'
ctx.contentType = "text/html"
span.getTags() >> [
span.getTags() >> TagMap.fromMap([
'http.url': 'https://localhost/a',
'http.status_code': 200i
]
])

when:
def flow = handler.apply(reqCtx, span)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import com.datadog.iast.Reporter
import com.datadog.iast.RequestEndedHandler
import com.datadog.iast.model.Vulnerability
import com.datadog.iast.model.VulnerabilityType
import datadog.trace.api.TagMap
import datadog.trace.api.gateway.Flow
import datadog.trace.api.iast.InstrumentationBridge
import datadog.trace.api.iast.sink.InsecureAuthProtocolModule
Expand Down Expand Up @@ -42,9 +43,9 @@ class InsecureAuthProtocolModuleTest extends IastModuleImplTestBase{
given:
final handler = new RequestEndedHandler(dependencies)
ctx.authorization = value
span.getTags() >> [
span.getTags() >> TagMap.fromMap([
'http.status_code': status_code
]
])

when:
def flow = handler.apply(reqCtx, span)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import com.datadog.iast.Reporter
import com.datadog.iast.RequestEndedHandler
import com.datadog.iast.model.Vulnerability
import com.datadog.iast.model.VulnerabilityType
import datadog.trace.api.TagMap
import datadog.trace.api.gateway.Flow
import datadog.trace.api.iast.InstrumentationBridge
import datadog.trace.api.internal.TraceSegment
Expand Down Expand Up @@ -33,9 +34,9 @@ class XContentTypeOptionsModuleTest extends IastModuleImplTestBase {
given:
final handler = new RequestEndedHandler(dependencies)
ctx.contentType = "text/html"
span.getTags() >> [
span.getTags() >> TagMap.fromMap([
'http.status_code': 200i
]
])

when:
def flow = handler.apply(reqCtx, span)
Expand All @@ -56,10 +57,10 @@ class XContentTypeOptionsModuleTest extends IastModuleImplTestBase {
final handler = new RequestEndedHandler(dependencies)
ctx.xForwardedProto = 'https'
ctx.contentType = "text/html"
span.getTags() >> [
span.getTags() >> TagMap.fromMap([
'http.url': url,
'http.status_code': status
]
])

when:
def flow = handler.apply(reqCtx, span)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import datadog.remoteconfig.Product
import datadog.remoteconfig.state.ConfigKey
import datadog.remoteconfig.state.ProductListener
import datadog.trace.api.Config
import datadog.trace.api.TagMap
import datadog.trace.api.gateway.Flow
import datadog.trace.api.gateway.IGSpanInfo
import datadog.trace.api.gateway.RequestContext
Expand Down Expand Up @@ -91,7 +92,7 @@ class AppSecSystemSpecification extends DDSpecification {
requestEndedCB.apply(requestContext, span)

then:
1 * span.getTags() >> ['http.client_ip':'1.1.1.1']
1 * span.getTags() >> TagMap.fromMap(['http.client_ip':'1.1.1.1'])
1 * subService.registerCallback(EVENTS.requestEnded(), _) >> { requestEndedCB = it[1]; null }
1 * requestContext.getData(RequestContextSlot.APPSEC) >> appSecReqCtx
1 * requestContext.traceSegment >> traceSegment
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import com.datadog.appsec.event.data.DataBundle
import com.datadog.appsec.event.data.KnownAddresses
import com.datadog.appsec.report.AppSecEvent
import com.datadog.appsec.report.AppSecEventWrapper
import datadog.trace.api.TagMap
import datadog.trace.api.ProductTraceSource
import datadog.trace.api.config.GeneralConfig
import static datadog.trace.api.config.IastConfig.IAST_DEDUPLICATION_ENABLED
Expand Down Expand Up @@ -173,7 +174,7 @@ class GatewayBridgeSpecification extends DDSpecification {
def flow = requestEndedCB.apply(mockCtx, spanInfo)

then:
1 * spanInfo.getTags() >> ['http.client_ip': '1.1.1.1']
1 * spanInfo.getTags() >> TagMap.fromMap(['http.client_ip': '1.1.1.1'])
1 * mockAppSecCtx.transferCollectedEvents() >> [event]
1 * mockAppSecCtx.peerAddress >> '2001::1'
1 * mockAppSecCtx.close()
Expand Down Expand Up @@ -212,7 +213,7 @@ class GatewayBridgeSpecification extends DDSpecification {

then:
1 * mockAppSecCtx.transferCollectedEvents() >> [Stub(AppSecEvent)]
1 * spanInfo.getTags() >> ['http.client_ip': '8.8.8.8']
1 * spanInfo.getTags() >> TagMap.fromMap(['http.client_ip': '8.8.8.8'])
1 * traceSegment.setTagTop('actor.ip', '8.8.8.8')
}

Expand Down Expand Up @@ -1008,7 +1009,7 @@ class GatewayBridgeSpecification extends DDSpecification {
getTraceSegment() >> traceSegment
}
final spanInfo = Mock(AgentSpan) {
getTags() >> ['http.route':'/']
getTags() >> TagMap.fromMap(['http.route':'/'])
}

when:
Expand Down Expand Up @@ -1196,7 +1197,7 @@ class GatewayBridgeSpecification extends DDSpecification {
def flow = requestEndedCB.apply(mockCtx, spanInfo)
then:
1 * mockAppSecCtx.transferCollectedEvents() >> []
1 * spanInfo.getTags() >> ['http.route': 'route']
1 * spanInfo.getTags() >> TagMap.fromMap(['http.route': 'route'])
1 * requestSampler.preSampleRequest(_) >> true
0 * traceSegment.setTagTop(Tags.ASM_KEEP, true)
0 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM)
Expand All @@ -1214,7 +1215,7 @@ class GatewayBridgeSpecification extends DDSpecification {
def flow = requestEndedCB.apply(mockCtx, spanInfo)
then:
1 * mockAppSecCtx.transferCollectedEvents() >> []
1 * spanInfo.getTags() >> ['http.route': 'route']
1 * spanInfo.getTags() >> TagMap.fromMap(['http.route': 'route'])
1 * requestSampler.preSampleRequest(_) >> false
0 * traceSegment.setTagTop(Tags.ASM_KEEP, true)
0 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM)
Expand All @@ -1233,7 +1234,7 @@ class GatewayBridgeSpecification extends DDSpecification {
def flow = requestEndedCB.apply(mockCtx, spanInfo)
then:
1 * mockAppSecCtx.transferCollectedEvents() >> []
1 * spanInfo.getTags() >> ['http.route': 'route']
1 * spanInfo.getTags() >> TagMap.fromMap(['http.route': 'route'])
1 * requestSampler.preSampleRequest(_) >> true
1 * traceSegment.setTagTop(Tags.ASM_KEEP, true)
1 * traceSegment.setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM)
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import datadog.trace.agent.test.AgentTestRunner
import datadog.trace.api.DDTags
import datadog.trace.api.TagMap
import datadog.trace.api.sampling.PrioritySampling
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags
Expand Down Expand Up @@ -584,7 +585,7 @@ class WebsocketTest extends AgentTestRunner {
clientHandshake.setSamplingPriority(PrioritySampling.SAMPLER_DROP) // simulate sampler drop
def serverHandshake = createHandshakeSpan("servlet.request", url,
new ExtractedContext(clientHandshake.context().getTraceId(), clientHandshake.context().getSpanId(), clientHandshake.context().getSamplingPriority(),
"test", 0, ["example_baggage": "test"], null, null, null, null, null)) // simulate server span
"test", 0, ["example_baggage": "test"], TagMap.EMPTY, null, null, null, null)) // simulate server span
def session = deployEndpointAndConnect(new Endpoints.TestEndpoint(new Endpoints.FullStringHandler()),
clientHandshake, serverHandshake, url)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,7 @@ public final class GeneralConfig {
public static final String APM_TRACING_ENABLED = "apm.tracing.enabled";
public static final String JDK_SOCKET_ENABLED = "jdk.socket.enabled";

public static final String OPTIMIZED_MAP_ENABLED = "optimized.map.enabled";
public static final String STACK_TRACE_LENGTH_LIMIT = "stack.trace.length.limit";

public static final String SSI_INJECTION_ENABLED = "injection.enabled";
Expand Down
Loading