Add methods to wrap Runnable and Callbacks and to run them. (#778)

* Add methods to wrap Runnable and Callbacks and to run them.

* Fix more comments.

* Fix use of Assert.fail().

* Improve javadocs for the start and wrap methods.

* Remove startSpanAndWrap API for the moment.

* Change wrap to withSpan.

* Fix examples in SpanBuilders.

6 files changed, 620 insertions, 38 deletions

diff --git a/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java b/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java
index 713c8816..539339bc 100644
--- a/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java
+++ b/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java
@@ -16,9 +16,11 @@
 
 package io.opencensus.trace;
 
+import com.google.common.base.Throwables;
 import io.grpc.Context;
 import io.opencensus.common.Scope;
 import io.opencensus.trace.unsafe.ContextUtils;
+import java.util.concurrent.Callable;
 
 /** Util methods/functionality to interact with the {@link Span} in the {@link io.grpc.Context}. */
 final class CurrentSpanUtils {
@@ -49,6 +51,30 @@ final class CurrentSpanUtils {
     return new ScopeInSpan(span, endSpan);
   }
 
+  /**
+   * Wraps a {@link Runnable} so that it executes with the {@code span} as the current {@code Span}.
+   *
+   * @param span the {@code Span} to be set as current.
+   * @param endSpan if {@code true} the returned {@code Runnable} will close the {@code Span}.
+   * @param runnable the {@code Runnable} to run in the {@code Span}.
+   * @return the wrapped {@code Runnable}.
+   */
+  static Runnable withSpan(Span span, boolean endSpan, Runnable runnable) {
+    return new RunnableInSpan(span, runnable, endSpan);
+  }
+
+  /**
+   * Wraps a {@link Callable} so that it executes with the {@code span} as the current {@code Span}.
+   *
+   * @param span the {@code Span} to be set as current.
+   * @param endSpan if {@code true} the returned {@code Runnable} will close the {@code Span}.
+   * @param callable the {@code Callable} to run in the {@code Span}.
+   * @return the wrapped {@code Callable}.
+   */
+  static <C> Callable<C> withSpan(Span span, boolean endSpan, Callable<C> callable) {
+    return new CallableInSpan<C>(span, callable, endSpan);
+  }
+
   // Defines an arbitrary scope of code as a traceable operation. Supports try-with-resources idiom.
   private static final class ScopeInSpan implements Scope {
     private final Context origContext;
@@ -60,7 +86,7 @@ final class CurrentSpanUtils {
      *
      * @param span is the {@code Span} to be added to the current {@code io.grpc.Context}.
      */
-    ScopeInSpan(Span span, boolean endSpan) {
+    private ScopeInSpan(Span span, boolean endSpan) {
       this.span = span;
       this.endSpan = endSpan;
       origContext = Context.current().withValue(ContextUtils.CONTEXT_SPAN_KEY, span).attach();
@@ -74,4 +100,74 @@ final class CurrentSpanUtils {
       }
     }
   }
+
+  private static final class RunnableInSpan implements Runnable {
+    // TODO(bdrutu): Investigate if the extra private visibility increases the generated bytecode.
+    private final Span span;
+    private final Runnable runnable;
+    private final boolean endSpan;
+
+    private RunnableInSpan(Span span, Runnable runnable, boolean endSpan) {
+      this.span = span;
+      this.runnable = runnable;
+      this.endSpan = endSpan;
+    }
+
+    @Override
+    public void run() {
+      Context origContext =
+          Context.current().withValue(ContextUtils.CONTEXT_SPAN_KEY, span).attach();
+      try {
+        runnable.run();
+      } catch (Throwable t) {
+        setErrorStatus(span, t);
+        Throwables.propagateIfPossible(t);
+        throw new RuntimeException("unexpected", t);
+      } finally {
+        Context.current().detach(origContext);
+        if (endSpan) {
+          span.end();
+        }
+      }
+    }
+  }
+
+  private static final class CallableInSpan<V> implements Callable<V> {
+    private final Span span;
+    private final Callable<V> callable;
+    private final boolean endSpan;
+
+    private CallableInSpan(Span span, Callable<V> callable, boolean endSpan) {
+      this.span = span;
+      this.callable = callable;
+      this.endSpan = endSpan;
+    }
+
+    @Override
+    public V call() throws Exception {
+      Context origContext =
+          Context.current().withValue(ContextUtils.CONTEXT_SPAN_KEY, span).attach();
+      try {
+        return callable.call();
+      } catch (Exception e) {
+        setErrorStatus(span, e);
+        throw e;
+      } catch (Throwable t) {
+        setErrorStatus(span, t);
+        Throwables.propagateIfPossible(t);
+        throw new RuntimeException("unexpected", t);
+      } finally {
+        Context.current().detach(origContext);
+        if (endSpan) {
+          span.end();
+        }
+      }
+    }
+  }
+
+  private static void setErrorStatus(Span span, Throwable t) {
+    span.setStatus(
+        Status.UNKNOWN.withDescription(
+            t.getMessage() == null ? t.getClass().getSimpleName() : t.getMessage()));
+  }
 }
diff --git a/api/src/main/java/io/opencensus/trace/SpanBuilder.java b/api/src/main/java/io/opencensus/trace/SpanBuilder.java
index 4c28147e..0a7c9177 100644
--- a/api/src/main/java/io/opencensus/trace/SpanBuilder.java
+++ b/api/src/main/java/io/opencensus/trace/SpanBuilder.java
@@ -21,6 +21,7 @@ import static com.google.common.base.Preconditions.checkNotNull;
 import com.google.errorprone.annotations.MustBeClosed;
 import io.opencensus.common.Scope;
 import java.util.List;
+import java.util.concurrent.Callable;
 import javax.annotation.Nullable;
 
 /**
@@ -223,6 +224,68 @@ public abstract class SpanBuilder {
     return CurrentSpanUtils.withSpan(startSpan(), true);
   }
 
+  /**
+   * Starts a new span and runs the given {@code Runnable} with the newly created {@code Span} as
+   * the current {@code Span}, and ends the {@code Span} after the {@code Runnable} is run.
+   *
+   * <p>Any error will end up as a {@link Status#UNKNOWN}.
+   *
+   * <pre><code>
+   * tracer.spanBuilder("MyRunnableSpan").startSpanAndRun(myRunnable);
+   * </code></pre>
+   *
+   * <p>It is equivalent with the following code:
+   *
+   * <pre><code>
+   * Span span = tracer.spanBuilder("MyRunnableSpan").startSpan();
+   * Runnable newRunnable = tracer.withSpan(span, myRunnable);
+   * try {
+   *   newRunnable.run();
+   * } finally {
+   *   span.end();
+   * }
+   * </code></pre>
+   *
+   * @param runnable the {@code Runnable} to run in the {@code Span}.
+   * @since 0.11.0
+   */
+  public final void startSpanAndRun(final Runnable runnable) {
+    final Span span = startSpan();
+    CurrentSpanUtils.withSpan(span, true, runnable).run();
+  }
+
+  /**
+   * Starts a new span and calls the given {@code Callable} with the newly created {@code Span} as
+   * the current {@code Span}, and ends the {@code Span} after the {@code Callable} is called.
+   *
+   * <p>Any error will end up as a {@link Status#UNKNOWN}.
+   *
+   * <pre><code>
+   * MyResult myResult = tracer.spanBuilder("MyCallableSpan").startSpanAndCall(myCallable);
+   * </code></pre>
+   *
+   * <p>It is equivalent with the following code:
+   *
+   * <pre><code>
+   * Span span = tracer.spanBuilder("MyCallableSpan").startSpan();
+   * {@code Callable<MyResult>} newCallable = tracer.withSpan(span, myCallable);
+   * MyResult myResult = null;
+   * try {
+   *   myResult = newCallable.call();
+   * } finally {
+   *   span.end();
+   * }
+   * );
+   * </code></pre>
+   *
+   * @param callable the {@code Callable} to run in the {@code Span}.
+   * @since 0.11.0
+   */
+  public final <V> V startSpanAndCall(Callable<V> callable) throws Exception {
+    final Span span = startSpan();
+    return CurrentSpanUtils.withSpan(span, true, callable).call();
+  }
+
   static final class NoopSpanBuilder extends SpanBuilder {
     static NoopSpanBuilder createWithParent(String spanName, @Nullable Span parent) {
       return new NoopSpanBuilder(spanName);
diff --git a/api/src/main/java/io/opencensus/trace/Tracer.java b/api/src/main/java/io/opencensus/trace/Tracer.java
index 8215f7b1..cd4fc822 100644
--- a/api/src/main/java/io/opencensus/trace/Tracer.java
+++ b/api/src/main/java/io/opencensus/trace/Tracer.java
@@ -21,6 +21,7 @@ import static com.google.common.base.Preconditions.checkNotNull;
 import com.google.errorprone.annotations.MustBeClosed;
 import io.opencensus.common.Scope;
 import io.opencensus.trace.SpanBuilder.NoopSpanBuilder;
+import java.util.concurrent.Callable;
 import javax.annotation.Nullable;
 
 /**
@@ -151,6 +152,136 @@ public abstract class Tracer {
   }
 
   /**
+   * Returns a {@link Runnable} that runs the given task with the given {@code Span} in the current
+   * context.
+   *
+   * <p>Users may consider to use {@link SpanBuilder#startSpanAndRun(Runnable)}.
+   *
+   * <p>Any error will end up as a {@link Status#UNKNOWN}.
+   *
+   * <p>IMPORTANT: Caller must manually propagate the entire {@code io.grpc.Context} when wraps a
+   * {@code Runnable}, see the examples.
+   *
+   * <p>IMPORTANT: Caller must manually end the {@code Span} within the {@code Runnable}, or after
+   * the {@code Runnable} is executed.
+   *
+   * <p>Example with Executor wrapped with {@link io.grpc.Context#currentContextExecutor}:
+   *
+   * <pre><code>
+   * class MyClass {
+   *   private static Tracer tracer = Tracing.getTracer();
+   *   void handleRequest(Executor executor) {
+   *     Span span = new Tracer.spanBuilder("MyRunnableSpan");
+   *     executor.execute(tracer.withSpan(span, new Runnable() {
+   *      {@literal @}Override
+   *       public void run() {
+   *         try {
+   *           sendResult();
+   *         } finally {
+   *           span.end();
+   *         }
+   *       }
+   *     }));
+   *   }
+   * }
+   * </code></pre>
+   *
+   * <p>Example without Executor wrapped with {@link io.grpc.Context#currentContextExecutor}:
+   *
+   * <pre><code>
+   * class MyClass {
+   *   private static Tracer tracer = Tracing.getTracer();
+   *   void handleRequest(Executor executor) {
+   *     Span span = new Tracer.spanBuilder("MyRunnableSpan");
+   *     executor.execute(Context.wrap(tracer.withSpan(span, new Runnable() {
+   *      {@literal @}Override
+   *       public void run() {
+   *         try {
+   *           sendResult();
+   *         } finally {
+   *           span.end();
+   *         }
+   *       }
+   *     })));
+   *   }
+   * }
+   * </code></pre>
+   *
+   * @param span the {@code Span} to be set as current.
+   * @param runnable the {@code Runnable} to withSpan in the {@code Span}.
+   * @return the {@code Runnable}.
+   * @since 0.11.0
+   */
+  public final Runnable withSpan(Span span, Runnable runnable) {
+    return CurrentSpanUtils.withSpan(span, false, runnable);
+  }
+
+  /**
+   * Returns a {@link Callable} that runs the given task with the given {@code Span} in the current
+   * context.
+   *
+   * <p>Users may consider to use {@link SpanBuilder#startSpanAndCall(Callable)}.
+   *
+   * <p>Any error will end up as a {@link Status#UNKNOWN}.
+   *
+   * <p>IMPORTANT: Caller must manually propagate the entire {@code io.grpc.Context} when wraps a
+   * {@code Callable}, see the examples.
+   *
+   * <p>IMPORTANT: Caller must manually end the {@code Span} within the {@code Callable}, or after
+   * the {@code Callable} is executed.
+   *
+   * <p>Example with Executor wrapped with {@link io.grpc.Context#currentContextExecutor}:
+   *
+   * <pre><code>
+   * class MyClass {
+   *   private static Tracer tracer = Tracing.getTracer();
+   *   void handleRequest(Executor executor) {
+   *     Span span = new Tracer.spanBuilder("MyRunnableSpan");
+   *     executor.execute(tracer.withSpan(span, {@code new Callable<MyResult>()} {
+   *      {@literal @}Override
+   *       public MyResult call() throws Exception {
+   *         try {
+   *           return sendResult();
+   *         } finally {
+   *           span.end();
+   *         }
+   *       }
+   *     }));
+   *   }
+   * }
+   * </code></pre>
+   *
+   * <p>Example without Executor wrapped with {@link io.grpc.Context#currentContextExecutor}:
+   *
+   * <pre><code>
+   * class MyClass {
+   *   private static Tracer tracer = Tracing.getTracer();
+   *   void handleRequest(Executor executor) {
+   *     Span span = new Tracer.spanBuilder("MyRunnableSpan");
+   *     executor.execute(Context.wrap(tracer.withSpan(span, {@code new Callable<MyResult>()} {
+   *      {@literal @}Override
+   *       public MyResult call() throws Exception {
+   *         try {
+   *           return sendResult();
+   *         } finally {
+   *           span.end();
+   *         }
+   *       }
+   *     })));
+   *   }
+   * }
+   * </code></pre>
+   *
+   * @param span the {@code Span} to be set as current.
+   * @param callable the {@code Callable} to run in the {@code Span}.
+   * @return the {@code Callable}.
+   * @since 0.11.0
+   */
+  public final <C> Callable<C> withSpan(Span span, final Callable<C> callable) {
+    return CurrentSpanUtils.withSpan(span, false, callable);
+  }
+
+  /**
    * Returns a {@link SpanBuilder} to create and start a new child {@link Span} as a child of to the
    * current {@code Span} if any, otherwise creates a root {@code Span}.
    *
diff --git a/api/src/test/java/io/opencensus/trace/CurrentSpanUtilsTest.java b/api/src/test/java/io/opencensus/trace/CurrentSpanUtilsTest.java
index 8ba910b0..6b16c3d0 100644
--- a/api/src/test/java/io/opencensus/trace/CurrentSpanUtilsTest.java
+++ b/api/src/test/java/io/opencensus/trace/CurrentSpanUtilsTest.java
@@ -24,6 +24,7 @@ import static org.mockito.Mockito.verifyZeroInteractions;
 import io.grpc.Context;
 import io.opencensus.common.Scope;
 import io.opencensus.trace.unsafe.ContextUtils;
+import java.util.concurrent.Callable;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -41,6 +42,30 @@ public class CurrentSpanUtilsTest {
     MockitoAnnotations.initMocks(this);
   }
 
+  // TODO(bdrutu): When update to junit 4.13 use assertThrows instead of this.
+  private void executeRunnableAndExpectError(Runnable runnable, Throwable error) {
+    boolean called = false;
+    try {
+      CurrentSpanUtils.withSpan(span, true, runnable).run();
+    } catch (Throwable e) {
+      assertThat(e).isEqualTo(error);
+      called = true;
+    }
+    assertThat(called).isTrue();
+  }
+
+  // TODO(bdrutu): When update to junit 4.13 use assertThrows instead of this.
+  private void executeCallableAndExpectError(Callable<Object> callable, Throwable error) {
+    boolean called = false;
+    try {
+      CurrentSpanUtils.withSpan(span, true, callable).call();
+    } catch (Throwable e) {
+      assertThat(e).isEqualTo(error);
+      called = true;
+    }
+    assertThat(called).isTrue();
+  }
+
   @Test
   public void getCurrentSpan_WhenNoContext() {
     assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
@@ -86,25 +111,184 @@ public class CurrentSpanUtilsTest {
   }
 
   @Test
-  public void propagationViaRunnable() {
-    Runnable runnable = null;
-    Scope ws = CurrentSpanUtils.withSpan(span, false);
-    try {
-      assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
-      runnable =
-          Context.current()
-              .wrap(
-                  new Runnable() {
-                    @Override
-                    public void run() {
-                      assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
-                    }
-                  });
-    } finally {
-      ws.close();
-    }
-    assertThat(CurrentSpanUtils.getCurrentSpan()).isNotSameAs(span);
-    // When we run the runnable we will have the span in the current Context.
-    runnable.run();
+  public void withSpanRunnable() {
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+    Runnable runnable =
+        new Runnable() {
+          @Override
+          public void run() {
+            // When we run the runnable we will have the span in the current Context.
+            assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
+          }
+        };
+    CurrentSpanUtils.withSpan(span, false, runnable).run();
+    verifyZeroInteractions(span);
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+  }
+
+  @Test
+  public void withSpanRunnable_EndSpan() {
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+    Runnable runnable =
+        new Runnable() {
+          @Override
+          public void run() {
+            // When we run the runnable we will have the span in the current Context.
+            assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
+          }
+        };
+    CurrentSpanUtils.withSpan(span, true, runnable).run();
+    verify(span).end(EndSpanOptions.DEFAULT);
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+  }
+
+  @Test
+  public void withSpanRunnable_WithError() {
+    final AssertionError error = new AssertionError("MyError");
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+    Runnable runnable =
+        new Runnable() {
+          @Override
+          public void run() {
+            // When we run the runnable we will have the span in the current Context.
+            assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
+            throw error;
+          }
+        };
+    executeRunnableAndExpectError(runnable, error);
+    verify(span).setStatus(Status.UNKNOWN.withDescription("MyError"));
+    verify(span).end(EndSpanOptions.DEFAULT);
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+  }
+
+  @Test
+  public void withSpanRunnable_WithErrorNoMessage() {
+    final AssertionError error = new AssertionError();
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+    Runnable runnable =
+        new Runnable() {
+          @Override
+          public void run() {
+            // When we run the runnable we will have the span in the current Context.
+            assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
+            throw error;
+          }
+        };
+    executeRunnableAndExpectError(runnable, error);
+    verify(span).setStatus(Status.UNKNOWN.withDescription("AssertionError"));
+    verify(span).end(EndSpanOptions.DEFAULT);
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+  }
+
+  @Test
+  public void withSpanCallable() throws Exception {
+    final Object ret = new Object();
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+    Callable<Object> callable =
+        new Callable<Object>() {
+          @Override
+          public Object call() throws Exception {
+            // When we run the runnable we will have the span in the current Context.
+            assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
+            return ret;
+          }
+        };
+    assertThat(CurrentSpanUtils.withSpan(span, false, callable).call()).isEqualTo(ret);
+    verifyZeroInteractions(span);
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+  }
+
+  @Test
+  public void withSpanCallable_EndSpan() throws Exception {
+    final Object ret = new Object();
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+    Callable<Object> callable =
+        new Callable<Object>() {
+          @Override
+          public Object call() throws Exception {
+            // When we run the runnable we will have the span in the current Context.
+            assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
+            return ret;
+          }
+        };
+    assertThat(CurrentSpanUtils.withSpan(span, true, callable).call()).isEqualTo(ret);
+    verify(span).end(EndSpanOptions.DEFAULT);
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+  }
+
+  @Test
+  public void withSpanCallable_WithException() {
+    final Exception exception = new Exception("MyException");
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+    Callable<Object> callable =
+        new Callable<Object>() {
+          @Override
+          public Object call() throws Exception {
+            // When we run the runnable we will have the span in the current Context.
+            assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
+            throw exception;
+          }
+        };
+    executeCallableAndExpectError(callable, exception);
+    verify(span).setStatus(Status.UNKNOWN.withDescription("MyException"));
+    verify(span).end(EndSpanOptions.DEFAULT);
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+  }
+
+  @Test
+  public void withSpanCallable_WithExceptionNoMessage() {
+    final Exception exception = new Exception();
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+    Callable<Object> callable =
+        new Callable<Object>() {
+          @Override
+          public Object call() throws Exception {
+            // When we run the runnable we will have the span in the current Context.
+            assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
+            throw exception;
+          }
+        };
+    executeCallableAndExpectError(callable, exception);
+    verify(span).setStatus(Status.UNKNOWN.withDescription("Exception"));
+    verify(span).end(EndSpanOptions.DEFAULT);
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+  }
+
+  @Test
+  public void withSpanCallable_WithError() {
+    final AssertionError error = new AssertionError("MyError");
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+    Callable<Object> callable =
+        new Callable<Object>() {
+          @Override
+          public Object call() throws Exception {
+            // When we run the runnable we will have the span in the current Context.
+            assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
+            throw error;
+          }
+        };
+    executeCallableAndExpectError(callable, error);
+    verify(span).setStatus(Status.UNKNOWN.withDescription("MyError"));
+    verify(span).end(EndSpanOptions.DEFAULT);
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+  }
+
+  @Test
+  public void withSpanCallable_WithErrorNoMessage() {
+    final AssertionError error = new AssertionError();
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+    Callable<Object> callable =
+        new Callable<Object>() {
+          @Override
+          public Object call() throws Exception {
+            // When we run the runnable we will have the span in the current Context.
+            assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
+            throw error;
+          }
+        };
+    executeCallableAndExpectError(callable, error);
+    verify(span).setStatus(Status.UNKNOWN.withDescription("AssertionError"));
+    verify(span).end(EndSpanOptions.DEFAULT);
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
   }
 }
diff --git a/api/src/test/java/io/opencensus/trace/SpanBuilderTest.java b/api/src/test/java/io/opencensus/trace/SpanBuilderTest.java
new file mode 100644
index 00000000..1403db15
--- /dev/null
+++ b/api/src/test/java/io/opencensus/trace/SpanBuilderTest.java
@@ -0,0 +1,91 @@
+/*
+ * Copyright 2017, OpenCensus 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 io.opencensus.trace;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import io.opencensus.common.Scope;
+import java.util.concurrent.Callable;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/** Unit tests for {@link SpanBuilder}. */
+@RunWith(JUnit4.class)
+// Need to suppress warnings for MustBeClosed because Java-6 does not support try-with-resources.
+@SuppressWarnings("MustBeClosedChecker")
+public class SpanBuilderTest {
+  private Tracer tracer = Tracing.getTracer();
+  @Mock private SpanBuilder spanBuilder;
+  @Mock private Span span;
+
+  @Before
+  public void setUp() {
+    MockitoAnnotations.initMocks(this);
+    when(spanBuilder.startSpan()).thenReturn(span);
+  }
+
+  @Test
+  public void startScopedSpan() {
+    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+    Scope scope = spanBuilder.startScopedSpan();
+    try {
+      assertThat(tracer.getCurrentSpan()).isSameAs(span);
+    } finally {
+      scope.close();
+    }
+    verify(span).end(EndSpanOptions.DEFAULT);
+    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+  }
+
+  @Test
+  public void startSpanAndRun() {
+    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+    spanBuilder.startSpanAndRun(
+        new Runnable() {
+          @Override
+          public void run() {
+            assertThat(tracer.getCurrentSpan()).isSameAs(span);
+          }
+        });
+    verify(span).end(EndSpanOptions.DEFAULT);
+    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+  }
+
+  @Test
+  public void startSpanAndCall() throws Exception {
+    final Object ret = new Object();
+    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+    assertThat(
+            spanBuilder.startSpanAndCall(
+                new Callable<Object>() {
+                  @Override
+                  public Object call() throws Exception {
+                    assertThat(tracer.getCurrentSpan()).isSameAs(span);
+                    return ret;
+                  }
+                }))
+        .isEqualTo(ret);
+    verify(span).end(EndSpanOptions.DEFAULT);
+    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+  }
+}
diff --git a/api/src/test/java/io/opencensus/trace/TracerTest.java b/api/src/test/java/io/opencensus/trace/TracerTest.java
index 8f0ac64f..58dd4bbe 100644
--- a/api/src/test/java/io/opencensus/trace/TracerTest.java
+++ b/api/src/test/java/io/opencensus/trace/TracerTest.java
@@ -18,10 +18,11 @@ package io.opencensus.trace;
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Matchers.same;
+import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
-import io.grpc.Context;
 import io.opencensus.common.Scope;
+import java.util.concurrent.Callable;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -68,26 +69,42 @@ public class TracerTest {
   }
 
   @Test
-  public void propagationViaRunnable() {
+  public void wrapRunnable() {
     Runnable runnable;
-    Scope ws = noopTracer.withSpan(span);
-    try {
-      assertThat(noopTracer.getCurrentSpan()).isSameAs(span);
-      runnable =
-          Context.current()
-              .wrap(
-                  new Runnable() {
-                    @Override
-                    public void run() {
-                      assertThat(noopTracer.getCurrentSpan()).isSameAs(span);
-                    }
-                  });
-    } finally {
-      ws.close();
-    }
     assertThat(noopTracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+    runnable =
+        tracer.withSpan(
+            span,
+            new Runnable() {
+              @Override
+              public void run() {
+                assertThat(noopTracer.getCurrentSpan()).isSameAs(span);
+              }
+            });
     // When we run the runnable we will have the span in the current Context.
     runnable.run();
+    verifyZeroInteractions(span);
+    assertThat(noopTracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+  }
+
+  @Test
+  public void wrapCallable() throws Exception {
+    final Object ret = new Object();
+    Callable<Object> callable;
+    assertThat(noopTracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
+    callable =
+        tracer.withSpan(
+            span,
+            new Callable<Object>() {
+              @Override
+              public Object call() throws Exception {
+                assertThat(noopTracer.getCurrentSpan()).isSameAs(span);
+                return ret;
+              }
+            });
+    // When we call the callable we will have the span in the current Context.
+    assertThat(callable.call()).isEqualTo(ret);
+    verifyZeroInteractions(span);
     assertThat(noopTracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
   }