contrib/net/http: add tracing for http.RoundTripper (#299)

contrib/net/http/option.go

@@ -1,5 +1,11 @@
 package http
 
+import (
+	"net/http"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
+)
+
 type muxConfig struct{ serviceName string }
 
 // MuxOption represents an option that can be passed to NewServeMux.
@@ -15,3 +21,36 @@ func WithServiceName(name string) MuxOption {
 		cfg.serviceName = name
 	}
 }
+
+// A RoundTripperBeforeFunc can be used to modify a span before an http
+// RoundTrip is made.
+type RoundTripperBeforeFunc func(*http.Request, ddtrace.Span)
+
+// A RoundTripperAfterFunc can be used to modify a span after an http
+// RoundTrip is made. It is possible for the http Response to be nil.
+type RoundTripperAfterFunc func(*http.Response, ddtrace.Span)
+
+type roundTripperConfig struct {
+	before RoundTripperBeforeFunc
+	after  RoundTripperAfterFunc
+}
+
+// A RoundTripperOption represents an option that can be passed to
+// WrapRoundTripper.
+type RoundTripperOption func(*roundTripperConfig)
+
+// WithBefore adds a RoundTripperBeforeFunc to the RoundTripper
+// config.
+func WithBefore(f RoundTripperBeforeFunc) RoundTripperOption {
+	return func(cfg *roundTripperConfig) {
+		cfg.before = f
+	}
+}
+
+// WithAfter adds a RoundTripperAfterFunc to the RoundTripper
+// config.
+func WithAfter(f RoundTripperAfterFunc) RoundTripperOption {
+	return func(cfg *roundTripperConfig) {
+		cfg.after = f
+	}
+}

contrib/net/http/roundtripper.go

@@ -0,0 +1,68 @@
+package http
+
+import (
+	"errors"
+	"fmt"
+	"net/http"
+	"os"
+	"strconv"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+)
+
+const defaultResourceName = "http.request"
+
+type roundTripper struct {
+	base http.RoundTripper
+	cfg  *roundTripperConfig
+}
+
+func (rt *roundTripper) RoundTrip(req *http.Request) (res *http.Response, err error) {
+	span, _ := tracer.StartSpanFromContext(req.Context(), "http.request",
+		tracer.SpanType(ext.SpanTypeHTTP),
+		tracer.ResourceName(defaultResourceName),
+		tracer.Tag(ext.HTTPMethod, req.Method),
+		tracer.Tag(ext.HTTPURL, req.URL.Path),
+	)
+	defer func() {
+		if rt.cfg.after != nil {
+			rt.cfg.after(res, span)
+		}
+		span.Finish(tracer.WithError(err))
+	}()
+	if rt.cfg.before != nil {
+		rt.cfg.before(req, span)
+	}
+	// inject the span context into the http request
+	err = tracer.Inject(span.Context(), tracer.HTTPHeadersCarrier(req.Header))
+	if err != nil {
+		// this should never happen
+		fmt.Fprintf(os.Stderr, "failed to inject http headers for round tripper: %v\n", err)
+	}
+	res, err = rt.base.RoundTrip(req)
+	if err != nil {
+		span.SetTag("http.errors", err.Error())
+	} else {
+		span.SetTag(ext.HTTPCode, strconv.Itoa(res.StatusCode))
+		// treat 5XX as errors
+		if res.StatusCode/100 == 5 {
+			span.SetTag("http.errors", res.Status)
+			err = errors.New(res.Status)
+		}
+	}
+	return res, err
+}
+
+// WrapRoundTripper returns a new RoundTripper which traces all requests sent
+// over the transport.
+func WrapRoundTripper(rt http.RoundTripper, opts ...RoundTripperOption) http.RoundTripper {
+	cfg := new(roundTripperConfig)
+	for _, opt := range opts {
+		opt(cfg)
+	}
+	return &roundTripper{
+		base: rt,
+		cfg:  cfg,
+	}
+}

contrib/net/http/roundtripper_test.go

@@ -0,0 +1,61 @@
+package http
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"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/mocktracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+)
+
+func TestRoundTripper(t *testing.T) {
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		spanctx, err := tracer.Extract(tracer.HTTPHeadersCarrier(r.Header))
+		assert.NoError(t, err)
+
+		span := tracer.StartSpan("test",
+			tracer.ChildOf(spanctx))
+		defer span.Finish()
+
+		w.Write([]byte("Hello World"))
+	}))
+	defer s.Close()
+
+	rt := WrapRoundTripper(http.DefaultTransport,
+		WithBefore(func(req *http.Request, span ddtrace.Span) {
+			span.SetTag("CalledBefore", true)
+		}),
+		WithAfter(func(res *http.Response, span ddtrace.Span) {
+			span.SetTag("CalledAfter", true)
+		}))
+
+	client := &http.Client{
+		Transport: rt,
+	}
+
+	client.Get(s.URL + "/hello/world")
+
+	spans := mt.FinishedSpans()
+	assert.Len(t, spans, 2)
+	assert.Equal(t, spans[0].TraceID(), spans[1].TraceID())
+
+	s0 := spans[0]
+	assert.Equal(t, "test", s0.OperationName())
+	assert.Equal(t, "test", s0.Tag(ext.ResourceName))
+
+	s1 := spans[1]
+	assert.Equal(t, "http.request", s1.OperationName())
+	assert.Equal(t, "http.request", s1.Tag(ext.ResourceName))
+	assert.Equal(t, "200", s1.Tag(ext.HTTPCode))
+	assert.Equal(t, "GET", s1.Tag(ext.HTTPMethod))
+	assert.Equal(t, "/hello/world", s1.Tag(ext.HTTPURL))
+	assert.Equal(t, true, s1.Tag("CalledBefore"))
+	assert.Equal(t, true, s1.Tag("CalledAfter"))
+}