Merge branch 'metrics-plug' of github.com:coredns/coredns into metrics-plug

plugin/metrics/handler.go

@@ -6,7 +6,6 @@ import (
 
 	"github.com/coredns/coredns/plugin"
 	"github.com/coredns/coredns/plugin/metrics/vars"
-	"github.com/coredns/coredns/plugin/pkg/dnstest"
 	"github.com/coredns/coredns/plugin/pkg/rcode"
 	"github.com/coredns/coredns/request"
 
@@ -24,7 +23,7 @@ func (m *Metrics) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg
 	}
 
 	// Record response to get status code and size of the reply.
-	rw := dnstest.NewRecorder(w)
+	rw := NewRecorder(w)
 	status, err := plugin.NextOrFailure(m.Name(), m.Next, ctx, rw, r)
 
 	rc := rw.Rcode

plugin/metrics/recorder.go

@@ -0,0 +1,31 @@
+package metrics
+
+import (
+	"runtime"
+
+	"github.com/coredns/coredns/plugin/pkg/dnstest"
+	"github.com/miekg/dns"
+)
+
+// Recorder is a dnstest.Recorder specific to the metrics plugin.
+type Recorder struct {
+	*dnstest.Recorder
+	// CallerN holds the string return value of the call to runtime.Caller(N)
+	Caller1 string
+	Caller2 string
+	Caller3 string
+}
+
+// NewRecorder makes and returns a new Recorder.
+func NewRecorder(w dns.ResponseWriter) *Recorder { return &Recorder{Recorder: dnstest.NewRecorder(w)} }
+
+// WriteMsg records the status code and calls the
+// underlying ResponseWriter's WriteMsg method.
+func (r *Recorder) WriteMsg(res *dns.Msg) error {
+	_, r.Caller1, _, _ = runtime.Caller(1)
+	_, r.Caller2, _, _ = runtime.Caller(2)
+	_, r.Caller3, _, _ = runtime.Caller(3)
+	r.Len += res.Len()
+	r.Msg = res
+	return r.ResponseWriter.WriteMsg(res)
+}