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.

3 files changed, 329 insertions, 37 deletions

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);
   }