Merge pull request #99 from remind101/aws-tracing-helper

Add a helper method to trace aws clients

Author: bmarini

tracing/contrib/aws/request.go

@@ -0,0 +1,58 @@
+package aws
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go/aws/request"
+	opentracing "github.com/opentracing/opentracing-go"
+)
+
+// Option is a hook for adding span tags.
+type Option func(opentracing.Span, *request.Request, error)
+
+// Tracer holds tagging options for request tracing.
+type Tracer struct {
+	Opts []Option
+}
+
+// New returns a new AWS request tracer with default tagging.
+func New(opts ...Option) *Tracer {
+	opts = append([]Option{defaultTags}, opts...)
+	return &Tracer{opts}
+}
+
+// Send wraps the aws request with an opentracing span.
+//
+// req, output := dynamoClient.PutItemRequest(input)
+// err := t.Send(ctx, req)
+func (t *Tracer) Send(ctx context.Context, r *request.Request, opts ...Option) error {
+	span, ctx := opentracing.StartSpanFromContext(ctx, "client.request")
+	defer span.Finish()
+
+	r.HTTPRequest = r.HTTPRequest.WithContext(ctx)
+	err := r.Send()
+
+	opts = append(t.Opts, opts...)
+	for _, fn := range opts {
+		fn(span, r, err)
+	}
+
+	return err
+}
+
+func defaultTags(span opentracing.Span, r *request.Request, err error) {
+	span.SetTag("resource.name", r.Operation.Name)
+	span.SetTag("http.method", r.Operation.HTTPMethod)
+	span.SetTag("http.url", r.ClientInfo.Endpoint+r.Operation.HTTPPath)
+	span.SetTag("out.host", r.ClientInfo.Endpoint)
+	span.SetTag("aws.operation", r.Operation.Name)
+	span.SetTag("aws.retry_count", r.RetryCount)
+
+	if r.HTTPResponse != nil {
+		span.SetTag("http.status_code", r.HTTPResponse.StatusCode)
+	}
+
+	if err != nil {
+		span.SetTag("error.error", err)
+	}
+}

tracing/contrib/aws/request_test.go

@@ -0,0 +1,72 @@
+package aws_test
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/request"
+	"github.com/aws/aws-sdk-go/awstesting/mock"
+	opentracing "github.com/opentracing/opentracing-go"
+	"github.com/opentracing/opentracing-go/mocktracer"
+	awsot "github.com/remind101/pkg/tracing/contrib/aws"
+)
+
+func TestSend(t *testing.T) {
+	tracer := mocktracer.New()
+	opentracing.SetGlobalTracer(tracer)
+
+	_, ctx := opentracing.StartSpanFromContext(context.Background(), "root")
+
+	client := mock.NewMockClient(&aws.Config{Region: aws.String("us-west-2")})
+	r := client.NewRequest(&request.Operation{
+		Name:       "object.get",
+		HTTPMethod: "GET",
+		HTTPPath:   "/foobar",
+	}, nil, nil)
+
+	awsTracer := awsot.New()
+
+	err := awsTracer.Send(ctx, r, func(s opentracing.Span, r *request.Request, err error) {
+		s.SetTag("span.type", "external")
+		s.SetTag("service.name", "aws.s3")
+	})
+
+	if err != nil {
+		t.Fatalf("expected no error; got %v", err)
+	}
+
+	spans := tracer.FinishedSpans()
+	if len(spans) != 1 {
+		t.Fatal("expected 1 finished span")
+	}
+	span := spans[0]
+	if got, want := span.OperationName, "client.request"; got != want {
+		t.Errorf("got: %+v; expected %+v", got, want)
+	}
+
+	tags := map[string]string{
+		"span.type":     "external",
+		"service.name":  "aws.s3",
+		"resource.name": "object.get",
+		"http.method":   "GET",
+		"http.url":      client.Endpoint + "/foobar",
+		"out.host":      client.Endpoint,
+		"aws.operation": "object.get",
+	}
+
+	for k, v := range tags {
+		if got, want := span.Tag(k).(string), v; got != want {
+			t.Errorf("span.Tag('%s'): %+v; expected %+v", k, got, want)
+		}
+	}
+
+	if got, want := span.Tag("http.status_code").(int), http.StatusOK; got != want {
+		t.Errorf("got: %+v; expected %+v", got, want)
+	}
+
+	if got, want := span.Tag("aws.retry_count").(int), 0; got != want {
+		t.Errorf("got: %+v; expected %+v", got, want)
+	}
+}