finagle-base-http: Add support for new "b3" tracing header

Problem

OpenZipkin has created a new set of header to carry
trace information. These headers are not compatible
with the old ones and are needed to interoperate
with services that use the new version of the header.

Solution

If a "b3" header exists in the request this extracts
all the new information and populates finagles
tracing header with the existing values.

#749
Signed-off-by: Ian Bennett <ibennett@twitter.com>

Differential Revision: https://phabricator.twitter.biz/D334419

CHANGELOG.rst

@@ -7,6 +7,11 @@ Note that ``PHAB_ID=#`` and ``RB_ID=#`` correspond to associated messages in com
 Unreleased
 ----------
 
+New Features
+~~~~~~~~~~~~
+
+* finagle-base-http: Add support for new "b3" tracing header. ``PHAB_ID=D334419``
+
 Runtime Behavior Changes
 ~~~~~~~~~~~~~~~~~~~~~~~~
 

finagle-base-http/src/main/scala/com/twitter/finagle/http/HttpTracing.scala

@@ -8,6 +8,7 @@ object HttpTracing {
    * See [[headers()]] for Java compatibility.
    */
   object Header {
+    val TraceContext = "b3"
     val TraceId = "X-B3-TraceId"
     val SpanId = "X-B3-SpanId"
     val ParentSpanId = "X-B3-ParentSpanId"

finagle-base-http/src/main/scala/com/twitter/finagle/http/TraceInfo.scala

@@ -11,7 +11,54 @@ private object TraceInfo {
     while (iter.hasNext) headers -= iter.next()
   }
 
+  def convertB3Trace(request: Request) =
+    if (request.headerMap.contains(Header.TraceContext)) {
+      def handleSampled(headers: HeaderMap, value: String): Unit =
+        value match {
+          case "0" => headers.set(Header.Flags, "0")
+          case "d" => headers.set(Header.Flags, "1")
+          case "1" => headers.set(Header.Sampled, "1")
+          case _ => ()
+        }
+      def handleTraceAndSpanIds(headers: HeaderMap, a: String, b: String): Unit = {
+        headers.set(Header.TraceId, a)
+        headers.set(Header.SpanId, b)
+      }
+
+      val _ = request.headerMap.get(Header.TraceContext).map(_.split("-").toList) match {
+        case Some(a) =>
+          val headers = request.headerMap
+          a.size match {
+            case 0 =>
+              // bogus
+              ()
+            case 1 =>
+              // either debug flag or sampled
+              handleSampled(headers, a(0))
+            case 2 =>
+              // this is required to be traceId, spanId
+              handleTraceAndSpanIds(headers, a(0), a(1))
+            case 3 =>
+              handleTraceAndSpanIds(headers, a(0), a(1))
+              handleSampled(headers, a(2))
+            case 4 =>
+              handleTraceAndSpanIds(headers, a(0), a(1))
+              handleSampled(headers, a(2))
+              headers.set(Header.ParentSpanId, a(3))
+          }
+        case None =>
+          ()
+      }
+      request.headerMap -= "b3"
+    }
+
   def letTraceIdFromRequestHeaders[R](request: Request)(f: => R): R = {
+    // rather than rewrite all this to handle reading and writing the
+    // new b3 trace header format this code sets up the request to
+    // allow the existing code to consume, but not produce, b3 header
+    // traces.
+    convertB3Trace(request)
+
     val id =
       if (Header.hasAllRequired(request.headerMap)) {
         val spanId = SpanId.fromString(request.headerMap(Header.SpanId))

finagle-base-http/src/test/scala/com/twitter/finagle/http/TraceInfoTest.scala

@@ -84,4 +84,34 @@ class TraceInfoTest extends FunSuite {
       )
     )
   }
+
+  test("b3-header is parsed into a proper TraceId, most basic header") {
+    val req = Request(Method.Get, "/")
+    req.headerMap.put("b3", "0000000000000abc-0000000000000def")
+    TraceInfo.convertB3Trace(req)
+    assert(req.headerMap.keys == Set("X-B3-TraceId", "X-B3-SpanId"))
+  }
+
+  test("b3 header, just sampled/debug flag") {
+    val req = Request(Method.Get, "/")
+    // flags == 0
+    req.headerMap.put("b3", "0")
+    TraceInfo.convertB3Trace(req)
+    assert(req.headerMap.keys == Set("X-B3-Flags"))
+
+    // debug == d
+    req.headerMap.clear
+    req.headerMap.put("b3", "1")
+    TraceInfo.convertB3Trace(req)
+
+    assert(req.headerMap.keys == Set("X-B3-Sampled"))
+  }
+
+  test("b3 header with all the fields") {
+    val req = Request(Method.Get, "/")
+    req.headerMap.put("b3", "0000000000000abc-0000000000000def-1-0000000000000def")
+    TraceInfo.convertB3Trace(req)
+    assert(
+      req.headerMap.keys == Set("X-B3-SpanId", "X-B3-Sampled", "X-B3-ParentSpanId", "X-B3-TraceId"))
+  }
 }