ProcessPipeline tracing and some tags

.gitignore

@@ -0,0 +1 @@
+.idea/*

README.md

@@ -1,16 +1,16 @@
 # opentracing go-redis
 
-[OpenTracing](http://opentracing.io/) instrumentation for [go-redis](https://github.com/go-redis/redis).
+[OpenTracing](http://opentracing.io/) instrumentation for [go-redis](https://github.com/go-redis/redis) with support of `process` and `processPipeline` commands.
 
 ## Install
 
 ```
-go get -u github.com/smacker/opentracing-go-redis
+go get -u github.com/ticketmaster/opentracing-go-redis
 ```
 
 ## Usage
 
-Clone redis client `c := otredis.WrapRedisClient(ctx, c)` with a span.
+Clone redis client `c := otredis.WrapWithOpenTracing(ctx, c)` with a span.
 
 Example:
 
@@ -22,7 +22,7 @@ func Handler(ctx context.Context) {
     defer span.Finish()
 
     // clone redis with proper context
-    client := otredis.WrapRedisClient(ctx, client)
+    client := otredis.WrapWithOpenTracing(ctx, client)
 
     // make requests to redis
     client.Get("foo")

otredis.go

@@ -2,34 +2,68 @@ package otredis
 
 import (
 	"context"
+	"fmt"
+	"strings"
 
 	"github.com/go-redis/redis"
-	opentracing "github.com/opentracing/opentracing-go"
+	"github.com/opentracing/opentracing-go"
 	"github.com/opentracing/opentracing-go/ext"
 )
 
-// WrapRedisClient adds opentracing measurements for commands and returns cloned client
-func WrapRedisClient(ctx context.Context, c *redis.Client) *redis.Client {
+func WrapWithOpenTracing(ctx context.Context, client *redis.Client) *redis.Client {
 	if ctx == nil {
-		return c
+		return client
 	}
 	parentSpan := opentracing.SpanFromContext(ctx)
 	if parentSpan == nil {
-		return c
+		return client
 	}
+	clientWithContext := client.WithContext(ctx)
+	opts := clientWithContext.Options()
+	clientWithContext.WrapProcess(process(parentSpan, opts))
+	clientWithContext.WrapProcessPipeline(processPipeline(parentSpan, opts))
+	return clientWithContext
+}
 
-	// clone using context
-	copy := c.WithContext(c.Context())
-	copy.WrapProcess(func(oldProcess func(cmd redis.Cmder) error) func(cmd redis.Cmder) error {
+func process(parentSpan opentracing.Span, opts *redis.Options) func(oldProcess func(cmd redis.Cmder) error) func(cmd redis.Cmder) error {
+	return func(oldProcess func(cmd redis.Cmder) error) func(cmd redis.Cmder) error {
 		return func(cmd redis.Cmder) error {
-			tr := parentSpan.Tracer()
-			sp := tr.StartSpan("redis", opentracing.ChildOf(parentSpan.Context()))
-			ext.DBType.Set(sp, "redis")
-			sp.SetTag("db.method", cmd.Name())
-			defer sp.Finish()
-
+			dbStatement := formatCommandAsDbStatement(cmd)
+			doSpan(parentSpan, opts, "redis-cmd", dbStatement)
 			return oldProcess(cmd)
 		}
-	})
-	return copy
+	}
+}
+
+func processPipeline(parentSpan opentracing.Span, opts *redis.Options) func(oldProcess func(cmds []redis.Cmder) error) func(cmds []redis.Cmder) error {
+	return func(oldProcess func(cmds []redis.Cmder) error) func(cmds []redis.Cmder) error {
+		return func(cmds []redis.Cmder) error {
+			dbStatement := formatCommandsAsDbStatement(cmds)
+			doSpan(parentSpan, opts, "redis-pipeline-cmd", dbStatement)
+			return oldProcess(cmds)
+		}
+	}
+}
+
+func formatCommandAsDbStatement(cmd redis.Cmder) string {
+	return fmt.Sprintf("%s", cmd)
+}
+
+func formatCommandsAsDbStatement(cmds []redis.Cmder) string {
+	cmdsAsDbStatements := make([]string, len(cmds))
+	for i, cmd := range cmds {
+		cmdAsDbStatement := formatCommandAsDbStatement(cmd)
+		cmdsAsDbStatements[i] = cmdAsDbStatement
+	}
+	return strings.Join(cmdsAsDbStatements, "\n")
+}
+
+func doSpan(parentSpan opentracing.Span, opts *redis.Options, operationName, dbStatement string) {
+	tr := parentSpan.Tracer()
+	span := tr.StartSpan(operationName, opentracing.ChildOf(parentSpan.Context()))
+	defer span.Finish()
+	ext.DBType.Set(span, "redis")
+	ext.DBStatement.Set(span, dbStatement)
+	ext.PeerAddress.Set(span, opts.Addr)
+	ext.SpanKind.Set(span, ext.SpanKindEnum("client"))
 }

otredis_test.go

@@ -2,117 +2,133 @@ package otredis
 
 import (
 	"context"
-	"os"
 	"testing"
 
 	"github.com/alicebob/miniredis"
 	"github.com/go-redis/redis"
-	opentracing "github.com/opentracing/opentracing-go"
+	"github.com/opentracing/opentracing-go"
+	"github.com/opentracing/opentracing-go/ext"
 	"github.com/opentracing/opentracing-go/mocktracer"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/suite"
 )
 
-var client *redis.Client
-var tracer *mocktracer.MockTracer
+type OpenTracingRedisTestSuite struct {
+	suite.Suite
+	miniRedis  *miniredis.Miniredis
+	client     *redis.Client
+	mockTracer *mocktracer.MockTracer
+}
 
-func init() {
-	tracer = mocktracer.New()
-	opentracing.SetGlobalTracer(tracer)
+func TestOpenTracingRedisTestSuite(t *testing.T) {
+	tests := new(OpenTracingRedisTestSuite)
+	suite.Run(t, tests)
 }
 
-func TestMain(m *testing.M) {
-	// in-memory redis
-	s, err := miniredis.Run()
-	if err != nil {
-		panic(err)
-	}
-	defer s.Close()
-	client = redis.NewClient(&redis.Options{
-		Addr: s.Addr(),
+func (ts *OpenTracingRedisTestSuite) SetupSuite() {
+	// Common
+	redisAddr := "127.0.0.1:6379"
+	// MiniRedis
+	ts.miniRedis = miniredis.NewMiniRedis()
+	ts.miniRedis.StartAddr(redisAddr)
+	// Client
+	ts.client = redis.NewClient(&redis.Options{
+		Addr: redisAddr,
 	})
+	// MockTracer
+	ts.mockTracer = mocktracer.New()
+	opentracing.SetGlobalTracer(ts.mockTracer)
+}
 
-	os.Exit(m.Run())
+func (ts *OpenTracingRedisTestSuite) TearDownSuite() {
+	// Client
+	ts.client.Close()
+	// MiniRedis
+	ts.miniRedis.Close()
 }
 
-func TestSet(t *testing.T) {
+func (ts *OpenTracingRedisTestSuite) Test_ProcessExecution() {
+
+	t := ts.T()
 	ctx := context.Background()
 
-	span, ctx := opentracing.StartSpanFromContext(ctx, "test-params")
-	ctxClient := WrapRedisClient(ctx, client)
-	callSet(t, ctxClient, "with span")
-	span.Finish()
+	hmSetKey, hmSetValues := buildHMSetCmdData("PROCESS_EXEC")
 
-	spans := tracer.FinishedSpans()
-	if len(spans) != 2 {
-		t.Fatalf("should be 2 finished spans but there are %d: %v", len(spans), spans)
-	}
+	span, ctx := opentracing.StartSpanFromContext(ctx, "test-process-execution")
 
-	redisSpan := spans[0]
-	if redisSpan.OperationName != "redis" {
-		t.Errorf("first span operation should be redis but it's '%s'", redisSpan.OperationName)
-	}
+	_, err := WrapWithOpenTracing(ctx, ts.client).HMSet(hmSetKey, hmSetValues).Result()
+	assert.Nil(t, err, "redis execution failed: %+v", err)
 
-	testTags(t, redisSpan, map[string]string{
-		"db.type":   "redis",
-		"db.method": "set",
-	})
+	span.Finish()
+
+	finishedSpans := ts.mockTracer.FinishedSpans()
+	expectedFinishedSpansNumber := 2
+	assert.Equal(t, expectedFinishedSpansNumber, len(finishedSpans), "the number of finished spans is invalid")
 
-	tracer.Reset()
+	redisSpan := finishedSpans[0]
+	expectedOperationName := "redis-cmd"
+	assert.Equal(t, expectedOperationName, redisSpan.OperationName)
+	expectedTags := buildExpectedTags("PROCESS_EXEC")
+	assertTags(t, redisSpan, expectedTags)
+
+	ts.mockTracer.Reset()
 }
 
-func TestGet(t *testing.T) {
+func (ts *OpenTracingRedisTestSuite) Test_ProcessPipelineExecution() {
+
+	t := ts.T()
 	ctx := context.Background()
 
-	span, ctx := opentracing.StartSpanFromContext(ctx, "test-params")
-	ctxClient := WrapRedisClient(ctx, client)
-	callGet(t, ctxClient)
-	span.Finish()
+	hmSetKey, hmSetValues := buildHMSetCmdData("PROCESS_PIPELINE_EXEC")
 
-	spans := tracer.FinishedSpans()
-	if len(spans) != 2 {
-		t.Fatalf("should be 2 finished spans but there are %d: %v", len(spans), spans)
-	}
+	span, ctx := opentracing.StartSpanFromContext(ctx, "test-process-pipeline-execution")
 
-	redisSpan := spans[0]
-	if redisSpan.OperationName != "redis" {
-		t.Errorf("first span operation should be redis but it's '%s'", redisSpan.OperationName)
-	}
+	pipeline := WrapWithOpenTracing(ctx, ts.client).TxPipeline()
+	pipeline.HMSet(hmSetKey, hmSetValues)
+	_, err := pipeline.Exec()
+	assert.Nil(t, err, "redis pipeline execution failed: %+v", err)
 
-	testTags(t, redisSpan, map[string]string{
-		"db.type":   "redis",
-		"db.method": "get",
-	})
+	span.Finish()
+
+	finishedSpans := ts.mockTracer.FinishedSpans()
+	expectedFinishedSpansNumber := 2
+	assert.Equal(t, expectedFinishedSpansNumber, len(finishedSpans), "the number of finished spans is invalid")
+
+	redisSpan := finishedSpans[0]
+	expectedOperationName := "redis-pipeline-cmd"
+	assert.Equal(t, expectedOperationName, redisSpan.OperationName)
+	expectedTags := buildExpectedTags("PROCESS_PIPELINE_EXEC")
+	assertTags(t, redisSpan, expectedTags)
 
-	tracer.Reset()
+	ts.mockTracer.Reset()
 }
 
-func callSet(t *testing.T, c *redis.Client, value string) {
-	_, err := c.Set("foo", value, 0).Result()
-	if err != nil {
-		t.Fatalf("Redis returned error: %v", err)
-	}
+func buildHMSetCmdData(testSuffix string) (string, map[string]interface{}) {
+
+	key := "key:TEST_" + testSuffix
+	values := make(map[string]interface{})
+	values["TEST_KEY_"+testSuffix] = "TEST_VALUE_" + testSuffix
+
+	return key, values
 }
 
-func callGet(t *testing.T, c *redis.Client) {
-	_, err := c.Get("foo").Result()
-	if err != nil {
-		t.Fatalf("Redis returned error: %v", err)
-	}
+func buildExpectedTags(testSuffix string) map[string]interface{} {
+	expectedTags := make(map[string]interface{})
+	expectedTags["db.type"] = "redis"
+	expectedTags["db.statement"] = "hmset key:TEST_" + testSuffix + " TEST_KEY_" + testSuffix + " TEST_VALUE_" + testSuffix + ": "
+	expectedTags["peer.address"] = "127.0.0.1:6379"
+	expectedTags["span.kind"] = ext.SpanKindEnum("client")
+	return expectedTags
 }
 
-func testTags(t *testing.T, redisSpan *mocktracer.MockSpan, expectedTags map[string]string) {
-	redisTags := redisSpan.Tags()
-	if len(redisTags) != len(expectedTags) {
-		t.Errorf("redis span should have %d tags but it has %d", len(expectedTags), len(redisTags))
-	}
+func assertTags(t *testing.T, redisSpan *mocktracer.MockSpan, expectedTags map[string]interface{}) {
+
+	actualTags := redisSpan.Tags()
+	assert.Equal(t, len(expectedTags), len(actualTags), "redis span tags number is invalid")
 
-	for name, expected := range expectedTags {
-		value, ok := redisTags[name]
-		if !ok {
-			t.Errorf("redis span doesn't have tag '%s'", name)
-			continue
-		}
-		if value != expected {
-			t.Errorf("redis span tag '%s' should have value '%s' but it has '%s'", name, expected, value)
-		}
+	for expectedTagKey, expectedTagValue := range expectedTags {
+		actualTag, ok := actualTags[expectedTagKey]
+		assert.True(t, ok, "redis span doesn't have tag '%s'", expectedTagKey)
+		assert.Equal(t, expectedTagValue, actualTag, "redis span tag '%s' is invalid", expectedTagKey)
 	}
 }