diff options
Diffstat (limited to 'core')
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); } } |