Grpc tracing (#544)

* checkpoint

* Pass context through ServeDNS, enable gRPC tracing

* Fix types and make tracer available to proxy. go fmt

* Fix imports

* Use the DoNotStartTrace option

* Change to SpanFilter from DoNotStartTrace

* Use new name (IncludeSpan)

* Final names

* Add tests; fix possible client/conn leaks in grpc

* go fmt

middleware/trace/setup.go

@@ -34,9 +34,9 @@ func setup(c *caddy.Controller) error {
 	return nil
 }
 
-func traceParse(c *caddy.Controller) (*Trace, error) {
+func traceParse(c *caddy.Controller) (*trace, error) {
 	var (
-		tr  = &Trace{Endpoint: defEP, EndpointType: defEpType, every: 1, serviceName: defServiceName}
+		tr  = &trace{Endpoint: defEP, EndpointType: defEpType, every: 1, serviceName: defServiceName}
 		err error
 	)
 

middleware/trace/trace.go

@@ -15,12 +15,17 @@ import (
 )
 
 // Trace holds the tracer and endpoint info
-type Trace struct {
+type Trace interface {
+	middleware.Handler
+	Tracer() ot.Tracer
+}
+
+type trace struct {
 	Next            middleware.Handler
 	ServiceEndpoint string
 	Endpoint        string
 	EndpointType    string
-	Tracer          ot.Tracer
+	tracer          ot.Tracer
 	serviceName     string
 	clientServer    bool
 	every           uint64
@@ -28,8 +33,12 @@ type Trace struct {
 	Once            sync.Once
 }
 
+func (t *trace) Tracer() ot.Tracer {
+	return t.tracer
+}
+
 // OnStartup sets up the tracer
-func (t *Trace) OnStartup() error {
+func (t *trace) OnStartup() error {
 	var err error
 	t.Once.Do(func() {
 		switch t.EndpointType {
@@ -42,7 +51,7 @@ func (t *Trace) OnStartup() error {
 	return err
 }
 
-func (t *Trace) setupZipkin() error {
+func (t *trace) setupZipkin() error {
 
 	collector, err := zipkin.NewHTTPCollector(t.Endpoint)
 	if err != nil {
@@ -50,7 +59,7 @@ func (t *Trace) setupZipkin() error {
 	}
 
 	recorder := zipkin.NewRecorder(collector, false, t.ServiceEndpoint, t.serviceName)
-	t.Tracer, err = zipkin.NewTracer(recorder, zipkin.ClientServerSameSpan(t.clientServer))
+	t.tracer, err = zipkin.NewTracer(recorder, zipkin.ClientServerSameSpan(t.clientServer))
 	if err != nil {
 		return err
 	}
@@ -58,12 +67,12 @@ func (t *Trace) setupZipkin() error {
 }
 
 // Name implements the Handler interface.
-func (t *Trace) Name() string {
+func (t *trace) Name() string {
 	return "trace"
 }
 
 // ServeDNS implements the middleware.Handle interface.
-func (t *Trace) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
+func (t *trace) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
 	trace := false
 	if t.every > 0 {
 		queryNr := atomic.AddUint64(&t.count, 1)
@@ -73,7 +82,7 @@ func (t *Trace) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg)
 		}
 	}
 	if span := ot.SpanFromContext(ctx); span == nil && trace {
-		span := t.Tracer.StartSpan("servedns")
+		span := t.Tracer().StartSpan("servedns")
 		defer span.Finish()
 		ctx = ot.ContextWithSpan(ctx, span)
 	}

middleware/trace/trace_test.go

@@ -0,0 +1,33 @@
+package trace
+
+import (
+	"testing"
+
+	"github.com/mholt/caddy"
+)
+
+// CreateTestTrace creates a trace middleware to be used in tests
+func CreateTestTrace(config string) (*caddy.Controller, *trace, error) {
+	c := caddy.NewTestController("dns", config)
+	m, err := traceParse(c)
+	return c, m, err
+}
+
+func TestTrace(t *testing.T) {
+	_, m, err := CreateTestTrace(`trace`)
+	if err != nil {
+		t.Errorf("Error parsing test input: %s", err)
+		return
+	}
+	if m.Name() != "trace" {
+		t.Errorf("Wrong name from GetName: %s", m.Name())
+	}
+	err = m.OnStartup()
+	if err != nil {
+		t.Errorf("Error starting tracing middleware: %s", err)
+		return
+	}
+	if m.Tracer() == nil {
+		t.Errorf("Error, no tracer created")
+	}
+}