From 003af1dd8ef1e3b4a8f1a9eff56a3c15c01cb2a9 Mon Sep 17 00:00:00 2001
From: Ivan Topolnjak <ivantopo@gmail.com>
Date: Mon, 27 Jan 2014 09:11:32 -0300
Subject: rename UOW to TraceToken and generate it by default when a Trace
 starts

---
 kamon-spray/src/main/resources/reference.conf                    | 6 ++++++
 kamon-spray/src/main/scala/kamon/spray/UowDirectives.scala       | 4 ++--
 .../src/main/scala/spray/can/server/ServerRequestTracing.scala   | 9 ++++++++-
 3 files changed, 16 insertions(+), 3 deletions(-)
 create mode 100644 kamon-spray/src/main/resources/reference.conf

(limited to 'kamon-spray/src/main')

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
-- 
cgit v1.2.3