Change the trace API to directly depend on the Context implementation. (#154)

* Change the trace API to directly depend on the Context implementation.

Other changes in this PR:
Enable maven tests for core (common and trace).
Make trace tests in core to compile with java6.

Reasons to change to directly use the io.grpc.Context API:

We would need to have a different API that allows users to pass the generic context between threads and via Executors (similar with io.grpc.Context).
io.grpc.Context is an API the Storage is decoupled and can be backed by any mechanism not necessary thread-local.
Support for Deadline and Cancellation is a nice to have.
Other products may already support propagation of the io.grpc.Context.
It would have been very hard to make binaries work correctly when they have libraries that are using io.grpc.Context and libraries that are using something different.

11 files changed, 339 insertions, 254 deletions

diff --git a/core/build.gradle b/core/build.gradle
index da14e524..067b34c1 100644
--- a/core/build.gradle
+++ b/core/build.gradle
@@ -2,9 +2,11 @@ description = 'Instrumentation: Core'
 
 dependencies {
     compile project(':shared'),
+            libraries.grpc_context,
             libraries.guava,
             libraries.errorprone,
             libraries.jsr305
 
-    testCompile project(':instrumentation-core-impl')
+    testCompile project(':instrumentation-core-impl'),
+            libraries.grpc_context
 }
diff --git a/core/pom.xml b/core/pom.xml
index 5a623297..b8bcf678 100644
--- a/core/pom.xml
+++ b/core/pom.xml
@@ -16,6 +16,15 @@
     <version>0.1.0-SNAPSHOT</version>
   </parent>
 
+  <dependencies>
+    <dependency>
+      <groupId>com.google.io</groupId>
+      <artifactId>shared</artifactId>
+      <version>${project.version}</version>
+      <scope>test</scope>
+    </dependency>
+  </dependencies>
+
   <build>
     <plugins>
       <plugin>
@@ -30,10 +39,12 @@
             </goals>
             <configuration>
               <!--
-              Don't build the tests, because they currently create a circular
+              Don't build the stats tests, because they currently create a circular
               dependency between the 'core' and 'core_impl' directories.
               -->
-              <skip>true</skip>
+              <testExcludes>
+                <exclude>**/stats/*.java</exclude>
+              </testExcludes>
             </configuration>
           </execution>
         </executions>
diff --git a/core/src/main/java/com/google/instrumentation/trace/ContextSpanHandler.java b/core/src/main/java/com/google/instrumentation/trace/ContextSpanHandler.java
deleted file mode 100644
index c7eac797..00000000
--- a/core/src/main/java/com/google/instrumentation/trace/ContextSpanHandler.java
+++ /dev/null
@@ -1,38 +0,0 @@
-/*
- * Copyright 2017, Google 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 com.google.instrumentation.trace;
-
-import com.google.instrumentation.common.NonThrowingCloseable;
-
-/** Class to assist in {@link Span}/Context interactions. */
-abstract class ContextSpanHandler {
-  /**
-   * Returns The {@link Span} from the current context.
-   *
-   * @return The {@code Span} from the current context.
-   */
-  abstract Span getCurrentSpan();
-
-  /**
-   * Enters the scope of code where the given {@link Span} is in the current context, and returns an
-   * object that represents that scope. The scope is exited when the returned object is closed.
-   *
-   * <p>Supports try-with-resource idiom.
-   *
-   * @param span The {@code Span} to be set to the current context.
-   * @return An object that defines a scope where the given {@code Span} will be set to the current
-   *     context.
-   */
-  abstract NonThrowingCloseable withSpan(Span span);
-}
diff --git a/core/src/main/java/com/google/instrumentation/trace/ContextUtils.java b/core/src/main/java/com/google/instrumentation/trace/ContextUtils.java
new file mode 100644
index 00000000..bab9c189
--- /dev/null
+++ b/core/src/main/java/com/google/instrumentation/trace/ContextUtils.java
@@ -0,0 +1,75 @@
+/*
+ * Copyright 2017, Google 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 com.google.instrumentation.trace;
+
+import com.google.instrumentation.common.NonThrowingCloseable;
+import io.grpc.Context;
+
+/**
+ * Util methods/functionality to interact with the {@link io.grpc.Context}.
+ *
+ * <p>Users must interact with the current Context via the public APIs in {@link Tracer} and avoid
+ * usages of the {@link #CONTEXT_SPAN_KEY} directly.
+ */
+public final class ContextUtils {
+  /** The {@link io.grpc.Context.Key} used to interact with {@link io.grpc.Context}. */
+  public static final Context.Key<Span> CONTEXT_SPAN_KEY = Context.key("instrumentation-trace-key");
+
+  // No instance of this class.
+  private ContextUtils() {}
+
+  /**
+   * Returns The {@link Span} from the current context.
+   *
+   * @return The {@code Span} from the current context.
+   */
+  static Span getCurrentSpan() {
+    return CONTEXT_SPAN_KEY.get(Context.current());
+  }
+
+  /**
+   * Enters the scope of code where the given {@link Span} is in the current context, and returns an
+   * object that represents that scope. The scope is exited when the returned object is closed.
+   *
+   * <p>Supports try-with-resource idiom.
+   *
+   * @param span The {@code Span} to be set to the current context.
+   * @return An object that defines a scope where the given {@code Span} is set to the current
+   *     context.
+   */
+  static NonThrowingCloseable withSpan(Span span) {
+    return new WithSpan(span, CONTEXT_SPAN_KEY);
+  }
+
+  // Defines an arbitrary scope of code as a traceable operation. Supports try-with-resources idiom.
+  private static final class WithSpan implements NonThrowingCloseable {
+    private final Context origContext;
+
+    /**
+     * Constructs a new {@link WithSpan}.
+     *
+     * @param span is the {@code Span} to be added to the current {@code io.grpc.Context}.
+     * @param contextKey is the {@code Context.Key} used to set/get {@code Span} from the {@code
+     *     Context}.
+     */
+    WithSpan(Span span, Context.Key<Span> contextKey) {
+      origContext = Context.current().withValue(contextKey, span).attach();
+    }
+
+    @Override
+    public void close() {
+      Context.current().detach(origContext);
+    }
+  }
+}
diff --git a/core/src/main/java/com/google/instrumentation/trace/ScopedSpanHandle.java b/core/src/main/java/com/google/instrumentation/trace/ScopedSpanHandle.java
index 152e17e3..2838071b 100644
--- a/core/src/main/java/com/google/instrumentation/trace/ScopedSpanHandle.java
+++ b/core/src/main/java/com/google/instrumentation/trace/ScopedSpanHandle.java
@@ -30,11 +30,10 @@ final class ScopedSpanHandle implements NonThrowingCloseable {
    * Creates a {@code ScopedSpanHandle}
    *
    * @param span The span that will be installed in the current context.
-   * @param contextSpanHandler The handler that is used to interact with the current context.
    */
-  ScopedSpanHandle(Span span, ContextSpanHandler contextSpanHandler) {
+  ScopedSpanHandle(Span span) {
     this.span = span;
-    this.withSpan = contextSpanHandler.withSpan(span);
+    this.withSpan = ContextUtils.withSpan(span);
   }
 
   /**
diff --git a/core/src/main/java/com/google/instrumentation/trace/SpanBuilder.java b/core/src/main/java/com/google/instrumentation/trace/SpanBuilder.java
index 69298e68..9341d8cf 100644
--- a/core/src/main/java/com/google/instrumentation/trace/SpanBuilder.java
+++ b/core/src/main/java/com/google/instrumentation/trace/SpanBuilder.java
@@ -22,25 +22,6 @@ import javax.annotation.Nullable;
  * {@link SpanBuilder} is used to construct {@link Span} instances which define arbitrary scopes of
  * code that are sampled for distributed tracing as a single atomic unit.
  *
- * <p>This is a simple example where all the work is being done within a single scope and the
- * Context is manually propagated:
- *
- * <pre>{@code
- * class MyClass {
- *   private static final Tracer tracer = Tracer.getTracer();
- *   void DoWork() {
- *     Span span = tracer.spanBuilder(null, "MyRootSpan").startSpan();
- *     span.addAnnotation("my annotation");
- *     try {
- *       doSomeWork(span); // Manually propagate the new span down the stack.
- *     } finally {
- *       // To make sure we end the span even in case of an exception.
- *       span.end();  // Manually end the span.
- *     }
- *   }
- * }
- * }</pre>
- *
  * <p>This is a simple example where all the work is being done within a single scope and a single
  * thread and the Context is automatically propagated:
  *
@@ -96,20 +77,37 @@ import javax.annotation.Nullable;
  * }
  * }</pre>
  *
+ * <p>This is a simple example where all the work is being done within a single scope and the
+ * Context is manually propagated:
+ *
+ * <pre>{@code
+ * class MyClass {
+ *   private static final Tracer tracer = Tracer.getTracer();
+ *   void DoWork() {
+ *     Span span = tracer.spanBuilder(null, "MyRootSpan").startSpan();
+ *     span.addAnnotation("my annotation");
+ *     try {
+ *       doSomeWork(span); // Manually propagate the new span down the stack.
+ *     } finally {
+ *       // To make sure we end the span even in case of an exception.
+ *       span.end();  // Manually end the span.
+ *     }
+ *   }
+ * }
+ * }</pre>
+ *
  * <p>If your Java version is less than Java SE 7, see {@link SpanBuilder#startSpan} and {@link
  * SpanBuilder#startScopedSpan} for usage examples.
  */
 public final class SpanBuilder {
   private final SpanFactory spanFactory;
-  private final ContextSpanHandler contextSpanHandler;
-  private SpanContext parentSpanContext;
-  private boolean hasRemoteParent;
   private final String name;
   private final StartSpanOptions startSpanOption = new StartSpanOptions();
+  private SpanContext parentSpanContext;
+  private boolean hasRemoteParent;
 
   SpanBuilder(
       SpanFactory spanFactory,
-      ContextSpanHandler contextSpanHandler,
       SpanContext parentSpanContext,
       boolean hasRemoteParent,
       String name) {
@@ -117,7 +115,6 @@ public final class SpanBuilder {
     this.hasRemoteParent = hasRemoteParent;
     this.name = name;
     this.spanFactory = spanFactory;
-    this.contextSpanHandler = contextSpanHandler;
   }
 
   /**
@@ -273,7 +270,7 @@ public final class SpanBuilder {
    *     current Context.
    */
   public NonThrowingCloseable startScopedSpan() {
-    return new ScopedSpanHandle(start(), contextSpanHandler);
+    return new ScopedSpanHandle(start());
   }
 
   // Utility method to start a Span.
diff --git a/core/src/main/java/com/google/instrumentation/trace/Tracer.java b/core/src/main/java/com/google/instrumentation/trace/Tracer.java
index 432e7db4..07cac993 100644
--- a/core/src/main/java/com/google/instrumentation/trace/Tracer.java
+++ b/core/src/main/java/com/google/instrumentation/trace/Tracer.java
@@ -31,14 +31,21 @@ import javax.annotation.Nullable;
  * <p>Users may choose to use manual or automatic Context propagation. Because of that this class
  * offers APIs to facilitate both usages.
  *
+ * <p>The automatic context propagation is done using {@link io.grpc.Context} which is a gRPC
+ * independent implementation for in-process Context propagation mechanism which can carry
+ * scoped-values across API boundaries and between threads. Users of the library must propagate the
+ * {@link io.grpc.Context} between different threads.
+ *
  * <p>Example usage with automatic context propagation:
  *
  * <pre>{@code
  * class MyClass {
  *   private static final Tracer tracer = Tracer.getTracer();
- *   void DoWork() {
+ *   void doWork() {
  *     try(NonThrowingCloseable ss = tracer.spanBuilder("MyClass.DoWork").startScopedSpan) {
- *       tracer.getCurrentSpan().addAnnotation("We did the work.");
+ *       tracer.getCurrentSpan().addAnnotation("Starting the work.");
+ *       doWorkInternal();
+ *       tracer.getCurrentSpan().addAnnotation("Finished working.");
  *     }
  *   }
  * }
@@ -49,9 +56,15 @@ import javax.annotation.Nullable;
  * <pre>{@code
  * class MyClass {
  *   private static final Tracer tracer = Tracer.getTracer();
- *   void DoWork(Span parent) {
- *     try(NonThrowingCloseable ss = tracer.spanBuilder(parent, "MyClass.DoWork").startScopedSpan) {
- *       tracer.getCurrentSpan().addAnnotation("We did the work.");
+ *   void doWork() {
+ *     Span span = tracer.spanBuilder(null, "MyRootSpan").startSpan();
+ *     span.addAnnotation("Starting the work.");
+ *     try {
+ *       doSomeWork(span); // Manually propagate the new span down the stack.
+ *     } finally {
+ *       span.addAnnotation("Finished working.");
+ *       // To make sure we end the span even in case of an exception.
+ *       span.end();  // Manually end the span.
  *     }
  *   }
  * }
@@ -60,25 +73,38 @@ import javax.annotation.Nullable;
 public final class Tracer {
   private static final Logger logger = Logger.getLogger(Tracer.class.getName());
   private static final Tracer INSTANCE =
-      new Tracer(
-          loadContextSpanHandler(Provider.getCorrectClassLoader(ContextSpanHandler.class)),
-          loadSpanFactory(Provider.getCorrectClassLoader(SpanFactory.class)));
-  private final ContextSpanHandler contextSpanHandler;
+      new Tracer(loadSpanFactory(Provider.getCorrectClassLoader(SpanFactory.class)));
   private final SpanFactory spanFactory;
 
+  @VisibleForTesting
+  Tracer(SpanFactory spanFactory) {
+    this.spanFactory = checkNotNull(spanFactory, "spanFactory");
+  }
+
   /**
-   * Returns the {@link Tracer} with the provided implementations for {@link ContextSpanHandler} and
-   * {@link SpanFactory}.
-   *
-   * <p>If no implementation is provided for any of the {@link Tracer} modules then no-op
-   * implementations will be used.
+   * Returns the {@link Tracer} with the provided implementations for {@link SpanFactory}. If no
+   * implementation is provided then no-op implementations will be used.
    *
-   * @return the {@link Tracer}.
+   * @return the {@code Tracer}.
    */
   public static Tracer getTracer() {
     return INSTANCE;
   }
 
+  // Any provider that may be used for SpanFactory can be added here.
+  @VisibleForTesting
+  static SpanFactory loadSpanFactory(ClassLoader classLoader) {
+    try {
+      // Call Class.forName with literal string name of the class to help shading tools.
+      return Provider.createInstance(
+          Class.forName("com.google.instrumentation.trace.SpanFactoryImpl", true, classLoader),
+          SpanFactory.class);
+    } catch (ClassNotFoundException e) {
+      logger.log(Level.FINE, "Using default implementation for SpanFactory.", e);
+    }
+    return new NoopSpanFactory();
+  }
+
   /**
    * Gets the current Span from the current Context.
    *
@@ -92,7 +118,7 @@ public final class Tracer {
    *     from the Context.
    */
   public Span getCurrentSpan() {
-    Span currentSpan = contextSpanHandler.getCurrentSpan();
+    Span currentSpan = ContextUtils.getCurrentSpan();
     return currentSpan != null ? currentSpan : BlankSpan.INSTANCE;
   }
 
@@ -146,7 +172,7 @@ public final class Tracer {
    * @throws NullPointerException if span is null.
    */
   public NonThrowingCloseable withSpan(Span span) {
-    return contextSpanHandler.withSpan(checkNotNull(span, "span"));
+    return ContextUtils.withSpan(checkNotNull(span, "span"));
   }
 
   /**
@@ -163,7 +189,7 @@ public final class Tracer {
    * @throws NullPointerException if name is null.
    */
   public SpanBuilder spanBuilder(String name) {
-    return spanBuilder(contextSpanHandler.getCurrentSpan(), name);
+    return spanBuilder(ContextUtils.getCurrentSpan(), name);
   }
 
   /**
@@ -183,7 +209,6 @@ public final class Tracer {
   public SpanBuilder spanBuilder(@Nullable Span parent, String name) {
     return new SpanBuilder(
         spanFactory,
-        contextSpanHandler,
         parent == null ? null : parent.getContext(),
         /* hasRemoteParent = */ false,
         checkNotNull(name, "name"));
@@ -205,39 +230,7 @@ public final class Tracer {
    */
   public SpanBuilder spanBuilderWithRemoteParent(@Nullable SpanContext remoteParent, String name) {
     return new SpanBuilder(
-        spanFactory,
-        contextSpanHandler,
-        remoteParent,
-        /* hasRemoteParent = */ true,
-        checkNotNull(name, "name"));
-  }
-
-  @VisibleForTesting
-  Tracer(ContextSpanHandler contextSpanHandler, SpanFactory spanFactory) {
-    this.contextSpanHandler = checkNotNull(contextSpanHandler, "contextSpanHandler");
-    this.spanFactory = checkNotNull(spanFactory, "spanFactory");
-  }
-
-  // No-op implementation of the ContextSpanHandler
-  private static final class NoopContextSpanHandler extends ContextSpanHandler {
-    private static final NonThrowingCloseable defaultWithSpan =
-        new NonThrowingCloseable() {
-          @Override
-          public void close() {
-            // Do nothing.
-          }
-        };
-
-    @Override
-    @Nullable
-    public Span getCurrentSpan() {
-      return null;
-    }
-
-    @Override
-    public NonThrowingCloseable withSpan(Span span) {
-      return defaultWithSpan;
-    }
+        spanFactory, remoteParent, /* hasRemoteParent = */ true, checkNotNull(name, "name"));
   }
 
   // No-op implementation of the SpanFactory
@@ -251,35 +244,4 @@ public final class Tracer {
       return BlankSpan.INSTANCE;
     }
   }
-
-  // Any provider that may be used for SpanFactory can be added here.
-  @VisibleForTesting
-  static SpanFactory loadSpanFactory(ClassLoader classLoader) {
-    try {
-      // Because of shading tools we must call Class.forName with the literal string name of the
-      // class.
-      return Provider.createInstance(
-          Class.forName("com.google.instrumentation.trace.SpanFactoryImpl", true, classLoader),
-          SpanFactory.class);
-    } catch (ClassNotFoundException e) {
-      logger.log(Level.FINE, "Using default implementation for SpanFactory.", e);
-    }
-    return new NoopSpanFactory();
-  }
-
-  // Any provider that may be used for ContextSpanHandler can be added here.
-  @VisibleForTesting
-  static ContextSpanHandler loadContextSpanHandler(ClassLoader classLoader) {
-    try {
-      // Because of shading tools we must call Class.forName with the literal string name of the
-      // class.
-      return Provider.createInstance(
-          Class.forName(
-              "com.google.instrumentation.trace.ContextSpanHandlerImpl", true, classLoader),
-          ContextSpanHandler.class);
-    } catch (ClassNotFoundException e) {
-      logger.log(Level.FINE, "Using default implementation for ContextSpanHandler.", e);
-    }
-    return new NoopContextSpanHandler();
-  }
 }
diff --git a/core/src/test/java/com/google/instrumentation/trace/ContextUtilsTest.java b/core/src/test/java/com/google/instrumentation/trace/ContextUtilsTest.java
new file mode 100644
index 00000000..425f6731
--- /dev/null
+++ b/core/src/test/java/com/google/instrumentation/trace/ContextUtilsTest.java
@@ -0,0 +1,87 @@
+/*
+ * Copyright 2017, Google 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 com.google.instrumentation.trace;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.instrumentation.common.NonThrowingCloseable;
+import io.grpc.Context;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/** Unit tests for {@link ContextUtils}. */
+public class ContextUtilsTest {
+  @Mock private Span span;
+
+  @Before
+  public void setUp() {
+    MockitoAnnotations.initMocks(this);
+  }
+
+  @Test
+  public void getCurrentSpan_WhenNoContext() {
+    assertThat(ContextUtils.getCurrentSpan()).isNull();
+  }
+
+  @Test
+  public void getCurrentSpan() {
+    assertThat(ContextUtils.getCurrentSpan()).isNull();
+    Context origContext = Context.current().withValue(ContextUtils.CONTEXT_SPAN_KEY, span).attach();
+    // Make sure context is detached even if test fails.
+    try {
+      assertThat(ContextUtils.getCurrentSpan()).isSameAs(span);
+    } finally {
+      Context.current().detach(origContext);
+    }
+    assertThat(ContextUtils.getCurrentSpan()).isNull();
+  }
+
+  @Test
+  public void withSpan() {
+    assertThat(ContextUtils.getCurrentSpan()).isNull();
+    NonThrowingCloseable ws = ContextUtils.withSpan(span);
+    try {
+      assertThat(ContextUtils.getCurrentSpan()).isSameAs(span);
+    } finally {
+      ws.close();
+    }
+    assertThat(ContextUtils.getCurrentSpan()).isNull();
+    ;
+  }
+
+  @Test
+  public void propagationViaRunnable() {
+    Runnable runnable = null;
+    NonThrowingCloseable ws = ContextUtils.withSpan(span);
+    try {
+      assertThat(ContextUtils.getCurrentSpan()).isSameAs(span);
+      runnable =
+          Context.current()
+              .wrap(
+                  new Runnable() {
+                    @Override
+                    public void run() {
+                      assertThat(ContextUtils.getCurrentSpan()).isSameAs(span);
+                    }
+                  });
+    } finally {
+      ws.close();
+    }
+    assertThat(ContextUtils.getCurrentSpan()).isNotSameAs(span);
+    // When we run the runnable we will have the span in the current Context.
+    runnable.run();
+  }
+}
diff --git a/core/src/test/java/com/google/instrumentation/trace/ScopedSpanHandleTest.java b/core/src/test/java/com/google/instrumentation/trace/ScopedSpanHandleTest.java
index 971904b1..ec17428b 100644
--- a/core/src/test/java/com/google/instrumentation/trace/ScopedSpanHandleTest.java
+++ b/core/src/test/java/com/google/instrumentation/trace/ScopedSpanHandleTest.java
@@ -13,9 +13,9 @@
 
 package com.google.instrumentation.trace;
 
+import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Matchers.same;
 import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
 
 import com.google.instrumentation.common.NonThrowingCloseable;
 import org.junit.Before;
@@ -28,12 +28,9 @@ import org.mockito.MockitoAnnotations;
 /** Unit tests for {@link ScopedSpanHandle}. */
 @RunWith(JUnit4.class)
 public class ScopedSpanHandleTest {
+  private static final Tracer tracer = Tracer.getTracer();
   @Mock private Span span;
 
-  @Mock private ContextSpanHandler csh;
-
-  @Mock private NonThrowingCloseable ntc;
-
   @Before
   public void setUp() {
     MockitoAnnotations.initMocks(this);
@@ -41,11 +38,14 @@ public class ScopedSpanHandleTest {
 
   @Test
   public void initAttachesSpan_CloseDetachesAndEndsSpan() {
-    when(csh.withSpan(same(span))).thenReturn(ntc);
-    try (NonThrowingCloseable ss = new ScopedSpanHandle(span, csh)) {
-      // Do nothing.
+    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+    NonThrowingCloseable ss = new ScopedSpanHandle(span);
+    try {
+      assertThat(tracer.getCurrentSpan()).isSameAs(span);
+    } finally {
+      ss.close();
     }
-    verify(ntc).close();
+    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
     verify(span).end(same(EndSpanOptions.DEFAULT));
   }
 }
diff --git a/core/src/test/java/com/google/instrumentation/trace/SpanBuilderTest.java b/core/src/test/java/com/google/instrumentation/trace/SpanBuilderTest.java
index 657f3fe7..32a977cc 100644
--- a/core/src/test/java/com/google/instrumentation/trace/SpanBuilderTest.java
+++ b/core/src/test/java/com/google/instrumentation/trace/SpanBuilderTest.java
@@ -35,12 +35,10 @@ import org.mockito.MockitoAnnotations;
 /** Unit tests for {@link Tracer}. */
 @RunWith(JUnit4.class)
 public class SpanBuilderTest {
-  @Mock private ContextSpanHandler contextSpanHandler;
+  private static final String SPAN_NAME = "MySpanName";
+  private static final Tracer tracer = Tracer.getTracer();
   @Mock private Span span;
   @Mock private SpanFactory spanFactory;
-  @Mock private NonThrowingCloseable withSpan;
-
-  private static final String SPAN_NAME = "MySpanName";
   private Random random;
   private SpanContext spanContext;
   private SpanBuilder spanBuilder;
@@ -54,9 +52,7 @@ public class SpanBuilderTest {
             TraceId.generateRandomId(random),
             SpanId.generateRandomId(random),
             TraceOptions.DEFAULT);
-    spanBuilder =
-        new SpanBuilder(
-            spanFactory, contextSpanHandler, spanContext, false /* hasRemoteParent */, SPAN_NAME);
+    spanBuilder = new SpanBuilder(spanFactory, spanContext, false /* hasRemoteParent */, SPAN_NAME);
   }
 
   @Test
@@ -64,9 +60,12 @@ public class SpanBuilderTest {
     when(spanFactory.startSpan(
             isNull(SpanContext.class), eq(false), same(SPAN_NAME), eq(new StartSpanOptions())))
         .thenReturn(span);
-    when(contextSpanHandler.withSpan(same(span))).thenReturn(withSpan);
-    try (NonThrowingCloseable ss = spanBuilder.becomeRoot().startScopedSpan()) {}
-    verify(withSpan).close();
+    NonThrowingCloseable ss = spanBuilder.becomeRoot().startScopedSpan();
+    try {
+      assertThat(tracer.getCurrentSpan()).isSameAs(span);
+    } finally {
+      ss.close();
+    }
     verify(span).end(same(EndSpanOptions.DEFAULT));
   }
 
@@ -78,14 +77,17 @@ public class SpanBuilderTest {
     when(spanFactory.startSpan(
             isNull(SpanContext.class), eq(false), same(SPAN_NAME), eq(startSpanOptions)))
         .thenReturn(span);
-    when(contextSpanHandler.withSpan(same(span))).thenReturn(withSpan);
-    try (NonThrowingCloseable ss =
+    NonThrowingCloseable ss =
         spanBuilder
             .becomeRoot()
             .setSampler(Samplers.neverSample())
             .setStartTime(Timestamp.fromMillis(1234567L))
-            .startScopedSpan()) {}
-    verify(withSpan).close();
+            .startScopedSpan();
+    try {
+      assertThat(tracer.getCurrentSpan()).isSameAs(span);
+    } finally {
+      ss.close();
+    }
     verify(span).end(same(EndSpanOptions.DEFAULT));
   }
 
@@ -102,9 +104,7 @@ public class SpanBuilderTest {
 
   @Test
   public void startSpan_WithNullParent() {
-    spanBuilder =
-        new SpanBuilder(
-            spanFactory, contextSpanHandler, null, false /* hasRemoteParent */, SPAN_NAME);
+    spanBuilder = new SpanBuilder(spanFactory, null, false /* hasRemoteParent */, SPAN_NAME);
     when(spanFactory.startSpan(
             isNull(SpanContext.class), eq(false), same(SPAN_NAME), eq(new StartSpanOptions())))
         .thenReturn(span);
@@ -161,9 +161,7 @@ public class SpanBuilderTest {
 
   @Test
   public void startSpanWitRemoteParent() {
-    spanBuilder =
-        new SpanBuilder(
-            spanFactory, contextSpanHandler, spanContext, true /* hasRemoteParent */, SPAN_NAME);
+    spanBuilder = new SpanBuilder(spanFactory, spanContext, true /* hasRemoteParent */, SPAN_NAME);
     when(spanFactory.startSpan(
             same(spanContext), eq(true), same(SPAN_NAME), eq(new StartSpanOptions())))
         .thenReturn(span);
@@ -175,9 +173,7 @@ public class SpanBuilderTest {
 
   @Test
   public void startSpanWitRemoteParent_WithNullParent() {
-    spanBuilder =
-        new SpanBuilder(
-            spanFactory, contextSpanHandler, null, true /* hasRemoteParent */, SPAN_NAME);
+    spanBuilder = new SpanBuilder(spanFactory, null, true /* hasRemoteParent */, SPAN_NAME);
     when(spanFactory.startSpan(
             isNull(SpanContext.class), eq(true), same(SPAN_NAME), eq(new StartSpanOptions())))
         .thenReturn(span);
diff --git a/core/src/test/java/com/google/instrumentation/trace/TracerTest.java b/core/src/test/java/com/google/instrumentation/trace/TracerTest.java
index cc60fe5c..69661177 100644
--- a/core/src/test/java/com/google/instrumentation/trace/TracerTest.java
+++ b/core/src/test/java/com/google/instrumentation/trace/TracerTest.java
@@ -21,7 +21,9 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import com.google.instrumentation.common.NonThrowingCloseable;
+import io.grpc.Context;
 import java.util.Random;
+import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
@@ -33,14 +35,15 @@ import org.mockito.MockitoAnnotations;
 /** Unit tests for {@link Tracer}. */
 @RunWith(JUnit4.class)
 public class TracerTest {
-  @Rule public ExpectedException thrown = ExpectedException.none();
-
   private static final Tracer tracer = Tracer.getTracer();
-
-  @Mock private ContextSpanHandler contextSpanHandler;
+  @Rule public ExpectedException thrown = ExpectedException.none();
   @Mock private SpanFactory spanFactory;
   @Mock private Span span;
-  @Mock private NonThrowingCloseable withSpan;
+
+  @Before
+  public void setUp() {
+    MockitoAnnotations.initMocks(this);
+  }
 
   @Test
   public void defaultGetCurrentSpan() {
@@ -49,14 +52,43 @@ public class TracerTest {
 
   @Test(expected = NullPointerException.class)
   public void withSpan_NullSpan() {
-    try (NonThrowingCloseable ws = tracer.withSpan(null)) {}
+    tracer.withSpan(null);
   }
 
   @Test
-  public void defaultWithSpan() {
-    try (NonThrowingCloseable ss = tracer.withSpan(BlankSpan.INSTANCE)) {
-      assertThat(tracer.getCurrentSpan()).isEqualTo(BlankSpan.INSTANCE);
+  public void getCurrentSpan_WithSpan() {
+    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+    NonThrowingCloseable ws = tracer.withSpan(span);
+    try {
+      assertThat(tracer.getCurrentSpan()).isSameAs(span);
+    } finally {
+      ws.close();
     }
+    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+  }
+
+  @Test
+  public void propagationViaRunnable() {
+    Runnable runnable = null;
+    NonThrowingCloseable ws = tracer.withSpan(span);
+    try {
+      assertThat(tracer.getCurrentSpan()).isSameAs(span);
+      runnable =
+          Context.current()
+              .wrap(
+                  new Runnable() {
+                    @Override
+                    public void run() {
+                      assertThat(tracer.getCurrentSpan()).isSameAs(span);
+                    }
+                  });
+    } finally {
+      ws.close();
+    }
+    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+    // When we run the runnable we will have the span in the current Context.
+    runnable.run();
+    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
   }
 
   @Test(expected = NullPointerException.class)
@@ -92,35 +124,6 @@ public class TracerTest {
   }
 
   @Test
-  public void loadContextSpanHandler_UsesProvidedClassLoader() {
-    final RuntimeException toThrow = new RuntimeException("UseClassLoader");
-    thrown.expect(RuntimeException.class);
-    thrown.expectMessage("UseClassLoader");
-    Tracer.loadContextSpanHandler(
-        new ClassLoader() {
-          @Override
-          public Class<?> loadClass(String name) {
-            throw toThrow;
-          }
-        });
-  }
-
-  @Test
-  public void loadContextSpanHandler_IgnoresMissingClasses() {
-    assertThat(
-            Tracer.loadContextSpanHandler(
-                    new ClassLoader() {
-                      @Override
-                      public Class<?> loadClass(String name) throws ClassNotFoundException {
-                        throw new ClassNotFoundException();
-                      }
-                    })
-                .getClass()
-                .getName())
-        .isEqualTo("com.google.instrumentation.trace.Tracer$NoopContextSpanHandler");
-  }
-
-  @Test
   public void loadSpanFactory_UsesProvidedClassLoader() {
     final RuntimeException toThrow = new RuntimeException("UseClassLoader");
     thrown.expect(RuntimeException.class);
@@ -150,47 +153,39 @@ public class TracerTest {
   }
 
   @Test
-  public void getCurrentSpan() {
-    Tracer mockTracer = newTracerWithMocks();
-    when(contextSpanHandler.getCurrentSpan()).thenReturn(span);
-    assertThat(mockTracer.getCurrentSpan()).isEqualTo(span);
-  }
-
-  @Test
-  public void withSpan() {
-    Tracer mockTracer = newTracerWithMocks();
-    when(contextSpanHandler.withSpan(same(span))).thenReturn(withSpan);
-    try (NonThrowingCloseable ss = mockTracer.withSpan(span)) {}
-    verify(withSpan).close();
-  }
-
-  @Test
   public void startScopedSpanRoot() {
     Tracer mockTracer = newTracerWithMocks();
-    when(contextSpanHandler.getCurrentSpan()).thenReturn(null);
     when(spanFactory.startSpan(
             isNull(SpanContext.class), eq(false), eq("MySpanName"), eq(new StartSpanOptions())))
         .thenReturn(span);
-    when(contextSpanHandler.withSpan(same(span))).thenReturn(withSpan);
-    try (NonThrowingCloseable ss =
-        mockTracer.spanBuilder("MySpanName").becomeRoot().startScopedSpan()) {}
-    verify(withSpan).close();
+    NonThrowingCloseable ss = mockTracer.spanBuilder("MySpanName").becomeRoot().startScopedSpan();
+    try {
+      assertThat(tracer.getCurrentSpan()).isSameAs(span);
+    } finally {
+      ss.close();
+    }
     verify(span).end(same(EndSpanOptions.DEFAULT));
   }
 
   @Test
   public void startScopedSpanChild() {
     Tracer mockTracer = newTracerWithMocks();
-    when(contextSpanHandler.getCurrentSpan()).thenReturn(BlankSpan.INSTANCE);
-    when(spanFactory.startSpan(
-            same(BlankSpan.INSTANCE.getContext()),
-            eq(false),
-            eq("MySpanName"),
-            eq(new StartSpanOptions())))
-        .thenReturn(span);
-    when(contextSpanHandler.withSpan(same(span))).thenReturn(withSpan);
-    try (NonThrowingCloseable ss = mockTracer.spanBuilder("MySpanName").startScopedSpan()) {}
-    verify(withSpan).close();
+    NonThrowingCloseable ws = mockTracer.withSpan(BlankSpan.INSTANCE);
+    try {
+      assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+      when(spanFactory.startSpan(
+              same(SpanContext.INVALID), eq(false), eq("MySpanName"), eq(new StartSpanOptions())))
+          .thenReturn(span);
+      NonThrowingCloseable ss = mockTracer.spanBuilder("MySpanName").startScopedSpan();
+      try {
+        assertThat(tracer.getCurrentSpan()).isSameAs(span);
+      } finally {
+        ss.close();
+      }
+      assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+    } finally {
+      ws.close();
+    }
     verify(span).end(same(EndSpanOptions.DEFAULT));
   }
 
@@ -242,7 +237,6 @@ public class TracerTest {
   }
 
   private Tracer newTracerWithMocks() {
-    MockitoAnnotations.initMocks(this);
-    return new Tracer(contextSpanHandler, spanFactory);
+    return new Tracer(spanFactory);
   }
 }