Merge pull request #20 from llinder/issue-19

issue-19: ensure only one trace context is in the registry.

src/main/java/ratpack/zipkin/internal/RatpackCurrentTraceContext.java

@@ -13,7 +13,7 @@
  */
 package ratpack.zipkin.internal;
 
-import brave.http.HttpTracing;
+
 import brave.propagation.CurrentTraceContext;
 import brave.propagation.TraceContext;
 import java.util.function.Supplier;
@@ -46,20 +46,27 @@ public TraceContext get() {
 
   @Override
   public Scope newScope(TraceContext current) {
-    final TraceContextHolder previous = registrySupplier.get()
+    final MutableRegistry registry = registrySupplier.get();
+
+    // get previous entry if one exists so we can re-add it when
+    // the scope is closed.
+    final TraceContextHolder previous = registry
         .maybeGet(TraceContextHolder.class)
         .orElse(TraceContextHolder.EMPTY);
 
+    removeAll(registry);
+
     if (current != null) {
-      registrySupplier.get().add(new TraceContextHolder(current));
+      registry.add(new TraceContextHolder(current));
       MDC.put(TRACE_ID_KEY, current.traceIdString());
     } else {
-      registrySupplier.get().add(TraceContextHolder.EMPTY);
+      registry.add(TraceContextHolder.EMPTY);
       MDC.remove(TRACE_ID_KEY);
     }
 
     return () -> {
-      registrySupplier.get().add(previous);
+      removeAll(registry);
+      registry.add(previous);
       if (previous.context != null) {
         MDC.put(TRACE_ID_KEY, previous.context.traceIdString());
       } else {
@@ -68,6 +75,12 @@ public Scope newScope(TraceContext current) {
     };
   }
 
+  private void removeAll(MutableRegistry registry) {
+    registry
+      .getAll(TraceContextHolder.class)
+      .forEach(tch -> registry.remove(TraceContextHolder.class));
+  }
+
   /**
    * Used by TracedParallelBatch where its used to wrap a TraceContext and puts it in the
    * registry for the forked execution.  This is marked deprecated as we prefer not to
@@ -99,16 +112,10 @@ public static class TracingPropagationExecInitializer implements ExecInitializer
 
     @Override
     public void init(Execution execution) {
-      execution.maybeParent().ifPresent(parent -> {
-        parent.maybeGet(TraceContextHolder.class).ifPresent(traceContextHolder -> {
-          TraceContext traceContext = traceContextHolder.context;
-          if (traceContext == null) {
-            execution.add(RatpackCurrentTraceContext.TraceContextHolder.EMPTY);
-          } else {
-            execution.add(new RatpackCurrentTraceContext.TraceContextHolder(traceContext));
-          }
-        });
-      });
+      execution
+        .maybeParent()
+        .flatMap(parent -> parent.maybeGet(TraceContextHolder.class))
+        .ifPresent(execution::add);
     }
   }
 

src/test/groovy/ratpack/zipkin/internal/RatpackCurrentTraceContextSpec.groovy

@@ -1,5 +1,5 @@
 /**
- * Copyright 2016-2018 The OpenZipkin Authors
+ * Copyright 2016-2019 The OpenZipkin 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
@@ -15,9 +15,12 @@ package ratpack.zipkin.internal
 
 import brave.propagation.CurrentTraceContext
 import brave.propagation.TraceContext
+import ratpack.exec.ExecSpec
+import ratpack.exec.Execution
 import ratpack.registry.MutableRegistry
 import ratpack.registry.Registry
 import spock.lang.Specification
+import ratpack.zipkin.internal.RatpackCurrentTraceContext.TraceContextHolder
 
 class RatpackCurrentTraceContextSpec extends Specification {
     MutableRegistry registry = Registry.mutable()
@@ -40,6 +43,7 @@ class RatpackCurrentTraceContextSpec extends Specification {
             def current = traceContext.get()
         expect:
             current == null
+            registry.getAll(TraceContextHolder).size() == 0
     }
 
     def 'When setting TraceContext span, should return same TraceContext'() {
@@ -51,6 +55,7 @@ class RatpackCurrentTraceContextSpec extends Specification {
             def result = traceContext.get()
         then:
             result == expected
+            registry.getAll(TraceContextHolder).size() == 1
     }
 
     def 'When closing a scope, trace context should revert back to previous'() {
@@ -65,6 +70,7 @@ class RatpackCurrentTraceContextSpec extends Specification {
             def traceContext = traceContext.get()
         then:
             expected ==  traceContext
+            registry.getAll(TraceContextHolder).size() == 1
     }
 
 
@@ -80,6 +86,7 @@ class RatpackCurrentTraceContextSpec extends Specification {
             def traceContext = traceContext.get()
         then:
             traceContext == null
+            registry.getAll(TraceContextHolder).size() == 1
     }
 
     def 'When TraceContext is null the context should be cleared'() {
@@ -91,10 +98,12 @@ class RatpackCurrentTraceContextSpec extends Specification {
             def result = traceContext.get()
         then:
             result == expected
+            registry.getAll(TraceContextHolder).size() == 1
         when:
             traceContext.newScope(null)
         then:
             traceContext.get() == null
+            registry.getAll(TraceContextHolder).size() == 1
     }
 
 }

src/test/groovy/ratpack/zipkin/internal/TracePropagationExecInitializerSpec.groovy

@@ -0,0 +1,73 @@
+/**
+ * Copyright 2016-2019 The OpenZipkin 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 ratpack.zipkin.internal
+
+import brave.propagation.TraceContext
+import ratpack.exec.Execution
+import ratpack.exec.ExecutionRef
+import spock.lang.Specification
+import ratpack.zipkin.internal.RatpackCurrentTraceContext.TraceContextHolder
+import ratpack.zipkin.internal.RatpackCurrentTraceContext.TracingPropagationExecInitializer
+
+class TracePropagationExecInitializerSpec extends Specification {
+
+    def dummyContext() {
+        TraceContext
+                .newBuilder()
+                .traceId(new Random().nextLong())
+                .spanId(new Random().nextLong())
+                .build()
+    }
+
+    def 'should copy context from parent when present'() {
+        given:
+            def initializer = new TracingPropagationExecInitializer()
+            def parent = Mock(ExecutionRef)
+            def execution = Mock(Execution)
+            def parentContext = dummyContext()
+            def parentContextHolder = new TraceContextHolder(parentContext)
+        when:
+            initializer.init(execution)
+        then:
+            1 * execution.maybeParent() >> Optional.of(parent)
+            1 * parent.maybeGet(TraceContextHolder.class) >> Optional.of(parentContextHolder)
+            1 * execution.add(parentContextHolder)
+            0 * _
+    }
+
+    def 'should not copy context when parent is not present'() {
+        given:
+            def initializer = new TracingPropagationExecInitializer()
+            def execution = Mock(Execution)
+        when:
+            initializer.init(execution)
+        then:
+            1 * execution.maybeParent() >> Optional.empty()
+            0 * _
+    }
+
+    def 'should not copy context when parent context is empty'() {
+        given:
+            def initializer = new TracingPropagationExecInitializer()
+            def parent = Mock(ExecutionRef)
+            def execution = Mock(Execution)
+        when:
+            initializer.init(execution)
+        then:
+            1 * execution.maybeParent() >> Optional.of(parent)
+            1 * parent.maybeGet(TraceContextHolder.class) >> Optional.empty()
+            0 * _
+    }
+
+}