Remove ScopedSpanHandle and make SpanInScope able to close the Span. (#848)

* Remove ScopedSpanHandle and make SpanInScope able to close the Span.

* Fix extra ";".

* Fix comments.

6 files changed, 36 insertions, 117 deletions

diff --git a/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java b/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java
index 291de53a..713c8816 100644
--- a/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java
+++ b/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java
@@ -41,31 +41,37 @@ final class CurrentSpanUtils {
    * <p>Supports try-with-resource idiom.
    *
    * @param span The {@code Span} to be set to the current context.
+   * @param endSpan if {@code true} the returned {@code Scope} will close the {@code Span}.
    * @return An object that defines a scope where the given {@code Span} is set to the current
    *     context.
    */
-  static Scope withSpan(Span span) {
-    return new WithSpan(span, ContextUtils.CONTEXT_SPAN_KEY);
+  static Scope withSpan(Span span, boolean endSpan) {
+    return new ScopeInSpan(span, endSpan);
   }
 
   // Defines an arbitrary scope of code as a traceable operation. Supports try-with-resources idiom.
-  private static final class WithSpan implements Scope {
+  private static final class ScopeInSpan implements Scope {
     private final Context origContext;
+    private final Span span;
+    private boolean endSpan;
 
     /**
-     * Constructs a new {@link WithSpan}.
+     * Constructs a new {@link ScopeInSpan}.
      *
      * @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();
+    ScopeInSpan(Span span, boolean endSpan) {
+      this.span = span;
+      this.endSpan = endSpan;
+      origContext = Context.current().withValue(ContextUtils.CONTEXT_SPAN_KEY, span).attach();
     }
 
     @Override
     public void close() {
       Context.current().detach(origContext);
+      if (endSpan) {
+        span.end();
+      }
     }
   }
 }
diff --git a/api/src/main/java/io/opencensus/trace/ScopedSpanHandle.java b/api/src/main/java/io/opencensus/trace/ScopedSpanHandle.java
deleted file mode 100644
index b5beb4aa..00000000
--- a/api/src/main/java/io/opencensus/trace/ScopedSpanHandle.java
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * 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 io.opencensus.common.Scope;
-
-/**
- * Defines a scope of code where the given {@link Span} is in the current context. The scope is
- * exited when the object is closed and the previous context is restored. When the scope exits the
- * given {@code Span} will be ended using {@link Span#end}.
- *
- * <p>Supports try-with-resource idiom.
- */
-final class ScopedSpanHandle implements Scope {
-  private final Span span;
-  private final Scope withSpan;
-
-  /**
-   * Creates a {@code ScopedSpanHandle}
-   *
-   * @param span The span that will be installed in the current context.
-   */
-  ScopedSpanHandle(Span span) {
-    this.span = span;
-    this.withSpan = CurrentSpanUtils.withSpan(span);
-  }
-
-  /**
-   * Uninstalls the {@code Span} from the current context and ends it by calling {@link Span#end()}.
-   */
-  @Override
-  public void close() {
-    withSpan.close();
-    span.end();
-  }
-}
diff --git a/api/src/main/java/io/opencensus/trace/SpanBuilder.java b/api/src/main/java/io/opencensus/trace/SpanBuilder.java
index 530af97e..4c28147e 100644
--- a/api/src/main/java/io/opencensus/trace/SpanBuilder.java
+++ b/api/src/main/java/io/opencensus/trace/SpanBuilder.java
@@ -220,7 +220,7 @@ public abstract class SpanBuilder {
    */
   @MustBeClosed
   public final Scope startScopedSpan() {
-    return new ScopedSpanHandle(startSpan());
+    return CurrentSpanUtils.withSpan(startSpan(), true);
   }
 
   static final class NoopSpanBuilder extends SpanBuilder {
diff --git a/api/src/main/java/io/opencensus/trace/Tracer.java b/api/src/main/java/io/opencensus/trace/Tracer.java
index 24853fbc..8215f7b1 100644
--- a/api/src/main/java/io/opencensus/trace/Tracer.java
+++ b/api/src/main/java/io/opencensus/trace/Tracer.java
@@ -147,7 +147,7 @@ public abstract class Tracer {
    */
   @MustBeClosed
   public final Scope withSpan(Span span) {
-    return CurrentSpanUtils.withSpan(checkNotNull(span, "span"));
+    return CurrentSpanUtils.withSpan(checkNotNull(span, "span"), false);
   }
 
   /**
diff --git a/api/src/test/java/io/opencensus/trace/CurrentSpanUtilsTest.java b/api/src/test/java/io/opencensus/trace/CurrentSpanUtilsTest.java
index afc7db63..8ba910b0 100644
--- a/api/src/test/java/io/opencensus/trace/CurrentSpanUtilsTest.java
+++ b/api/src/test/java/io/opencensus/trace/CurrentSpanUtilsTest.java
@@ -17,6 +17,9 @@
 package io.opencensus.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.verifyZeroInteractions;
 
 import io.grpc.Context;
 import io.opencensus.common.Scope;
@@ -57,21 +60,35 @@ public class CurrentSpanUtilsTest {
   }
 
   @Test
-  public void withSpan() {
+  public void withSpan_CloseDetaches() {
     assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
-    Scope ws = CurrentSpanUtils.withSpan(span);
+    Scope ws = CurrentSpanUtils.withSpan(span, false);
     try {
       assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
     } finally {
       ws.close();
     }
     assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+    verifyZeroInteractions(span);
+  }
+
+  @Test
+  public void withSpan_CloseDetachesAndEndsSpan() {
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+    Scope ss = CurrentSpanUtils.withSpan(span, true);
+    try {
+      assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
+    } finally {
+      ss.close();
+    }
+    assertThat(CurrentSpanUtils.getCurrentSpan()).isNull();
+    verify(span).end(same(EndSpanOptions.DEFAULT));
   }
 
   @Test
   public void propagationViaRunnable() {
     Runnable runnable = null;
-    Scope ws = CurrentSpanUtils.withSpan(span);
+    Scope ws = CurrentSpanUtils.withSpan(span, false);
     try {
       assertThat(CurrentSpanUtils.getCurrentSpan()).isSameAs(span);
       runnable =
diff --git a/api/src/test/java/io/opencensus/trace/ScopedSpanHandleTest.java b/api/src/test/java/io/opencensus/trace/ScopedSpanHandleTest.java
deleted file mode 100644
index a8c18819..00000000
--- a/api/src/test/java/io/opencensus/trace/ScopedSpanHandleTest.java
+++ /dev/null
@@ -1,54 +0,0 @@
-/*
- * 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.Matchers.same;
-import static org.mockito.Mockito.verify;
-
-import io.opencensus.common.Scope;
-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 ScopedSpanHandle}. */
-@RunWith(JUnit4.class)
-public class ScopedSpanHandleTest {
-  private static final Tracer tracer = Tracer.getNoopTracer();
-  @Mock private Span span;
-
-  @Before
-  public void setUp() {
-    MockitoAnnotations.initMocks(this);
-  }
-
-  @Test
-  public void initAttachesSpan_CloseDetachesAndEndsSpan() {
-    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
-    Scope ss = new ScopedSpanHandle(span);
-    try {
-      assertThat(tracer.getCurrentSpan()).isSameAs(span);
-    } finally {
-      ss.close();
-    }
-    assertThat(tracer.getCurrentSpan()).isSameAs(BlankSpan.INSTANCE);
-    verify(span).end(same(EndSpanOptions.DEFAULT));
-  }
-}