Various trace improvements (#527)

4 files changed, 130 insertions, 46 deletions

diff --git a/middleware/trace/README.md b/middleware/trace/README.md
index 394de67f1..aa157e1e2 100644
--- a/middleware/trace/README.md
+++ b/middleware/trace/README.md
@@ -5,26 +5,42 @@ middleware chain.
 
 ## Syntax
 
+The simplest form is just:
+
 ~~~
 trace [ENDPOINT-TYPE] [ENDPOINT]
 ~~~
 
-For each server you which to trace.
+* **ENDPOINT-TYPE** is the type of tracing destination. Currently only `zipkin` is supported
+  and that is what it defaults to.
+* **ENDPOINT** is the tracing destination, and defaults to `localhost:9411`. For Zipkin, if
+  ENDPOINT does not begin with `http`, then it will be transformed to `http://ENDPOINT/api/v1/spans`.
+
+With this form, all queries will be traced.
+
+Additional features can be enabled with this syntax:
+
+~~~
+trace [ENDPOINT-TYPE] [ENDPOINT] {
+	every AMOUNT
+	service NAME
+	client_server
+}
+~~~
 
-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.
+* `every` **AMOUNT** will only trace one query of each AMOUNT queries. For example, to trace 1 in every
+  100 queries, use AMOUNT of 100. The default is 1.
+* `service` **NAME** allows you to specify the service name reported to the tracing server.
+  Default is `coredns`.
+* `client_server` will enable the `ClientServerSameSpan` OpenTracing feature.
 
-The only ENDPOINT-TYPE supported so far is `zipkin`. You can run Zipkin on a Docker host
-like this:
+## 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:
@@ -45,3 +61,13 @@ the standard Zipkin URL you can do something like:
 ~~~
 trace http://tracinghost:9411/zipkin/api/v1/spans
 ~~~
+
+Trace one query every 10000 queries, rename the service, and enable same span:
+
+~~~
+trace tracinghost:9411 {
+	every 10000
+	service dnsproxy
+	client_server
+}
+~~~
diff --git a/middleware/trace/setup.go b/middleware/trace/setup.go
index 302d43d4d..57531d34c 100644
--- a/middleware/trace/setup.go
+++ b/middleware/trace/setup.go
@@ -2,8 +2,8 @@ package trace
 
 import (
 	"fmt"
+	"strconv"
 	"strings"
-	"sync"
 
 	"github.com/miekg/coredns/core/dnsserver"
 	"github.com/miekg/coredns/middleware"
@@ -29,38 +29,65 @@ func setup(c *caddy.Controller) error {
 		return t
 	})
 
-	traceOnce.Do(func() {
-		c.OnStartup(t.OnStartup)
-	})
+	c.OnStartup(t.OnStartup)
 
 	return nil
 }
 
 func traceParse(c *caddy.Controller) (*Trace, error) {
 	var (
-		tr  = &Trace{Endpoint: defEP, EndpointType: defEpType}
+		tr  = &Trace{Endpoint: defEP, EndpointType: defEpType, every: 1, serviceName: defServiceName}
 		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
+	for c.Next() { // 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
+		}
+                for c.NextBlock() {
+                        switch c.Val() {
+                        case "every":
+                                args := c.RemainingArgs()
+                                if len(args) != 1 {
+                                        return nil, c.ArgErr()
+                                }
+                                tr.every, err = strconv.ParseUint(args[0], 10, 64)
+                                if err != nil {
+                                        return nil, err
+                                }
+			case "service":
+				args := c.RemainingArgs()
+                                if len(args) != 1 {
+                                        return nil, c.ArgErr()
+                                }
+				tr.serviceName = args[0]
+			case "client_server":
+				args := c.RemainingArgs()
+				if len(args) > 1 {
+					return nil, c.ArgErr()
+				}
+				tr.clientServer = true
+				if len(args) == 1 {
+					tr.clientServer, err = strconv.ParseBool(args[0])
+				}
+				if err != nil {
+					return nil, err
+				}
 			}
 		}
 	}
@@ -79,9 +106,8 @@ func normalizeEndpoint(epType, ep string) (string, error) {
 	}
 }
 
-var traceOnce sync.Once
-
 const (
 	defEP     = "localhost:9411"
 	defEpType = "zipkin"
+	defServiceName = "coredns"
 )
diff --git a/middleware/trace/setup_test.go b/middleware/trace/setup_test.go
index 0b9ac8c42..b565cd5d1 100644
--- a/middleware/trace/setup_test.go
+++ b/middleware/trace/setup_test.go
@@ -11,15 +11,23 @@ func TestTraceParse(t *testing.T) {
 		input     string
 		shouldErr bool
 		endpoint  string
+		every	uint64
+		serviceName	string
+		clientServer	bool
 	}{
 		// 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"},
+		{`trace`, false, "http://localhost:9411/api/v1/spans", 1, `coredns`, false},
+		{`trace localhost:1234`, false, "http://localhost:1234/api/v1/spans", 1, `coredns`, false},
+		{`trace http://localhost:1234/somewhere/else`, false, "http://localhost:1234/somewhere/else", 1, `coredns`, false},
+		{`trace zipkin localhost:1234`, false, "http://localhost:1234/api/v1/spans", 1, `coredns`, false},
+		{`trace zipkin http://localhost:1234/somewhere/else`, false, "http://localhost:1234/somewhere/else", 1, `coredns`, false},
+		{"trace {\n every 100\n}", false, "http://localhost:9411/api/v1/spans", 100, `coredns`, false},
+		{"trace {\n every 100\n service foobar\nclient_server\n}", false, "http://localhost:9411/api/v1/spans", 100, `foobar`, true},
+		{"trace {\n every 2\n client_server true\n}", false, "http://localhost:9411/api/v1/spans", 2, `coredns`, true},
+		{"trace {\n client_server false\n}", false, "http://localhost:9411/api/v1/spans", 1, `coredns`, false},
 		// fails
-		{`trace footype localhost:4321`, true, ""},
+		{`trace footype localhost:4321`, true, "", 1, "", false},
+		{"trace {\n every 2\n client_server junk\n}", true, "", 1, "", false},
 	}
 	for i, test := range tests {
 		c := caddy.NewTestController("dns", test.input)
@@ -39,5 +47,14 @@ func TestTraceParse(t *testing.T) {
 		if test.endpoint != m.Endpoint {
 			t.Errorf("Test %v: Expected endpoint %s but found: %s", i, test.endpoint, m.Endpoint)
 		}
+		if test.every != m.every {
+			t.Errorf("Test %v: Expected every %d but found: %d", i, test.every, m.every)
+		}
+		if test.serviceName != m.serviceName {
+			t.Errorf("Test %v: Expected service name %s but found: %s", i, test.serviceName, m.serviceName)
+		}
+		if test.clientServer != m.clientServer {
+			t.Errorf("Test %v: Expected client_server %t but found: %t", i, test.clientServer, m.clientServer)
+		}
 	}
 }
diff --git a/middleware/trace/trace.go b/middleware/trace/trace.go
index 22b174016..b7ce3c27c 100644
--- a/middleware/trace/trace.go
+++ b/middleware/trace/trace.go
@@ -4,13 +4,14 @@ package trace
 import (
 	"fmt"
 	"sync"
-
-	"golang.org/x/net/context"
+	"sync/atomic"
 
 	"github.com/miekg/coredns/middleware"
 	"github.com/miekg/dns"
 	ot "github.com/opentracing/opentracing-go"
 	zipkin "github.com/openzipkin/zipkin-go-opentracing"
+
+	"golang.org/x/net/context"
 )
 
 // Trace holds the tracer and endpoint info
@@ -20,6 +21,10 @@ type Trace struct {
 	Endpoint        string
 	EndpointType    string
 	Tracer          ot.Tracer
+	serviceName	string
+	clientServer	bool
+	every		uint64
+	count		uint64
 	Once            sync.Once
 }
 
@@ -44,8 +49,8 @@ func (t *Trace) setupZipkin() error {
 		return err
 	}
 
-	recorder := zipkin.NewRecorder(collector, false, t.ServiceEndpoint, "coredns")
-	t.Tracer, err = zipkin.NewTracer(recorder, zipkin.ClientServerSameSpan(false))
+	recorder := zipkin.NewRecorder(collector, false, t.ServiceEndpoint, t.serviceName)
+	t.Tracer, err = zipkin.NewTracer(recorder, zipkin.ClientServerSameSpan(t.clientServer))
 	if err != nil {
 		return err
 	}
@@ -59,8 +64,18 @@ func (t *Trace) Name() string {
 
 // ServeDNS implements the middleware.Handle interface.
 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)
+	trace := false
+	if t.every > 0 {
+		queryNr := atomic.AddUint64(&t.count, 1)
+
+		if queryNr%t.every == 0 {
+			trace = true
+		}
+	}
+	if span := ot.SpanFromContext(ctx); span == nil && trace {
+		span := t.Tracer.StartSpan("servedns")
+		defer span.Finish()
+		ctx = ot.ContextWithSpan(ctx, span)
+	}
 	return middleware.NextOrFailure(t.Name(), t.Next, ctx, w, r)
 }