Replace ThreadContextSpan#set with withSpan

lucidsoftware · Mar 6, 2017 · 4e625de · 4e625de
1 parent 9fa2b5a
commit 4e625de
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 10 deletions.
diff --git 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
@@ -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<Span> 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();
+        });
     }
 
 }