Adds better default impl from https://github.com/DistributedTracing/c…

…ontinuation-local-storage-jvm/pull/1/files

brave/README.md

@@ -374,6 +374,10 @@ least concerns this way. This approach is only possible due to the Recorder
 architecture, which allows different span instances with the same trace
 context to mutate the same shared state.
 
+The first design of `DefaultSpanScoper` included code from Eirik
+Sletteberg's [ContinuationLocal](https://github.com/DistributedTracing/continuation-local-storage-jvm/pull/1)
+draft, which is a dependency-free port of `com.twitter.util.Local`.
+
 ### Public namespace
 Brave 4's pubic namespace is more defensive that the past, using a package
 accessor design from [OkHttp](https://github.com/square/okhttp).

brave/src/main/java/brave/Tracer.java

@@ -10,6 +10,7 @@
 import brave.propagation.TraceContext;
 import brave.propagation.TraceContextOrSamplingFlags;
 import brave.sampler.Sampler;
+import java.util.concurrent.atomic.AtomicInteger;
 import zipkin.Endpoint;
 import zipkin.reporter.AsyncReporter;
 import zipkin.reporter.Reporter;
@@ -317,19 +318,51 @@ public SpanScoper.SpanInScope withSpan(Span span) {
     return currentContext != null ? toSpan(currentContext) : null;
   }
 
-  /** Default implementation which is backed by a thread local */
+  /**
+   * Default implementation which is backed by a thread local
+   *
+   * <p>This is based on com.twitter.util.Local port named ContinuationLocal written by Eirik
+   * Sletteberg
+   */
   static final class DefaultSpanScoper implements SpanScoper {
-    // TODO: make this implementation more realistic
-    final ThreadLocal<TraceContext> localSpan = new ThreadLocal<>();
+    static final InheritableThreadLocal<Object[]> localCtx = new InheritableThreadLocal<>();
+    static final AtomicInteger size = new AtomicInteger();
+
+    final int me = size.getAndIncrement();
 
     @Override public TraceContext currentSpan() {
-      return localSpan.get();
+      return (TraceContext) get(me);
     }
 
     @Override public SpanInScope open(TraceContext currentSpan) {
-      final TraceContext previous = localSpan.get();
-      localSpan.set(currentSpan);
-      return () -> localSpan.set(previous);
+      TraceContext saved = currentSpan();
+      set(me, currentSpan);
+      return () -> set(me, saved);
+    }
+
+    static Object get(int i) {
+      Object[] ctx = localCtx.get();
+      if (ctx == null || ctx.length <= i) {
+        return null;
+      }
+      return ctx[i];
+    }
+
+    static void set(int i, Object v) {
+      int currentSize = size.get();
+      assert i < currentSize;
+      Object[] ctx = localCtx.get();
+
+      if (ctx == null) {
+        ctx = new Object[currentSize];
+      } else {
+        Object[] oldCtx = ctx;
+        ctx = new Object[currentSize];
+        System.arraycopy(oldCtx, 0, ctx, 0, oldCtx.length);
+      }
+
+      ctx[i] = v;
+      localCtx.set(ctx);
     }
   }
 }

brave/src/test/java/brave/DefaultSpanScoperTest.java

@@ -0,0 +1,63 @@
+package brave;
+
+import brave.Tracer.DefaultSpanScoper;
+import brave.propagation.SpanScoper;
+import brave.propagation.TraceContext;
+import org.junit.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class DefaultSpanScoperTest {
+  DefaultSpanScoper scoper = new DefaultSpanScoper();
+  Tracer tracer = Tracer.newBuilder().build();
+  TraceContext context = tracer.newTrace().context();
+  TraceContext context2 = tracer.newTrace().context();
+
+  @Test public void currentSpan_defaultsToNull() {
+    assertThat(scoper.currentSpan()).isNull();
+  }
+
+  @Test public void scope_retainsContext() {
+    try (SpanScoper.SpanInScope scope = scoper.open(context)) {
+      assertThat(scoper.currentSpan())
+          .isEqualTo(context);
+    }
+  }
+
+  @Test public void scope_isDefinedPerThread() throws InterruptedException {
+    final TraceContext[] threadValue = new TraceContext[1];
+
+    try (SpanScoper.SpanInScope scope = scoper.open(context)) {
+      Thread t = new Thread(() -> { // inheritable thread local
+        assertThat(scoper.currentSpan())
+            .isEqualTo(context);
+
+        try (SpanScoper.SpanInScope scope2 = scoper.open(context2)) {
+          assertThat(scoper.currentSpan())
+              .isEqualTo(context2);
+          threadValue[0] = context2;
+        }
+      });
+
+      t.start();
+      t.join();
+      assertThat(scoper.currentSpan())
+          .isEqualTo(context);
+      assertThat(threadValue[0])
+          .isEqualTo(context2);
+    }
+  }
+
+  @Test public void instancesAreIndependent() {
+    DefaultSpanScoper scoper2 = new DefaultSpanScoper();
+
+    try (SpanScoper.SpanInScope scope1 = scoper.open(context)) {
+      assertThat(scoper2.currentSpan()).isNull();
+
+      try (SpanScoper.SpanInScope scope2 = scoper2.open(context2)) {
+        assertThat(scoper.currentSpan()).isEqualTo(context);
+        assertThat(scoper2.currentSpan()).isEqualTo(context2);
+      }
+    }
+  }
+}