remake of trace context and allow different tracing levels

7 files changed, 87 insertions, 94 deletions

diff --git a/kamon-core/src/test/scala/akka/testkit/TestProbeTracing.scala b/kamon-core/src/test/scala/akka/testkit/TestProbeTracing.scala
index c681e921..a050145a 100644
--- a/kamon-core/src/test/scala/akka/testkit/TestProbeTracing.scala
+++ b/kamon-core/src/test/scala/akka/testkit/TestProbeTracing.scala
@@ -46,7 +46,7 @@ class TestProbeTracing {
       case _                ⇒ None
     }
 
-    TraceRecorder.withContext(traceContext) {
+    TraceRecorder.withTraceContext(traceContext) {
       pjp.proceed
     }
   }
diff --git a/kamon-core/src/test/scala/kamon/trace/instrumentation/ActorLoggingSpec.scala b/kamon-core/src/test/scala/kamon/trace/instrumentation/ActorLoggingSpec.scala
index d6648cef..a0d8e933 100644
--- a/kamon-core/src/test/scala/kamon/trace/instrumentation/ActorLoggingSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/instrumentation/ActorLoggingSpec.scala
@@ -18,24 +18,24 @@ package kamon.trace.instrumentation
 import akka.testkit.TestKit
 import org.scalatest.{ Inspectors, Matchers, WordSpecLike }
 import akka.actor.{ Props, ActorLogging, Actor, ActorSystem }
-import akka.event.Logging.{ LogEvent }
-import kamon.trace.{ ContextAware, TraceContext, Trace }
+import akka.event.Logging.LogEvent
+import kamon.trace.{TraceContextAware, TraceRecorder}
 
 class ActorLoggingSpec extends TestKit(ActorSystem("actor-logging-spec")) with WordSpecLike with Matchers with Inspectors {
 
   "the ActorLogging instrumentation" should {
     "attach the TraceContext (if available) to log events" in {
-      val testTraceContext = Some(TraceContext(Actor.noSender, 1, "test", "test-1"))
       val loggerActor = system.actorOf(Props[LoggerActor])
       system.eventStream.subscribe(testActor, classOf[LogEvent])
 
-      Trace.withContext(testTraceContext) {
+      val testTraceContext = TraceRecorder.withNewTraceContext("logging") {
         loggerActor ! "info"
+        TraceRecorder.currentContext
       }
 
       fishForMessage() {
         case event: LogEvent if event.message.toString contains "TraceContext =>" ⇒
-          val ctxInEvent = event.asInstanceOf[ContextAware].traceContext
+          val ctxInEvent = event.asInstanceOf[TraceContextAware].traceContext
           ctxInEvent === testTraceContext
 
         case event: LogEvent ⇒ false
@@ -46,6 +46,6 @@ class ActorLoggingSpec extends TestKit(ActorSystem("actor-logging-spec")) with W
 
 class LoggerActor extends Actor with ActorLogging {
   def receive = {
-    case "info" ⇒ log.info("TraceContext => {}", Trace.context())
+    case "info" ⇒ log.info("TraceContext => {}", TraceRecorder.currentContext)
   }
 }
diff --git a/kamon-core/src/test/scala/kamon/trace/instrumentation/ActorMessagePassingTracingSpec.scala b/kamon-core/src/test/scala/kamon/trace/instrumentation/ActorMessagePassingTracingSpec.scala
index f32623b9..acc939fb 100644
--- a/kamon-core/src/test/scala/kamon/trace/instrumentation/ActorMessagePassingTracingSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/instrumentation/ActorMessagePassingTracingSpec.scala
@@ -15,70 +15,71 @@
  * ========================================================== */
 package kamon.trace.instrumentation
 
-import org.scalatest.{ WordSpecLike, Matchers }
-import akka.actor.{ ActorRef, Actor, Props, ActorSystem }
+import org.scalatest.WordSpecLike
+import akka.actor.{ Actor, Props, ActorSystem }
 
 import akka.testkit.{ ImplicitSender, TestKit }
-import kamon.trace.Trace
+import kamon.trace.TraceRecorder
 import akka.pattern.{ pipe, ask }
 import akka.util.Timeout
 import scala.concurrent.duration._
-import scala.concurrent.{ Await, Future }
 import akka.routing.RoundRobinRouter
-import kamon.trace.TraceContext
 
 class ActorMessagePassingTracingSpec extends TestKit(ActorSystem("actor-message-passing-tracing-spec")) with WordSpecLike with ImplicitSender {
   implicit val executionContext = system.dispatcher
 
   "the message passing instrumentation" should {
-    "propagate the TraceContext using bang" in new TraceContextEchoFixture {
-      Trace.withContext(testTraceContext) {
+    "propagate the TraceContext using bang" in new EchoActorFixture {
+      val testTraceContext = TraceRecorder.withNewTraceContext("bang-reply") {
         ctxEchoActor ! "test"
+        TraceRecorder.currentContext
       }
 
       expectMsg(testTraceContext)
     }
 
-    "propagate the TraceContext using tell" in new TraceContextEchoFixture {
-      Trace.withContext(testTraceContext) {
+    "propagate the TraceContext using tell" in new EchoActorFixture {
+      val testTraceContext = TraceRecorder.withNewTraceContext("tell-reply") {
         ctxEchoActor.tell("test", testActor)
+        TraceRecorder.currentContext
       }
 
       expectMsg(testTraceContext)
     }
 
-    "propagate the TraceContext using ask" in new TraceContextEchoFixture {
+    "propagate the TraceContext using ask" in new EchoActorFixture {
       implicit val timeout = Timeout(1 seconds)
-      Trace.withContext(testTraceContext) {
+      val testTraceContext = TraceRecorder.withNewTraceContext("ask-reply") {
         // The pipe pattern use Futures internally, so FutureTracing test should cover the underpinnings of it.
         (ctxEchoActor ? "test") pipeTo (testActor)
+        TraceRecorder.currentContext
       }
 
       expectMsg(testTraceContext)
     }
 
-    "propagate the TraceContext to actors behind a router" in new RoutedTraceContextEchoFixture {
-      Trace.withContext(testTraceContext) {
+    "propagate the TraceContext to actors behind a router" in new RoutedEchoActorFixture {
+      val testTraceContext = TraceRecorder.withNewTraceContext("router-reply") {
         ctxEchoActor ! "test"
+        TraceRecorder.currentContext
       }
 
       expectMsg(testTraceContext)
     }
   }
 
-  trait TraceContextEchoFixture {
-    val testTraceContext = Some(Trace.newTraceContext("test", "test-1"))
+  trait EchoActorFixture {
     val ctxEchoActor = system.actorOf(Props[TraceContextEcho])
   }
 
-  trait RoutedTraceContextEchoFixture extends TraceContextEchoFixture {
+  trait RoutedEchoActorFixture extends EchoActorFixture {
     override val ctxEchoActor = system.actorOf(Props[TraceContextEcho].withRouter(RoundRobinRouter(nrOfInstances = 1)))
   }
 }
 
 class TraceContextEcho extends Actor {
   def receive = {
-    case msg: String ⇒ sender ! Trace.context()
+    case msg: String ⇒ sender ! TraceRecorder.currentContext
   }
 }
 
diff --git a/kamon-core/src/test/scala/kamon/trace/instrumentation/ActorSystemMessagePassingInstrumentationSpec.scala b/kamon-core/src/test/scala/kamon/trace/instrumentation/ActorSystemMessagePassingInstrumentationSpec.scala
index 7d539370..00ecae79 100644
--- a/kamon-core/src/test/scala/kamon/trace/instrumentation/ActorSystemMessagePassingInstrumentationSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/instrumentation/ActorSystemMessagePassingInstrumentationSpec.scala
@@ -3,7 +3,7 @@ package kamon.trace.instrumentation
 import akka.testkit.{ ImplicitSender, TestKit }
 import akka.actor._
 import org.scalatest.WordSpecLike
-import kamon.trace.Trace
+import kamon.trace.TraceRecorder
 import scala.util.control.NonFatal
 import akka.actor.SupervisorStrategy.{ Escalate, Stop, Restart, Resume }
 import scala.concurrent.duration._
@@ -12,43 +12,44 @@ class ActorSystemMessagePassingInstrumentationSpec extends TestKit(ActorSystem("
   implicit val executionContext = system.dispatcher
 
   "the system message passing instrumentation" should {
-    "keep the TraceContext while processing the Create message in top level actors" in new TraceContextFixture {
-      Trace.withContext(testTraceContext) {
+    "keep the TraceContext while processing the Create message in top level actors" in {
+      val testTraceContext = TraceRecorder.withNewTraceContext("creating-top-level-actor") {
         system.actorOf(Props(new Actor {
-
-          testActor ! Trace.context()
-
+          testActor ! TraceRecorder.currentContext
           def receive: Actor.Receive = { case any ⇒ }
         }))
+
+        TraceRecorder.currentContext
       }
 
       expectMsg(testTraceContext)
     }
 
-    "keep the TraceContext while processing the Create message in non top level actors" in new TraceContextFixture {
-      Trace.withContext(testTraceContext) {
+    "keep the TraceContext while processing the Create message in non top level actors" in  {
+      val testTraceContext = TraceRecorder.withNewTraceContext("creating-non-top-level-actor") {
         system.actorOf(Props(new Actor {
           def receive: Actor.Receive = {
             case any ⇒
               context.actorOf(Props(new Actor {
-
-                testActor ! Trace.context()
-
+                testActor ! TraceRecorder.currentContext
                 def receive: Actor.Receive = { case any ⇒ }
               }))
           }
         })) ! "any"
+
+        TraceRecorder.currentContext
       }
 
       expectMsg(testTraceContext)
     }
 
     "keep the TraceContext in the supervision cycle" when {
-      "the actor is resumed" in new TraceContextFixture {
+      "the actor is resumed" in {
         val supervisor = supervisorWithDirective(Resume)
 
-        Trace.withContext(testTraceContext) {
+        val testTraceContext = TraceRecorder.withNewTraceContext("fail-and-resume") {
           supervisor ! "fail"
+          TraceRecorder.currentContext
         }
 
         expectMsg(testTraceContext) // From the parent executing the supervision strategy
@@ -58,11 +59,12 @@ class ActorSystemMessagePassingInstrumentationSpec extends TestKit(ActorSystem("
         expectMsg(None)
       }
 
-      "the actor is restarted" in new TraceContextFixture {
+      "the actor is restarted" in {
         val supervisor = supervisorWithDirective(Restart, sendPreRestart = true, sendPostRestart = true)
 
-        Trace.withContext(testTraceContext) {
+        val testTraceContext = TraceRecorder.withNewTraceContext("fail-and-restart") {
           supervisor ! "fail"
+          TraceRecorder.currentContext
         }
 
         expectMsg(testTraceContext) // From the parent executing the supervision strategy
@@ -74,11 +76,12 @@ class ActorSystemMessagePassingInstrumentationSpec extends TestKit(ActorSystem("
         expectMsg(None)
       }
 
-      "the actor is stopped" in new TraceContextFixture {
+      "the actor is stopped" in {
         val supervisor = supervisorWithDirective(Stop, sendPostStop = true)
 
-        Trace.withContext(testTraceContext) {
+        val testTraceContext = TraceRecorder.withNewTraceContext("fail-and-stop") {
           supervisor ! "fail"
+          TraceRecorder.currentContext
         }
 
         expectMsg(testTraceContext) // From the parent executing the supervision strategy
@@ -86,11 +89,12 @@ class ActorSystemMessagePassingInstrumentationSpec extends TestKit(ActorSystem("
         expectNoMsg(1 second)
       }
 
-      "the failure is escalated" in new TraceContextFixture {
+      "the failure is escalated" in {
         val supervisor = supervisorWithDirective(Escalate, sendPostStop = true)
 
-        Trace.withContext(testTraceContext) {
+        val testTraceContext = TraceRecorder.withNewTraceContext("fail-and-escalate") {
           supervisor ! "fail"
+          TraceRecorder.currentContext
         }
 
         expectMsg(testTraceContext) // From the parent executing the supervision strategy
@@ -108,7 +112,7 @@ class ActorSystemMessagePassingInstrumentationSpec extends TestKit(ActorSystem("
       val child = context.actorOf(Props(new Parent))
 
       override def supervisorStrategy: SupervisorStrategy = OneForOneStrategy() {
-        case NonFatal(throwable) ⇒ testActor ! Trace.context(); Stop
+        case NonFatal(throwable) ⇒ testActor ! TraceRecorder.currentContext; Stop
       }
 
       def receive = {
@@ -120,7 +124,7 @@ class ActorSystemMessagePassingInstrumentationSpec extends TestKit(ActorSystem("
       val child = context.actorOf(Props(new Child))
 
       override def supervisorStrategy: SupervisorStrategy = OneForOneStrategy() {
-        case NonFatal(throwable) ⇒ testActor ! Trace.context(); directive
+        case NonFatal(throwable) ⇒ testActor ! TraceRecorder.currentContext; directive
       }
 
       def receive: Actor.Receive = {
@@ -128,7 +132,7 @@ class ActorSystemMessagePassingInstrumentationSpec extends TestKit(ActorSystem("
       }
 
       override def postStop(): Unit = {
-        if (sendPostStop) testActor ! Trace.context()
+        if (sendPostStop) testActor ! TraceRecorder.currentContext
         super.postStop()
       }
     }
@@ -136,26 +140,26 @@ class ActorSystemMessagePassingInstrumentationSpec extends TestKit(ActorSystem("
     class Child extends Actor {
       def receive = {
         case "fail"    ⇒ 1 / 0
-        case "context" ⇒ sender ! Trace.context()
+        case "context" ⇒ sender ! TraceRecorder.currentContext
       }
 
       override def preRestart(reason: Throwable, message: Option[Any]): Unit = {
-        if (sendPreRestart) testActor ! Trace.context()
+        if (sendPreRestart) testActor ! TraceRecorder.currentContext
         super.preRestart(reason, message)
       }
 
       override def postRestart(reason: Throwable): Unit = {
-        if (sendPostRestart) testActor ! Trace.context()
+        if (sendPostRestart) testActor ! TraceRecorder.currentContext
         super.postRestart(reason)
       }
 
       override def postStop(): Unit = {
-        if (sendPostStop) testActor ! Trace.context()
+        if (sendPostStop) testActor ! TraceRecorder.currentContext
         super.postStop()
       }
 
       override def preStart(): Unit = {
-        if (sendPreStart) testActor ! Trace.context()
+        if (sendPreStart) testActor ! TraceRecorder.currentContext
         super.preStart()
       }
     }
diff --git a/kamon-core/src/test/scala/kamon/trace/instrumentation/AskPatternTracingSpec.scala b/kamon-core/src/test/scala/kamon/trace/instrumentation/AskPatternTracingSpec.scala
index 9df67391..0387386c 100644
--- a/kamon-core/src/test/scala/kamon/trace/instrumentation/AskPatternTracingSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/instrumentation/AskPatternTracingSpec.scala
@@ -22,32 +22,30 @@ import akka.event.Logging.Warning
 import scala.concurrent.duration._
 import akka.pattern.ask
 import akka.util.Timeout
-import kamon.trace.{ Trace, ContextAware }
+import kamon.trace.{TraceContextAware, TraceRecorder}
 import org.scalatest.OptionValues._
 
 class AskPatternTracingSpec extends TestKit(ActorSystem("ask-pattern-tracing-spec")) with WordSpecLike with Matchers {
 
   "the AskPatternTracing" should {
-    "log a warning with a stack trace and TraceContext taken from the moment the ask was triggered" in new TraceContextFixture {
+    "log a warning with a stack trace and TraceContext taken from the moment the ask was triggered" in {
       implicit val ec = system.dispatcher
       implicit val timeout = Timeout(10 milliseconds)
       val noReply = system.actorOf(Props[NoReply])
       system.eventStream.subscribe(testActor, classOf[Warning])
 
-      within(500 milliseconds) {
-        val initialCtx = Trace.withContext(testTraceContext) {
-          noReply ? "hello"
-          Trace.context()
-        }
-
-        val warn = expectMsgPF() {
-          case warn: Warning if warn.message.toString.contains("Timeout triggered for ask pattern") ⇒ warn
-        }
-        val capturedCtx = warn.asInstanceOf[ContextAware].traceContext
+      val testTraceContext = TraceRecorder.withNewTraceContext("ask-timeout-warning") {
+        noReply ? "hello"
+        TraceRecorder.currentContext
+      }
 
-        capturedCtx should be('defined)
-        capturedCtx should equal(initialCtx)
+      val warn = expectMsgPF() {
+        case warn: Warning if warn.message.toString.contains("Timeout triggered for ask pattern") ⇒ warn
       }
+      val capturedCtx = warn.asInstanceOf[TraceContextAware].traceContext
+
+      capturedCtx should be('defined)
+      capturedCtx should equal(testTraceContext)
     }
   }
 }
diff --git a/kamon-core/src/test/scala/kamon/trace/instrumentation/FutureTracingSpec.scala b/kamon-core/src/test/scala/kamon/trace/instrumentation/FutureTracingSpec.scala
index a5554836..e6797148 100644
--- a/kamon-core/src/test/scala/kamon/trace/instrumentation/FutureTracingSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/instrumentation/FutureTracingSpec.scala
@@ -15,42 +15,42 @@
  * ========================================================== */
 package kamon.trace.instrumentation
 
-import scala.concurrent.{ ExecutionContext, Await, Promise, Future }
-import org.scalatest.{ Matchers, OptionValues, WordSpec }
+import scala.concurrent.{ ExecutionContext, Future }
+import org.scalatest.{ Matchers, OptionValues, WordSpecLike }
 import org.scalatest.concurrent.{ ScalaFutures, PatienceConfiguration }
-import java.util.UUID
-import scala.util.{ Random, Success }
-import scala.concurrent.duration._
-import java.util.concurrent.TimeUnit
-import akka.actor.{ Actor, ActorSystem }
-import kamon.trace.{ Trace, TraceContext }
+import kamon.trace.TraceRecorder
+import akka.testkit.TestKit
+import akka.actor.ActorSystem
 
-class FutureTracingSpec extends WordSpec with Matchers with ScalaFutures with PatienceConfiguration with OptionValues {
+class FutureTracingSpec extends TestKit(ActorSystem("actor-message-passing-tracing-spec")) with WordSpecLike with Matchers
+  with ScalaFutures with PatienceConfiguration with OptionValues {
 
-  implicit val execContext = ExecutionContext.Implicits.global
+  implicit val execContext = system.dispatcher
 
   "a Future created with FutureTracing" should {
     "capture the TraceContext available when created" which {
-      "must be available when executing the future's body" in new TraceContextFixture {
-        var future: Future[Option[TraceContext]] = _
+      "must be available when executing the future's body" in {
 
-        Trace.withContext(testTraceContext) {
-          future = Future(Trace.context)
+        val (future, testTraceContext) = TraceRecorder.withNewTraceContext("future-body") {
+          val future = Future(TraceRecorder.currentContext)
+
+          (future, TraceRecorder.currentContext)
         }
 
         whenReady(future)(ctxInFuture ⇒
           ctxInFuture should equal(testTraceContext))
       }
 
-      "must be available when executing callbacks on the future" in new TraceContextFixture {
-        var future: Future[Option[TraceContext]] = _
+      "must be available when executing callbacks on the future" in  {
 
-        Trace.withContext(testTraceContext) {
-          future = Future("Hello Kamon!")
+        val (future, testTraceContext) = TraceRecorder.withNewTraceContext("future-body") {
+          val future = Future("Hello Kamon!")
             // The TraceContext is expected to be available during all intermediate processing.
             .map(_.length)
             .flatMap(len ⇒ Future(len.toString))
-            .map(s ⇒ Trace.context())
+            .map(s ⇒ TraceRecorder.currentContext)
+
+          (future, TraceRecorder.currentContext)
         }
 
         whenReady(future)(ctxInFuture ⇒
diff --git a/kamon-core/src/test/scala/kamon/trace/instrumentation/TraceContextFixture.scala b/kamon-core/src/test/scala/kamon/trace/instrumentation/TraceContextFixture.scala
deleted file mode 100644
index 2df95d09..00000000
--- a/kamon-core/src/test/scala/kamon/trace/instrumentation/TraceContextFixture.scala
+++ /dev/null
@@ -1,10 +0,0 @@
-package kamon.trace.instrumentation
-
-import scala.util.Random
-import kamon.trace.TraceContext
-import akka.actor.Actor
-
-trait TraceContextFixture {
-  val random = new Random(System.nanoTime)
-  val testTraceContext = Some(TraceContext(Actor.noSender, random.nextInt, "test", "test-1"))
-}
\ No newline at end of file