pprof middleware (#138)

Add pprof middleware, enabled by pprof directive.

core/directives.go

@@ -46,6 +46,7 @@ var directiveOrder = []directive{
 	{"bind", setup.BindHost},
 	{"tls", https.Setup},
 	{"health", setup.Health},
+	{"pprof", setup.PProf},
 
 	// Other directives that don't create HTTP handlers
 	{"startup", setup.Startup},

core/setup/metrics.go

@@ -9,7 +9,7 @@ import (
 
 const addr = "localhost:9135" // 9153 is occupied by bind_exporter
 
-var once sync.Once
+var metricsOnce sync.Once
 
 func Prometheus(c *Controller) (middleware.Middleware, error) {
 	met, err := parsePrometheus(c)
@@ -17,7 +17,7 @@ func Prometheus(c *Controller) (middleware.Middleware, error) {
 		return nil, err
 	}
 
-	once.Do(func() {
+	metricsOnce.Do(func() {
 		c.Startup = append(c.Startup, met.Start)
 	})
 

core/setup/pprof.go

@@ -0,0 +1,36 @@
+package setup
+
+import (
+	"sync"
+
+	"github.com/miekg/coredns/middleware"
+	"github.com/miekg/coredns/middleware/pprof"
+)
+
+var pprofOnce sync.Once
+
+// PProf returns a new instance of a pprof handler. It accepts no arguments or options.
+func PProf(c *Controller) (middleware.Middleware, error) {
+	found := false
+	for c.Next() {
+		if found {
+			return nil, c.Err("pprof can only be specified once")
+		}
+		if len(c.RemainingArgs()) != 0 {
+			return nil, c.ArgErr()
+		}
+		if c.NextBlock() {
+			return nil, c.ArgErr()
+		}
+		found = true
+	}
+	handler := &pprof.Handler{}
+	pprofOnce.Do(func() {
+		c.Startup = append(c.Startup, handler.Start)
+	})
+
+	return func(next middleware.Handler) middleware.Handler {
+		handler.Next = next
+		return handler
+	}, nil
+}

core/setup/pprof_test.go

@@ -0,0 +1,28 @@
+package setup
+
+import "testing"
+
+func TestPProf(t *testing.T) {
+	tests := []struct {
+		input     string
+		shouldErr bool
+	}{
+		{`pprof`, false},
+		{`pprof {}`, true},
+		{`pprof /foo`, true},
+		{`pprof {
+            a b
+        }`, true},
+		{`pprof
+          pprof`, true},
+	}
+	for i, test := range tests {
+		c := NewTestController(test.input)
+		_, err := PProf(c)
+		if test.shouldErr && err == nil {
+			t.Errorf("Test %v: Expected error but found nil", i)
+		} else if !test.shouldErr && err != nil {
+			t.Errorf("Test %v: Expected no error but found error: %v", i, err)
+		}
+	}
+}

middleware/health/README.md

@@ -19,3 +19,7 @@ will just return "OK", when CoreDNS is healthy.
 This middleware only needs to be enabled once.
 
 ## Examples
+
+~~~
+health localhost:8091
+~~~

middleware/metrics/handler.go

@@ -3,10 +3,10 @@ package metrics
 import (
 	"time"
 
-	"golang.org/x/net/context"
-
 	"github.com/miekg/coredns/middleware"
+
 	"github.com/miekg/dns"
+	"golang.org/x/net/context"
 )
 
 func (m Metrics) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {

middleware/pprof/README.md

@@ -0,0 +1,25 @@
+# pprof
+
+pprof publishes runtime profiling data at endpoints under /debug/pprof. You can visit /debug/pprof
+on your site for an index of the available endpoints. By default it will listen on localhost:8053.
+
+> This is a debugging tool. Certain requests (such as collecting execution traces) can be slow. If
+> you use pprof on a live site, consider restricting access or enabling it only temporarily.
+
+For more information, please see [Go's pprof
+documentation](https://golang.org/pkg/net/http/pprof/s://golang.org/pkg/net/http/pprof/) and read
+[Profiling Go Programs](https://blog.golang.org/profiling-go-programs).
+
+## Syntax
+
+~~~
+pprof
+~~~
+
+## Examples
+
+Enable pprof endpoints:
+
+~~~
+pprof
+~~~

middleware/pprof/pprof.go

@@ -0,0 +1,32 @@
+package pprof
+
+import (
+	"log"
+	"net/http"
+	_ "net/http/pprof"
+
+	"github.com/miekg/coredns/middleware"
+
+	"github.com/miekg/dns"
+	"golang.org/x/net/context"
+)
+
+const addr = "localhost:8053"
+
+type Handler struct {
+	Next middleware.Handler
+}
+
+// ServeDNS passes all other requests up the chain.
+func (h *Handler) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
+	return h.Next.ServeDNS(ctx, w, r)
+}
+
+func (h *Handler) Start() error {
+	go func() {
+		if err := http.ListenAndServe(addr, nil); err != nil {
+			log.Printf("[ERROR] Failed to start pprof handler: %s", err)
+		}
+	}()
+	return nil
+}