rename UOW to TraceToken and generate it by default when a Trace starts

5 files changed, 43 insertions, 15 deletions

diff --git a/kamon-spray/src/main/resources/reference.conf b/kamon-spray/src/main/resources/reference.conf
new file mode 100644
index 00000000..88bd8fb8
--- /dev/null
+++ b/kamon-spray/src/main/resources/reference.conf
@@ -0,0 +1,6 @@
+kamon {
+  spray {
+    include-trace-token-header = true
+    trace-token-header-name = "X-Trace-Token"
+  }
+}
\ No newline at end of file
diff --git a/kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala b/kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala
index 44db55e7..78a5b336 100644
--- a/kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala
+++ b/kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala
@@ -28,10 +28,10 @@ trait UowDirectives extends BasicDirectives {
     val uowHeader = request.headers.find(_.name == "X-UOW")
 
     val generatedUow = uowHeader.map(_.value).getOrElse(UowDirectives.newUow)
-    Trace.transformContext(_.copy(uow = generatedUow))
+    Trace.transformContext(_.copy(token = generatedUow))
     request
   }
-  def respondWithUow = mapHttpResponseHeaders(headers ⇒ Trace.context().map(ctx ⇒ RawHeader("X-UOW", ctx.uow) :: headers).getOrElse(headers))
+  def respondWithUow = mapHttpResponseHeaders(headers ⇒ Trace.context().map(ctx ⇒ RawHeader("X-UOW", ctx.token) :: headers).getOrElse(headers))
 }
 
 object UowDirectives {
diff --git a/kamon-spray/src/main/scala/spray/can/server/ServerRequestTracing.scala b/kamon-spray/src/main/scala/spray/can/server/ServerRequestTracing.scala
index 91dba87b..a7d48ac8 100644
--- a/kamon-spray/src/main/scala/spray/can/server/ServerRequestTracing.scala
+++ b/kamon-spray/src/main/scala/spray/can/server/ServerRequestTracing.scala
@@ -37,9 +37,16 @@ class ServerRequestTracing {
   @After("openRequestInit(openRequest, request)")
   def afterInit(openRequest: ContextAware, request: HttpRequest): Unit = {
     val system: ActorSystem = openRequest.asInstanceOf[OpenRequest].context.actorContext.system
+    val config = system.settings.config.getConfig("kamon.spray")
+
+    val token = if(config.getBoolean("include-trace-token-header")) {
+      val traceTokenHeader = config.getString("trace-token-header-name")
+      request.headers.find(_.name == traceTokenHeader).map(_.value)
+    } else None
+
     val defaultTraceName: String = request.method.value + ": " + request.uri.path
 
-    Trace.start(defaultTraceName)(system)
+    Trace.start(defaultTraceName, token)(system)
 
     // Necessary to force initialization of traceContext when initiating the request.
     openRequest.traceContext
diff --git a/kamon-spray/src/test/scala/kamon/spray/ClientRequestTracingSpec.scala b/kamon-spray/src/test/scala/kamon/spray/ClientRequestTracingSpec.scala
index c7ec2a5f..eae0c992 100644
--- a/kamon-spray/src/test/scala/kamon/spray/ClientRequestTracingSpec.scala
+++ b/kamon-spray/src/test/scala/kamon/spray/ClientRequestTracingSpec.scala
@@ -14,7 +14,7 @@ class ClientRequestTracingSpec extends TestKit(ActorSystem("server-request-traci
   "the client instrumentation" should {
     "record segments for a client http request" in {
 
-      Trace.start("record-segments")(system)
+      Trace.start("record-segments", None)(system)
 
       send {
         Get(s"http://127.0.0.1:$port/ok")
diff --git a/kamon-spray/src/test/scala/kamon/spray/ServerRequestTracingSpec.scala b/kamon-spray/src/test/scala/kamon/spray/ServerRequestTracingSpec.scala
index 44ed3baf..2a2865e1 100644
--- a/kamon-spray/src/test/scala/kamon/spray/ServerRequestTracingSpec.scala
+++ b/kamon-spray/src/test/scala/kamon/spray/ServerRequestTracingSpec.scala
@@ -19,25 +19,33 @@ import _root_.spray.httpx.RequestBuilding
 import _root_.spray.routing.SimpleRoutingApp
 import akka.testkit.TestKit
 import akka.actor.ActorSystem
-import org.scalatest.WordSpecLike
+import org.scalatest.{ Matchers, WordSpecLike }
 import scala.concurrent.Await
 import scala.concurrent.duration._
 import _root_.spray.client.pipelining._
 import akka.util.Timeout
 import kamon.trace.{ UowTrace, Trace }
 import kamon.Kamon
+import org.scalatest.concurrent.{ PatienceConfiguration, ScalaFutures }
+import spray.http.HttpHeaders.RawHeader
+import spray.http.HttpRequest
+import spray.http.HttpHeaders.Host
 
-class ServerRequestTracingSpec extends TestKit(ActorSystem("server-request-tracing-spec")) with WordSpecLike with RequestBuilding with TestServer {
+class ServerRequestTracingSpec extends TestKit(ActorSystem("server-request-tracing-spec")) with WordSpecLike with Matchers with RequestBuilding with ScalaFutures with PatienceConfiguration with TestServer {
 
   "the spray server request tracing instrumentation" should {
     "trace a request start/finish sequence when proper TraceContext is received" in {
-      send {
-        Get(s"http://127.0.0.1:$port/ok")
+      val response = send {
+        Get("/ok")
       }
 
       within(5 seconds) {
         fishForNamedTrace("ok")
       }
+
+      whenReady(response) { rsp ⇒
+        rsp.headers should contain(RawHeader("X-Trace-Token", ""))
+      }
     }
 
     "finish a request even if no TraceContext is received in the response" in {
@@ -60,6 +68,10 @@ class ServerRequestTracingSpec extends TestKit(ActorSystem("server-request-traci
       }
     }
   }
+  /*
+  - si no llega uow, crear una
+  - si llega con uow hay que propagarla
+   */
 
   def fishForNamedTrace(traceName: String) = fishForMessage() {
     case trace: UowTrace if trace.name.contains(traceName) ⇒ true
@@ -79,15 +91,18 @@ trait TestServer extends SimpleRoutingApp {
         path("ok") {
           complete("ok")
         } ~
-          path("clearcontext") {
-            complete {
-              println("The Context in the route is: " + Trace.context)
-              Trace.clear
-              "ok"
-            }
+        path("clearcontext") {
+          complete {
+            Trace.clear
+            "ok"
           }
+        }
       }), timeout.duration).localAddress.getPort
 
-  val send = sendReceive(system, system.dispatcher, timeout)
+  val send = includeHost("127.0.0.1", port) ~> sendReceive(system, system.dispatcher, timeout)
+
+  def includeHost(host: String, port: Int) = { request: HttpRequest ⇒
+    request.withEffectiveUri(port == 443, Host(host, port))
+  }
 
 }