Migrate OtelContext wrapper to new Context API (#8645)

Author: mcculls

dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java

@@ -1,170 +1,109 @@
 package datadog.opentelemetry.shim.context;
 
 import datadog.opentelemetry.shim.trace.OtelSpan;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper;
-import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.ContextKey;
 import io.opentelemetry.context.Scope;
-import java.util.Arrays;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import javax.annotation.Nullable;
 import javax.annotation.ParametersAreNonnullByDefault;
 
+@SuppressWarnings({"rawtypes", "unchecked"})
 @ParametersAreNonnullByDefault
 public class OtelContext implements Context {
-  private static final Object[] NO_ENTRIES = {};
 
   /** Overridden root context. */
-  public static final Context ROOT = new OtelContext(OtelSpan.invalid(), OtelSpan.invalid());
+  public static final Context ROOT = new OtelContext(datadog.context.Context.root());
 
   private static final String OTEL_CONTEXT_SPAN_KEY = "opentelemetry-trace-span-key";
   private static final String OTEL_CONTEXT_ROOT_SPAN_KEY = "opentelemetry-traces-local-root-span";
 
-  /** Keep track of propagated context that has not been captured on the scope stack. */
-  private static final ThreadLocal<OtelContext> lastPropagated = new ThreadLocal<>();
+  /** Records the keys needed to access the delegate context, mapped by key name. */
+  private static final Map<ContextKey<?>, datadog.context.ContextKey<?>> DELEGATE_KEYS =
+      new ConcurrentHashMap<>();
 
-  private final Span currentSpan;
-  private final Span rootSpan;
+  private final datadog.context.Context delegate;
 
-  private final Object[] entries;
+  public OtelContext(datadog.context.Context delegate) {
+    this.delegate = delegate;
+  }
 
-  public OtelContext(Span currentSpan, Span rootSpan) {
-    this(currentSpan, rootSpan, NO_ENTRIES);
+  public static Context current() {
+    return new OtelContext(datadog.context.Context.current());
   }
 
-  public OtelContext(Span currentSpan, Span rootSpan, Object[] entries) {
-    this.currentSpan = currentSpan;
-    this.rootSpan = rootSpan;
-    this.entries = entries;
+  @Override
+  public Scope makeCurrent() {
+    return new OtelScope(Context.super.makeCurrent(), delegate.attach());
   }
 
   @Nullable
   @Override
-  @SuppressWarnings("unchecked")
   public <V> V get(ContextKey<V> key) {
     if (OTEL_CONTEXT_SPAN_KEY.equals(key.toString())) {
-      return (V) this.currentSpan;
+      AgentSpan span = AgentSpan.fromContext(delegate);
+      if (span != null) {
+        return (V) toOtelSpan(span);
+      }
+      // fall-through and check for non-datadog span data
     } else if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(key.toString())) {
-      return (V) this.rootSpan;
-    }
-    for (int i = 0; i < this.entries.length; i += 2) {
-      if (this.entries[i] == key) {
-        return (V) this.entries[i + 1];
+      AgentSpan span = AgentSpan.fromContext(delegate);
+      if (span != null) {
+        return (V) toOtelSpan(span.getLocalRootSpan());
       }
+      // fall-through and check for non-datadog span data
     }
-    return null;
+    return (V) delegate.get(delegateKey(key));
   }
 
   @Override
   public <V> Context with(ContextKey<V> key, V value) {
     if (OTEL_CONTEXT_SPAN_KEY.equals(key.toString())) {
-      return new OtelContext((Span) value, this.rootSpan, this.entries);
-    } else if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(key.toString())) {
-      return new OtelContext(this.currentSpan, (Span) value, this.entries);
-    }
-    Object[] newEntries = null;
-    int oldEntriesLength = this.entries.length;
-    for (int i = 0; i < oldEntriesLength; i += 2) {
-      if (this.entries[i] == key) {
-        if (this.entries[i + 1] == value) {
-          return this;
-        }
-        newEntries = this.entries.clone();
-        newEntries[i + 1] = value;
-        break;
+      if (value instanceof OtelSpan) {
+        AgentSpan span = ((OtelSpan) value).asAgentSpan();
+        return new OtelContext(delegate.with(span));
       }
+      // fall-through and store as non-datadog span data
     }
-    if (null == newEntries) {
-      newEntries = Arrays.copyOf(this.entries, oldEntriesLength + 2);
-      newEntries[oldEntriesLength] = key;
-      newEntries[oldEntriesLength + 1] = value;
-    }
-    return new OtelContext(this.currentSpan, this.rootSpan, newEntries);
+    return new OtelContext(delegate.with(delegateKey(key), value));
   }
 
   @Override
-  public Scope makeCurrent() {
-    final Scope scope = Context.super.makeCurrent();
-    if (this.currentSpan instanceof OtelSpan) {
-      // only keep propagated context until next span activation
-      lastPropagated.remove();
-      AgentScope agentScope = ((OtelSpan) this.currentSpan).activate();
-      return new OtelScope(scope, agentScope, this.entries);
-    } else {
-      // propagated context not on the scope stack, capture it here
-      lastPropagated.set(this);
-      return new Scope() {
-        @Override
-        public void close() {
-          lastPropagated.remove();
-          scope.close();
-        }
-      };
-    }
-  }
-
-  public static Context current() {
-    // Check for propagated context not on the scope stack
-    Context context = lastPropagated.get();
-    if (null != context) {
-      return context;
-    }
-    // Check empty context
-    AgentScope agentCurrentScope = AgentTracer.activeScope();
-    if (null == agentCurrentScope) {
-      return OtelContext.ROOT;
-    }
-    // Get OTel current span
-    Span otelCurrentSpan = null;
-    AgentSpan agentCurrentSpan = agentCurrentScope.span();
-    if (agentCurrentSpan instanceof AttachableWrapper) {
-      Object wrapper = ((AttachableWrapper) agentCurrentSpan).getWrapper();
-      if (wrapper instanceof OtelSpan) {
-        otelCurrentSpan = (OtelSpan) wrapper;
-      }
+  public boolean equals(Object o) {
+    if (o == null || getClass() != o.getClass()) {
+      return false;
     }
-    if (otelCurrentSpan == null) {
-      otelCurrentSpan = new OtelSpan(agentCurrentSpan);
-    }
-    // Get OTel root span
-    Span otelRootSpan = null;
-    AgentSpan agentRootSpan = agentCurrentSpan.getLocalRootSpan();
-    if (agentRootSpan instanceof AttachableWrapper) {
-      Object wrapper = ((AttachableWrapper) agentRootSpan).getWrapper();
-      if (wrapper instanceof OtelSpan) {
-        otelRootSpan = (OtelSpan) wrapper;
-      }
-    }
-    if (otelRootSpan == null) {
-      otelRootSpan = new OtelSpan(agentRootSpan);
-    }
-    // Get OTel custom context entries
-    Object[] contextEntries = NO_ENTRIES;
-    if (agentCurrentScope instanceof AttachableWrapper) {
-      Object wrapper = ((AttachableWrapper) agentCurrentScope).getWrapper();
-      if (wrapper instanceof OtelScope) {
-        contextEntries = ((OtelScope) wrapper).contextEntries();
-      }
-    }
-    return new OtelContext(otelCurrentSpan, otelRootSpan, contextEntries);
+    return delegate.equals(((OtelContext) o).delegate);
   }
 
-  /** Last propagated context not on the scope stack; {@code null} if there's no such context. */
-  @Nullable
-  public static Context lastPropagated() {
-    return lastPropagated.get();
+  @Override
+  public int hashCode() {
+    return delegate.hashCode();
   }
 
   @Override
   public String toString() {
-    return "OtelContext{"
-        + "currentSpan="
-        + this.currentSpan.getSpanContext()
-        + ", rootSpan="
-        + this.rootSpan.getSpanContext()
-        + '}';
+    return "OtelContext{" + "delegate=" + delegate + '}';
+  }
+
+  private static datadog.context.ContextKey delegateKey(ContextKey key) {
+    return DELEGATE_KEYS.computeIfAbsent(key, OtelContext::mapByKeyName);
+  }
+
+  private static datadog.context.ContextKey mapByKeyName(ContextKey key) {
+    return datadog.context.ContextKey.named(key.toString());
+  }
+
+  private static OtelSpan toOtelSpan(AgentSpan span) {
+    if (span instanceof AttachableWrapper) {
+      Object wrapper = ((AttachableWrapper) span).getWrapper();
+      if (wrapper instanceof OtelSpan) {
+        return (OtelSpan) wrapper;
+      }
+    }
+    return new OtelSpan(span);
   }
 }

dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelScope.java

@@ -1,31 +1,20 @@
 package datadog.opentelemetry.shim.context;
 
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
-import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper;
+import datadog.context.ContextScope;
 import io.opentelemetry.context.Scope;
 
 public class OtelScope implements Scope {
   private final Scope scope;
-  private final AgentScope delegate;
-  private final Object[] contextEntries;
+  private final ContextScope delegate;
 
-  public OtelScope(Scope scope, AgentScope delegate, Object[] contextEntries) {
+  public OtelScope(Scope scope, ContextScope delegate) {
     this.scope = scope;
     this.delegate = delegate;
-    this.contextEntries = contextEntries;
-    if (delegate instanceof AttachableWrapper) {
-      ((AttachableWrapper) delegate).attachWrapper(this);
-    }
-  }
-
-  /** Context entries from {@link OtelContext}, captured when the context was made current. */
-  Object[] contextEntries() {
-    return contextEntries;
   }
 
   @Override
   public void close() {
-    this.delegate.close();
-    this.scope.close();
+    delegate.close();
+    scope.close();
   }
 }

dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java

@@ -7,7 +7,6 @@
 
 import datadog.opentelemetry.shim.context.OtelContext;
 import datadog.opentelemetry.shim.trace.OtelExtractedContext;
-import datadog.opentelemetry.shim.trace.OtelSpan;
 import datadog.trace.api.TracePropagationStyle;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
@@ -59,7 +58,7 @@ public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C
     } else {
       TraceState traceState = extractTraceState(extracted, carrier, getter);
       SpanContext spanContext = fromRemote(extracted, traceState);
-      return new OtelContext(Span.wrap(spanContext), OtelSpan.invalid());
+      return Span.wrap(spanContext).storeInContext(OtelContext.ROOT);
     }
   }
 

dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/trace/OtelSpanBuilder.java

@@ -11,7 +11,6 @@
 import static java.lang.Boolean.parseBoolean;
 import static java.util.Locale.ROOT;
 
-import datadog.opentelemetry.shim.context.OtelContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -202,11 +201,7 @@ public Span startSpan() {
       setSpanKind(INTERNAL);
     }
     if (!this.ignoreActiveSpan) {
-      // support automatic parenting from propagated context not on the scope stack
-      Context context = OtelContext.lastPropagated();
-      if (null != context) {
-        setParent(context);
-      }
+      setParent(Context.current());
     }
     AgentSpan delegate = this.delegate.start();
     // Apply overrides

dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/main/java/datadog/trace/instrumentation/opentelemetry14/OpenTelemetryInstrumentation.java

@@ -58,7 +58,6 @@ public boolean onlyMatchKnownTypes() {
   public String[] helperClassNames() {
     return new String[] {
       "datadog.opentelemetry.shim.context.OtelContext",
-      "datadog.opentelemetry.shim.context.OtelContext$1",
       "datadog.opentelemetry.shim.context.OtelScope",
       "datadog.opentelemetry.shim.context.propagation.AgentTextMapPropagator",
       "datadog.opentelemetry.shim.context.propagation.OtelContextPropagators",

dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/main/java/datadog/trace/instrumentation/opentelemetry14/context/OpenTelemetryContextInstrumentation.java

@@ -55,7 +55,6 @@ public boolean onlyMatchKnownTypes() {
   public String[] helperClassNames() {
     return new String[] {
       "datadog.opentelemetry.shim.context.OtelContext",
-      "datadog.opentelemetry.shim.context.OtelContext$1",
       "datadog.opentelemetry.shim.context.OtelScope",
       "datadog.opentelemetry.shim.trace.OtelExtractedContext",
       "datadog.opentelemetry.shim.trace.OtelConventions",

dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/main/java/datadog/trace/instrumentation/opentelemetry14/context/OpenTelemetryContextStorageInstrumentation.java

@@ -56,7 +56,6 @@ public boolean onlyMatchKnownTypes() {
   public String[] helperClassNames() {
     return new String[] {
       "datadog.opentelemetry.shim.context.OtelContext",
-      "datadog.opentelemetry.shim.context.OtelContext$1",
       "datadog.opentelemetry.shim.context.OtelScope",
       "datadog.opentelemetry.shim.trace.OtelExtractedContext",
       "datadog.opentelemetry.shim.trace.OtelConventions",

dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextTest.groovy renamed to dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/ContextForkedTest.groovy

@@ -14,7 +14,7 @@ import static datadog.opentelemetry.shim.context.OtelContext.OTEL_CONTEXT_ROOT_S
 import static datadog.opentelemetry.shim.context.OtelContext.OTEL_CONTEXT_SPAN_KEY
 import static datadog.opentelemetry.shim.trace.OtelConventions.SPAN_KIND_INTERNAL
 
-class ContextTest extends AgentTestRunner {
+class ContextForkedTest extends AgentTestRunner {
   @Subject
   def tracer = GlobalOpenTelemetry.get().tracerProvider.get("context-instrumentation")
 
@@ -169,6 +169,10 @@ class ContextTest extends AgentTestRunner {
     then:
     currentSpan != null
     !currentSpan.spanContext.isValid()
+
+    cleanup:
+    ddScope.close()
+    ddSpan.finish()
   }
 
   def "test clearing context"() {
@@ -238,6 +242,14 @@ class ContextTest extends AgentTestRunner {
         }
       }
     }
+
+    cleanup:
+    otelGrandChildScope?.close()
+    otelGrandChildSpan?.end()
+    ddChildScope?.close()
+    ddChildSpan?.finish()
+    otelParentScope.close()
+    otelParentSpan.end()
   }
 
   def "test context spans retrieval"() {