[Go] Opentracing hook for go-redis v7+ example

hook.go

package redistracing

import (
	"context"
	"fmt"
	"strings"

	"github.com/go-redis/redis/v7"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
)

const (
	redisDBType = "redis"
)

// Hook is a hook for tracing requests in Redis.
// When using *redis.Client without calling WithContext(), it is thread unsafe.
// https://github.com/go-redis/redis/issues/1139#issuecomment-529530118
type Hook struct {
	Opts *redis.Options
}

func (h Hook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (context.Context, error) {
	if ctx == nil {
		return ctx, nil
	}

	parentSpan := opentracing.SpanFromContext(ctx)
	if parentSpan == nil {
		return ctx, nil
	}

	tracer := parentSpan.Tracer()
	span := tracer.StartSpan(getOperationName(cmd), opentracing.ChildOf(parentSpan.Context()))
	h.setDefaultTags(span, cmd)

	return opentracing.ContextWithSpan(ctx, span), nil
}

func (h Hook) AfterProcess(ctx context.Context, _ redis.Cmder) error {
	if ctx == nil {
		return nil
	}

	span := opentracing.SpanFromContext(ctx)
	if span != nil {
		span.Finish()
	}
	return nil
}

func (h Hook) BeforeProcessPipeline(ctx context.Context, cmds []redis.Cmder) (context.Context, error) {
	if ctx == nil {
		return ctx, nil
	}

	parentSpan := opentracing.SpanFromContext(ctx)
	if parentSpan == nil {
		return ctx, nil
	}

	tracer := parentSpan.Tracer()
	pipelineSpan := tracer.StartSpan("PIPELINE", opentracing.ChildOf(parentSpan.Context()))
	h.setDefaultTags(pipelineSpan, nil)

	for _, cmd := range cmds {
		span := tracer.StartSpan(getOperationName(cmd), opentracing.ChildOf(pipelineSpan.Context()))
		h.setDefaultTags(span, cmd)
		span.Finish()
	}
	return opentracing.ContextWithSpan(ctx, pipelineSpan), nil
}

func (h Hook) AfterProcessPipeline(ctx context.Context, _ []redis.Cmder) error {
	if ctx == nil {
		return nil
	}

	span := opentracing.SpanFromContext(ctx)
	if span != nil {
		span.Finish()
	}
	return nil
}

func (h Hook) setDefaultTags(span opentracing.Span, cmd redis.Cmder) {
	ext.DBType.Set(span, redisDBType)
	ext.PeerAddress.Set(span, h.Opts.Addr)
	ext.SpanKind.Set(span, "client")

	if cmd != nil {
		ext.DBStatement.Set(span, fmt.Sprintf("%v", cmd.Args()))
	}
}

func getOperationName(cmd redis.Cmder) string {
	return strings.ToUpper(cmd.Name())
}