Fixes bug where we didn't merge extra fields from implicit context (#605

)

While unusual, and likely a leak, it is possible to have a span in scope
when extracting a remote context. This fixes the code to merge the extra
fields in the current span with extracted fields and backfills tests.

Thanks to @aldex32 and @aukevanleeuwen for hunting this down

brave/src/main/java/brave/propagation/ExtraFieldPropagation.java

@@ -167,7 +167,7 @@ public Factory build() {
    *
    * <p>Prefer {@link #set(TraceContext, String, String)} if you have a reference to a span.
    */
-   public static void set(String name, String value) {
+  public static void set(String name, String value) {
     TraceContext context = currentTraceContext();
     if (context != null) set(context, name, value);
   }
@@ -252,32 +252,55 @@ static final class Factory extends Propagation.Factory {
 
     @Override public TraceContext decorate(TraceContext context) {
       TraceContext result = delegate.decorate(context);
-      int extraIndex = indexOfExtra(result.extra());
-      if (extraIndex != -1) {
-        Extra extra = (Extra) result.extra().get(extraIndex);
-        if (!Arrays.equals(extra.fieldNames, fieldNames)) {
-          throw new IllegalStateException(
-              String.format("Mixed name configuration unsupported: found %s, expected %s",
-                  Arrays.asList(extra.fieldNames), Arrays.asList(fieldNames))
-          );
-        }
 
-        // If this extra is unassociated (due to remote extraction),
-        // or it is the same span ID, re-use the instance.
-        if (extra.tryAssociate(context)) return context;
+      // If there's an implicit context while extracting only extra fields, we will have two extras!
+      List<Object> extras = result.extra();
+      int thisExtraIndex = -1, parentExtraIndex = -1;
+      for (int i = 0, length = extras.size(); i < length; i++) {
+        if (extras.get(i) instanceof Extra) {
+          Extra extra = (Extra) result.extra().get(i);
+          if (!Arrays.equals(extra.fieldNames, fieldNames)) {
+            throw new IllegalStateException(
+                String.format("Mixed name configuration unsupported: found %s, expected %s",
+                    Arrays.toString(extra.fieldNames), Arrays.toString(fieldNames))
+            );
+          }
+          // If this extra is unassociated (due to remote extraction),
+          // or it is the same span ID, re-use the instance.
+          if (extra.tryAssociate(context)) {
+            thisExtraIndex = i;
+          } else {
+            parentExtraIndex = i;
+          }
+        }
       }
-      // otherwise, we are creating a new instance of
+
+      if (thisExtraIndex != -1 && parentExtraIndex == -1) return context;
+
+      // otherwise, we are creating a new instance
       List<Object> copyOfExtra = new ArrayList<>(result.extra());
       Extra extra;
-      if (extraIndex != -1) {
-        extra = ((Extra) copyOfExtra.get(extraIndex)).clone();
-        copyOfExtra.set(extraIndex, extra);
-      } else {
+      if (thisExtraIndex == -1 && parentExtraIndex != -1) { // clone then parent (for copy-on-write)
+        extra = ((Extra) copyOfExtra.get(parentExtraIndex)).clone();
+        copyOfExtra.set(parentExtraIndex, extra);
+      } else if (thisExtraIndex != -1 && parentExtraIndex != -1) { // merge with the parent
+        extra = ((Extra) copyOfExtra.get(thisExtraIndex));
+        Extra parent = (Extra) copyOfExtra.remove(parentExtraIndex); // ensures only one extra
+        if (parent.values != null) { // then values were added to our parent
+          for (int i = 0; i < parent.values.length; i++) {
+            if (parent.values[i] != null && extra.get(i) == null) { // extracted wins vs parent
+              extra.set(i, parent.values[i]);
+            }
+          }
+        }
+      } else { // no fields were extracted and the parent also had no fields. create a new copy
         extra = new Extra(fieldNames);
         copyOfExtra.add(extra);
       }
-      extra.context = context;
-      return result.toBuilder().extra(Collections.unmodifiableList(copyOfExtra)).build();
+      TraceContext resultContext =
+          result.toBuilder().extra(Collections.unmodifiableList(copyOfExtra)).build();
+      extra.context = resultContext; // associate this with the new context
+      return resultContext;
     }
   }
 
@@ -467,7 +490,7 @@ static Map<String, String> getAll(List<Object> extraList) {
     if (elements == null) return Collections.emptyMap();
 
     Map<String, String> result = new LinkedHashMap<>();
-    for (int i = 0, length = elements.length; i<length; i++) {
+    for (int i = 0, length = elements.length; i < length; i++) {
       String value = elements[i];
       if (value != null) result.put(extra.fieldNames[i], value);
     }

brave/src/test/java/brave/propagation/ExtraFieldPropagationTest.java

@@ -1,5 +1,7 @@
 package brave.propagation;
 
+import brave.Span;
+import brave.Tracer;
 import brave.Tracing;
 import java.util.Arrays;
 import java.util.Collections;
@@ -88,6 +90,74 @@ public class ExtraFieldPropagationTest {
     }
   }
 
+  /**
+   * This scenario is possible, albeit rare. {@code tracer.nextSpan(extracted} } is called when
+   * there is an implicit parent. For example, you have a trace in progress when extracting trace
+   * state from an incoming message. Another example is where there is a span in scope due to a leak
+   * such as from using {@link CurrentTraceContext.Default#inheritable()}.
+   *
+   * <p>When we are only extracting extra fields, the state should merge as opposed to creating
+   * duplicate copies of {@link ExtraFieldPropagation.Extra}.
+   */
+  @Test public void nextSpanMergesExtraWithImplicitParent_hasFields() {
+    try (Tracing tracing = Tracing.newBuilder().propagationFactory(factory).build()) {
+      ExtraFieldPropagation.set(context, "x-vcap-request-id", "foo");
+      Span span = tracing.tracer().toSpan(context);
+
+      try (Tracer.SpanInScope ws = tracing.tracer().withSpanInScope(span)) {
+        carrier.put("x-amzn-trace-id", awsTraceId);
+        carrier.put("x-vcap-request-id", "bar"); // extracted should win!
+        TraceContext context1 = tracing.tracer().nextSpan(extractor.extract(carrier)).context();
+
+        assertThat(context1.extra()) // merged
+            .hasSize(1);
+        ExtraFieldPropagation.Extra extra = ((ExtraFieldPropagation.Extra) context1.extra().get(0));
+        assertThat(extra.values)
+            .containsExactly("bar", awsTraceId);
+        assertThat(extra.context)
+            .isSameAs(context1);
+      }
+    }
+  }
+
+  @Test public void nextSpanExtraWithImplicitParent_butNoImplicitExtraFields() {
+    try (Tracing tracing = Tracing.newBuilder().propagationFactory(factory).build()) {
+      Span span = tracing.tracer().toSpan(context);
+
+      try (Tracer.SpanInScope ws = tracing.tracer().withSpanInScope(span)) {
+        carrier.put("x-amzn-trace-id", awsTraceId);
+        TraceContext context1 = tracing.tracer().nextSpan(extractor.extract(carrier)).context();
+
+        assertThat(context1.extra()) // merged
+            .hasSize(1);
+        ExtraFieldPropagation.Extra extra = ((ExtraFieldPropagation.Extra) context1.extra().get(0));
+        assertThat(extra.values)
+            .containsExactly(null, awsTraceId);
+        assertThat(extra.context)
+            .isSameAs(context1);
+      }
+    }
+  }
+
+  @Test public void nextSpanExtraWithImplicitParent_butNoExtractedExtraFields() {
+    try (Tracing tracing = Tracing.newBuilder().propagationFactory(factory).build()) {
+      ExtraFieldPropagation.set(context, "x-vcap-request-id", "foo");
+      Span span = tracing.tracer().toSpan(context);
+
+      try (Tracer.SpanInScope ws = tracing.tracer().withSpanInScope(span)) {
+        TraceContext context1 = tracing.tracer().nextSpan(extractor.extract(carrier)).context();
+
+        assertThat(context1.extra()) // merged
+            .hasSize(1);
+        ExtraFieldPropagation.Extra extra = ((ExtraFieldPropagation.Extra) context1.extra().get(0));
+        assertThat(extra.values)
+            .containsExactly("foo", null);
+        assertThat(extra.context)
+            .isSameAs(context1);
+      }
+    }
+  }
+
   @Test public void downcasesNames() {
     ExtraFieldPropagation.Factory factory =
         (ExtraFieldPropagation.Factory) ExtraFieldPropagation.newFactory(B3Propagation.FACTORY,
@@ -168,6 +238,20 @@ public class ExtraFieldPropagationTest {
     }
   }
 
+  @Test public void toSpan_selfLinksContext() {
+    try (Tracing t = Tracing.newBuilder().propagationFactory(factory).build()) {
+      ExtraFieldPropagation.set(context, "x-amzn-trace-id", awsTraceId);
+
+      Span span = t.tracer().toSpan(context);
+
+      ExtraFieldPropagation.Extra extra =
+          (ExtraFieldPropagation.Extra) span.context().extra().get(0);
+
+      assertThat(extra.context)
+          .isSameAs(span.context());
+    }
+  }
+
   @Test public void current_set_noop_if_no_current_context() {
     try (Tracing t = Tracing.newBuilder().propagationFactory(factory).build()) {
       ExtraFieldPropagation.set("x-amzn-trace-id", awsTraceId); // doesn't throw

pom.xml

@@ -66,8 +66,8 @@
     <jetty-servlet25.version>7.6.21.v20160908</jetty-servlet25.version>
     <!-- Note: 3.1.x requires Java 8; 3.0.20.Final is broken -->
     <resteasy.version>3.1.4.Final</resteasy.version>
-    <zipkin.version>2.4.5</zipkin.version>
-    <zipkin-reporter2.version>2.3.1</zipkin-reporter2.version>
+    <zipkin.version>2.4.6</zipkin.version>
+    <zipkin-reporter2.version>2.3.2</zipkin-reporter2.version>
     <zipkin-reporter.version>1.1.2</zipkin-reporter.version>
     <finagle.version>18.1.0</finagle.version>
     <log4j.version>2.10.0</log4j.version>