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)
+		}
+	}
+}