wip

1 files changed, 7 insertions, 3 deletions

diff --git a/src/main/scala/kamon/Kamon.scala b/src/main/scala/kamon/Kamon.scala
index b5998f81..4cba2342 100644
--- a/src/main/scala/kamon/Kamon.scala
+++ b/src/main/scala/kamon/Kamon.scala
@@ -6,15 +6,19 @@ object Kamon {
 
   implicit val actorSystem = ActorSystem("kamon")
 
-  private val ctx = new ThreadLocal[Option[TraceContext]] {
+  val ctx = new ThreadLocal[Option[TraceContext]] {
     override def initialValue() = None
   }
 
-  def context = ctx.get()
+  def context() = ctx.get()
   def clear = ctx.remove()
   def set(traceContext: TraceContext) = ctx.set(Some(traceContext))
 
-  def start: Unit = set(newTraceContext)
+  def start = set(newTraceContext)
+  def stop = ctx.get match {
+    case Some(context) => context.close
+    case None =>
+  }
 
   def newTraceContext(): TraceContext = TraceContext()