diff --git a/api/src/main/java/io/opencensus/trace/ContextManager.java b/api/src/main/java/io/opencensus/trace/ContextManager.java new file mode 100644 index 0000000000..42f4c7c07d --- /dev/null +++ b/api/src/main/java/io/opencensus/trace/ContextManager.java @@ -0,0 +1,7 @@ +package io.opencensus.trace; + +public interface ContextManager { + Ctx currentContext(); + Ctx withValue(Ctx ctx, @javax.annotation.Nullable Span span); + Span getValue(Ctx ctx); +} diff --git a/api/src/main/java/io/opencensus/trace/Ctx.java b/api/src/main/java/io/opencensus/trace/Ctx.java new file mode 100644 index 0000000000..70f36dd742 --- /dev/null +++ b/api/src/main/java/io/opencensus/trace/Ctx.java @@ -0,0 +1,6 @@ +package io.opencensus.trace; + +public interface Ctx { + Ctx attach(); + void detach(Ctx ctx); +} diff --git a/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java b/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java index 19d6ac9978..a9df2c908e 100644 --- a/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java +++ b/api/src/main/java/io/opencensus/trace/CurrentSpanUtils.java @@ -16,9 +16,8 @@ package io.opencensus.trace; -import io.grpc.Context; import io.opencensus.common.Scope; -import io.opencensus.trace.unsafe.ContextUtils; +import io.opencensus.trace.unsafe.CtxUtils; import java.util.concurrent.Callable; import javax.annotation.Nullable; @@ -34,7 +33,7 @@ private CurrentSpanUtils() {} */ @Nullable static Span getCurrentSpan() { - return ContextUtils.getValue(Context.current()); + return CtxUtils.getValue(CtxUtils.currentContext()); } /** @@ -78,7 +77,7 @@ static Callable withSpan(Span span, boolean endSpan, Callable callable // 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; + private final Ctx origContext; private final Span span; private final boolean endSpan; @@ -90,12 +89,12 @@ private static final class ScopeInSpan implements Scope { private ScopeInSpan(Span span, boolean endSpan) { this.span = span; this.endSpan = endSpan; - origContext = ContextUtils.withValue(Context.current(), span).attach(); + origContext = CtxUtils.withValue(CtxUtils.currentContext(), span).attach(); } @Override public void close() { - Context.current().detach(origContext); + CtxUtils.currentContext().detach(origContext); if (endSpan) { span.end(); } @@ -116,7 +115,7 @@ private RunnableInSpan(Span span, Runnable runnable, boolean endSpan) { @Override public void run() { - Context origContext = ContextUtils.withValue(Context.current(), span).attach(); + Ctx origContext = CtxUtils.withValue(CtxUtils.currentContext(), span).attach(); try { runnable.run(); } catch (Throwable t) { @@ -128,7 +127,7 @@ public void run() { } throw new RuntimeException("unexpected", t); } finally { - Context.current().detach(origContext); + CtxUtils.currentContext().detach(origContext); if (endSpan) { span.end(); } @@ -149,7 +148,7 @@ private CallableInSpan(Span span, Callable callable, boolean endSpan) { @Override public V call() throws Exception { - Context origContext = ContextUtils.withValue(Context.current(), span).attach(); + Ctx origContext = CtxUtils.withValue(CtxUtils.currentContext(), span).attach(); try { return callable.call(); } catch (Exception e) { @@ -162,7 +161,7 @@ public V call() throws Exception { } throw new RuntimeException("unexpected", t); } finally { - Context.current().detach(origContext); + CtxUtils.currentContext().detach(origContext); if (endSpan) { span.end(); } diff --git a/api/src/main/java/io/opencensus/trace/unsafe/ContextManagerImpl.java b/api/src/main/java/io/opencensus/trace/unsafe/ContextManagerImpl.java new file mode 100644 index 0000000000..91b494a734 --- /dev/null +++ b/api/src/main/java/io/opencensus/trace/unsafe/ContextManagerImpl.java @@ -0,0 +1,39 @@ +package io.opencensus.trace.unsafe; + +import io.grpc.Context; +import io.opencensus.trace.ContextManager; +import io.opencensus.trace.Ctx; +import io.opencensus.trace.Span; +import javax.annotation.Nullable; + +/** + * Default {@code ContextManager} implementation using {@see io.grpc.Context} + */ +public class ContextManagerImpl implements ContextManager { + + @Override + public Ctx currentContext() { + return wrapContext(Context.current()); + } + + @Override + public Ctx withValue(Ctx ctx, @Nullable Span span) { + return wrapContext(ContextUtils.withValue(unwrapContext(ctx), span)); + } + + @Override + public Span getValue(Ctx ctx) { + return ContextUtils.getValue(unwrapContext(ctx)); + } + + private static Ctx wrapContext(Context context) { + return new CtxImpl(context); + } + + private static Context unwrapContext(Ctx ctx) { + return ((CtxImpl)ctx).getContext(); + } + + protected ContextManagerImpl() { + } +} diff --git a/api/src/main/java/io/opencensus/trace/unsafe/ContextUtils.java b/api/src/main/java/io/opencensus/trace/unsafe/ContextUtils.java index 1026512eec..46ce2856ee 100644 --- a/api/src/main/java/io/opencensus/trace/unsafe/ContextUtils.java +++ b/api/src/main/java/io/opencensus/trace/unsafe/ContextUtils.java @@ -33,7 +33,7 @@ * * @since 0.5 */ -public final class ContextUtils { +final class ContextUtils { // No instance of this class. private ContextUtils() {} diff --git a/api/src/main/java/io/opencensus/trace/unsafe/CtxImpl.java b/api/src/main/java/io/opencensus/trace/unsafe/CtxImpl.java new file mode 100644 index 0000000000..a10a863d2b --- /dev/null +++ b/api/src/main/java/io/opencensus/trace/unsafe/CtxImpl.java @@ -0,0 +1,30 @@ +package io.opencensus.trace.unsafe; + +import io.grpc.Context; +import io.opencensus.trace.Ctx; + +/** + * {@code Ctx} implementation using {@see io.grpc.Context} + */ +class CtxImpl implements Ctx { + private final Context context; + + public CtxImpl(Context context) { + this.context = context; + } + + Context getContext() { + return context; + } + + @Override + public Ctx attach() { + return new CtxImpl(context.attach()); + } + + @Override + public void detach(Ctx ctx) { + CtxImpl impl = (CtxImpl) ctx; + context.detach(impl.context); + } +} diff --git a/api/src/main/java/io/opencensus/trace/unsafe/CtxUtils.java b/api/src/main/java/io/opencensus/trace/unsafe/CtxUtils.java new file mode 100644 index 0000000000..4425523146 --- /dev/null +++ b/api/src/main/java/io/opencensus/trace/unsafe/CtxUtils.java @@ -0,0 +1,46 @@ +package io.opencensus.trace.unsafe; + +import io.opencensus.trace.ContextManager; +import io.opencensus.trace.Ctx; +import io.opencensus.trace.Span; + +public class CtxUtils { + // No instance of this class. + private CtxUtils() {} + + private static final ContextManager DEFAULT_CONTEXT_MANAGER = new ContextManagerImpl(); + private static ContextManager contextManager = DEFAULT_CONTEXT_MANAGER; + + /** + * Overrides context manager with a custom implementation + * @param cm custom {@code ContextManager} to be used instead of a default one. + */ + public static void setContextManager(ContextManager cm) { + contextManager = cm; + } + + public static Ctx currentContext() { + return contextManager.currentContext(); + } + + /** + * Creates a new {@code Ctx} with the given value set. + * + * @param context the parent {@code Ctx}. + * @param span the value to be set. + * @return a new context with the given value set. + */ + public static Ctx withValue(Ctx context, @javax.annotation.Nullable Span span) { + return contextManager.withValue(context, span); + } + + /** + * Returns the value from the specified {@code Ctx}. + * + * @param context the specified {@code Ctx}. + * @return the value from the specified {@code Ctx}. + */ + public static Span getValue(Ctx context) { + return contextManager.getValue(context); + } +} diff --git a/api/src/test/java/io/opencensus/trace/CurrentSpanUtilsTest.java b/api/src/test/java/io/opencensus/trace/CurrentSpanUtilsTest.java index 4cc6fccbc1..cd56f344ce 100644 --- a/api/src/test/java/io/opencensus/trace/CurrentSpanUtilsTest.java +++ b/api/src/test/java/io/opencensus/trace/CurrentSpanUtilsTest.java @@ -21,9 +21,8 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyZeroInteractions; -import io.grpc.Context; import io.opencensus.common.Scope; -import io.opencensus.trace.unsafe.ContextUtils; +import io.opencensus.trace.unsafe.CtxUtils; import java.util.concurrent.Callable; import org.junit.Before; import org.junit.Test; @@ -74,12 +73,12 @@ public void getCurrentSpan_WhenNoContext() { @Test public void getCurrentSpan() { assertThat(CurrentSpanUtils.getCurrentSpan()).isEqualTo(BlankSpan.INSTANCE); - Context origContext = ContextUtils.withValue(Context.current(), span).attach(); + Ctx origContext = CtxUtils.withValue(CtxUtils.currentContext(), span).attach(); // Make sure context is detached even if test fails. try { assertThat(CurrentSpanUtils.getCurrentSpan()).isSameInstanceAs(span); } finally { - Context.current().detach(origContext); + CtxUtils.currentContext().detach(origContext); } assertThat(CurrentSpanUtils.getCurrentSpan()).isEqualTo(BlankSpan.INSTANCE); } diff --git a/api/src/test/java/io/opencensus/trace/unsafe/ContextUtilsTest.java b/api/src/test/java/io/opencensus/trace/unsafe/ContextUtilsTest.java index 0ad9feba3d..53cb0d023a 100644 --- a/api/src/test/java/io/opencensus/trace/unsafe/ContextUtilsTest.java +++ b/api/src/test/java/io/opencensus/trace/unsafe/ContextUtilsTest.java @@ -18,8 +18,8 @@ import static com.google.common.truth.Truth.assertThat; -import io.grpc.Context; import io.opencensus.trace.BlankSpan; +import io.opencensus.trace.Ctx; import io.opencensus.trace.Span; import org.junit.Test; import org.junit.runner.RunWith; @@ -31,19 +31,19 @@ public class ContextUtilsTest { @Test public void testGetCurrentSpan_DefaultContext() { - Span span = ContextUtils.getValue(Context.current()); + Span span = CtxUtils.getValue(CtxUtils.currentContext()); assertThat(span).isEqualTo(BlankSpan.INSTANCE); } @Test public void testGetCurrentSpan_ContextSetToNull() { - Context orig = ContextUtils.withValue(Context.current(), null).attach(); + Ctx orig = CtxUtils.withValue(CtxUtils.currentContext(), null).attach(); try { - Span span = ContextUtils.getValue(Context.current()); + Span span = CtxUtils.getValue(CtxUtils.currentContext()); // ContextUtils.getValue always returns non-null. assertThat(span).isEqualTo(BlankSpan.INSTANCE); } finally { - Context.current().detach(orig); + CtxUtils.currentContext().detach(orig); } } } diff --git a/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/ContextDataUtils.java b/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/ContextDataUtils.java index d8bc5819ab..d182f53492 100644 --- a/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/ContextDataUtils.java +++ b/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/ContextDataUtils.java @@ -16,10 +16,9 @@ package io.opencensus.contrib.logcorrelation.log4j2; -import io.grpc.Context; import io.opencensus.trace.Span; import io.opencensus.trace.SpanContext; -import io.opencensus.trace.unsafe.ContextUtils; +import io.opencensus.trace.unsafe.CtxUtils; import java.util.Collection; import java.util.List; import javax.annotation.Nullable; @@ -81,7 +80,7 @@ static StringMap getContextAndTracingData() { } private static SpanContext getCurrentSpanContext() { - Span span = ContextUtils.getValue(Context.current()); + Span span = CtxUtils.getValue(CtxUtils.currentContext()); return span == null ? SpanContext.INVALID : span.getContext(); } } diff --git a/contrib/log_correlation/stackdriver/src/main/java/io/opencensus/contrib/logcorrelation/stackdriver/OpenCensusTraceLoggingEnhancer.java b/contrib/log_correlation/stackdriver/src/main/java/io/opencensus/contrib/logcorrelation/stackdriver/OpenCensusTraceLoggingEnhancer.java index bc3f2b2451..186c14d552 100644 --- a/contrib/log_correlation/stackdriver/src/main/java/io/opencensus/contrib/logcorrelation/stackdriver/OpenCensusTraceLoggingEnhancer.java +++ b/contrib/log_correlation/stackdriver/src/main/java/io/opencensus/contrib/logcorrelation/stackdriver/OpenCensusTraceLoggingEnhancer.java @@ -19,11 +19,10 @@ import com.google.cloud.ServiceOptions; import com.google.cloud.logging.LogEntry; import com.google.cloud.logging.LoggingEnhancer; -import io.grpc.Context; import io.opencensus.trace.Span; import io.opencensus.trace.SpanContext; import io.opencensus.trace.TraceId; -import io.opencensus.trace.unsafe.ContextUtils; +import io.opencensus.trace.unsafe.CtxUtils; import java.util.logging.LogManager; import javax.annotation.Nullable; @@ -99,7 +98,7 @@ public void enhanceLogEntry(LogEntry.Builder builder) { } private static SpanContext getCurrentSpanContext() { - Span span = ContextUtils.getValue(Context.current()); + Span span = CtxUtils.getValue(CtxUtils.currentContext()); return span == null ? SpanContext.INVALID : span.getContext(); } diff --git a/contrib/spring_sleuth_v1x/src/main/java/io/opencensus/contrib/spring/sleuth/v1x/OpenCensusSleuthSpanContextHolder.java b/contrib/spring_sleuth_v1x/src/main/java/io/opencensus/contrib/spring/sleuth/v1x/OpenCensusSleuthSpanContextHolder.java index 332058d02a..5b77470629 100644 --- a/contrib/spring_sleuth_v1x/src/main/java/io/opencensus/contrib/spring/sleuth/v1x/OpenCensusSleuthSpanContextHolder.java +++ b/contrib/spring_sleuth_v1x/src/main/java/io/opencensus/contrib/spring/sleuth/v1x/OpenCensusSleuthSpanContextHolder.java @@ -18,7 +18,8 @@ import io.grpc.Context; import io.opencensus.common.ExperimentalApi; -import io.opencensus.trace.unsafe.ContextUtils; +import io.opencensus.trace.Ctx; +import io.opencensus.trace.unsafe.CtxUtils; import org.apache.commons.logging.Log; import org.springframework.cloud.sleuth.Span; import org.springframework.core.NamedThreadLocal; @@ -136,14 +137,14 @@ private static class SpanContext { final boolean autoClose; @javax.annotation.Nullable final SpanContext parent; final OpenCensusSleuthSpan ocSpan; - final Context ocCurrentContext; + final Ctx ocCurrentContext; private SpanContext(Span span, boolean autoClose) { this.span = span; this.autoClose = autoClose; this.parent = CURRENT_SPAN.get(); this.ocSpan = new OpenCensusSleuthSpan(span); - this.ocCurrentContext = ContextUtils.withValue(Context.current(), this.ocSpan); + this.ocCurrentContext = CtxUtils.withValue(CtxUtils.currentContext(), this.ocSpan); } }