Incorporate PR feedback

* Style cleanup
* Cleaner construction of 128 bit trace id
* Remove and comment on lack of legacy trace id support in TReq
* Loop on spanId=0 in Trace#nextId

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

@@ -1,6 +1,6 @@
 package com.twitter.finagle.http
 
-import com.twitter.finagle.tracing.{Flags, SpanId, Trace, TraceId}
+import com.twitter.finagle.tracing.{Flags, SpanId, Trace, TraceId, TraceId128}
 import java.lang.{Boolean => JBoolean, Long => JLong}
 
 private object TraceInfo {
@@ -19,7 +19,7 @@ private object TraceInfo {
         spanId match {
           case None => None
           case Some(sid) =>
-            val (highTraceId, lowTraceId) = TraceId.mk128BitTraceId(request.headerMap(Header.TraceId))
+            val trace128Bit = TraceId128(request.headerMap(Header.TraceId))
 
             val parentSpanId =
               if (request.headerMap.contains(Header.ParentSpanId))
@@ -37,7 +37,7 @@ private object TraceInfo {
             }
 
             val flags = getFlags(request)
-            Some(TraceId(lowTraceId, parentSpanId, sid, sampled, flags, highTraceId))
+            Some(TraceId(trace128Bit.low, parentSpanId, sid, sampled, flags, trace128Bit.high))
         }
       } else if (request.headerMap.contains(Header.Flags)) {
         // even if there are no id headers we want to get the debug flag

finagle-core/src/main/scala/com/twitter/finagle/tracing/Id.scala

@@ -7,10 +7,10 @@ import java.lang.{Boolean => JBool}
 import scala.util.control.NonFatal
 
 /**
-  * Defines trace identifiers.  Span IDs name a particular (unique)
-  * span, while TraceIds contain a span ID as well as context (parentId
-  * and traceId).
-  */
+ * Defines trace identifiers.  Span IDs name a particular (unique)
+ * span, while TraceIds contain a span ID as well as context (parentId
+ * and traceId).
+ */
 final class SpanId(val self: Long) extends Proxy {
   def toLong = self
 
@@ -26,7 +26,7 @@ object SpanId {
       val s = "%02x".format(bb)
       Array(s(0), s(1))
     }
-    ).toArray
+   ).toArray
 
   private def byteToChars(b: Byte): Array[Char] = lut(b + 128)
 
@@ -50,7 +50,7 @@ object SpanId {
     try {
       // Tolerates 128 bit X-B3-TraceId by reading the right-most 16 hex
       // characters (as opposed to overflowing a U64 and starting a new trace).
-      // For TraceId, prefer TraceId#make128BitTraceId.
+      // For TraceId, prefer TraceId128#apply.
       val length = spanId.length()
       val lower64Bits = if (length <= 16) spanId else spanId.substring(length - 16)
       Some(SpanId(new RichU64String(lower64Bits).toU64Long))
@@ -59,6 +59,29 @@ object SpanId {
     }
 }
 
+case class TraceId128(low: Option[SpanId], high: Option[SpanId])
+object TraceId128 {
+  val empty = TraceId128(None, None)
+
+  /**
+   * Extracts the high 64bits (if set and valid) and low 64bits (if valid) from a B3 TraceID's string representation.
+   *
+   * @param spanId A 64bit or 128bit Trace ID.
+   */
+  def apply(spanId : String) : TraceId128 = {
+    try {
+      val length = spanId.length()
+      val lower64Bits = if (length <= 16) spanId else spanId.substring(length - 16)
+      val low = Some(SpanId(new RichU64String(lower64Bits).toU64Long))
+
+      if (length == 32) TraceId128(low, Some(SpanId(new RichU64String(spanId.substring(0, 16)).toU64Long)))
+      else TraceId128(low, None)
+    } catch {
+      case NonFatal(_) => empty
+    }
+  }
+}
+
 object TraceId {
 
   /**
@@ -73,8 +96,8 @@ object TraceId {
     TraceId(traceId, parentId, spanId, sampled, Flags(), None)
 
   /**
-    * Creates a 64bit TraceID. See case class for more info.
-    */
+   * Creates a 64bit TraceID. See case class for more info.
+   */
   def apply(
     traceId: Option[SpanId],
     parentId: Option[SpanId],
@@ -97,20 +120,21 @@ object TraceId {
         traceId.flags.setFlag(Flags.SamplingKnown)
     }
 
-    val bytes = new Array[Byte](40)
-    ByteArrays.put64be(bytes, 0, traceId.spanId.self)
-    ByteArrays.put64be(bytes, 8, traceId.parentId.self)
-    ByteArrays.put64be(bytes, 16, traceId.traceId.self)
+    // For backward compatibility for TraceID: 40 bytes if 128bit, 32 bytes if 64bit
+    val bytes = new Array[Byte](if(traceId.traceIdHigh.isDefined) 40 else 32)
+    ByteArrays.put64be(bytes, 0, traceId.spanId.toLong)
+    ByteArrays.put64be(bytes, 8, traceId.parentId.toLong)
+    ByteArrays.put64be(bytes, 16, traceId.traceId.toLong)
     ByteArrays.put64be(bytes, 24, flags.toLong)
-    if(traceId.traceIdHigh.isDefined) ByteArrays.put64be(bytes, 32, traceId.traceIdHigh.get.self)
+    if(traceId.traceIdHigh.isDefined) ByteArrays.put64be(bytes, 32, traceId.traceIdHigh.get.toLong)
     bytes
   }
 
   /**
    * Deserialize a TraceId from an array of bytes.
+   * Allows for 64-bit or 128-bit trace identifiers.
    */
   def deserialize(bytes: Array[Byte]): Try[TraceId] = {
-    // Allows for 64-bit or 128-bit trace identifiers
     if (bytes.length != 32 && bytes.length != 40) {
       Throw(new IllegalArgumentException("Expected 32 or 40 bytes"))
     } else {
@@ -119,7 +143,7 @@ object TraceId {
       val trace64 = ByteArrays.get64be(bytes, 16)
       val flags64 = ByteArrays.get64be(bytes, 24)
 
-      val traceIdHigh = if(bytes.length == 40) ByteArrays.get64be(bytes, 32) else 0L
+      val traceIdHigh = if(bytes.length == 40) Some(SpanId(ByteArrays.get64be(bytes, 32))) else None
 
       val flags = Flags(flags64)
       val sampled = if (flags.isFlagSet(Flags.SamplingKnown)) {
@@ -132,35 +156,11 @@ object TraceId {
         SpanId(span64),
         sampled,
         flags,
-        Some(SpanId(traceIdHigh))
+        traceIdHigh
       )
       Return(traceId)
     }
   }
-
-  /**
-   * Extracts the high 64bits (if set) and low 64bits from a B3 TraceID's string representation.
-   *
-   * @param spanId A 64bit or 128bit Trace ID.
-   *
-   * @return An <code>(Option[SpanId], Option[SpanId])</code> of 64bit or 128bit TraceID. Structure is (high, low).
-   */
-  def mk128BitTraceId(spanId: String): (Option[SpanId], Option[SpanId]) = {
-    try {
-      val length = spanId.length()
-      val lower64Bits = if (length <= 16) spanId else spanId.substring(length - 16)
-      val upper64bits = if (length == 32) Some(spanId.substring(0, 16)) else None
-
-      upper64bits match {
-        case Some(high) =>
-          (Some(SpanId(new RichU64String(high).toU64Long)), Some(SpanId(new RichU64String(lower64Bits).toU64Long)))
-        case None =>
-          (None, Some(SpanId(new RichU64String(lower64Bits).toU64Long)))
-      }
-    } catch {
-      case NonFatal(_) => (None, None)
-    }
-  }
 }
 
 /**

finagle-core/src/main/scala/com/twitter/finagle/tracing/Trace.scala

@@ -6,7 +6,6 @@ import com.twitter.finagle.util.ByteArrays
 import com.twitter.io.Buf
 import com.twitter.util._
 import java.net.InetSocketAddress
-
 import scala.util.Random
 
 /**
@@ -145,7 +144,13 @@ object Trace {
    * Create a derived id from the current TraceId.
    */
   def nextId: TraceId = {
-    val spanId = SpanId(rng.nextLong())
+    var nextLong = rng.nextLong()
+    if (nextLong == 0) {
+      // NOTE: spanId of 0 is invalid, so guard against that here.
+      do nextLong = rng.nextLong() while (nextLong == 0)
+    }
+
+    val spanId = SpanId(nextLong)
     idOption match {
       case Some(id) =>
         TraceId(Some(id.traceId), Some(id.spanId), spanId, id.sampled, id.flags, id.traceIdHigh)

finagle-core/src/test/scala/com/twitter/finagle/tracing/IdTest.scala

@@ -63,23 +63,27 @@ class IdTest extends FunSuite {
   test("extract 64bit only ids") {
     val low = 5208512171318403364L
     val spanId = hex(low)
-    val (highBits, lowBits) = TraceId.mk128BitTraceId(spanId)
+    val traceId = TraceId128(spanId)
 
-    assert(highBits.isDefined == false)
-    assert(lowBits.isDefined)
-    assert(lowBits.get.self == low)
+    assert(traceId.high.isDefined == false)
+    assert(traceId.low.isDefined)
+    assert(traceId.low.get.self == low)
   }
 
   test("extract 128bit ids") {
     val low = 5208512171318403364L
     val high = 5060571933882717101L
     val spanId =  hex(high) + hex(low)
-    val (highBits, lowBits) = TraceId.mk128BitTraceId(spanId)
+    val traceId = TraceId128(spanId)
 
-    assert(highBits.isDefined)
-    assert(highBits.get.self == high)
-    assert(lowBits.isDefined)
-    assert(lowBits.get.self == low)
+    assert(traceId.high.isDefined)
+    assert(traceId.high.get.self == high)
+    assert(traceId.low.isDefined)
+    assert(traceId.low.get.self == low)
+  }
+
+  test("extract to TraceId128Bit.empty when span id is invalid") {
+    assert(TraceId128("invalid") == TraceId128.empty)
   }
 
   test("drops invalid traceIdHigh = 0L") {

finagle-mux/src/main/scala/com/twitter/finagle/mux/transport/Message.scala

@@ -168,6 +168,8 @@ private[finagle] object Message {
    *
    * Note, Treq messages are deprecated in favor of [[Tdispatch]] and will likely
    * be removed in a future version of mux.
+   *
+   * Treq does not support 128bit TraceId.
    */
   case class Treq(tag: Int, traceId: Option[TraceId], req: Buf) extends Message {
     import Treq._
@@ -190,8 +192,6 @@ private[finagle] object Message {
           hd.writeByte(1) // key 1 size
           hd.writeByte(traceId.flags.toLong.toByte)
 
-          // TODO: support traceIdHigh
-
           hd.owned()
 
         case None =>
@@ -478,7 +478,7 @@ private[finagle] object Message {
       // TODO: technically we should probably check for duplicate
       // keys, but for now, just pick the latest one.
       key match {
-        // TODO: support traceIdHigh
+        // NOTE: Treq is deprecated and therefore won't support 128bit TraceID. see Tdispatch/Rdispatch.
         case Treq.Keys.TraceId =>
           if (vsize != 24)
             throwBadMessageException(s"bad traceid size $vsize")