From 90c48ee1d8ba04eadca7100848bd0aad088a3328 Mon Sep 17 00:00:00 2001
From: Gabriel Aszalos <gabriel.aszalos@gmail.com>
Date: Thu, 23 Aug 2018 17:29:36 +0300
Subject: [PATCH] gorilla/mux: allow setting custom span options (#320)

---
 contrib/gorilla/mux/mux.go      |  1 +
 contrib/gorilla/mux/mux_test.go | 16 ++++++++++++++++
 contrib/gorilla/mux/option.go   | 15 ++++++++++++++-
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/contrib/gorilla/mux/mux.go b/contrib/gorilla/mux/mux.go
index 76824a20df..e12887f974 100644
--- a/contrib/gorilla/mux/mux.go
+++ b/contrib/gorilla/mux/mux.go
@@ -97,6 +97,7 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 			spanopts = append(spanopts, tracer.Tag("mux.host", h))
 		}
 	}
+	spanopts = append(spanopts, r.config.spanOpts...)
 	resource := req.Method + " " + route
 	httputil.TraceAndServe(r.Router, w, req, r.config.serviceName, resource, spanopts...)
 }
diff --git a/contrib/gorilla/mux/mux_test.go b/contrib/gorilla/mux/mux_test.go
index 1958c3aba6..cc8c0fd000 100644
--- a/contrib/gorilla/mux/mux_test.go
+++ b/contrib/gorilla/mux/mux_test.go
@@ -9,6 +9,7 @@ import (
 
 	"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"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -92,6 +93,21 @@ func TestDomain(t *testing.T) {
 	assert.Equal("localhost", spans[0].Tag("mux.host"))
 }
 
+func TestSpanOptions(t *testing.T) {
+	assert := assert.New(t)
+	mt := mocktracer.Start()
+	defer mt.Stop()
+	mux := NewRouter(WithSpanOptions(tracer.Tag(ext.SamplingPriority, 2)))
+	mux.Handle("/200", okHandler()).Host("localhost")
+	r := httptest.NewRequest("GET", "http://localhost/200", nil)
+	w := httptest.NewRecorder()
+	mux.ServeHTTP(w, r)
+
+	spans := mt.FinishedSpans()
+	assert.Equal(1, len(spans))
+	assert.Equal(2, spans[0].Tag(ext.SamplingPriority))
+}
+
 // TestImplementingMethods is a regression tests asserting that all the mux.Router methods
 // returning the router will return the modified traced version of it and not the original
 // router.
diff --git a/contrib/gorilla/mux/option.go b/contrib/gorilla/mux/option.go
index e491840d7a..0b96d26f2e 100644
--- a/contrib/gorilla/mux/option.go
+++ b/contrib/gorilla/mux/option.go
@@ -1,6 +1,11 @@
 package mux
 
-type routerConfig struct{ serviceName string }
+import "gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
+
+type routerConfig struct {
+	serviceName string
+	spanOpts    []ddtrace.StartSpanOption // additional span options to be applied
+}
 
 // RouterOption represents an option that can be passed to NewRouter.
 type RouterOption func(*routerConfig)
@@ -15,3 +20,11 @@ func WithServiceName(name string) RouterOption {
 		cfg.serviceName = name
 	}
 }
+
+// WithSpanOptions applies the given set of options to the spans started
+// by the router.
+func WithSpanOptions(opts ...ddtrace.StartSpanOption) RouterOption {
+	return func(cfg *routerConfig) {
+		cfg.spanOpts = opts
+	}
+}