Store SpanShim/Baggage in the Context.

opentracing-shim/src/main/java/io/opentelemetry/opentracingshim/ScopeManagerShim.java

@@ -5,42 +5,50 @@
 
 package io.opentelemetry.opentracingshim;
 
+import io.opentelemetry.context.Context;
 import io.opentracing.Scope;
 import io.opentracing.ScopeManager;
 import io.opentracing.Span;
 
 final class ScopeManagerShim extends BaseShimObject implements ScopeManager {
-
   public ScopeManagerShim(TelemetryInfo telemetryInfo) {
     super(telemetryInfo);
   }
 
   @Override
   @SuppressWarnings("ReturnMissingNullable")
   public Span activeSpan() {
+    SpanShim spanShim = SpanShim.current();
+    io.opentelemetry.api.trace.Span span = null;
+    if (spanShim == null) {
+      span = io.opentelemetry.api.trace.Span.current();
+    } else {
+      span = spanShim.getSpan();
+    }
+
     // As OpenTracing simply returns null when no active instance is available,
     // we need to do map an invalid OpenTelemetry span to null here.
-    io.opentelemetry.api.trace.Span span = io.opentelemetry.api.trace.Span.current();
     if (!span.getSpanContext().isValid()) {
       return null;
     }
 
-    // TODO: Properly include the bagagge/distributedContext.
+    // If there's a SpanShim for the *actual* active Span, simply return it.
+    if (spanShim != null && span == io.opentelemetry.api.trace.Span.current()) {
+      return spanShim;
+    }
+
+    // Span was activated from outside the Shim layer unfortunately.
     return new SpanShim(telemetryInfo(), span);
   }
 
   @Override
   @SuppressWarnings("MustBeClosedChecker")
   public Scope activate(Span span) {
-    io.opentelemetry.api.trace.Span actualSpan = getActualSpan(span);
-    return new ScopeShim(actualSpan.makeCurrent());
-  }
-
-  static io.opentelemetry.api.trace.Span getActualSpan(Span span) {
     if (!(span instanceof SpanShim)) {
       throw new IllegalArgumentException("span is not a valid SpanShim object");
     }
 
-    return ((SpanShim) span).getSpan();
+    SpanShim spanShim = (SpanShim) span;
+    return new ScopeShim(Context.current().with(spanShim).makeCurrent());
   }
 }

opentracing-shim/src/main/java/io/opentelemetry/opentracingshim/SpanShim.java

@@ -14,6 +14,9 @@
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
+import io.opentelemetry.context.ImplicitContextKeyed;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import io.opentracing.Span;
 import io.opentracing.SpanContext;
@@ -34,9 +37,11 @@
  * Calling context() or setBaggageItem() will effectively force the creation
  * of SpanContextShim object if none existed yet.
  */
-final class SpanShim extends BaseShimObject implements Span {
+final class SpanShim extends BaseShimObject implements Span, ImplicitContextKeyed {
   private static final String DEFAULT_EVENT_NAME = "log";
   private static final String ERROR = "error";
+  private static final ContextKey<SpanShim> SPAN_SHIM_KEY =
+      ContextKey.named("opentracing-shim-key");
 
   private final io.opentelemetry.api.trace.Span span;
 
@@ -49,6 +54,22 @@ io.opentelemetry.api.trace.Span getSpan() {
     return span;
   }
 
+  public static SpanShim current() {
+    return Context.current().get(SPAN_SHIM_KEY);
+  }
+
+  @Override
+  public Context storeInContext(Context context) {
+    context = context.with(SPAN_SHIM_KEY, this).with(span);
+
+    SpanContextShim spanContextShim = spanContextTable().get(this);
+    if (spanContextShim != null) {
+      context = context.with(spanContextShim.getBaggage());
+    }
+
+    return context;
+  }
+
   @Override
   public SpanContext context() {
     /* Read the value using the read lock first. */

opentracing-shim/src/test/java/io/opentelemetry/opentracingshim/TracerShimTest.java

@@ -28,6 +28,11 @@
 
 class TracerShimTest {
 
+  static final io.opentelemetry.api.trace.Span INVALID_SPAN =
+      io.opentelemetry.api.trace.Span.getInvalid();
+  static final io.opentelemetry.api.baggage.Baggage EMPTY_BAGGAGE =
+      io.opentelemetry.api.baggage.Baggage.empty();
+
   @RegisterExtension public OpenTelemetryExtension otelTesting = OpenTelemetryExtension.create();
 
   TracerShim tracerShim;
@@ -51,16 +56,145 @@ void defaultTracer() {
   @Test
   void activateSpan() {
     Span otSpan = tracerShim.buildSpan("one").start();
-    io.opentelemetry.api.trace.Span span = ((SpanShim) otSpan).getSpan();
+    io.opentelemetry.api.trace.Span actualSpan = ((SpanShim) otSpan).getSpan();
 
     assertThat(tracerShim.activeSpan()).isNull();
     assertThat(tracerShim.scopeManager().activeSpan()).isNull();
+    assertThat(io.opentelemetry.api.trace.Span.current()).isSameAs(INVALID_SPAN);
+    assertThat(io.opentelemetry.api.baggage.Baggage.current()).isSameAs(EMPTY_BAGGAGE);
 
     try (Scope scope = tracerShim.activateSpan(otSpan)) {
       assertThat(tracerShim.activeSpan()).isNotNull();
       assertThat(tracerShim.scopeManager().activeSpan()).isNotNull();
-      assertThat(((SpanShim) tracerShim.activeSpan()).getSpan()).isEqualTo(span);
-      assertThat(((SpanShim) tracerShim.scopeManager().activeSpan()).getSpan()).isEqualTo(span);
+      assertThat(((SpanShim) tracerShim.activeSpan()).getSpan()).isSameAs(actualSpan);
+      assertThat(((SpanShim) tracerShim.scopeManager().activeSpan()).getSpan())
+          .isSameAs(actualSpan);
+
+      assertThat(io.opentelemetry.api.trace.Span.current()).isSameAs(actualSpan);
+      assertThat(io.opentelemetry.api.baggage.Baggage.current()).isSameAs(EMPTY_BAGGAGE);
+    }
+
+    assertThat(tracerShim.activeSpan()).isNull();
+    assertThat(tracerShim.scopeManager().activeSpan()).isNull();
+    assertThat(io.opentelemetry.api.trace.Span.current()).isSameAs(INVALID_SPAN);
+    assertThat(io.opentelemetry.api.baggage.Baggage.current()).isSameAs(EMPTY_BAGGAGE);
+  }
+
+  @Test
+  void activateSpan_withBaggage() {
+    Span otSpan = tracerShim.buildSpan("one").start();
+    otSpan.setBaggageItem("foo", "bar");
+    otSpan.setBaggageItem("hello", "world");
+
+    io.opentelemetry.api.trace.Span actualSpan = ((SpanShim) otSpan).getSpan();
+    io.opentelemetry.api.baggage.Baggage actualBaggage =
+        ((SpanContextShim) otSpan.context()).getBaggage();
+    assertThat(
+            io.opentelemetry.api.baggage.Baggage.builder()
+                .put("foo", "bar")
+                .put("hello", "world")
+                .build())
+        .isEqualTo(actualBaggage);
+
+    assertThat(tracerShim.activeSpan()).isNull();
+    assertThat(tracerShim.scopeManager().activeSpan()).isNull();
+    assertThat(io.opentelemetry.api.trace.Span.current()).isSameAs(INVALID_SPAN);
+    assertThat(io.opentelemetry.api.baggage.Baggage.current()).isSameAs(EMPTY_BAGGAGE);
+
+    try (Scope scope = tracerShim.activateSpan(otSpan)) {
+      assertThat(tracerShim.activeSpan()).isNotNull();
+      assertThat(tracerShim.scopeManager().activeSpan()).isNotNull();
+      assertThat(((SpanShim) tracerShim.activeSpan()).getSpan()).isSameAs(actualSpan);
+      assertThat(((SpanShim) tracerShim.scopeManager().activeSpan()).getSpan())
+          .isSameAs(actualSpan);
+
+      assertThat(io.opentelemetry.api.trace.Span.current()).isSameAs(actualSpan);
+      assertThat(io.opentelemetry.api.baggage.Baggage.current()).isSameAs(actualBaggage);
+    }
+
+    assertThat(tracerShim.activeSpan()).isNull();
+    assertThat(tracerShim.scopeManager().activeSpan()).isNull();
+    assertThat(io.opentelemetry.api.trace.Span.current()).isSameAs(INVALID_SPAN);
+    assertThat(io.opentelemetry.api.baggage.Baggage.current()).isSameAs(EMPTY_BAGGAGE);
+  }
+
+  /* If the Span has no baggage, allow any previously active Baggage go through */
+  @Test
+  public void activateSpan_withExistingBaggage() {
+    Span otSpan = tracerShim.buildSpan("one").start();
+    io.opentelemetry.api.trace.Span actualSpan = ((SpanShim) otSpan).getSpan();
+
+    io.opentelemetry.api.baggage.Baggage otherBaggage =
+        io.opentelemetry.api.baggage.Baggage.builder()
+            .put("foo", "bar")
+            .put("hello", "world")
+            .build();
+
+    assertThat(tracerShim.activeSpan()).isNull();
+    assertThat(tracerShim.scopeManager().activeSpan()).isNull();
+    assertThat(io.opentelemetry.api.trace.Span.current()).isSameAs(INVALID_SPAN);
+    assertThat(io.opentelemetry.api.baggage.Baggage.current()).isSameAs(EMPTY_BAGGAGE);
+
+    try (io.opentelemetry.context.Scope scope1 = otherBaggage.makeCurrent()) {
+      try (Scope scope2 = tracerShim.activateSpan(otSpan)) {
+        assertThat(tracerShim.activeSpan()).isNotNull();
+        assertThat(tracerShim.scopeManager().activeSpan()).isNotNull();
+        assertThat(((SpanShim) tracerShim.activeSpan()).getSpan()).isSameAs(actualSpan);
+        assertThat(((SpanShim) tracerShim.scopeManager().activeSpan()).getSpan())
+            .isSameAs(actualSpan);
+
+        assertThat(io.opentelemetry.api.trace.Span.current()).isSameAs(actualSpan);
+        assertThat(io.opentelemetry.api.baggage.Baggage.current()).isSameAs(otherBaggage);
+      }
+    }
+
+    assertThat(tracerShim.activeSpan()).isNull();
+    assertThat(tracerShim.scopeManager().activeSpan()).isNull();
+    assertThat(io.opentelemetry.api.trace.Span.current()).isSameAs(INVALID_SPAN);
+    assertThat(io.opentelemetry.api.baggage.Baggage.current()).isSameAs(EMPTY_BAGGAGE);
+  }
+
+  /*
+   * Make sure that, upon activation from the Trace Shim layer, the
+   * returned active Span shim is cached/stored in the context, so
+   * we always get the *same* instance.
+   */
+  @Test
+  void activateSpan_cacheSpanShim() {
+    Span otSpan = tracerShim.buildSpan("one").start();
+    io.opentelemetry.api.trace.Span actualSpan = ((SpanShim) otSpan).getSpan();
+
+    assertThat(tracerShim.activeSpan()).isNull();
+    assertThat(tracerShim.scopeManager().activeSpan()).isNull();
+    assertThat(io.opentelemetry.api.trace.Span.current()).isSameAs(INVALID_SPAN);
+    assertThat(io.opentelemetry.api.baggage.Baggage.current()).isSameAs(EMPTY_BAGGAGE);
+
+    try (Scope scope = tracerShim.activateSpan(otSpan)) {
+      assertThat(tracerShim.activeSpan()).isSameAs(tracerShim.activeSpan());
+      assertThat(tracerShim.activeSpan()).isSameAs(tracerShim.scopeManager().activeSpan());
+
+      assertThat(((SpanShim) tracerShim.activeSpan()).getSpan()).isSameAs(actualSpan);
+    }
+
+    assertThat(tracerShim.activeSpan()).isNull();
+    assertThat(tracerShim.scopeManager().activeSpan()).isNull();
+    assertThat(io.opentelemetry.api.trace.Span.current()).isSameAs(INVALID_SPAN);
+    assertThat(io.opentelemetry.api.baggage.Baggage.current()).isSameAs(EMPTY_BAGGAGE);
+  }
+
+  /*
+   * Create and activate a Span without the Shim layer, in order to verify
+   * we keep on working as expected (even if not as efficiently).
+   */
+  @Test
+  void activateSpan_withoutShim() {
+    io.opentelemetry.api.trace.Span span =
+        otelTesting.getOpenTelemetry().getTracer("opentracingshim").spanBuilder("one").startSpan();
+    try (io.opentelemetry.context.Scope scope = span.makeCurrent()) {
+      assertThat(tracerShim.activeSpan()).isNotNull();
+      assertThat(tracerShim.scopeManager().activeSpan()).isNotNull();
+      assertThat(((SpanShim) tracerShim.activeSpan()).getSpan()).isSameAs(span);
+      assertThat(((SpanShim) tracerShim.scopeManager().activeSpan()).getSpan()).isSameAs(span);
     }
 
     assertThat(tracerShim.activeSpan()).isNull();