contrib/net/http: add WithSpanOptions (#429)

This change adds the `WithSpanOption` to `WrapHandler` and `ServeMux` and type aliases `MuxOption` to `Option` in order to make it more generic for all components to use.

contrib/net/http/http.go

@@ -5,21 +5,20 @@ import (
 	"net/http"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/contrib/internal/httputil"
-	"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/tracer"
 )
 
 // ServeMux is an HTTP request multiplexer that traces all the incoming requests.
 type ServeMux struct {
 	*http.ServeMux
-	cfg *muxConfig
+	cfg *config
 }
 
 // NewServeMux allocates and returns an http.ServeMux augmented with the
 // global tracer.
-func NewServeMux(opts ...MuxOption) *ServeMux {
-	cfg := new(muxConfig)
+func NewServeMux(opts ...Option) *ServeMux {
+	cfg := new(config)
 	defaults(cfg)
 	for _, fn := range opts {
 		fn(cfg)
@@ -38,16 +37,21 @@ func (mux *ServeMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	// get the resource associated to this request
 	_, route := mux.Handler(r)
 	resource := r.Method + " " + route
-	var opts []ddtrace.StartSpanOption
+	opts := mux.cfg.spanOpts
 	if mux.cfg.analyticsRate > 0 {
 		opts = append(opts, tracer.Tag(ext.EventSampleRate, mux.cfg.analyticsRate))
 	}
 	httputil.TraceAndServe(mux.ServeMux, w, r, mux.cfg.serviceName, resource, opts...)
 }
 
 // WrapHandler wraps an http.Handler with tracing using the given service and resource.
-func WrapHandler(h http.Handler, service, resource string) http.Handler {
+func WrapHandler(h http.Handler, service, resource string, opts ...Option) http.Handler {
+	cfg := new(config)
+	defaults(cfg)
+	for _, fn := range opts {
+		fn(cfg)
+	}
 	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
-		httputil.TraceAndServe(h, w, req, service, resource)
+		httputil.TraceAndServe(h, w, req, service, resource, cfg.spanOpts...)
 	})
 }

contrib/net/http/http_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"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"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig"
 )
 
@@ -35,6 +36,7 @@ func TestHttpTracer200(t *testing.T) {
 	assert.Equal("GET", s.Tag(ext.HTTPMethod))
 	assert.Equal(url, s.Tag(ext.HTTPURL))
 	assert.Equal(nil, s.Tag(ext.Error))
+	assert.Equal("bar", s.Tag("foo"))
 }
 
 func TestHttpTracer500(t *testing.T) {
@@ -62,14 +64,16 @@ func TestHttpTracer500(t *testing.T) {
 	assert.Equal("GET", s.Tag(ext.HTTPMethod))
 	assert.Equal(url, s.Tag(ext.HTTPURL))
 	assert.Equal("500: Internal Server Error", s.Tag(ext.Error).(error).Error())
+	assert.Equal("bar", s.Tag("foo"))
 }
 
 func TestWrapHandler200(t *testing.T) {
 	mt := mocktracer.Start()
 	defer mt.Stop()
 	assert := assert.New(t)
 
-	handler := WrapHandler(http.HandlerFunc(handler200), "my-service", "my-resource")
+	handler := WrapHandler(http.HandlerFunc(handler200), "my-service", "my-resource",
+		WithSpanOptions(tracer.Tag("foo", "bar")))
 
 	url := "/"
 	r := httptest.NewRequest("GET", url, nil)
@@ -89,10 +93,11 @@ func TestWrapHandler200(t *testing.T) {
 	assert.Equal("GET", s.Tag(ext.HTTPMethod))
 	assert.Equal(url, s.Tag(ext.HTTPURL))
 	assert.Equal(nil, s.Tag(ext.Error))
+	assert.Equal("bar", s.Tag("foo"))
 }
 
 func TestAnalyticsSettings(t *testing.T) {
-	assertRate := func(t *testing.T, mt mocktracer.Tracer, rate interface{}, opts ...MuxOption) {
+	assertRate := func(t *testing.T, mt mocktracer.Tracer, rate interface{}, opts ...Option) {
 		mux := NewServeMux(opts...)
 		mux.HandleFunc("/200", handler200)
 		r := httptest.NewRequest("GET", "/200", nil)
@@ -150,7 +155,7 @@ func TestAnalyticsSettings(t *testing.T) {
 }
 
 func router() http.Handler {
-	mux := NewServeMux(WithServiceName("my-service"))
+	mux := NewServeMux(WithServiceName("my-service"), WithSpanOptions(tracer.Tag("foo", "bar")))
 	mux.HandleFunc("/200", handler200)
 	mux.HandleFunc("/500", handler500)
 	return mux

contrib/net/http/option.go

@@ -7,22 +7,26 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig"
 )
 
-type muxConfig struct {
+type config struct {
 	serviceName   string
 	analyticsRate float64
+	spanOpts      []ddtrace.StartSpanOption
 }
 
-// MuxOption represents an option that can be passed to NewServeMux.
-type MuxOption func(*muxConfig)
+// MuxOption has been deprecated in favor of Option.
+type MuxOption func(*config)
 
-func defaults(cfg *muxConfig) {
+// Option represents an option that can be passed to NewServeMux or WrapHandler.
+type Option = MuxOption
+
+func defaults(cfg *config) {
 	cfg.analyticsRate = globalconfig.AnalyticsRate()
 	cfg.serviceName = "http.router"
 }
 
 // WithServiceName sets the given service name for the returned ServeMux.
 func WithServiceName(name string) MuxOption {
-	return func(cfg *muxConfig) {
+	return func(cfg *config) {
 		cfg.serviceName = name
 	}
 }
@@ -38,11 +42,19 @@ func WithAnalytics(on bool) MuxOption {
 // WithAnalyticsRate sets the sampling rate for Trace Analytics events
 // correlated to started spans.
 func WithAnalyticsRate(rate float64) MuxOption {
-	return func(cfg *muxConfig) {
+	return func(cfg *config) {
 		cfg.analyticsRate = rate
 	}
 }
 
+// WithSpanOptions defines a set of additional ddtrace.StartSpanOption to be added
+// to spans started by the integration.
+func WithSpanOptions(opts ...ddtrace.StartSpanOption) Option {
+	return func(cfg *config) {
+		cfg.spanOpts = append(cfg.spanOpts, opts...)
+	}
+}
+
 // A RoundTripperBeforeFunc can be used to modify a span before an http
 // RoundTrip is made.
 type RoundTripperBeforeFunc func(*http.Request, ddtrace.Span)