simple instrumentation for keeping a trace context within Futures callbacks

1 files changed, 4 insertions, 14 deletions

diff --git a/src/main/scala/kamon/TraceContext.scala b/src/main/scala/kamon/TraceContext.scala
index 1fbedf86..18a91145 100644
--- a/src/main/scala/kamon/TraceContext.scala
+++ b/src/main/scala/kamon/TraceContext.scala
@@ -36,21 +36,11 @@ case class CodeBlockExecutionTime(blockName: String, begin: Long, end: Long) ext
 
 
 trait TraceSupport {
-  import TraceContext.current
-
-
-  def trace[T](blockName: String)(f: => T): T = {
-    val before = System.currentTimeMillis
-
-    val result = f
-
-    val after = System.currentTimeMillis
-    //swapContext(current.get().withEntry(CodeBlockExecutionTime(blockName, before, after)))
+  def withContext[Out](func: => Any => Out, ctx: TraceContext) = {
+    TraceContext.set(ctx)
+    val result = func
+    TraceContext.clear
 
     result
   }
-
-  def swapContext(newContext: TraceContext) {
-    //current.set(newContext)
-  }
 }