Support custom OpenTelemetry context values (#7118)

Author: mcculls

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

@@ -9,11 +9,14 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.ContextKey;
 import io.opentelemetry.context.Scope;
+import java.util.Arrays;
 import javax.annotation.Nullable;
 import javax.annotation.ParametersAreNonnullByDefault;
 
 @ParametersAreNonnullByDefault
 public class OtelContext implements Context {
+  private static final Object[] NO_ENTRIES = {};
+
   /** Overridden root context. */
   public static final OtelContext ROOT = new OtelContext(OtelSpan.invalid(), OtelSpan.invalid());
 
@@ -26,30 +29,60 @@ public class OtelContext implements Context {
   private final Span currentSpan;
   private final Span rootSpan;
 
+  private final Object[] entries;
+
   public OtelContext(Span currentSpan, Span rootSpan) {
+    this(currentSpan, rootSpan, NO_ENTRIES);
+  }
+
+  public OtelContext(Span currentSpan, Span rootSpan, Object[] entries) {
     this.currentSpan = currentSpan;
     this.rootSpan = rootSpan;
+    this.entries = entries;
   }
 
   @Nullable
   @Override
+  @SuppressWarnings("unchecked")
   public <V> V get(ContextKey<V> key) {
     if (OTEL_CONTEXT_SPAN_KEY.equals(key.toString())) {
       return (V) this.currentSpan;
     } 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];
+      }
+    }
     return null;
   }
 
   @Override
-  public <V> Context with(ContextKey<V> k1, V v1) {
-    if (OTEL_CONTEXT_SPAN_KEY.equals(k1.toString())) {
-      return new OtelContext((Span) v1, this.rootSpan);
-    } else if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(k1.toString())) {
-      return new OtelContext(this.currentSpan, (Span) v1);
+  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;
+      }
     }
-    return this;
+    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);
   }
 
   @Override
@@ -59,7 +92,7 @@ public Scope makeCurrent() {
       // only keep propagated context until next span activation
       lastPropagated.remove();
       AgentScope agentScope = ((OtelSpan) this.currentSpan).activate();
-      return new OtelScope(scope, agentScope);
+      return new OtelScope(scope, agentScope, this.entries);
     } else {
       // propagated context not on the scope stack, capture it here
       lastPropagated.set(this);
@@ -80,12 +113,13 @@ public static Context current() {
       return context;
     }
     // Check empty context
-    AgentSpan agentCurrentSpan = AgentTracer.activeSpan();
-    if (null == agentCurrentSpan) {
+    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) {
@@ -107,7 +141,15 @@ public static Context current() {
     if (otelRootSpan == null) {
       otelRootSpan = new OtelSpan(agentRootSpan);
     }
-    return new OtelContext(otelCurrentSpan, otelRootSpan);
+    // 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);
   }
 
   /** Last propagated context not on the scope stack; {@code null} if there's no such context. */

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

@@ -7,15 +7,22 @@
 public class OtelScope implements Scope {
   private final Scope scope;
   private final AgentScope delegate;
+  private final Object[] contextEntries;
 
-  public OtelScope(Scope scope, AgentScope delegate) {
+  public OtelScope(Scope scope, AgentScope delegate, Object[] contextEntries) {
     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();

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

@@ -8,7 +8,6 @@ import io.opentelemetry.context.Context
 import io.opentelemetry.context.ContextKey
 import io.opentelemetry.context.ImplicitContextKeyed
 import io.opentelemetry.context.ThreadLocalContextStorage
-import spock.lang.Ignore
 import spock.lang.Subject
 
 import static datadog.trace.bootstrap.instrumentation.api.ScopeSource.MANUAL
@@ -279,7 +278,6 @@ class ContextTest extends AgentTestRunner {
     parentSpan.end()
   }
 
-  @Ignore("Not supported")
   def "test custom object storage"() {
     setup:
     def context = Context.root()