From 4f31ab556117398c3df39052b65b5c8898431b09 Mon Sep 17 00:00:00 2001
From: Santiago Mola <santiago.mola@datadoghq.com>
Date: Wed, 22 May 2024 09:42:30 +0200
Subject: [PATCH] Prevent AppSec context from being closed more than once on
 partial flush 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 teams, and
 earlier than expected.

---
 .../java/datadog/trace/core/CoreTracer.java   | 33 ++++++++++++-------
 .../java/datadog/trace/core/PendingTrace.java |  6 +++-
 2 files changed, 27 insertions(+), 12 deletions(-)

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
index b415222b080..75123095c2b 100644
--- 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,9 +240,31 @@ 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 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 inthe top span.
+    RequestContext requestContext = root.getRequestContext();
+    if (requestContext != null) {
+      try {
+        requestContext.close();
+      } catch (IOException e) {
+        log.warn("Error closing request context data", e);
+      }
+    }
   }
 
   @Override
@@ -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
index 088aab42d66..5f8c4c9813c 100644
--- 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