From 53cf8096d13ec4caaec38e7f4662f6b8a70d8460 Mon Sep 17 00:00:00 2001 From: ValentinZakharov Date: Tue, 17 Dec 2024 19:03:23 +0200 Subject: [PATCH] Refactored requires post-processing flag (#8104) --- .../com/datadog/appsec/gateway/GatewayBridge.java | 11 +---------- .../trace/common/writer/TraceProcessingWorker.java | 4 +--- .../src/main/java/datadog/trace/core/DDSpan.java | 5 +++++ .../java/datadog/trace/api/gateway/IGSpanInfo.java | 4 ++++ .../bootstrap/instrumentation/api/AgentSpan.java | 2 -- .../bootstrap/instrumentation/api/AgentTracer.java | 3 +++ 6 files changed, 14 insertions(+), 15 deletions(-) diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java index e27848c39f0..6b98dad668c 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java @@ -33,7 +33,6 @@ import datadog.trace.api.internal.TraceSegment; import datadog.trace.api.telemetry.RuleType; import datadog.trace.api.telemetry.WafMetricCollector; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter; import datadog.trace.util.stacktrace.StackTraceEvent; @@ -618,18 +617,10 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) { } } - ctx.close(requiresPostProcessing(spanInfo)); + ctx.close(spanInfo.isRequiresPostProcessing()); return NoopFlow.INSTANCE; } - private boolean requiresPostProcessing(final IGSpanInfo spanInfo) { - if (!(spanInfo instanceof AgentSpan)) { - return true; // be conservative - } - final AgentSpan span = (AgentSpan) spanInfo; - return span.isRequiresPostProcessing(); - } - private Flow onRequestHeadersDone(RequestContext ctx_) { AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC); if (ctx == null || ctx.isReqDataPublished()) { diff --git a/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java b/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java index 3dbccb71bad..d2ed59f3f8a 100644 --- a/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java +++ b/dd-trace-core/src/main/java/datadog/trace/common/writer/TraceProcessingWorker.java @@ -13,7 +13,6 @@ import datadog.trace.common.writer.ddagent.PrioritizationStrategy; import datadog.trace.core.CoreSpan; import datadog.trace.core.DDSpan; -import datadog.trace.core.DDSpanContext; import datadog.trace.core.monitor.HealthMetrics; import datadog.trace.core.postprocessor.SpanPostProcessor; import java.util.ArrayList; @@ -266,8 +265,7 @@ private void maybeTracePostProcessing(List trace) { // Filter spans that need post-processing List spansToPostProcess = null; for (DDSpan span : trace) { - DDSpanContext context = span.context(); - if (context != null && context.isRequiresPostProcessing()) { + if (span.isRequiresPostProcessing()) { if (spansToPostProcess == null) { spansToPostProcess = new ArrayList<>(); } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java index 180889600d3..e56f9a5206c 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java @@ -841,6 +841,11 @@ public boolean isRequiresPostProcessing() { return context.isRequiresPostProcessing(); } + @Override + public void setRequiresPostProcessing(boolean requiresPostProcessing) { + context.setRequiresPostProcessing(requiresPostProcessing); + } + // to be accessible in Spock spies, which the field wouldn't otherwise be public long getStartTimeNano() { return startTimeNano; diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/IGSpanInfo.java b/internal-api/src/main/java/datadog/trace/api/gateway/IGSpanInfo.java index 60bca47f184..28ddbe27ad3 100644 --- a/internal-api/src/main/java/datadog/trace/api/gateway/IGSpanInfo.java +++ b/internal-api/src/main/java/datadog/trace/api/gateway/IGSpanInfo.java @@ -16,4 +16,8 @@ public interface IGSpanInfo { void setRequestBlockingAction(Flow.Action.RequestBlockingAction rba); Flow.Action.RequestBlockingAction getRequestBlockingAction(); + + boolean isRequiresPostProcessing(); + + void setRequiresPostProcessing(boolean requiresPostProcessing); } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java index f4e63553114..8b113565f9a 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java @@ -141,8 +141,6 @@ public interface AgentSpan extends MutableSpan, IGSpanInfo, ImplicitContextKeyed void addLink(AgentSpanLink link); - boolean isRequiresPostProcessing(); - @Override default ScopedContext storeInto(ScopedContext context) { return context.with(ScopedContextKey.SPAN_KEY, this); diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java index 0518c92563b..630d91ad7cf 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java @@ -865,6 +865,9 @@ public boolean isOutbound() { public boolean isRequiresPostProcessing() { return false; } + + @Override + public void setRequiresPostProcessing(boolean requiresPostProcessing) {} } public static final class NoopAgentScope implements AgentScope {