Supports reading multiple spans per Kafka message

zipkin-receiver-kafka/src/main/scala/com/twitter/zipkin/receiver/kafka/KafkaSpanReceiver.scala

@@ -50,7 +50,7 @@ trait KafkaSpanReceiverFactory { self: App =>
     process: Seq[ThriftSpan] => Future[Unit],
     stats: StatsReceiver = DefaultStatsReceiver.scope("KafkaSpanReceiver"),
     keyDecoder: Decoder[T] = KafkaProcessor.defaultKeyDecoder,
-    valueDecoder: KafkaProcessor.KafkaDecoder = new SpanCodec()
+    valueDecoder: KafkaProcessor.KafkaDecoder = new SpanDecoder()
   ): SpanReceiver = new SpanReceiver {
 
 

zipkin-receiver-kafka/src/main/scala/com/twitter/zipkin/receiver/kafka/SpanDecoder.scala

@@ -0,0 +1,23 @@
+package com.twitter.zipkin.receiver.kafka
+
+import com.twitter.scrooge.BinaryThriftStructSerializer
+import com.twitter.zipkin.conversions.thrift
+import com.twitter.zipkin.thriftscala.{Span => ThriftSpan}
+import org.apache.thrift.protocol.TType
+
+class SpanDecoder extends KafkaProcessor.KafkaDecoder {
+  val deserializer = new BinaryThriftStructSerializer[ThriftSpan] {
+    def codec = ThriftSpan
+  }
+
+  // Given the thrift encoding is TBinaryProtocol..
+  // .. When serializing a Span (Struct), the first byte will be the type of a field
+  // .. When serializing a List[ThriftSpan], the first byte is the member type, TType.STRUCT
+  // Span has no STRUCT fields: we assume that if the first byte is TType.STRUCT is a list.
+  def fromBytes(bytes: Array[Byte]) =
+    if (bytes(0) == TType.STRUCT) {
+      thrift.thriftListToThriftSpans(bytes)
+    } else {
+      List(deserializer.fromBytes(bytes))
+    }
+}

zipkin-receiver-kafka/src/test/scala/com/twitter/zipkin/receiver/kafka/KafkaProcessorSpec.scala

@@ -3,11 +3,14 @@ package com.twitter.zipkin.receiver.kafka
 import java.util.concurrent.LinkedBlockingQueue
 
 import com.github.charithe.kafka.KafkaJunitRule
+import com.twitter.io.Buf
 import com.twitter.util.{Await, Future, Promise}
 import com.twitter.zipkin.common._
 import com.twitter.zipkin.conversions.thrift._
 import com.twitter.zipkin.thriftscala.{Span => ThriftSpan}
 import kafka.producer._
+import org.apache.thrift.protocol.{TType, TList, TBinaryProtocol}
+import org.apache.thrift.transport.TMemoryBuffer
 import org.junit.{ClassRule, Test}
 import org.scalatest.junit.JUnitSuite
 
@@ -26,7 +29,7 @@ class KafkaProcessorSpec extends JUnitSuite {
   val annotation = Annotation(1, "value", Some(Endpoint(1, 2, "service")))
   // Intentionally leaving timestamp and duration unset, as legacy instrumentation don't set this.
   val span = Span(1234, "methodname", 4567, annotations = List(annotation))
-  val codec = new SpanCodec()
+  val codec = new SpanDecoder()
 
   @Test def messageWithSingleSpan() {
     val topic = "single_span"
@@ -38,14 +41,33 @@ class KafkaProcessorSpec extends JUnitSuite {
     }, codec, codec)
 
     val producer = new Producer[Array[Byte], Array[Byte]](kafkaRule.producerConfigWithDefaultEncoder())
-    producer.send(new KeyedMessage(topic, codec.encode(span)))
+    producer.send(new KeyedMessage(topic, encode(span)))
     producer.close()
 
     assert(Await.result(recvdSpan) == Seq(span.toThrift))
 
     Await.result(service.close())
   }
 
+  @Test def messageWithMultipleSpans() {
+    val topic = "multiple_spans"
+    val recvdSpan = new Promise[Seq[ThriftSpan]]
+
+    val service = KafkaProcessor(Map(topic -> 1), kafkaRule.consumerConfig(), { s =>
+      recvdSpan.setValue(s)
+      Future.value(true)
+    }, codec, codec)
+
+    val producer = new Producer[Array[Byte], Array[Byte]](kafkaRule.producerConfigWithDefaultEncoder())
+    producer.send(new KeyedMessage(topic, encode(Seq(span, span)))) // 2 spans in one message
+    producer.close()
+
+    // make sure we decoded both spans from the same message
+    assert(Await.result(recvdSpan) == Seq(span.toThrift, span.toThrift))
+
+    Await.result(service.close())
+  }
+
   @Test def skipsMalformedData() {
     val topic = "malformed"
     val recvdSpans = new LinkedBlockingQueue[Seq[ThriftSpan]](3)
@@ -57,14 +79,32 @@ class KafkaProcessorSpec extends JUnitSuite {
 
     val producer = new Producer[Array[Byte], Array[Byte]](kafkaRule.producerConfigWithDefaultEncoder())
 
-    producer.send(new KeyedMessage(topic, codec.encode(span)))
+    producer.send(new KeyedMessage(topic, encode(span)))
     producer.send(new KeyedMessage(topic, "malformed".getBytes()))
-    producer.send(new KeyedMessage(topic, codec.encode(span)))
+    producer.send(new KeyedMessage(topic, encode(span)))
     producer.close()
 
     for (elem <- 1 until 2)
       assert(recvdSpans.take() == Seq(span.toThrift))
 
     Await.result(service.close())
   }
+
+  def encode(span: Span) = {
+    val transport = new TMemoryBuffer(0)
+    val oproto = new TBinaryProtocol(transport)
+    val tspan = spanToThriftSpan(span)
+    tspan.toThrift.write(oproto)
+    transport.getArray()
+  }
+
+  def encode(spans: Seq[Span]) = {
+    // serialize all spans as a thrift list
+    val transport = new TMemoryBuffer(0)
+    val oproto = new TBinaryProtocol(transport)
+    oproto.writeListBegin(new TList(TType.STRUCT, spans.size))
+    spans.map(spanToThriftSpan).foreach(_.toThrift.write(oproto))
+    oproto.writeListEnd()
+    transport.getArray()
+  }
 }

zipkin-scrooge/src/main/scala/com/twitter/zipkin/conversions/thrift.scala

@@ -132,21 +132,23 @@ object thrift {
   implicit def spanToThriftSpan(s: Span) = new ThriftSpan(s)
   implicit def thriftSpanToSpan(s: thriftscala.Span) = new WrappedSpan(s)
 
-  def thriftListToSpans(bytes: Array[Byte]) = {
+  def thriftListToThriftSpans(bytes: Array[Byte]) = {
     val proto = new TBinaryProtocol(TArrayByteTransport(bytes))
     val _list = proto.readListBegin()
     if (_list.size > 10000) {
       throw new IllegalArgumentException(_list.size + " > 10000: possibly malformed thrift")
     }
-    val result = new ArrayBuffer[Span](_list.size)
+    val result = new ArrayBuffer[thriftscala.Span](_list.size)
     for (i <- 1 to _list.size) {
-      val thrift = thriftscala.Span.decode(proto)
-      result += thriftSpanToSpan(thrift).toSpan
+      result += thriftscala.Span.decode(proto)
     }
     proto.readListEnd()
     result.toList
   }
 
+  def thriftListToSpans(bytes: Array[Byte]) =
+    thriftListToThriftSpans(bytes).map(thriftSpanToSpan(_).toSpan)
+
   class WrappedDependencyLink(dl: DependencyLink) {
     lazy val toThrift = thriftscala.DependencyLink(dl.parent, dl.child, dl.callCount)
   }