Fix sampling when jaeger-baggage header is given

CHANGELOG.md

@@ -4,7 +4,7 @@ Changes by Version
 0.31.1 (unreleased)
 -------------------
 
-- Nothing yet
+- Fix sampling when `jaeger-baggage` header is given ([#542](https://github.com/jaegertracing/jaeger-client-java/pull/542), [@yurishkuro](https://github.com/yurishkuro))
 
 0.31.0 (2018-08-28)
 -------------------

jaeger-core/src/main/java/io/jaegertracing/internal/JaegerSpanContext.java

@@ -125,16 +125,15 @@ public JaegerSpanContext withFlags(byte flags) {
   }
 
   /**
-   * @return true when the instance of the context is only used to return the debug/correlation ID
-   * from extract() method. This happens in the situation when "jaeger-debug-id" header is passed in
-   * the carrier to the extract() method, but the request otherwise has no span context in it.
-   * Previously this would've returned null from the extract method, but now it returns a dummy
-   * context with only debugId filled in.
+   * @return true when the instance of the context contains a non-zero trace and span ID,
+   * indicating a valid trace. It may return false if the context was created with only
+   * a debugId or baggage passed via special ad-hoc headers.
    *
    * @see Constants#DEBUG_ID_HEADER_KEY
+   * @see Constants#BAGGAGE_HEADER_KEY
    */
-  boolean isDebugIdContainerOnly() {
-    return traceId == 0 && debugId != null;
+  boolean hasTrace() {
+    return traceId != 0 && spanId != 0;
   }
 
   /**

jaeger-core/src/main/java/io/jaegertracing/internal/JaegerTracer.java

@@ -289,7 +289,8 @@ public JaegerTracer.SpanBuilder withStartTimestamp(long microseconds) {
       return this;
     }
 
-    private JaegerSpanContext createNewContext(String debugId) {
+    private JaegerSpanContext createNewContext() {
+      String debugId = getDebugId();
       long id = Utils.uniqueId();
 
       byte flags = 0;
@@ -298,7 +299,7 @@ private JaegerSpanContext createNewContext(String debugId) {
         tags.put(Constants.DEBUG_ID_HEADER_KEY, debugId);
         metrics.traceStartedSampled.inc(1);
       } else {
-        //TODO(prithvi): Don't assume operationName is set on creation
+        // TODO: (prithvi) Don't assume operationName is set on creation
         SamplingStatus samplingStatus = sampler.sample(operationName, id);
         if (samplingStatus.isSampled()) {
           flags |= JaegerSpanContext.flagSampled;
@@ -314,11 +315,11 @@ private JaegerSpanContext createNewContext(String debugId) {
           id,
           0,
           flags,
-          Collections.<String, String>emptyMap(),
+          getBaggage(),
           debugId);
     }
 
-    private Map<String, String> createChildBaggage() {
+    private Map<String, String> getBaggage() {
       Map<String, String> baggage = null;
 
       // optimization for 99% use cases, when there is only one parent
@@ -360,7 +361,7 @@ private JaegerSpanContext createChildContext() {
           preferredReference.getSpanId(),
           // should we do OR across passed references?
           preferredReference.getFlags(),
-          createChildBaggage(),
+          getBaggage(),
           null);
     }
 
@@ -393,11 +394,11 @@ private boolean isSampled() {
       return false;
     }
 
-    private String debugId() {
-      if (references.size() == 1 && references.get(0).getSpanContext().isDebugIdContainerOnly()) {
-        return references.get(0).getSpanContext().getDebugId();
+    private String getDebugId() {
+      if (references.isEmpty()) {
+        return null;
       }
-      return null;
+      return references.get(0).getSpanContext().getDebugId();
     }
 
     @Override
@@ -409,9 +410,8 @@ public JaegerSpan start() {
         asChildOf(scopeManager.active().span());
       }
 
-      String debugId = debugId();
-      if (references.isEmpty() || debugId != null) {
-        context = createNewContext(debugId);
+      if (references.isEmpty() || !references.get(0).getSpanContext().hasTrace()) {
+        context = createNewContext();
       } else {
         context = createChildContext();
       }

jaeger-core/src/test/java/io/jaegertracing/internal/PropagationTest.java → jaeger-core/src/test/java/io/jaegertracing/internal/ActiveSpanTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016, Uber Technologies, Inc
+ * Copyright (c) 2018, The Jaeger 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
@@ -16,68 +16,50 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;
 
-import io.jaegertracing.internal.Constants;
-import io.jaegertracing.internal.JaegerSpan;
-import io.jaegertracing.internal.JaegerSpanContext;
-import io.jaegertracing.internal.JaegerTracer;
 import io.jaegertracing.internal.reporters.InMemoryReporter;
 import io.jaegertracing.internal.samplers.ConstSampler;
 import io.opentracing.References;
 import io.opentracing.Scope;
 import io.opentracing.ScopeManager;
 import io.opentracing.Span;
-import io.opentracing.propagation.Format;
-import io.opentracing.propagation.TextMap;
-import io.opentracing.propagation.TextMapExtractAdapter;
-import java.util.Collections;
-import java.util.Map;
 
+import org.junit.Before;
 import org.junit.Test;
+import org.mockito.Mockito;
 
-public class PropagationTest {
-  @Test
-  public void testDebugCorrelationId() {
-    JaegerTracer tracer = new JaegerTracer.Builder("test")
-            .withReporter(new InMemoryReporter())
+public class ActiveSpanTest {
+  InMemoryReporter reporter;
+  JaegerTracer tracer;
+
+  @Before
+  public void setUp() {
+    reporter = new InMemoryReporter();
+    tracer =
+        new JaegerTracer.Builder("TracerTestService")
+            .withReporter(reporter)
             .withSampler(new ConstSampler(true))
             .build();
-    Map<String, String> headers = Collections.singletonMap(Constants.DEBUG_ID_HEADER_KEY, "Coraline");
-    TextMap carrier = new TextMapExtractAdapter(headers);
-
-    JaegerSpanContext jaegerSpanContext = tracer.extract(Format.Builtin.TEXT_MAP, carrier);
-    assertNotNull(jaegerSpanContext);
-    assertTrue(jaegerSpanContext.isDebugIdContainerOnly());
-    assertEquals("Coraline", jaegerSpanContext.getDebugId());
-
-    JaegerSpan span = tracer.buildSpan("span").asChildOf(jaegerSpanContext).start();
-    jaegerSpanContext = span.context();
-    assertTrue(jaegerSpanContext.isSampled());
-    assertTrue(jaegerSpanContext.isDebug());
-    assertEquals("Coraline", span.getTags().get(Constants.DEBUG_ID_HEADER_KEY));
+  }
+
+  @Test
+  public void testActiveSpan() {
+    JaegerSpan mockSpan = Mockito.mock(JaegerSpan.class);
+    tracer.scopeManager().activate(mockSpan, true);
+    assertEquals(mockSpan, tracer.activeSpan());
   }
 
   @Test
   public void testActiveSpanPropagation() {
-    JaegerTracer tracer = new JaegerTracer.Builder("test")
-            .withReporter(new InMemoryReporter())
-            .withSampler(new ConstSampler(true))
-            .build();
     try (Scope parent = tracer.buildSpan("parent").startActive(true)) {
       assertEquals(parent, tracer.scopeManager().active());
     }
   }
 
   @Test
   public void testActiveSpanAutoReference() {
-    InMemoryReporter reporter = new InMemoryReporter();
-    JaegerTracer tracer = new JaegerTracer.Builder("test")
-            .withReporter(reporter)
-            .withSampler(new ConstSampler(true))
-            .build();
     try (Scope ignored = tracer.buildSpan("parent").startActive(true)) {
       tracer.buildSpan("child").startActive(true).close();
     }
@@ -101,22 +83,12 @@ public void testActiveSpanAutoReference() {
 
   @Test
   public void testActiveSpanAutoFinishOnClose() {
-    InMemoryReporter reporter = new InMemoryReporter();
-    JaegerTracer tracer = new JaegerTracer.Builder("test")
-            .withReporter(reporter)
-            .withSampler(new ConstSampler(true))
-            .build();
     tracer.buildSpan("parent").startActive(true).close();
     assertEquals(1, reporter.getSpans().size());
   }
 
   @Test
   public void testActiveSpanNotAutoFinishOnClose() {
-    InMemoryReporter reporter = new InMemoryReporter();
-    JaegerTracer tracer = new JaegerTracer.Builder("test")
-            .withReporter(reporter)
-            .withSampler(new ConstSampler(true))
-            .build();
     Scope scope = tracer.buildSpan("parent").startActive(false);
     Span span = scope.span();
     scope.close();
@@ -127,11 +99,6 @@ public void testActiveSpanNotAutoFinishOnClose() {
 
   @Test
   public void testIgnoreActiveSpan() {
-    InMemoryReporter reporter = new InMemoryReporter();
-    JaegerTracer tracer = new JaegerTracer.Builder("test")
-            .withReporter(reporter)
-            .withSampler(new ConstSampler(true))
-            .build();
     try (Scope ignored = tracer.buildSpan("parent").startActive(true)) {
       tracer.buildSpan("child").ignoreActiveSpan().startActive(true).close();
     }
@@ -150,12 +117,6 @@ public void testIgnoreActiveSpan() {
 
   @Test
   public void testNoAutoRefWithExistingRefs() {
-    InMemoryReporter reporter = new InMemoryReporter();
-    JaegerTracer tracer = new JaegerTracer.Builder("test")
-        .withReporter(reporter)
-        .withSampler(new ConstSampler(true))
-        .build();
-
     JaegerSpan initialSpan = tracer.buildSpan("initial").start();
 
     try (Scope ignored = tracer.buildSpan("parent").startActive(true)) {
@@ -204,4 +165,18 @@ public Scope active() {
         }).build();
     assertEquals(scope, tracer.scopeManager().active());
   }
+
+
+  @Test
+  public void testCustomSpanOnSpanManager() {
+    // prepare
+    Span activeSpan = mock(Span.class);
+    ScopeManager scopeManager = tracer.scopeManager();
+
+    // test
+    scopeManager.activate(activeSpan, false);
+
+    // check
+    assertEquals(activeSpan, tracer.activeSpan());
+  }
 }

jaeger-core/src/test/java/io/jaegertracing/internal/AdhocHeadersTest.java

@@ -0,0 +1,93 @@
+/*
+ * Copyright (c) 2018, The Jaeger Authors.
+ * Copyright (c) 2016-2017, Uber Technologies, Inc
+ *
+ * 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.jaegertracing.internal;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import io.jaegertracing.internal.reporters.InMemoryReporter;
+import io.jaegertracing.internal.samplers.ConstSampler;
+import io.opentracing.Span;
+import io.opentracing.propagation.Format;
+import io.opentracing.propagation.TextMap;
+import io.opentracing.propagation.TextMapExtractAdapter;
+import io.opentracing.propagation.TextMapInjectAdapter;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.junit.Before;
+import org.junit.Test;
+
+public class AdhocHeadersTest {
+  JaegerTracer tracer;
+
+  @Before
+  public void setUp() {
+    tracer =
+        new JaegerTracer.Builder("TracerTestService")
+            .withReporter(new InMemoryReporter())
+            .withSampler(new ConstSampler(true))
+            .build();
+  }
+
+  @Test
+  public void testDebugCorrelationId() {
+    Map<String, String> headers = Collections.singletonMap(Constants.DEBUG_ID_HEADER_KEY, "Coraline");
+    TextMap carrier = new TextMapExtractAdapter(headers);
+
+    JaegerSpanContext inboundSpanContext = tracer.extract(Format.Builtin.TEXT_MAP, carrier);
+    assertNotNull(inboundSpanContext);
+    assertFalse(inboundSpanContext.hasTrace());
+    assertEquals("Coraline", inboundSpanContext.getDebugId());
+
+    JaegerSpan span = tracer.buildSpan("span").asChildOf(inboundSpanContext).start();
+    JaegerSpanContext serverSpanContext = span.context();
+    assertTrue(serverSpanContext.isSampled());
+    assertTrue(serverSpanContext.isDebug());
+    assertEquals("Coraline", span.getTags().get(Constants.DEBUG_ID_HEADER_KEY));
+  }
+
+
+  @Test
+  public void testStartTraceWithAdhocBaggage() {
+    traceWithAdhocBaggage(new HashMap<String, String>());
+  }
+
+  @Test
+  public void testJoinTraceWithAdhocBaggage() {
+    Span span = tracer.buildSpan("test").start();
+    Map<String, String> headers = new HashMap<String, String>();
+    tracer.inject(span.context(), Format.Builtin.HTTP_HEADERS, new TextMapInjectAdapter(headers));
+    assertEquals(1, headers.size());
+
+    traceWithAdhocBaggage(headers);
+  }
+
+  private void traceWithAdhocBaggage(Map<String, String> headers) {
+    headers.put("jaeger-baggage", "k1=v1, k2 = v2");
+
+    JaegerSpanContext parent = tracer.extract(Format.Builtin.HTTP_HEADERS, new TextMapExtractAdapter(headers));
+    JaegerSpan span = tracer.buildSpan("test").asChildOf(parent).start();
+
+    assertTrue(span.context().isSampled());
+    assertEquals("must have baggage", "v1", span.getBaggageItem("k1"));
+    assertEquals("must have baggage", "v2", span.getBaggageItem("k2"));
+  }
+}

jaeger-core/src/test/java/io/jaegertracing/internal/JaegerSpanTest.java

@@ -237,11 +237,8 @@ public void testLogWithTimestamp() {
     long expectedTimestamp = 2222;
     final String expectedLog = "some-log";
     final String expectedEvent = "event";
-    Map<String, String> expectedFields = new HashMap<String, String>() {
-      {
-        put(expectedEvent, expectedLog);
-      }
-    };
+    Map<String, String> expectedFields = new HashMap<String, String>();
+    expectedFields.put(expectedEvent, expectedLog);
 
     jaegerSpan.log(expectedTimestamp, expectedEvent);
     jaegerSpan.log(expectedTimestamp, expectedFields);

jaeger-core/src/test/java/io/jaegertracing/internal/JaegerSubclassTest.java

@@ -17,7 +17,6 @@
 import io.jaegertracing.Configuration;
 import io.opentracing.Scope;
 
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;