Merge pull request #560 from ivantopo/trace-identifiers-in-response

include trace identifiers in HTTP responses, fixes #558

4 files changed, 48 insertions, 6 deletions

diff --git a/kamon-core-tests/src/test/resources/reference.conf b/kamon-core-tests/src/test/resources/reference.conf
index c092af08..60852722 100644
--- a/kamon-core-tests/src/test/resources/reference.conf
+++ b/kamon-core-tests/src/test/resources/reference.conf
@@ -21,8 +21,14 @@ kamon {
   instrumentation {
     http-server {
       default {
-        tracing.preferred-trace-id-tag = "correlation-id"
-        tracing.tags.from-context.peer = span
+        tracing {
+          preferred-trace-id-tag = "correlation-id"
+          tags.from-context.peer = span
+          response-headers {
+            trace-id = "x-trace-id"
+            span-id = "x-span-id"
+          }
+        }
       }
 
       no-span-metrics {
diff --git a/kamon-core-tests/src/test/scala/kamon/instrumentation/HttpServerInstrumentationSpec.scala b/kamon-core-tests/src/test/scala/kamon/instrumentation/HttpServerInstrumentationSpec.scala
index c3c5f131..62eae45b 100644
--- a/kamon-core-tests/src/test/scala/kamon/instrumentation/HttpServerInstrumentationSpec.scala
+++ b/kamon-core-tests/src/test/scala/kamon/instrumentation/HttpServerInstrumentationSpec.scala
@@ -220,6 +220,18 @@ class HttpServerInstrumentationSpec extends WordSpec with Matchers with SpanInsp
         val span = inspect(handler.span)
         span.tag("peer").value shouldBe "superservice"
       }
+
+      "write trace identifiers on the responses" in {
+        val handler = httpServer().receive(fakeRequest("http://localhost:8080/", "/", "GET", Map(
+          "x-correlation-id" -> "0011223344556677"
+        )))
+
+        val responseHeaders = mutable.Map.empty[String, String]
+        handler.send(fakeResponse(200, responseHeaders), handler.context)
+
+        responseHeaders.get("x-trace-id").value shouldBe "0011223344556677"
+        responseHeaders.get("x-span-id") shouldBe defined
+      }
     }
 
     "all capabilities are disabled" should {
diff --git a/kamon-core/src/main/resources/reference.conf b/kamon-core/src/main/resources/reference.conf
index 0fb3ce0a..c6d7a6e8 100644
--- a/kamon-core/src/main/resources/reference.conf
+++ b/kamon-core/src/main/resources/reference.conf
@@ -203,6 +203,7 @@ kamon {
           # Specify mappings between Context keys and the Propagation.EntryReader[HeaderReader] implementation in charge
           # of reading them from the incoming HTTP request into the Context.
           incoming {
+
             # kamon.trace.SpanPropagation$B3 for default header format or kamon.trace.SpanPropagation$B3Simple for 'b3 single' header format.
             span = "kamon.trace.SpanPropagation$B3"
           }
@@ -210,6 +211,7 @@ kamon {
           # Specify mappings betwen Context keys and the Propagation.EntryWriter[HeaderWriter] implementation in charge
           # of writing them to outgoing HTTP requests.
           outgoing {
+
             # kamon.trace.SpanPropagation$B3 for default header format or kamon.trace.SpanPropagation$B3Simple for 'b3 single' header format.
             span = "kamon.trace.SpanPropagation$B3"
           }
@@ -218,6 +220,7 @@ kamon {
     }
 
     binary {
+
       # Default HTTP propagation. Unless specified otherwise, all instrumentation will use the configuration on
       # this section for HTTP context propagation.
       #
@@ -338,6 +341,18 @@ kamon {
             }
           }
 
+          # Controls writing trace and span identifiers to HTTP response headers sent by the instrumented servers. The
+          # configuration can be set to either "none" to disable writing the identifiers on the response headers or to
+          # the header name to be used when writing the identifiers.
+          response-headers {
+
+            # HTTP response header name for the trace identifier, or "none" to disable it.
+            trace-id = "trace-id"
+
+            # HTTP response header name for the server span identifier, or "none" to disable it.
+            span-id = none
+          }
+
           # Custom mappings between routes and operation names.
           operations {
 
diff --git a/kamon-core/src/main/scala/kamon/instrumentation/HttpServer.scala b/kamon-core/src/main/scala/kamon/instrumentation/HttpServer.scala
index 72828424..659da8aa 100644
--- a/kamon-core/src/main/scala/kamon/instrumentation/HttpServer.scala
+++ b/kamon-core/src/main/scala/kamon/instrumentation/HttpServer.scala
@@ -294,14 +294,15 @@ object HttpServer {
             case TagMode.Off =>
           }
 
-          if(settings.enableContextPropagation) {
-            _propagation.write(context, response)
-          }
-
           _metrics.foreach { httpServerMetrics =>
             httpServerMetrics.countCompletedRequest(response.statusCode)
           }
 
+          if(span.nonEmpty()) {
+            settings.traceIDResponseHeader.foreach(traceIDHeader => response.write(traceIDHeader, span.context().traceID.string))
+            settings.spanIDResponseHeader.foreach(spanIDHeader => response.write(spanIDHeader, span.context().spanID.string))
+          }
+
           addResponseTag("http.status_code", response.statusCode.toString, settings.statusCodeTagMode)
           response.build()
         }
@@ -390,6 +391,8 @@ object HttpServer {
     methodTagMode: TagMode,
     statusCodeTagMode: TagMode,
     contextTags: Map[String, TagMode],
+    traceIDResponseHeader: Option[String],
+    spanIDResponseHeader: Option[String],
     unhandledOperationName: String,
     operationMappings: Map[GlobPathFilter, String]
   )
@@ -410,6 +413,7 @@ object HttpServer {
     }
 
     def from(config: Config): Settings = {
+      def optionalString(value: String): Option[String] = if(value.equalsIgnoreCase("none")) None else Some(value)
 
       // Context propagation settings
       val enablePropagation = config.getBoolean("propagation.enabled")
@@ -429,6 +433,9 @@ object HttpServer {
         case (tagName, mode) => (tagName, TagMode.from(mode))
       }
 
+      val traceIDResponseHeader = optionalString(config.getString("tracing.response-headers.trace-id"))
+      val spanIDResponseHeader = optionalString(config.getString("tracing.response-headers.span-id"))
+
       val unhandledOperationName = config.getString("tracing.operations.unhandled")
       val operationMappings = config.getConfig("tracing.operations.mappings").pairs.map {
         case (pattern, operationName) => (new GlobPathFilter(pattern), operationName)
@@ -445,6 +452,8 @@ object HttpServer {
         methodTagMode,
         statusCodeTagMode,
         contextTags,
+        traceIDResponseHeader,
+        spanIDResponseHeader,
         unhandledOperationName,
         operationMappings
       )