diff --git a/plugin/metrics/handler.go b/plugin/metrics/handler.go index 318254c6b..88fd2bcdb 100644 --- a/plugin/metrics/handler.go +++ b/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 diff --git a/plugin/metrics/recorder.go b/plugin/metrics/recorder.go new file mode 100644 index 000000000..fdfa1c2c4 --- /dev/null +++ b/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) +}