1 files changed, 9 insertions, 6 deletions

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 75320fdd..dce75596 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
@@ -20,10 +20,10 @@ import com.google.cloud.ServiceOptions;
 import com.google.cloud.logging.LogEntry;
 import com.google.cloud.logging.LoggingEnhancer;
 import io.opencensus.common.ExperimentalApi;
+import io.opencensus.trace.Span;
 import io.opencensus.trace.SpanContext;
 import io.opencensus.trace.TraceId;
-import io.opencensus.trace.Tracer;
-import io.opencensus.trace.Tracing;
+import io.opencensus.trace.unsafe.ContextUtils;
 import java.util.logging.LogManager;
 import javax.annotation.Nullable;
 
@@ -57,8 +57,6 @@ public final class OpenCensusTraceLoggingEnhancer implements LoggingEnhancer {
       "io.opencensus.contrib.logcorrelation.stackdriver."
           + "OpenCensusTraceLoggingEnhancer.spanSelection";
 
-  private static final Tracer tracer = Tracing.getTracer();
-
   private final String projectId;
   private final SpanSelection spanSelection;
 
@@ -184,18 +182,23 @@ public final class OpenCensusTraceLoggingEnhancer implements LoggingEnhancer {
       case NO_SPANS:
         return;
       case SAMPLED_SPANS:
-        SpanContext span = tracer.getCurrentSpan().getContext();
+        SpanContext span = getCurrentSpanContext();
         if (span.getTraceOptions().isSampled()) {
           addTracingData(tracePrefix, span, builder);
         }
         return;
       case ALL_SPANS:
-        addTracingData(tracePrefix, tracer.getCurrentSpan().getContext(), builder);
+        addTracingData(tracePrefix, getCurrentSpanContext(), builder);
         return;
     }
     throw new AssertionError("Unknown spanSelection: " + spanSelection);
   }
 
+  private static SpanContext getCurrentSpanContext() {
+    Span span = ContextUtils.CONTEXT_SPAN_KEY.get();
+    return span == null ? SpanContext.INVALID : span.getContext();
+  }
+
   private static void addTracingData(
       String tracePrefix, SpanContext span, LogEntry.Builder builder) {
     builder.setTrace(formatTraceId(tracePrefix, span.getTraceId()));