Allow passing in a custom HTTPServer (#341)

* Avoid opening browser when not in terminal

This also reverts e82ee9a, though only in spirit since the revert
did not apply cleanly.

This change may look cosmetic, but note that `plugin.UI` can be
passed in from the outside. This is important for embedding pprof
into an application for easy access to the web ui: with this change,
such an application can pass a non-Terminal `profile.UI` and avoids
the browser code.

* Allow passing in a custom HTTPServer

This allows embedding the pprof ui into an existing application.
See https://github.com/cockroachdb/cockroach/pull/24145 for an example.

driver/driver.go

@@ -41,24 +41,36 @@ func (o *Options) internalOptions() *plugin.Options {
 	if o.Sym != nil {
 		sym = &internalSymbolizer{o.Sym}
 	}
+	var httpServer func(args *plugin.HTTPServerArgs) error
+	if o.HTTPServer != nil {
+		httpServer = func(args *plugin.HTTPServerArgs) error {
+			return o.HTTPServer(((*HTTPServerArgs)(args)))
+		}
+	}
 	return &plugin.Options{
-		Writer:  o.Writer,
-		Flagset: o.Flagset,
-		Fetch:   o.Fetch,
-		Sym:     sym,
-		Obj:     obj,
-		UI:      o.UI,
+		Writer:     o.Writer,
+		Flagset:    o.Flagset,
+		Fetch:      o.Fetch,
+		Sym:        sym,
+		Obj:        obj,
+		UI:         o.UI,
+		HTTPServer: httpServer,
 	}
 }
 
+// HTTPServerArgs contains arguments needed by an HTTP server that
+// is exporting a pprof web interface.
+type HTTPServerArgs plugin.HTTPServerArgs
+
 // Options groups all the optional plugins into pprof.
 type Options struct {
-	Writer  Writer
-	Flagset FlagSet
-	Fetch   Fetcher
-	Sym     Symbolizer
-	Obj     ObjTool
-	UI      UI
+	Writer     Writer
+	Flagset    FlagSet
+	Fetch      Fetcher
+	Sym        Symbolizer
+	Obj        ObjTool
+	UI         UI
+	HTTPServer func(*HTTPServerArgs) error
 }
 
 // Writer provides a mechanism to write data under a certain name,

internal/driver/driver.go

@@ -54,7 +54,7 @@ func PProf(eo *plugin.Options) error {
 	}
 
 	if src.HTTPHostport != "" {
-		return serveWebInterface(src.HTTPHostport, p, o, true)
+		return serveWebInterface(src.HTTPHostport, p, o)
 	}
 	return interactive(p, o)
 }

internal/driver/webui.go

@@ -82,7 +82,7 @@ type webArgs struct {
 	FlameGraph template.JS
 }
 
-func serveWebInterface(hostport string, p *profile.Profile, o *plugin.Options, wantBrowser bool) error {
+func serveWebInterface(hostport string, p *profile.Profile, o *plugin.Options) error {
 	host, port, err := getHostAndPort(hostport)
 	if err != nil {
 		return err
@@ -117,7 +117,7 @@ func serveWebInterface(hostport string, p *profile.Profile, o *plugin.Options, w
 		},
 	}
 
-	if wantBrowser {
+	if o.UI.IsTerminal() {
 		go openBrowser("http://"+args.Hostport, o)
 	}
 	return server(args)

internal/driver/webui_test.go

@@ -31,6 +31,12 @@ import (
 	"github.com/google/pprof/profile"
 )
 
+type nonTerminalUI struct {
+	plugin.UI
+}
+
+func (*nonTerminalUI) IsTerminal() bool { return false }
+
 func TestWebInterface(t *testing.T) {
 	if runtime.GOOS == "nacl" {
 		t.Skip("test assumes tcp available")
@@ -55,9 +61,9 @@ func TestWebInterface(t *testing.T) {
 	// Start server and wait for it to be initialized
 	go serveWebInterface("unused:1234", prof, &plugin.Options{
 		Obj:        fakeObjTool{},
-		UI:         &stdUI{},
+		UI:         &nonTerminalUI{&stdUI{}}, // don't open browser
 		HTTPServer: creator,
-	}, false)
+	})
 	<-serverCreated
 	defer server.Close()