Merge pull request #248 from axw/v2-redis

module/apmredigo: introduce redigo instrumentation

docs/instrumenting.asciidoc

@@ -352,6 +352,46 @@ func main() {
 }
 ----
 
+[[builtin-modules-apmredigo]]
+===== module/apmredigo
+Package apmredigo provides a means of instrumenting https://github.com/gomodule/redigo[Redigo]
+so that Redis commands are reported as spans within the current transaction.
+
+To report Redis commands, you can use the top-level `Do` or `DoWithTimeout` functions.
+These functions have the same signature as the `redis.Conn` equivalents apart from an
+initial `context.Context` parameter. If the context passed in contains a sampled
+transaction, a span will be reported for the Redis command.
+
+Another top-level function, `Wrap`, is provided to wrap a `redis.Conn` such that its
+`Do` and `DoWithTimeout` methods call the above mentioned functions. Initially, the
+wrapped connection will be associated with the background context; its `WithContext`
+method may be used to obtain a shallow copy with another context. For example, in an
+HTTP middleware you might bind a connection to the request context, which would
+associate spans with the request's APM transaction.
+
+[source,go]
+----
+import (
+	"net/http"
+
+	"github.com/gomodule/redigo/redis"
+
+	"github.com/elastic/apm-agent-go/module/apmredigo"
+)
+
+var redisPool *redis.Pool // initialized at program startup
+
+func handleRequest(w http.ResponseWriter, req *http.Request) {
+	// Wrap and bind redis.Conn to request context. If the HTTP
+	// server is instrumented with Elastic APM (e.g. with apmhttp),
+	// Redis commands will be reported as spans within the request's
+	// transaction.
+	conn := apmredigo.Wrap(redisPool.Get()).WithContext(req.Context())
+	defer conn.Close()
+	...
+}
+----
+
 [[custom-instrumentation]]
 ==== Custom instrumentation
 

docs/supported-tech.asciidoc

@@ -117,6 +117,16 @@ is created for each query in the batch.
 See <<builtin-modules-apmgocql, module/apmgocql>> for more information
 about GoCQL instrumentation.
 
+[float]
+==== Redis (gomodule/redigo)
+
+We support https://github.com/gomodule/redigo[Redigo],
+https://github.com/gomodule/redigo/tree/v2.0.0[v2.0.0] and greater.
+We provide helper functions for reporting Redis commands as spans.
+
+See <<builtin-modules-apmredigo, module/apmredigo>> for more information
+about Redigo instrumentation.
+
 [float]
 [[supported-tech-rpc]]
 === RPC Frameworks

module/apmredigo/conn.go

@@ -0,0 +1,93 @@
+package apmredigo
+
+import (
+	"context"
+	"strings"
+	"time"
+
+	"github.com/gomodule/redigo/redis"
+
+	"github.com/elastic/apm-agent-go"
+)
+
+// Conn is the interface returned by ContextConn.
+//
+// Conn's Do method reports spans using the bound context.
+type Conn interface {
+	redis.Conn
+
+	// WithContext returns a shallow copy of the connection with
+	// its context changed to ctx.
+	//
+	// To report commands as spans, ctx must contain a transaction or span.
+	WithContext(ctx context.Context) Conn
+}
+
+// Wrap wraps conn such that its Do method calls apmredigo.Do with
+// context.Background(). The context can be changed using Conn.WithContext.
+//
+// If conn implements redis.ConnWithTimeout, then the DoWithTimeout method
+// will similarly call apmredigo.DoWithTimeout.
+//
+// Send and Receive calls are not currently captured.
+func Wrap(conn redis.Conn) Conn {
+	ctx := context.Background()
+	if cwt, ok := conn.(redis.ConnWithTimeout); ok {
+		return contextConnWithTimeout{ConnWithTimeout: cwt, ctx: ctx}
+	}
+	return contextConn{Conn: conn, ctx: ctx}
+}
+
+type contextConnWithTimeout struct {
+	redis.ConnWithTimeout
+	ctx context.Context
+}
+
+func (c contextConnWithTimeout) WithContext(ctx context.Context) Conn {
+	c.ctx = ctx
+	return c
+}
+
+func (c contextConnWithTimeout) Do(commandName string, args ...interface{}) (reply interface{}, err error) {
+	return Do(c.ctx, c.ConnWithTimeout, commandName, args...)
+}
+
+func (c contextConnWithTimeout) DoWithTimeout(timeout time.Duration, commandName string, args ...interface{}) (reply interface{}, err error) {
+	return DoWithTimeout(c.ctx, c.ConnWithTimeout, timeout, commandName, args...)
+}
+
+type contextConn struct {
+	redis.Conn
+	ctx context.Context
+}
+
+func (c contextConn) WithContext(ctx context.Context) Conn {
+	c.ctx = ctx
+	return c
+}
+
+func (c contextConn) Do(commandName string, args ...interface{}) (reply interface{}, err error) {
+	return Do(c.ctx, c.Conn, commandName, args...)
+}
+
+// Do calls conn.Do(commandName, args...), and also reports the operation as a span to Elastic APM.
+func Do(ctx context.Context, conn redis.Conn, commandName string, args ...interface{}) (interface{}, error) {
+	spanName := strings.ToUpper(commandName)
+	if spanName == "" {
+		spanName = "(flush pipeline)"
+	}
+	span, _ := elasticapm.StartSpan(ctx, spanName, "cache.redis")
+	defer span.End()
+	return conn.Do(commandName, args...)
+}
+
+// DoWithTimeout calls redis.DoWithTimeout(conn, timeout, commandName, args...), and also reports the operation as a span to Elastic APM.
+func DoWithTimeout(ctx context.Context, conn redis.Conn, timeout time.Duration, commandName string, args ...interface{}) (interface{}, error) {
+	spanName := strings.ToUpper(commandName)
+	if spanName == "" {
+		spanName = "(flush pipeline)"
+	}
+	span, _ := elasticapm.StartSpan(ctx, spanName, "cache.redis")
+	defer span.End()
+	return redis.DoWithTimeout(conn, timeout, commandName, args...)
+}

module/apmredigo/conn_test.go

@@ -0,0 +1,107 @@
+package apmredigo_test
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/gomodule/redigo/redis"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/elastic/apm-agent-go"
+	"github.com/elastic/apm-agent-go/apmtest"
+	"github.com/elastic/apm-agent-go/module/apmredigo"
+)
+
+func TestWrap(t *testing.T) {
+	var conn mockConn
+	_, spans, _ := apmtest.WithTransaction(func(ctx context.Context) {
+		conn := apmredigo.Wrap(conn).WithContext(ctx)
+		conn.Do("PING", "hello, world!")
+	})
+	require.Len(t, spans, 1)
+	assert.Equal(t, "PING", spans[0].Name)
+	assert.Equal(t, "cache.redis", spans[0].Type)
+}
+
+func TestWithContext(t *testing.T) {
+	ping := func(ctx context.Context, conn apmredigo.Conn) {
+		span, ctx := elasticapm.StartSpan(ctx, "ping", "custom")
+		defer span.End()
+
+		// bind conn to the ctx containing the span above
+		conn = conn.WithContext(ctx)
+		conn.Do("PING", "hello, world!")
+	}
+
+	var conn mockConn
+	_, spans, _ := apmtest.WithTransaction(func(ctx context.Context) {
+		conn := apmredigo.Wrap(conn)
+		ping(ctx, conn)
+	})
+	require.Len(t, spans, 2)
+	assert.Equal(t, "PING", spans[0].Name)
+	assert.Equal(t, "ping", spans[1].Name)
+	assert.Equal(t, spans[1].ID, spans[0].ParentID)
+}
+
+func TestConnWithTimeout(t *testing.T) {
+	var conn mockConnWithTimeout
+	_, spans, _ := apmtest.WithTransaction(func(ctx context.Context) {
+		conn := apmredigo.Wrap(conn).WithContext(ctx)
+		redis.DoWithTimeout(conn, time.Second, "PING", "hello, world!")
+	})
+	require.Len(t, spans, 1)
+	assert.Equal(t, "PING", spans[0].Name)
+	assert.Equal(t, "cache.redis", spans[0].Type)
+}
+
+func TestWrapPipeline(t *testing.T) {
+	var conn mockConnWithTimeout
+	_, spans, _ := apmtest.WithTransaction(func(ctx context.Context) {
+		conn := apmredigo.Wrap(conn).WithContext(ctx)
+		conn.Do("")
+		redis.DoWithTimeout(conn, time.Second, "")
+	})
+	require.Len(t, spans, 2)
+	assert.Equal(t, "(flush pipeline)", spans[0].Name)
+	assert.Equal(t, "(flush pipeline)", spans[1].Name)
+}
+
+type mockConnWithTimeout struct{ mockConn }
+
+func (mockConnWithTimeout) DoWithTimeout(timeout time.Duration, commandName string, args ...interface{}) (reply interface{}, err error) {
+	return []byte("Done"), errors.New("DoWithTimeout failed")
+}
+
+func (mockConnWithTimeout) ReceiveWithTimeout(timeout time.Duration) (reply interface{}, err error) {
+	return []byte("REceived"), errors.New("ReceiveWithTimeout failed")
+}
+
+type mockConn struct{}
+
+func (mockConn) Close() error {
+	panic("Close not implemented")
+}
+
+func (mockConn) Err() error {
+	panic("Err not implemented")
+}
+
+func (mockConn) Flush() error {
+	panic("Flush not implemented")
+}
+
+func (mockConn) Do(commandName string, args ...interface{}) (reply interface{}, err error) {
+	return []byte("Done"), errors.New("Do failed")
+}
+
+func (mockConn) Send(commandName string, args ...interface{}) error {
+	return errors.New("Send failed")
+}
+
+func (mockConn) Receive() (reply interface{}, err error) {
+	return []byte("Received"), errors.New("Receive failed")
+}

module/apmredigo/doc.go

@@ -0,0 +1,2 @@
+// Package apmredigo provides helpers for tracing github.com/gomodule/redigo/redis client operations as spans.
+package apmredigo