Prevent AppSec context from being closed more than once on partial flush (#7059)

smola · web-flow · commit ac3ca0988f3e · 2024-05-28T11:57:30.000+02:00
We had some hooks in `CoreTracer.write` meant to be run whenever a root
span is finished. However, they were effectively called not just when
the root span finished, but also whenever a partial flush on a child
span was performed. This ended up calling the hooks multiple times, and
earlier than expected.
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java
@@ -240,11 +240,33 @@ PropagationTags.Factory getPropagationTagsFactory() {
     return propagationTagsFactory;
   }
 
+  /**
+   * Called when a root span is finished before it is serialized. This is might be called multiple
+   * times per root span. If a child span is part of a partial flush, this method will be called for
+   * its root even if not finished.
+   */
   @Override
   public void onRootSpanFinished(AgentSpan root, EndpointTracker tracker) {
     profilingContextIntegration.onRootSpanFinished(root, tracker);
   }
 
+  /**
+   * Called when a root span is finished before it is serialized. This is guaranteed to be called
+   * exactly once per root span.
+   */
+  void onRootSpanPublished(final AgentSpan root) {
+    // Request context is propagated to contexts in child spans.
+    // Assume here that if present it will be so starting in the top span.
+    RequestContext requestContext = root.getRequestContext();
+    if (requestContext != null) {
+      try {
+        requestContext.close();
+      } catch (IOException e) {
+        log.warn("Error closing request context data", e);
+      }
+    }
+  }
+
   @Override
   public EndpointTracker onRootSpanStarted(AgentSpan root) {
     return profilingContextIntegration.onRootSpanStarted(root);
@@ -971,17 +993,6 @@ void write(final List<DDSpan> trace) {
     }
     if (null != rootSpan) {
       onRootSpanFinished(rootSpan, rootSpan.getEndpointTracker());
-
-      // request context is propagated to contexts in child spans
-      // Assume here that if present it will be so starting in the top span
-      RequestContext requestContext = rootSpan.getRequestContext();
-      if (requestContext != null) {
-        try {
-          requestContext.close();
-        } catch (IOException e) {
-          log.warn("Error closing request context data", e);
-        }
-      }
     }
   }
 
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/PendingTrace.java b/dd-trace-core/src/main/java/datadog/trace/core/PendingTrace.java
@@ -253,7 +253,11 @@ PublishState onPublish(final DDSpan span) {
     // write method.
     healthMetrics.onFinishSpan();
     COMPLETED_SPAN_COUNT.incrementAndGet(this);
-    return decrementRefAndMaybeWrite(span == getRootSpan());
+    final DDSpan rootSpan = getRootSpan();
+    if (span == rootSpan) {
+      tracer.onRootSpanPublished(rootSpan);
+    }
+    return decrementRefAndMaybeWrite(span == rootSpan);
   }
 
   @Override
diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/PendingTraceBufferTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/PendingTraceBufferTest.groovy
@@ -91,6 +91,7 @@ class PendingTraceBufferTest extends DDSpecification {
     1 * bufferSpy.enqueue(trace)
     _ * tracer.getPartialFlushMinSpans() >> 10
     1 * tracer.getTimeWithNanoTicks(_)
+    1 * tracer.onRootSpanPublished(span)
     0 * _
 
     when:
@@ -123,6 +124,7 @@ class PendingTraceBufferTest extends DDSpecification {
     _ * bufferSpy.longRunningSpansEnabled()
     _ * tracer.getPartialFlushMinSpans() >> 10
     1 * tracer.getTimeWithNanoTicks(_)
+    1 * tracer.onRootSpanPublished(parent)
     0 * _
 
     when:
@@ -180,6 +182,7 @@ class PendingTraceBufferTest extends DDSpecification {
     _ * tracer.getPartialFlushMinSpans() >> 10
     _ * traceConfig.getServiceMapping() >> [:]
     _ * tracer.getTimeWithNanoTicks(_)
+    buffer.queue.capacity() * tracer.onRootSpanPublished(_)
     0 * _
 
     when:
@@ -195,6 +198,7 @@ class PendingTraceBufferTest extends DDSpecification {
     _ * tracer.getPartialFlushMinSpans() >> 10
     _ * traceConfig.getServiceMapping() >> [:]
     2 * tracer.getTimeWithNanoTicks(_)
+    1 * tracer.onRootSpanPublished(_)
     0 * _
     pendingTrace.isEnqueued == 0
   }
@@ -216,14 +220,14 @@ class PendingTraceBufferTest extends DDSpecification {
     _ * tracer.getPartialFlushMinSpans() >> 10
     _ * traceConfig.getServiceMapping() >> [:]
     _ * tracer.getTimeWithNanoTicks(_)
+    buffer.queue.capacity() * tracer.onRootSpanPublished(_)
     0 * _
 
     when:
     def pendingTrace = factory.create(DDTraceId.ONE)
     pendingTrace.longRunningTrackedState = LongRunningTracesTracker.TO_TRACK
     addContinuation(newSpanOf(pendingTrace)).finish()
 
-    then:
     then:
     1 * tracer.captureTraceConfig() >> traceConfig
     1 * bufferSpy.enqueue(_)
@@ -234,6 +238,7 @@ class PendingTraceBufferTest extends DDSpecification {
     _ * tracer.getPartialFlushMinSpans() >> 10
     _ * traceConfig.getServiceMapping() >> [:]
     _ * tracer.getTimeWithNanoTicks(_)
+    1 * tracer.onRootSpanPublished(_)
     0 * _
 
     pendingTrace.isEnqueued == 0
@@ -261,6 +266,7 @@ class PendingTraceBufferTest extends DDSpecification {
     1 * bufferSpy.enqueue(trace)
     _ * tracer.getPartialFlushMinSpans() >> 10
     1 * tracer.getTimeWithNanoTicks(_)
+    1 * tracer.onRootSpanPublished(parent)
     0 * _
 
     when:
@@ -314,6 +320,7 @@ class PendingTraceBufferTest extends DDSpecification {
     }
     _ * tracer.getPartialFlushMinSpans() >> 10
     1 * tracer.getTimeWithNanoTicks(_)
+    1 * tracer.onRootSpanPublished(parent)
     0 * _
 
     when:
@@ -409,6 +416,7 @@ class PendingTraceBufferTest extends DDSpecification {
     1 * tracer.getPartialFlushMinSpans() >> 10000
     1 * traceConfig.getServiceMapping() >> [:]
     2 * tracer.getTimeWithNanoTicks(_)
+    1 * tracer.onRootSpanPublished(_)
     0 * _
 
     when: "fail to fill the buffer"
diff --git a/internal-api/src/main/java/datadog/trace/api/EndpointCheckpointer.java b/internal-api/src/main/java/datadog/trace/api/EndpointCheckpointer.java
@@ -4,7 +4,9 @@
 
 public interface EndpointCheckpointer {
   /**
-   * Callback to be called when a root span is written (together with the trace)
+   * Callback to be called when a root span is finished (together with the trace). With partial
+   * flushes, this may be called multiple times when any of the root span's children are finished
+   * even if the root span is not.
    *
    * @param rootSpan the local root span of the trace
    * @param tracker the endpoint tracker
