contrib/Shopify/sarama: add support for tracing the sarama kafka pack…

…age (#296)

contrib/Shopify/sarama/example_test.go

@@ -0,0 +1,77 @@
+package sarama_test
+
+import (
+	"log"
+
+	saramatrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/Shopify/sarama"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	sarama "gopkg.in/Shopify/sarama.v1"
+)
+
+func Example_asyncProducer() {
+	cfg := sarama.NewConfig()
+
+	producer, err := sarama.NewAsyncProducer([]string{"localhost:9092"}, cfg)
+	if err != nil {
+		panic(err)
+	}
+	defer producer.Close()
+
+	producer = saramatrace.WrapAsyncProducer(cfg, producer)
+
+	msg := &sarama.ProducerMessage{
+		Topic: "some-topic",
+		Value: sarama.StringEncoder("Hello World"),
+	}
+	producer.Input() <- msg
+}
+
+func Example_syncProducer() {
+	cfg := sarama.NewConfig()
+	cfg.Producer.Return.Successes = true
+
+	producer, err := sarama.NewSyncProducer([]string{"localhost:9092"}, cfg)
+	if err != nil {
+		panic(err)
+	}
+	defer producer.Close()
+
+	producer = saramatrace.WrapSyncProducer(cfg, producer)
+
+	msg := &sarama.ProducerMessage{
+		Topic: "some-topic",
+		Value: sarama.StringEncoder("Hello World"),
+	}
+	_, _, err = producer.SendMessage(msg)
+	if err != nil {
+		panic(err)
+	}
+}
+
+func Example_consumer() {
+	consumer, err := sarama.NewConsumer([]string{"localhost:9092"}, nil)
+	if err != nil {
+		panic(err)
+	}
+	defer consumer.Close()
+
+	consumer = saramatrace.WrapConsumer(consumer)
+
+	partitionConsumer, err := consumer.ConsumePartition("some-topic", 0, sarama.OffsetNewest)
+	if err != nil {
+		panic(err)
+	}
+	defer partitionConsumer.Close()
+
+	consumed := 0
+	for msg := range partitionConsumer.Messages() {
+		// if you want to use the kafka message as a parent span:
+		if spanctx, err := tracer.Extract(saramatrace.NewConsumerMessageCarrier(msg)); err == nil {
+			// you can create a span using ChildOf(spanctx)
+			_ = spanctx
+		}
+
+		log.Printf("Consumed message offset %d\n", msg.Offset)
+		consumed++
+	}
+}

contrib/Shopify/sarama/headers.go

@@ -0,0 +1,90 @@
+package sarama
+
+import (
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	sarama "gopkg.in/Shopify/sarama.v1"
+)
+
+// A ProducerMessageCarrier injects and extracts traces from a sarama.ProducerMessage.
+type ProducerMessageCarrier struct {
+	msg *sarama.ProducerMessage
+}
+
+var _ interface {
+	tracer.TextMapReader
+	tracer.TextMapWriter
+} = (*ProducerMessageCarrier)(nil)
+
+// ForeachKey iterates over every header.
+func (c ProducerMessageCarrier) ForeachKey(handler func(key, val string) error) error {
+	for _, h := range c.msg.Headers {
+		err := handler(string(h.Key), string(h.Value))
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// Set sets a header.
+func (c ProducerMessageCarrier) Set(key, val string) {
+	// ensure uniqueness of keys
+	for i := 0; i < len(c.msg.Headers); i++ {
+		if string(c.msg.Headers[i].Key) == key {
+			c.msg.Headers = append(c.msg.Headers[:i], c.msg.Headers[i+1:]...)
+			i--
+		}
+	}
+	c.msg.Headers = append(c.msg.Headers, sarama.RecordHeader{
+		Key:   []byte(key),
+		Value: []byte(val),
+	})
+}
+
+// NewProducerMessageCarrier creates a new ProducerMessageCarrier.
+func NewProducerMessageCarrier(msg *sarama.ProducerMessage) ProducerMessageCarrier {
+	return ProducerMessageCarrier{msg}
+}
+
+// A ConsumerMessageCarrier injects and extracts traces from a sarama.ConsumerMessage.
+type ConsumerMessageCarrier struct {
+	msg *sarama.ConsumerMessage
+}
+
+var _ interface {
+	tracer.TextMapReader
+	tracer.TextMapWriter
+} = (*ConsumerMessageCarrier)(nil)
+
+// NewConsumerMessageCarrier creates a new ConsumerMessageCarrier.
+func NewConsumerMessageCarrier(msg *sarama.ConsumerMessage) ConsumerMessageCarrier {
+	return ConsumerMessageCarrier{msg}
+}
+
+// ForeachKey iterates over every header.
+func (c ConsumerMessageCarrier) ForeachKey(handler func(key, val string) error) error {
+	for _, h := range c.msg.Headers {
+		if h != nil {
+			err := handler(string(h.Key), string(h.Value))
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+// Set sets a header.
+func (c ConsumerMessageCarrier) Set(key, val string) {
+	// ensure uniqueness of keys
+	for i := 0; i < len(c.msg.Headers); i++ {
+		if c.msg.Headers[i] != nil && string(c.msg.Headers[i].Key) == key {
+			c.msg.Headers = append(c.msg.Headers[:i], c.msg.Headers[i+1:]...)
+			i--
+		}
+	}
+	c.msg.Headers = append(c.msg.Headers, &sarama.RecordHeader{
+		Key:   []byte(key),
+		Value: []byte(val),
+	})
+}

contrib/Shopify/sarama/option.go

@@ -0,0 +1,19 @@
+package sarama
+
+type config struct {
+	serviceName string
+}
+
+func defaults(cfg *config) {
+	cfg.serviceName = "kafka"
+}
+
+// An Option is used to customize the config for the sarama tracer.
+type Option func(cfg *config)
+
+// WithServiceName sets the given service name for the intercepted client.
+func WithServiceName(name string) Option {
+	return func(cfg *config) {
+		cfg.serviceName = name
+	}
+}