Inherit the parent OT SpanContext baggage/distContext. (#648)

opentracing_shim/src/main/java/io/opentelemetry/opentracingshim/BaseShimObject.java

@@ -37,4 +37,8 @@ abstract class BaseShimObject {
   DistributedContextManager contextManager() {
     return telemetryInfo.contextManager();
   }
+
+  SpanContextShimTable spanContextTable() {
+    return telemetryInfo.spanContextTable();
+  }
 }

opentracing_shim/src/main/java/io/opentelemetry/opentracingshim/SpanBuilderShim.java

@@ -31,8 +31,9 @@ final class SpanBuilderShim extends BaseShimObject implements SpanBuilder {
   private final String spanName;
 
   // The parent will be either a Span or a SpanContext.
-  private io.opentelemetry.trace.Span parentSpan;
+  // Inherited baggage is supported only for the main parent.
-  private io.opentelemetry.trace.SpanContext parentSpanContext;
+  private SpanShim parentSpan;
+  private SpanContextShim parentSpanContext;
   private boolean ignoreActiveSpan;
 
   private final List<io.opentelemetry.trace.SpanContext> parentLinks = new ArrayList<>();
@@ -53,12 +54,12 @@ final class SpanBuilderShim extends BaseShimObject implements SpanBuilder {
     }
 
     // TODO - Verify we handle a no-op Span
-    io.opentelemetry.trace.Span actualParent = getActualSpan(parent);
+    SpanShim spanShim = getSpanShim(parent);
 
     if (parentSpan == null && parentSpanContext == null) {
-      parentSpan = actualParent;
+      parentSpan = spanShim;
     } else {
-      parentLinks.add(actualParent.getContext());
+      parentLinks.add(spanShim.getSpan().getContext());
     }
 
     return this;
@@ -76,12 +77,12 @@ final class SpanBuilderShim extends BaseShimObject implements SpanBuilder {
     }
 
     // TODO - Use referenceType
-    io.opentelemetry.trace.SpanContext actualContext = getActualContext(referencedContext);
+    SpanContextShim contextShim = getContextShim(referencedContext);
 
     if (parentSpan == null && parentSpanContext == null) {
-      parentSpanContext = actualContext;
+      parentSpanContext = contextShim;
     } else {
-      parentLinks.add(actualContext);
+      parentLinks.add(contextShim.getSpanContext());
     }
 
     return this;
@@ -175,14 +176,17 @@ final class SpanBuilderShim extends BaseShimObject implements SpanBuilder {
 
   @Override
   public Span start() {
+    io.opentelemetry.distributedcontext.DistributedContext distContext = null;
     io.opentelemetry.trace.Span.Builder builder = tracer().spanBuilder(spanName);
 
     if (ignoreActiveSpan && parentSpan == null && parentSpanContext == null) {
       builder.setNoParent();
     } else if (parentSpan != null) {
-      builder.setParent(parentSpan);
+      builder.setParent(parentSpan.getSpan());
+      distContext = spanContextTable().get(parentSpan).getDistributedContext();
     } else if (parentSpanContext != null) {
-      builder.setParent(parentSpanContext);
+      builder.setParent(parentSpanContext.getSpanContext());
+      distContext = parentSpanContext.getDistributedContext();
     }
 
     for (io.opentelemetry.trace.SpanContext link : parentLinks) {
@@ -204,22 +208,28 @@ final class SpanBuilderShim extends BaseShimObject implements SpanBuilder {
       span.setStatus(Status.UNKNOWN);
     }
 
-    return new SpanShim(telemetryInfo(), span);
+    SpanShim spanShim = new SpanShim(telemetryInfo(), span);
+
+    if (distContext != null && distContext != telemetryInfo().emptyDistributedContext()) {
+      spanContextTable().create(spanShim, distContext);
     }
 
-  private static io.opentelemetry.trace.Span getActualSpan(Span span) {
+    return spanShim;
+  }
+
+  private static SpanShim getSpanShim(Span span) {
     if (!(span instanceof SpanShim)) {
       throw new IllegalArgumentException("span is not a valid SpanShim object");
     }
 
-    return ((SpanShim) span).getSpan();
+    return (SpanShim) span;
   }
 
-  private static io.opentelemetry.trace.SpanContext getActualContext(SpanContext context) {
+  private static SpanContextShim getContextShim(SpanContext context) {
     if (!(context instanceof SpanContextShim)) {
       throw new IllegalArgumentException("context is not a valid SpanContextShim object");
     }
 
-    return ((SpanContextShim) context).getSpanContext();
+    return (SpanContextShim) context;
   }
 }

opentracing_shim/src/main/java/io/opentelemetry/opentracingshim/SpanContextShimTable.java

@@ -16,6 +16,7 @@
 
 package io.opentelemetry.opentracingshim;
 
+import io.opentelemetry.distributedcontext.DistributedContext;
 import io.opentelemetry.trace.Span;
 import java.util.Map;
 import java.util.WeakHashMap;
@@ -84,6 +85,10 @@ final class SpanContextShimTable {
   }
 
   public SpanContextShim create(SpanShim spanShim) {
+    return create(spanShim, spanShim.telemetryInfo().emptyDistributedContext());
+  }
+
+  public SpanContextShim create(SpanShim spanShim, DistributedContext distContext) {
     lock.writeLock().lock();
     try {
       SpanContextShim contextShim = shimsMap.get(spanShim.getSpan());
@@ -91,7 +96,9 @@ final class SpanContextShimTable {
         return contextShim;
       }
 
-      contextShim = new SpanContextShim(spanShim);
+      contextShim =
+          new SpanContextShim(
+              spanShim.telemetryInfo(), spanShim.getSpan().getContext(), distContext);
       shimsMap.put(spanShim.getSpan(), contextShim);
       return contextShim;
 

opentracing_shim/src/main/java/io/opentelemetry/opentracingshim/SpanShim.java

@@ -54,13 +54,13 @@ final class SpanShim extends BaseShimObject implements Span {
   @Override
   public SpanContext context() {
     /* Read the value using the read lock first. */
-    SpanContextShim contextShim = telemetryInfo().spanContextShimTable().get(this);
+    SpanContextShim contextShim = spanContextTable().get(this);
 
     /* Switch to the write lock *only* for the relatively exceptional case
      * of no context being created.
      * (as we cannot upgrade read->write lock sadly).*/
     if (contextShim == null) {
-      contextShim = telemetryInfo.spanContextShimTable().create(this);
+      contextShim = spanContextTable().create(this);
     }
 
     return contextShim;
@@ -147,7 +147,7 @@ final class SpanShim extends BaseShimObject implements Span {
       return this;
     }
 
-    telemetryInfo.spanContextShimTable().setBaggageItem(this, key, value);
+    spanContextTable().setBaggageItem(this, key, value);
 
     return this;
   }
@@ -159,7 +159,7 @@ final class SpanShim extends BaseShimObject implements Span {
       return null;
     }
 
-    return telemetryInfo.spanContextShimTable().getBaggageItem(this, key);
+    return spanContextTable().getBaggageItem(this, key);
   }
 
   @Override

opentracing_shim/src/main/java/io/opentelemetry/opentracingshim/TelemetryInfo.java

@@ -28,13 +28,13 @@ final class TelemetryInfo {
   private final Tracer tracer;
   private final DistributedContextManager contextManager;
   private final DistributedContext emptyDistributedContext;
-  private final SpanContextShimTable spanContextShimTable;
+  private final SpanContextShimTable spanContextTable;
 
   TelemetryInfo(Tracer tracer, DistributedContextManager contextManager) {
     this.tracer = tracer;
     this.contextManager = contextManager;
     this.emptyDistributedContext = contextManager.contextBuilder().build();
-    this.spanContextShimTable = new SpanContextShimTable();
+    this.spanContextTable = new SpanContextShimTable();
   }
 
   Tracer tracer() {
@@ -45,8 +45,8 @@ final class TelemetryInfo {
     return contextManager;
   }
 
-  SpanContextShimTable spanContextShimTable() {
+  SpanContextShimTable spanContextTable() {
-    return spanContextShimTable;
+    return spanContextTable;
   }
 
   DistributedContext emptyDistributedContext() {

opentracing_shim/src/test/java/io/opentelemetry/opentracingshim/SpanBuilderShimTest.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2019, OpenTelemetry Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.opentelemetry.opentracingshim;
+
+import static io.opentelemetry.opentracingshim.TestUtils.getBaggageMap;
+import static org.junit.Assert.assertEquals;
+
+import io.opentelemetry.sdk.distributedcontext.DistributedContextManagerSdk;
+import io.opentelemetry.sdk.trace.TracerSdkFactory;
+import io.opentelemetry.trace.Tracer;
+import org.junit.Test;
+
+public class SpanBuilderShimTest {
+  private final TracerSdkFactory tracerSdkFactory = TracerSdkFactory.create();
+  private final Tracer tracer = tracerSdkFactory.get("SpanShimTest");
+  private final TelemetryInfo telemetryInfo =
+      new TelemetryInfo(tracer, new DistributedContextManagerSdk());
+
+  private static final String SPAN_NAME = "Span";
+
+  @Test
+  public void baggage_parent() {
+    SpanShim parentSpan = (SpanShim) new SpanBuilderShim(telemetryInfo, SPAN_NAME).start();
+    try {
+      parentSpan.setBaggageItem("key1", "value1");
+
+      SpanShim childSpan =
+          (SpanShim) new SpanBuilderShim(telemetryInfo, SPAN_NAME).asChildOf(parentSpan).start();
+      try {
+        assertEquals(childSpan.getBaggageItem("key1"), "value1");
+        assertEquals(
+            getBaggageMap(childSpan.context().baggageItems()),
+            getBaggageMap(parentSpan.context().baggageItems()));
+      } finally {
+        childSpan.finish();
+      }
+    } finally {
+      parentSpan.finish();
+    }
+  }
+
+  @Test
+  public void baggage_parentContext() {
+    SpanShim parentSpan = (SpanShim) new SpanBuilderShim(telemetryInfo, SPAN_NAME).start();
+    try {
+      parentSpan.setBaggageItem("key1", "value1");
+
+      SpanShim childSpan =
+          (SpanShim)
+              new SpanBuilderShim(telemetryInfo, SPAN_NAME).asChildOf(parentSpan.context()).start();
+      try {
+        assertEquals(childSpan.getBaggageItem("key1"), "value1");
+        assertEquals(
+            getBaggageMap(childSpan.context().baggageItems()),
+            getBaggageMap(parentSpan.context().baggageItems()));
+      } finally {
+        childSpan.finish();
+      }
+    } finally {
+      parentSpan.finish();
+    }
+  }
+}

opentracing_shim/src/test/java/io/opentelemetry/opentracingshim/SpanShimTest.java

@@ -16,6 +16,7 @@
 
 package io.opentelemetry.opentracingshim;
 
+import static io.opentelemetry.opentracingshim.TestUtils.getBaggageMap;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
@@ -25,7 +26,6 @@ import static org.junit.Assert.assertTrue;
 import io.opentelemetry.sdk.distributedcontext.DistributedContextManagerSdk;
 import io.opentelemetry.sdk.trace.TracerSdkFactory;
 import io.opentelemetry.trace.Tracer;
-import java.util.HashMap;
 import java.util.Map;
 import org.junit.After;
 import org.junit.Before;
@@ -106,13 +106,4 @@ public class SpanShimTest {
         getBaggageMap(spanShim1.context().baggageItems()),
         getBaggageMap(spanShim2.context().baggageItems()));
   }
-
-  static Map<String, String> getBaggageMap(Iterable<Map.Entry<String, String>> baggage) {
-    Map<String, String> baggageMap = new HashMap<>();
-    for (Map.Entry<String, String> entry : baggage) {
-      baggageMap.put(entry.getKey(), entry.getValue());
-    }
-
-    return baggageMap;
-  }
 }

opentracing_shim/src/test/java/io/opentelemetry/opentracingshim/TestUtils.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2019, OpenTelemetry Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.opentelemetry.opentracingshim;
+
+import java.util.HashMap;
+import java.util.Map;
+
+final class TestUtils {
+  private TestUtils() {}
+
+  static Map<String, String> getBaggageMap(Iterable<Map.Entry<String, String>> baggage) {
+    Map<String, String> baggageMap = new HashMap<>();
+    for (Map.Entry<String, String> entry : baggage) {
+      baggageMap.put(entry.getKey(), entry.getValue());
+    }
+
+    return baggageMap;
+  }
+}

opentracing_shim/src/test/java/io/opentelemetry/opentracingshim/testbed/baggagehandling/BaggageHandlingTest.java

@@ -49,6 +49,15 @@ public final class BaggageHandlingTest {
                 /* add a new baggage item... */
                 span.setBaggageItem("newkey", "newvalue");
 
+                /* have a child that updates its own baggage
+                 * (should not be reflected in the original Span). */
+                Span childSpan = tracer.buildSpan("child").start();
+                try {
+                  childSpan.setBaggageItem("key1", "childvalue");
+                } finally {
+                  childSpan.finish();
+                }
+
                 /* and finish the Span. */
                 span.finish();
               }
@@ -57,7 +66,7 @@ public final class BaggageHandlingTest {
     /* Single call, no need to use await() */
     f.get(5, TimeUnit.SECONDS);
 
-    assertEquals(1, exporter.getFinishedSpanItems().size());
+    assertEquals(2, exporter.getFinishedSpanItems().size());
     assertEquals(span.getBaggageItem("key1"), "value2");
     assertEquals(span.getBaggageItem("newkey"), "newvalue");
   }