diff options
Diffstat (limited to 'middleware/pprof')
-rw-r--r-- | middleware/pprof/pprof.go | 2 | ||||
-rw-r--r-- | middleware/pprof/setup.go | 40 | ||||
-rw-r--r-- | middleware/pprof/setup_test.go | 32 |
3 files changed, 73 insertions, 1 deletions
diff --git a/middleware/pprof/pprof.go b/middleware/pprof/pprof.go index 42102d198..f538b3091 100644 --- a/middleware/pprof/pprof.go +++ b/middleware/pprof/pprof.go @@ -12,7 +12,7 @@ type Handler struct { mux *http.ServeMux } -func (h *Handler) Start() error { +func (h *Handler) Startup() error { if ln, err := net.Listen("tcp", addr); err != nil { log.Printf("[ERROR] Failed to start pprof handler: %s", err) return err diff --git a/middleware/pprof/setup.go b/middleware/pprof/setup.go new file mode 100644 index 000000000..2b6701f35 --- /dev/null +++ b/middleware/pprof/setup.go @@ -0,0 +1,40 @@ +package pprof + +import ( + "sync" + + "github.com/mholt/caddy" +) + +func init() { + caddy.RegisterPlugin("pprof", caddy.Plugin{ + ServerType: "dns", + Action: setup, + }) +} + +func setup(c *caddy.Controller) error { + found := false + for c.Next() { + if found { + return c.Err("pprof can only be specified once") + } + if len(c.RemainingArgs()) != 0 { + return c.ArgErr() + } + if c.NextBlock() { + return c.ArgErr() + } + found = true + } + + handler := &Handler{} + pprofOnce.Do(func() { + c.OnStartup(handler.Startup) + c.OnShutdown(handler.Shutdown) + }) + + return nil +} + +var pprofOnce sync.Once diff --git a/middleware/pprof/setup_test.go b/middleware/pprof/setup_test.go new file mode 100644 index 000000000..af46fd415 --- /dev/null +++ b/middleware/pprof/setup_test.go @@ -0,0 +1,32 @@ +package pprof + +import ( + "testing" + + "github.com/mholt/caddy" +) + +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 := caddy.NewTestController("dns", test.input) + err := setup(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) + } + } +} |