Add request and response context to traces (#2162)

Automatically submitted.

2 files changed, 104 insertions, 13 deletions

diff --git a/plugin/trace/trace.go b/plugin/trace/trace.go
index a0a2071ea..5421836a5 100644
--- a/plugin/trace/trace.go
+++ b/plugin/trace/trace.go
@@ -10,8 +10,11 @@ import (
 
 	"github.com/coredns/coredns/plugin"
 	"github.com/coredns/coredns/plugin/metrics"
+	"github.com/coredns/coredns/plugin/pkg/dnstest"
+	"github.com/coredns/coredns/plugin/pkg/rcode"
 	// Plugin the trace package.
 	_ "github.com/coredns/coredns/plugin/pkg/trace"
+	"github.com/coredns/coredns/request"
 
 	ddtrace "github.com/DataDog/dd-trace-go/opentracing"
 	"github.com/miekg/dns"
@@ -19,6 +22,12 @@ import (
 	zipkin "github.com/openzipkin/zipkin-go-opentracing"
 )
 
+const (
+	tagName  = "coredns.io/name"
+	tagType  = "coredns.io/type"
+	tagRcode = "coredns.io/rcode"
+)
+
 type trace struct {
 	Next            plugin.Handler
 	Endpoint        string
@@ -94,10 +103,26 @@ func (t *trace) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg)
 			trace = true
 		}
 	}
-	if span := ot.SpanFromContext(ctx); span == nil && trace {
-		span := t.Tracer().StartSpan("servedns:" + metrics.WithServer(ctx))
-		defer span.Finish()
-		ctx = ot.ContextWithSpan(ctx, span)
+	span := ot.SpanFromContext(ctx)
+	if !trace || span != nil {
+		return plugin.NextOrFailure(t.Name(), t.Next, ctx, w, r)
 	}
-	return plugin.NextOrFailure(t.Name(), t.Next, ctx, w, r)
+
+	req := request.Request{W: w, Req: r}
+	span = t.Tracer().StartSpan(spanName(ctx, req))
+	defer span.Finish()
+
+	rw := dnstest.NewRecorder(w)
+	ctx = ot.ContextWithSpan(ctx, span)
+	status, err := plugin.NextOrFailure(t.Name(), t.Next, ctx, rw, r)
+
+	span.SetTag(tagName, req.Name())
+	span.SetTag(tagType, req.Type())
+	span.SetTag(tagRcode, rcode.ToString(rw.Rcode))
+
+	return status, err
+}
+
+func spanName(ctx context.Context, req request.Request) string {
+	return "servedns:" + metrics.WithServer(ctx) + " " + req.Name()
 }
diff --git a/plugin/trace/trace_test.go b/plugin/trace/trace_test.go
index b006009c3..7fb1cb85f 100644
--- a/plugin/trace/trace_test.go
+++ b/plugin/trace/trace_test.go
@@ -1,20 +1,24 @@
 package trace
 
 import (
+	"context"
 	"testing"
 
+	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/plugin/pkg/dnstest"
+	"github.com/coredns/coredns/plugin/pkg/rcode"
+	"github.com/coredns/coredns/plugin/test"
+	"github.com/coredns/coredns/request"
+
 	"github.com/mholt/caddy"
+	"github.com/miekg/dns"
+	"github.com/opentracing/opentracing-go/mocktracer"
 )
 
-// createTestTrace creates a trace plugin 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
-}
+const server = "coolServer"
 
-func TestTrace(t *testing.T) {
-	_, m, err := createTestTrace(`trace`)
+func TestStartup(t *testing.T) {
+	m, err := traceParse(caddy.NewTestController("dns", `trace`))
 	if err != nil {
 		t.Errorf("Error parsing test input: %s", err)
 		return
@@ -31,3 +35,65 @@ func TestTrace(t *testing.T) {
 		t.Errorf("Error, no tracer created")
 	}
 }
+
+func TestTrace(t *testing.T) {
+	cases := []struct {
+		name     string
+		rcode    int
+		question *dns.Msg
+		server   string
+	}{
+		{
+			name:     "NXDOMAIN",
+			rcode:    dns.RcodeNameError,
+			question: new(dns.Msg).SetQuestion("example.org.", dns.TypeA),
+		},
+		{
+			name:     "NOERROR",
+			rcode:    dns.RcodeSuccess,
+			question: new(dns.Msg).SetQuestion("example.net.", dns.TypeCNAME),
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			w := dnstest.NewRecorder(&test.ResponseWriter{})
+			m := mocktracer.New()
+			tr := &trace{
+				Next: test.HandlerFunc(func(_ context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
+					m := new(dns.Msg)
+					m.SetRcode(r, tc.rcode)
+					w.WriteMsg(m)
+					return tc.rcode, nil
+				}),
+				every:  1,
+				tracer: m,
+			}
+			ctx := context.WithValue(context.TODO(), plugin.ServerCtx{}, server)
+			if _, err := tr.ServeDNS(ctx, w, tc.question); err != nil {
+				t.Fatalf("Error during tr.ServeDNS(ctx, w, %v): %v", tc.question, err)
+			}
+
+			fs := m.FinishedSpans()
+			// Each trace consists of two spans; the root and the Next function.
+			if len(fs) != 2 {
+				t.Fatalf("Unexpected span count: len(fs): want 2, got %v", len(fs))
+			}
+
+			rootSpan := fs[1]
+			req := request.Request{W: w, Req: tc.question}
+			if rootSpan.OperationName != spanName(ctx, req) {
+				t.Errorf("Unexpected span name: rootSpan.Name: want %v, got %v", spanName(ctx, req), rootSpan.OperationName)
+			}
+			if rootSpan.Tag(tagName) != req.Name() {
+				t.Errorf("Unexpected span tag: rootSpan.Tag(%v): want %v, got %v", tagName, req.Name(), rootSpan.Tag(tagName))
+			}
+			if rootSpan.Tag(tagType) != req.Type() {
+				t.Errorf("Unexpected span tag: rootSpan.Tag(%v): want %v, got %v", tagType, req.Type(), rootSpan.Tag(tagType))
+			}
+			if rootSpan.Tag(tagRcode) != rcode.ToString(tc.rcode) {
+				t.Errorf("Unexpected span tag: rootSpan.Tag(%v): want %v, got %v", tagRcode, rcode.ToString(tc.rcode), rootSpan.Tag(tagRcode))
+			}
+		})
+	}
+}