contrib/bradfitz/gomemcache: add memcache tracer (DataDog#291)

.circleci/config.yml

@@ -30,6 +30,7 @@ jobs:
         DD_APM_ENABLED: "true"
         DD_BIND_HOST: "0.0.0.0"
         DD_API_KEY: invalid_key_but_this_is_fine
+    - image: memcached:1.5.9
 
     steps:
     - checkout

contrib/bradfitz/gomemcache/memcache/example_test.go

@@ -0,0 +1,22 @@
+package memcache_test
+
+import (
+	"context"
+
+	"github.com/bradfitz/gomemcache/memcache"
+	memcachetrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/bradfitz/gomemcache/memcache"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+)
+
+func Example() {
+	span, ctx := tracer.StartSpanFromContext(context.Background(), "parent.request",
+		tracer.ServiceName("web"),
+		tracer.ResourceName("/home"),
+	)
+	defer span.Finish()
+
+	mc := memcachetrace.WrapClient(memcache.New("127.0.0.1:11211"))
+	// you can use WithContext to set the parent span
+	mc.WithContext(ctx).Set(&memcache.Item{Key: "my key", Value: []byte("my value")})
+
+}

contrib/bradfitz/gomemcache/memcache/memcache.go

@@ -0,0 +1,162 @@
+// Package memcache provides functions to trace the bradfitz/gomemcache package (https://github.com/bradfitz/gomemcache).
+//
+// `WrapClient` will wrap a memcache `Client` and return a new struct with all
+// the same methods, so should be seamless for existing applications. It also
+// has an additional `WithContext` method which can be used to connect a span
+// to an existing trace.
+package memcache // import "gopkg.in/DataDog/dd-trace-go.v1/contrib/bradfitz/gomemcache/memcache"
+
+import (
+	"context"
+
+	"github.com/bradfitz/gomemcache/memcache"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+)
+
+// WrapClient wraps a memcache.Client so that all requests are traced using the
+// default tracer with the service name "memcached".
+func WrapClient(client *memcache.Client, opts ...ClientOption) *Client {
+	cfg := new(clientConfig)
+	defaults(cfg)
+	for _, opt := range opts {
+		opt(cfg)
+	}
+	return &Client{
+		Client:  client,
+		cfg:     cfg,
+		context: context.Background(),
+	}
+}
+
+// A Client is used to trace requests to the memcached server.
+type Client struct {
+	*memcache.Client
+	cfg     *clientConfig
+	context context.Context
+}
+
+// WithContext creates a copy of the Client with the given context.
+func (c *Client) WithContext(ctx context.Context) *Client {
+	// the existing memcache client doesn't support context, but may in the
+	// future, so we do a runtime check to detect this
+	mc := c.Client
+	if wc, ok := (interface{})(c.Client).(interface {
+		WithContext(context.Context) *memcache.Client
+	}); ok {
+		mc = wc.WithContext(ctx)
+	}
+	return &Client{
+		Client:  mc,
+		cfg:     c.cfg,
+		context: ctx,
+	}
+}
+
+// startSpan starts a span from the context set with WithContext.
+func (c *Client) startSpan(resourceName string) ddtrace.Span {
+	span, _ := tracer.StartSpanFromContext(c.context, operationName,
+		tracer.SpanType(ext.SpanTypeMemcached),
+		tracer.ServiceName(c.cfg.serviceName),
+		tracer.ResourceName(resourceName))
+	return span
+}
+
+// wrapped methods:
+
+// Add invokes and traces Client.Add.
+func (c *Client) Add(item *memcache.Item) error {
+	span := c.startSpan("Add")
+	err := c.Client.Add(item)
+	span.Finish(tracer.WithError(err))
+	return err
+}
+
+// CompareAndSwap invokes and traces Client.CompareAndSwap.
+func (c *Client) CompareAndSwap(item *memcache.Item) error {
+	span := c.startSpan("CompareAndSwap")
+	err := c.Client.CompareAndSwap(item)
+	span.Finish(tracer.WithError(err))
+	return err
+}
+
+// Decrement invokes and traces Client.Decrement.
+func (c *Client) Decrement(key string, delta uint64) (newValue uint64, err error) {
+	span := c.startSpan("Decrement")
+	newValue, err = c.Client.Decrement(key, delta)
+	span.Finish(tracer.WithError(err))
+	return newValue, err
+}
+
+// Delete invokes and traces Client.Delete.
+func (c *Client) Delete(key string) error {
+	span := c.startSpan("Delete")
+	err := c.Client.Delete(key)
+	span.Finish(tracer.WithError(err))
+	return err
+}
+
+// DeleteAll invokes and traces Client.DeleteAll.
+func (c *Client) DeleteAll() error {
+	span := c.startSpan("DeleteAll")
+	err := c.Client.DeleteAll()
+	span.Finish(tracer.WithError(err))
+	return err
+}
+
+// FlushAll invokes and traces Client.FlushAll.
+func (c *Client) FlushAll() error {
+	span := c.startSpan("FlushAll")
+	err := c.Client.FlushAll()
+	span.Finish(tracer.WithError(err))
+	return err
+}
+
+// Get invokes and traces Client.Get.
+func (c *Client) Get(key string) (item *memcache.Item, err error) {
+	span := c.startSpan("Get")
+	item, err = c.Client.Get(key)
+	span.Finish(tracer.WithError(err))
+	return item, err
+}
+
+// GetMulti invokes and traces Client.GetMulti.
+func (c *Client) GetMulti(keys []string) (map[string]*memcache.Item, error) {
+	span := c.startSpan("GetMulti")
+	items, err := c.Client.GetMulti(keys)
+	span.Finish(tracer.WithError(err))
+	return items, err
+}
+
+// Increment invokes and traces Client.Increment.
+func (c *Client) Increment(key string, delta uint64) (newValue uint64, err error) {
+	span := c.startSpan("Increment")
+	newValue, err = c.Client.Increment(key, delta)
+	span.Finish(tracer.WithError(err))
+	return newValue, err
+}
+
+// Replace invokes and traces Client.Replace.
+func (c *Client) Replace(item *memcache.Item) error {
+	span := c.startSpan("Replace")
+	err := c.Client.Replace(item)
+	span.Finish(tracer.WithError(err))
+	return err
+}
+
+// Set invokes and traces Client.Set.
+func (c *Client) Set(item *memcache.Item) error {
+	span := c.startSpan("Set")
+	err := c.Client.Set(item)
+	span.Finish(tracer.WithError(err))
+	return err
+}
+
+// Touch invokes and traces Client.Touch.
+func (c *Client) Touch(key string, seconds int32) error {
+	span := c.startSpan("Touch")
+	err := c.Client.Touch(key, seconds)
+	span.Finish(tracer.WithError(err))
+	return err
+}

contrib/bradfitz/gomemcache/memcache/memcache_test.go

@@ -0,0 +1,146 @@
+package memcache
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"net"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/bradfitz/gomemcache/memcache"
+	"github.com/stretchr/testify/assert"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+)
+
+func TestMemcache(t *testing.T) {
+	li := makeFakeServer(t)
+	defer li.Close()
+
+	testMemcache(t, li.Addr().String())
+}
+
+func TestMemcacheIntegration(t *testing.T) {
+	if _, ok := os.LookupEnv("INTEGRATION"); !ok {
+		t.Skip("to enable integration test, set the INTEGRATION environment variable")
+	}
+
+	testMemcache(t, "localhost:11211")
+}
+
+func testMemcache(t *testing.T, addr string) {
+	client := WrapClient(memcache.New(addr), WithServiceName("test-memcache"))
+	defer client.DeleteAll()
+
+	validateMemcacheSpan := func(t *testing.T, span mocktracer.Span, resourceName string) {
+		assert.Equal(t, "test-memcache", span.Tag(ext.ServiceName),
+			"service name should be set to test-memcache")
+		assert.Equal(t, "memcached.query", span.OperationName(),
+			"operation name should be set to memcached.query")
+		assert.Equal(t, resourceName, span.Tag(ext.ResourceName),
+			"resource name should be set to the memcache command")
+	}
+
+	t.Run("traces without context", func(t *testing.T) {
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		err := client.
+			Add(&memcache.Item{
+				Key:   "key1",
+				Value: []byte("value1"),
+			})
+		assert.Nil(t, err)
+
+		spans := mt.FinishedSpans()
+		assert.Len(t, spans, 1)
+		validateMemcacheSpan(t, spans[0], "Add")
+	})
+
+	t.Run("traces with context", func(t *testing.T) {
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		ctx := context.Background()
+		span, ctx := tracer.StartSpanFromContext(ctx, "parent")
+
+		err := client.
+			WithContext(ctx).
+			Add(&memcache.Item{
+				Key:   "key2",
+				Value: []byte("value2"),
+			})
+		assert.Nil(t, err)
+
+		span.Finish()
+
+		spans := mt.FinishedSpans()
+		assert.Len(t, spans, 2)
+		validateMemcacheSpan(t, spans[0], "Add")
+		assert.Equal(t, span, spans[1])
+		assert.Equal(t, spans[1].TraceID(), spans[0].TraceID(),
+			"memcache span should be part of the parent trace")
+	})
+}
+
+func TestFakeServer(t *testing.T) {
+	li := makeFakeServer(t)
+	defer li.Close()
+
+	conn, err := net.Dial("tcp", li.Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+
+	fmt.Fprintf(conn, "add %s\r\n%s\r\n", "key", "value")
+	s := bufio.NewScanner(conn)
+	assert.True(t, s.Scan())
+	assert.Equal(t, "STORED", s.Text())
+}
+
+func makeFakeServer(t *testing.T) net.Listener {
+	li, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	go func() {
+		for {
+			c, err := li.Accept()
+			if err != nil {
+				break
+			}
+			go func() {
+				defer c.Close()
+
+				// the memcache textual protocol is line-oriented with each
+				// command being space separated:
+				//
+				//    command1 arg1 arg2
+				//    command2 arg1 arg2
+				//    ...
+				//
+				s := bufio.NewScanner(c)
+				for s.Scan() {
+					args := strings.Split(s.Text(), " ")
+					switch args[0] {
+					case "add":
+						if !s.Scan() {
+							return
+						}
+						fmt.Fprintf(c, "STORED\r\n")
+					default:
+						fmt.Fprintf(c, "SERVER ERROR unknown command: %v \r\n", args[0])
+						return
+					}
+				}
+			}()
+		}
+	}()
+
+	return li
+}

contrib/bradfitz/gomemcache/memcache/option.go

@@ -0,0 +1,22 @@
+package memcache
+
+const (
+	serviceName   = "memcached"
+	operationName = "memcached.query"
+)
+
+type clientConfig struct{ serviceName string }
+
+// ClientOption represents an option that can be passed to Dial.
+type ClientOption func(*clientConfig)
+
+func defaults(cfg *clientConfig) {
+	cfg.serviceName = serviceName
+}
+
+// WithServiceName sets the given service name for the dialled connection.
+func WithServiceName(name string) ClientOption {
+	return func(cfg *clientConfig) {
+		cfg.serviceName = name
+	}
+}