1 files changed, 24 insertions, 25 deletions
diff --git a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
index 48e56153..85b7396f 100644
--- a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
@@ -18,7 +18,9 @@ package kamon.trace
 
 import java.io.ObjectStreamException
 import kamon.trace.TraceContextAware.DefaultTraceContextAware
-import kamon.util.RelativeNanoTimestamp
+import kamon.util.{ SameThreadExecutionContext, Supplier, Function, RelativeNanoTimestamp }
+
+import scala.concurrent.Future
 
 trait TraceContext {
   def name: String
@@ -35,34 +37,31 @@ trait TraceContext {
   def addMetadata(key: String, value: String)
 
   def startTimestamp: RelativeNanoTimestamp
-}
-
-object TraceContext {
-  private[kamon] val _traceContextStorage = new ThreadLocal[TraceContext] {
-    override def initialValue(): TraceContext = EmptyTraceContext
-  }
-
-  def currentContext: TraceContext =
-    _traceContextStorage.get()
 
-  def setCurrentContext(context: TraceContext): Unit =
-    _traceContextStorage.set(context)
-
-  def clearCurrentContext: Unit =
-    _traceContextStorage.remove()
+  def collect[T](f: TraceContext ⇒ T): Option[T] =
+    if (nonEmpty)
+      Some(f(this))
+    else None
 
-  def withContext[T](context: TraceContext)(code: ⇒ T): T = {
-    val oldContext = _traceContextStorage.get()
-    _traceContextStorage.set(context)
+  def collect[T](f: Function[TraceContext, T]): Option[T] =
+    if (nonEmpty)
+      Some(f(this))
+    else None
 
-    try code finally _traceContextStorage.set(oldContext)
+  def withNewSegment[T](segmentName: String, category: String, library: String)(code: ⇒ T): T = {
+    val segment = startSegment(segmentName, category, library)
+    try code finally segment.finish()
   }
 
-  def map[T](f: TraceContext ⇒ T): Option[T] = {
-    val current = currentContext
-    if (current.nonEmpty)
-      Some(f(current))
-    else None
+  // Java variant.
+  def withNewSegment[T](segmentName: String, category: String, library: String, code: Supplier[T]): T =
+    withNewSegment(segmentName, category, library)(code.get)
+
+  def withNewAsyncSegment[T](segmentName: String, category: String, library: String)(code: ⇒ Future[T]): Future[T] = {
+    val segment = startSegment(segmentName, category, library)
+    val result = code
+    code.onComplete(_ ⇒ segment.finish())(SameThreadExecutionContext)
+    result
   }
 
 }
@@ -132,7 +131,7 @@ object TraceContextAware {
   def default: TraceContextAware = new DefaultTraceContextAware
 
   class DefaultTraceContextAware extends TraceContextAware {
-    @transient val traceContext = TraceContext.currentContext
+    @transient val traceContext = Tracer.currentContext
 
     //
     // Beware of this hack, it might bite us in the future!