Add tracing option (#487)

Adds a middleware to enable tracing with OpenTracing/OpenZipkin.
Enabling tracing will have a large impact on performance so it is
not advisable in production.

core/coredns.go

@@ -26,5 +26,6 @@ import (
 	_ "github.com/miekg/coredns/middleware/rewrite"
 	_ "github.com/miekg/coredns/middleware/root"
 	_ "github.com/miekg/coredns/middleware/secondary"
+	_ "github.com/miekg/coredns/middleware/trace"
 	_ "github.com/miekg/coredns/middleware/whoami"
 )

core/dnsserver/directives.go

@@ -75,6 +75,7 @@ func RegisterDevDirective(name, before string) {
 var directives = []string{
 	"root",
 	"bind",
+	"trace",
 	"health",
 	"pprof",
 

middleware/middleware.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 
 	"github.com/miekg/dns"
+        ot "github.com/opentracing/opentracing-go"
 	"golang.org/x/net/context"
 )
 
@@ -70,6 +71,11 @@ func Error(name string, err error) error { return fmt.Errorf("%s/%s: %s", "middl
 // and a nil error.
 func NextOrFailure(name string, next Handler, ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
 	if next != nil {
+		if span := ot.SpanFromContext(ctx); span != nil {
+			child := span.Tracer().StartSpan(next.Name(), ot.ChildOf(span.Context()))
+			defer child.Finish()
+			ctx = ot.ContextWithSpan(ctx, child)
+		}
 		return next.ServeDNS(ctx, w, r)
 	}
 

middleware/proxy/proxy.go

@@ -10,6 +10,7 @@ import (
 	"github.com/miekg/coredns/request"
 
 	"github.com/miekg/dns"
+        ot "github.com/opentracing/opentracing-go"
 	"golang.org/x/net/context"
 )
 
@@ -70,6 +71,8 @@ var tryDuration = 60 * time.Second
 
 // ServeDNS satisfies the middleware.Handler interface.
 func (p Proxy) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
+	var span, child ot.Span
+	span = ot.SpanFromContext(ctx)
 	state := request.Request{W: w, Req: r}
 	for _, upstream := range p.Upstreams {
 		start := time.Now()
@@ -85,12 +88,21 @@ func (p Proxy) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (
 				return dns.RcodeServerFailure, errUnreachable
 			}
 
+			if span != nil {
+				child = span.Tracer().StartSpan("exchange", ot.ChildOf(span.Context()))
+				ctx = ot.ContextWithSpan(ctx, child)
+			}
+
 			atomic.AddInt64(&host.Conns, 1)
 
 			reply, backendErr := host.Exchange(state)
 
 			atomic.AddInt64(&host.Conns, -1)
 
+			if child != nil {
+				child.Finish()
+			}
+
 			if backendErr == nil {
 				w.WriteMsg(reply)
 

middleware/trace/README.md

@@ -0,0 +1,47 @@
+# trace
+
+This module enables OpenTracing-based tracing of DNS requests as they go through the
+middleware chain.
+
+## Syntax
+
+~~~
+trace [ENDPOINT-TYPE] [ENDPOINT]
+~~~
+
+For each server you which to trace.
+
+It optionally takes the ENDPOINT-TYPE and ENDPOINT. The ENDPOINT-TYPE defaults to
+`zipkin` and the ENDPOINT to `localhost:9411`. A single argument will be interpreted as
+a Zipkin ENDPOINT.
+
+The only ENDPOINT-TYPE supported so far is `zipkin`. You can run Zipkin on a Docker host
+like this:
+
+```
+docker run -d -p 9411:9411 openzipkin/zipkin
+```
+
+For Zipkin, if ENDPOINT does not begin with `http`, then it will be transformed to
+`http://ENDPOINT/api/v1/spans`.
+
+## Examples
+
+Use an alternative Zipkin address:
+
+~~~
+trace tracinghost:9253
+~~~
+
+or
+
+~~~
+trace zipkin tracinghost:9253
+~~~
+
+If for some reason you are using an API reverse proxy or something and need to remap
+the standard Zipkin URL you can do something like:
+
+~~~
+trace http://tracinghost:9411/zipkin/api/v1/spans
+~~~

middleware/trace/setup.go

@@ -0,0 +1,87 @@
+package trace
+
+import (
+	"fmt"
+	"strings"
+	"sync"
+
+	"github.com/miekg/coredns/core/dnsserver"
+	"github.com/miekg/coredns/middleware"
+
+	"github.com/mholt/caddy"
+)
+
+func init() {
+	caddy.RegisterPlugin("trace", caddy.Plugin{
+		ServerType: "dns",
+		Action:     setup,
+	})
+}
+
+func setup(c *caddy.Controller) error {
+	t, err := traceParse(c)
+	if err != nil {
+		return middleware.Error("trace", err)
+	}
+
+	dnsserver.GetConfig(c).AddMiddleware(func(next middleware.Handler) middleware.Handler {
+		t.Next = next
+		return t
+	})
+
+	traceOnce.Do(func() {
+		c.OnStartup(t.OnStartup)
+	})
+
+	return nil
+}
+
+func traceParse(c *caddy.Controller) (*Trace, error) {
+	var (
+		tr = &Trace{Endpoint: defEP, EndpointType: defEpType}
+		err error
+	)
+
+	cfg := dnsserver.GetConfig(c)
+	tr.ServiceEndpoint = cfg.ListenHost + ":" + cfg.Port
+	for c.Next() {
+		if c.Val() == "trace" {
+			var err error
+			args := c.RemainingArgs()
+			switch len(args) {
+			case 0:
+				tr.Endpoint, err = normalizeEndpoint(tr.EndpointType, defEP)
+			case 1:
+				tr.Endpoint, err = normalizeEndpoint(defEpType, args[0])
+			case 2:
+				tr.EndpointType = strings.ToLower(args[0])
+				tr.Endpoint, err = normalizeEndpoint(tr.EndpointType, args[1])
+			default:
+				err = c.ArgErr()
+			}
+			if err != nil {
+				return tr, err
+			}
+		}
+	}
+	return tr, err
+}
+
+func normalizeEndpoint(epType, ep string) (string, error) {
+	switch epType {
+	case "zipkin":
+		if strings.Index(ep, "http") == -1 {
+			ep = "http://" + ep + "/api/v1/spans"
+		}
+		return ep, nil
+	default:
+		return "", fmt.Errorf("Tracing endpoint type '%s' is not supported.", epType)
+	}
+}
+
+var traceOnce sync.Once
+
+const (
+	defEP = "localhost:9411"
+	defEpType = "zipkin"
+)

middleware/trace/setup_test.go

@@ -0,0 +1,43 @@
+package trace
+
+import (
+	"testing"
+
+	"github.com/mholt/caddy"
+)
+
+func TestTraceParse(t *testing.T) {
+	tests := []struct {
+		input     string
+		shouldErr bool
+		endpoint      string
+	}{
+		// oks
+		{`trace`, false, "http://localhost:9411/api/v1/spans"},
+		{`trace localhost:1234`, false, "http://localhost:1234/api/v1/spans"},
+		{`trace http://localhost:1234/somewhere/else`, false, "http://localhost:1234/somewhere/else"},
+		{`trace zipkin localhost:1234`, false, "http://localhost:1234/api/v1/spans"},
+		{`trace zipkin http://localhost:1234/somewhere/else`, false, "http://localhost:1234/somewhere/else"},
+		// fails
+		{`trace footype localhost:4321`, true, ""},
+	}
+	for i, test := range tests {
+		c := caddy.NewTestController("dns", test.input)
+		m, err := traceParse(c)
+		if test.shouldErr && err == nil {
+			t.Errorf("Test %v: Expected error but found nil", i)
+			continue
+		} else if !test.shouldErr && err != nil {
+			t.Errorf("Test %v: Expected no error but found error: %v", i, err)
+			continue
+		}
+
+		if test.shouldErr {
+			continue
+		}
+
+		if test.endpoint != m.Endpoint {
+			t.Errorf("Test %v: Expected endpoint %s but found: %s", i, test.endpoint, m.Endpoint)
+		}
+	}
+}

middleware/trace/trace.go

@@ -0,0 +1,64 @@
+// Package trace implements OpenTracing-based tracing
+package trace
+
+import (
+	"fmt"
+	"sync"
+
+        "golang.org/x/net/context"
+
+	"github.com/miekg/coredns/middleware"
+        "github.com/miekg/dns"
+        ot "github.com/opentracing/opentracing-go"
+        zipkin "github.com/openzipkin/zipkin-go-opentracing"
+)
+
+// Trace holds the tracer and endpoint info
+type Trace struct {
+	Next middleware.Handler
+	ServiceEndpoint string
+	Endpoint string
+	EndpointType string
+	Tracer ot.Tracer
+	Once sync.Once
+}
+
+// OnStartup sets up the tracer
+func (t *Trace) OnStartup() error {
+	var err error
+	t.Once.Do(func() {
+		switch t.EndpointType {
+		case "zipkin":
+			err = t.setupZipkin()
+		default:
+			err = fmt.Errorf("Unknown endpoint type: %s", t.EndpointType)
+		}
+	})
+	return err
+}
+
+func (t *Trace) setupZipkin() error {
+
+	collector, err := zipkin.NewHTTPCollector(t.Endpoint)
+	if err != nil {
+		return err
+	}
+
+	recorder := zipkin.NewRecorder(collector, false, t.ServiceEndpoint, "coredns")
+	t.Tracer, err = zipkin.NewTracer(recorder, zipkin.ClientServerSameSpan(false))
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (t *Trace) Name() (string) {
+	return "trace"
+}
+
+func (t *Trace) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
+	span := t.Tracer.StartSpan("servedns")
+	defer span.Finish()
+	ctx = ot.ContextWithSpan(ctx, span)
+	return middleware.NextOrFailure(t.Name(), t.Next, ctx, w, r)
+}