Combine play spans with surrounding framework context before attaching them (#10554)

Add AgentSpan.attachWithCurrent() which attaches a combination of the span with the current context to the current execution unit

Combine play spans with surrounding custom context before attaching them

Clarify difference between AgentSpan.attach and AgentSpan.attachWithCurrent

Review feedback

Merge remote-tracking branch 'origin/master' into mcculls/merge-play-spans-with-framework-context

Co-authored-by: stuart.mcculloch <stuart.mcculloch@datadoghq.com>

Author: mcculls

dd-java-agent/instrumentation/play/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java

@@ -34,7 +34,7 @@ public static ContextScope onEnter(@Advice.Argument(0) final Request<?> req) {
       // Do not extract the context.
       span = startSpan("play", PLAY_REQUEST);
       span.setMeasured(true);
-      scope = span.attach();
+      scope = span.attachWithCurrent();
     }
 
     DECORATE.afterStart(span);

dd-java-agent/instrumentation/play/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java

@@ -40,7 +40,7 @@ public static ContextScope onEnter(@Advice.Argument(value = 0, readOnly = false)
       // Do not extract the context.
       span = startSpan("play", PLAY_REQUEST);
       span.setMeasured(true);
-      scope = span.attach();
+      scope = span.attachWithCurrent();
     }
 
     DECORATE.afterStart(span);

dd-java-agent/instrumentation/play/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java

@@ -41,7 +41,7 @@ public static ContextScope onEnter(@Advice.Argument(value = 0, readOnly = false)
       // Do not extract the context.
       parentContext = getRootContext();
       span = startSpan("play", PLAY_REQUEST);
-      scope = span.attach();
+      scope = span.attachWithCurrent();
     }
 
     span.setMeasured(true);

internal-api/src/main/java/datadog/trace/api/datastreams/DataStreamsContext.java

@@ -6,6 +6,7 @@
 import datadog.context.Context;
 import datadog.context.ContextKey;
 import datadog.context.ImplicitContextKeyed;
+import javax.annotation.Nonnull;
 
 public class DataStreamsContext implements ImplicitContextKeyed {
   private static final ContextKey<DataStreamsContext> CONTEXT_KEY =
@@ -102,7 +103,7 @@ public boolean sendCheckpoint() {
   }
 
   @Override
-  public Context storeInto(Context context) {
+  public Context storeInto(@Nonnull Context context) {
     return context.with(CONTEXT_KEY, this);
   }
 }

internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java

@@ -18,6 +18,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import javax.annotation.Nonnull;
 
 public class InferredProxySpan implements ImplicitContextKeyed {
   private static final ContextKey<InferredProxySpan> CONTEXT_KEY = named("inferred-proxy-key");
@@ -131,7 +132,7 @@ public void finish() {
   }
 
   @Override
-  public Context storeInto(Context context) {
+  public Context storeInto(@Nonnull Context context) {
     return context.with(CONTEXT_KEY, this);
   }
 }

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java

@@ -4,6 +4,7 @@
 
 import datadog.context.Context;
 import datadog.context.ContextKey;
+import datadog.context.ContextScope;
 import datadog.context.ImplicitContextKeyed;
 import datadog.trace.api.DDSpanId;
 import datadog.trace.api.DDTraceId;
@@ -209,8 +210,8 @@ default void copyPropagationAndBaggage(final AgentSpan source) {
   }
 
   @Override
-  default Context storeInto(Context context) {
-    return context.with(SPAN_KEY, this);
+  default Context storeInto(@Nonnull Context context) {
+    return context == Context.root() ? this : context.with(SPAN_KEY, this);
   }
 
   @Nullable
@@ -224,4 +225,25 @@ default <T> T get(@Nonnull ContextKey<T> key) {
   default <T> Context with(@Nonnull ContextKey<T> key, @Nullable T value) {
     return SPAN_KEY == key ? (Context) value : Context.root().with(SPAN_KEY, this, key, value);
   }
+
+  /**
+   * Attaches a context containing just the span to the current execution unit. Use this when you
+   * want to temporarily suppress any surrounding custom context during the span's scope.
+   *
+   * @return a scope to be closed when the span context is invalid.
+   */
+  @Override
+  default ContextScope attach() {
+    return Context.super.attach();
+  }
+
+  /**
+   * Attaches a context combining the span with the current context to the current execution unit.
+   * Use this when you want to maintain any surrounding custom context during the span's scope.
+   *
+   * @return a scope to be closed when the combined context is invalid.
+   */
+  default ContextScope attachWithCurrent() {
+    return storeInto(Context.current()).attach();
+  }
 }