More withSpan methods

lucidsoftware · Mar 6, 2017 · 6216fe9 · 6216fe9
1 parent 4e625de
commit 6216fe9
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/thread-context/build.sbt b/thread-context/build.sbt
@@ -1,5 +1,5 @@
 libraryDependencies ++= Seq(
-  "com.lucidchart" % "thread-context" % "0.4",
+  "com.lucidchart" % "thread-context" % "0.5",
   "io.opentracing" % "opentracing-api" % "0.20.7",
   "io.opentracing" % "opentracing-noop" % "0.20.7"
 )

diff --git a/thread-context/src/main/java/io/opentracing/threadcontext/ThreadContextSpan.java b/thread-context/src/main/java/io/opentracing/threadcontext/ThreadContextSpan.java
@@ -5,6 +5,7 @@
 import com.github.threadcontext.ThreadLocalSaver;
 import io.opentracing.NoopSpan;
 import io.opentracing.Span;
+import java.util.concurrent.Callable;
 import java.util.function.Supplier;
 
 public final class ThreadContextSpan {
@@ -35,4 +36,18 @@ public static void withSpan(Span span, Runnable runnable) {
         });
     }
 
+    public static <T> T withSpan(Span span, Callable<T> callable) throws Exception {
+        return saver.runAndRestore(() -> {
+            ThreadContextSpan.span.set(span);
+            return callable.call();
+        });
+    }
+
+    public static <T> T withSpan(Span span, Supplier<T> supplier) throws Exception {
+        return saver.runAndRestore(() -> {
+            ThreadContextSpan.span.set(span);
+            return supplier.get();
+        });
+    }
+
 }