Refactored requires post-processing flag (#8104)

DataDog · Dec 17, 2024 · 53cf809 · 53cf809
1 parent 772d178
commit 53cf809
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 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
@@ -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<Void> 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
@@ -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<DDSpan> trace) {
       // Filter spans that need post-processing
       List<DDSpan> 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
@@ -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
@@ -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
@@ -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
@@ -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 {