diff --git a/README.md b/README.md index 670acc2..a7dfab9 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# Opentracing thread-local propogation +# Opentracing thread-local propagation [![Build Status](https://travis-ci.org/lucidsoftware/opentracing-thread-context.svg?branch=master)](https://travis-ci.org/lucidsoftware/opentracing-thread-context) ![Maven Version](https://img.shields.io/maven-central/v/com.lucidchart/opentracing-thread-context.svg) @@ -12,6 +12,7 @@ import io.opentracing.Span; import io.opentracing.threadcontext.ThreadContextSpan; Span span = ... -ThreadContextSpan.set(span) -ThreadContextSpan.get(); +ThreadContextSpan.withSpan(span, () -> { + ThreadContextSpan.get(); // current Span +}); ``` diff --git a/thread-context/src/main/java/io/opentracing/threadcontext/ThreadContextSpan.java b/thread-context/src/main/java/io/opentracing/threadcontext/ThreadContextSpan.java index 91326e1..0ab4282 100644 --- a/thread-context/src/main/java/io/opentracing/threadcontext/ThreadContextSpan.java +++ b/thread-context/src/main/java/io/opentracing/threadcontext/ThreadContextSpan.java @@ -1,5 +1,6 @@ package io.opentracing.threadcontext; +import com.github.threadcontext.Saver; import com.github.threadcontext.ThreadContext; import com.github.threadcontext.ThreadLocalSaver; import io.opentracing.NoopSpan; @@ -13,8 +14,9 @@ protected Span initialValue() { return defaultSpan.get(); } }; + private static final Saver saver = new ThreadLocalSaver<>(span); static { - ThreadContext.savers.add(new ThreadLocalSaver<>(span)); + ThreadContext.savers.add(saver); } public static Supplier defaultSpan = () -> NoopSpan.INSTANCE; @@ -26,12 +28,11 @@ public static Span get() { return span.get(); } - public static void set(Span span) { - ThreadContextSpan.span.set(span); - } - - public static void clear() { - span.remove(); + public static void withSpan(Span span, Runnable runnable) { + saver.runAndRestore(() -> { + ThreadContextSpan.span.set(span); + runnable.run(); + }); } }