Add support for custom context storage

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

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

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 <C> Callable<C> withSpan(Span span, boolean endSpan, Callable<C> 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<V> 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();
         }

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() {
+  }
+}

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() {}
 

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

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

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

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

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

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

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