Add an OpenTelemetry context mechanism. (#1658)

* Add an OpenTelemetry context mechanism.

* Moar

* Extracted interfaces

* More context

* Cleanup / tests

* Move and cleanups

* Brackets

* Add example for brave context interop

* Brave in OTel

* Spotless

* Missing folder

* Another

* Spotless

build.gradle

@@ -354,9 +354,19 @@ configure(opentelemetryProjects) {
         }
     }
 
-    test {
+    tasks.withType(Test) {
         systemProperties project.properties.subMap(["enable.docker.tests"])
         useJUnitPlatform()
+
+        // At a test failure, log the stack trace to the console so that we don't
+        // have to open the HTML in a browser.
+        testLogging {
+            exceptionFormat = 'full'
+            showExceptions = true
+            showCauses = true
+            showStackTraces = true
+        }
+        maxHeapSize = '1500m'
     }
 
     javadoc.options {
@@ -381,18 +391,6 @@ configure(opentelemetryProjects) {
         sign configurations.archives
     }
 
-    // At a test failure, log the stack trace to the console so that we don't
-    // have to open the HTML in a browser.
-    test {
-        testLogging {
-            exceptionFormat = 'full'
-            showExceptions = true
-            showCauses = true
-            showStackTraces = true
-        }
-        maxHeapSize = '1500m'
-    }
-
     plugins.withId("ru.vyarus.animalsniffer") {
         animalsnifferTest {
             enabled = false

context/build.gradle

@@ -0,0 +1,32 @@
+plugins {
+    id "java"
+
+    id "org.unbroken-dome.test-sets"
+    id "ru.vyarus.animalsniffer"
+}
+
+description = 'OpenTelemetry Context (Incubator)'
+ext.moduleName = "io.opentelemetry.context"
+
+testSets {
+    grpcInOtelTest
+    otelInGrpcTest
+    otelAsGrpcTest
+
+    braveInOtelTest
+    otelAsBraveTest
+}
+
+dependencies {
+    grpcInOtelTestImplementation libraries.grpc_context
+    otelAsGrpcTestImplementation libraries.grpc_context
+    otelInGrpcTestImplementation libraries.grpc_context
+
+    braveInOtelTestImplementation "io.zipkin.brave:brave:5.12.6"
+    otelAsBraveTestImplementation "io.zipkin.brave:brave:5.12.6"
+
+    testImplementation libraries.awaitility
+
+    signature "org.codehaus.mojo.signature:java18:1.0@signature"
+    signature "net.sf.androidscents.signature:android-api-level-24:7.0_r2@signature"
+}

context/src/braveInOtelTest/java/io/opentelemetry/context/BraveInOtelTest.java

@@ -0,0 +1,125 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.context;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import brave.Tracing;
+import brave.propagation.CurrentTraceContext;
+import brave.propagation.TraceContext;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicReference;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+class BraveInOtelTest {
+
+  private static final ContextKey<String> ANIMAL = ContextKey.named("animal");
+
+  private static final Tracing TRACING =
+      Tracing.newBuilder().currentTraceContext(new OpenTelemetryCurrentTraceContext()).build();
+
+  private static final TraceContext TRACE_CONTEXT =
+      TraceContext.newBuilder().traceId(1).spanId(1).addExtra("japan").build();
+
+  private static ExecutorService otherThread;
+
+  @BeforeAll
+  static void setUp() {
+    otherThread = Executors.newSingleThreadExecutor();
+  }
+
+  @AfterAll
+  static void tearDown() {
+    otherThread.shutdown();
+  }
+
+  @Test
+  void braveOtelMix() {
+    try (CurrentTraceContext.Scope ignored =
+        TRACING.currentTraceContext().newScope(TRACE_CONTEXT)) {
+      assertThat(Tracing.current().currentTraceContext().get().extra()).contains("japan");
+      try (Scope ignored2 = Context.current().withValues(ANIMAL, "cat").makeCurrent()) {
+        assertThat(Tracing.current().currentTraceContext().get().extra()).contains("japan");
+        assertThat(Context.current().getValue(ANIMAL)).isEqualTo("cat");
+
+        TraceContext context2 =
+            Tracing.current().currentTraceContext().get().toBuilder().addExtra("cheese").build();
+        try (CurrentTraceContext.Scope ignored3 =
+            TRACING.currentTraceContext().newScope(context2)) {
+          assertThat(Tracing.current().currentTraceContext().get().extra()).contains("japan");
+          assertThat(Tracing.current().currentTraceContext().get().extra()).contains("cheese");
+          assertThat(Context.current().getValue(ANIMAL)).isEqualTo("cat");
+        }
+      }
+    }
+  }
+
+  @Test
+  void braveWrap() throws Exception {
+    try (CurrentTraceContext.Scope ignored =
+        TRACING.currentTraceContext().newScope(TRACE_CONTEXT)) {
+      try (Scope ignored2 = Context.current().withValues(ANIMAL, "cat").makeCurrent()) {
+        assertThat(Tracing.current().currentTraceContext().get().extra()).contains("japan");
+        assertThat(Context.current().getValue(ANIMAL)).isEqualTo("cat");
+        AtomicReference<Boolean> braveContainsJapan = new AtomicReference<>();
+        AtomicReference<String> otelValue = new AtomicReference<>();
+        Runnable runnable =
+            () -> {
+              TraceContext traceContext = Tracing.current().currentTraceContext().get();
+              if (traceContext != null && traceContext.extra().contains("japan")) {
+                braveContainsJapan.set(true);
+              } else {
+                braveContainsJapan.set(false);
+              }
+              otelValue.set(Context.current().getValue(ANIMAL));
+            };
+        otherThread.submit(runnable).get();
+        assertThat(braveContainsJapan).hasValue(false);
+        assertThat(otelValue).hasValue(null);
+
+        otherThread.submit(TRACING.currentTraceContext().wrap(runnable)).get();
+        assertThat(braveContainsJapan).hasValue(true);
+
+        // Since Brave context is inside the OTel context, propagating the Brave context does not
+        // propagate the OTel context.
+        assertThat(otelValue).hasValue(null);
+      }
+    }
+  }
+
+  @Test
+  void otelWrap() throws Exception {
+    try (CurrentTraceContext.Scope ignored =
+        TRACING.currentTraceContext().newScope(TRACE_CONTEXT)) {
+      try (Scope ignored2 = Context.current().withValues(ANIMAL, "cat").makeCurrent()) {
+        assertThat(Tracing.current().currentTraceContext().get().extra()).contains("japan");
+        assertThat(Context.current().getValue(ANIMAL)).isEqualTo("cat");
+        AtomicReference<Boolean> braveContainsJapan = new AtomicReference<>(false);
+        AtomicReference<String> otelValue = new AtomicReference<>();
+        Runnable runnable =
+            () -> {
+              TraceContext traceContext = Tracing.current().currentTraceContext().get();
+              if (traceContext != null && traceContext.extra().contains("japan")) {
+                braveContainsJapan.set(true);
+              } else {
+                braveContainsJapan.set(false);
+              }
+              otelValue.set(Context.current().getValue(ANIMAL));
+            };
+        otherThread.submit(runnable).get();
+        assertThat(braveContainsJapan).hasValue(false);
+        assertThat(otelValue).hasValue(null);
+
+        otherThread.submit(Context.current().wrap(runnable)).get();
+        assertThat(braveContainsJapan).hasValue(true);
+        assertThat(otelValue).hasValue("cat");
+      }
+    }
+  }
+}

context/src/braveInOtelTest/java/io/opentelemetry/context/OpenTelemetryCurrentTraceContext.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.context;
+
+import brave.propagation.CurrentTraceContext;
+import brave.propagation.TraceContext;
+
+public class OpenTelemetryCurrentTraceContext extends CurrentTraceContext {
+
+  private static final ContextKey<TraceContext> TRACE_CONTEXT_KEY =
+      ContextKey.named("brave-tracecontext");
+
+  @Override
+  public TraceContext get() {
+    return Context.current().getValue(TRACE_CONTEXT_KEY);
+  }
+
+  @SuppressWarnings("ReferenceEquality")
+  @Override
+  public Scope newScope(TraceContext context) {
+    Context currentOtel = Context.current();
+    TraceContext currentBrave = currentOtel.getValue(TRACE_CONTEXT_KEY);
+    if (currentBrave == context) {
+      return Scope.NOOP;
+    }
+
+    Context newOtel = currentOtel.withValues(TRACE_CONTEXT_KEY, context);
+    io.opentelemetry.context.Scope otelScope = newOtel.makeCurrent();
+    return otelScope::close;
+  }
+}

context/src/grpcInOtelTest/java/io/grpc/override/ContextStorageOverride.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.grpc.override;
+
+import io.grpc.Context;
+import io.opentelemetry.context.ContextKey;
+import io.opentelemetry.context.Scope;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+// This exact package / class name indicates to gRPC to use this override.
+public class ContextStorageOverride extends Context.Storage {
+
+  private static final Logger log = Logger.getLogger(ContextStorageOverride.class.getName());
+
+  private static final ContextKey<Context> GRPC_CONTEXT = ContextKey.named("grpc-context");
+  private static final Context.Key<Scope> OTEL_SCOPE = Context.key("otel-scope");
+
+  @Override
+  public Context doAttach(Context toAttach) {
+    io.opentelemetry.context.Context otelContext = io.opentelemetry.context.Context.current();
+    Context current = otelContext.getValue(GRPC_CONTEXT);
+
+    if (current == toAttach) {
+      return toAttach;
+    }
+
+    if (current == null) {
+      current = Context.ROOT;
+    }
+
+    io.opentelemetry.context.Context newOtelContext =
+        otelContext.withValues(GRPC_CONTEXT, toAttach);
+    Scope scope = newOtelContext.makeCurrent();
+    return current.withValue(OTEL_SCOPE, scope);
+  }
+
+  @Override
+  public void detach(Context toDetach, Context toRestore) {
+    if (current() != toDetach) {
+      // Log a severe message instead of throwing an exception as the context to attach is assumed
+      // to be the correct one and the unbalanced state represents a coding mistake in a lower
+      // layer in the stack that cannot be recovered from here.
+      log.log(
+          Level.SEVERE,
+          "Context was not attached when detaching",
+          new Throwable().fillInStackTrace());
+    }
+
+    Scope otelScope = OTEL_SCOPE.get(toRestore);
+    otelScope.close();
+  }
+
+  @Override
+  public Context current() {
+    return io.opentelemetry.context.Context.current().getValue(GRPC_CONTEXT);
+  }
+}