Pārlūkot izejas kodu

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.
Tobias Schottdorf 7 gadus atpakaļ
vecāks
revīzija
f0154354e2

+ 24
- 12
driver/driver.go Parādīt failu

41
 	if o.Sym != nil {
41
 	if o.Sym != nil {
42
 		sym = &internalSymbolizer{o.Sym}
42
 		sym = &internalSymbolizer{o.Sym}
43
 	}
43
 	}
44
+	var httpServer func(args *plugin.HTTPServerArgs) error
45
+	if o.HTTPServer != nil {
46
+		httpServer = func(args *plugin.HTTPServerArgs) error {
47
+			return o.HTTPServer(((*HTTPServerArgs)(args)))
48
+		}
49
+	}
44
 	return &plugin.Options{
50
 	return &plugin.Options{
45
-		Writer:  o.Writer,
46
-		Flagset: o.Flagset,
47
-		Fetch:   o.Fetch,
48
-		Sym:     sym,
49
-		Obj:     obj,
50
-		UI:      o.UI,
51
+		Writer:     o.Writer,
52
+		Flagset:    o.Flagset,
53
+		Fetch:      o.Fetch,
54
+		Sym:        sym,
55
+		Obj:        obj,
56
+		UI:         o.UI,
57
+		HTTPServer: httpServer,
51
 	}
58
 	}
52
 }
59
 }
53
 
60
 
61
+// HTTPServerArgs contains arguments needed by an HTTP server that
62
+// is exporting a pprof web interface.
63
+type HTTPServerArgs plugin.HTTPServerArgs
64
+
54
 // Options groups all the optional plugins into pprof.
65
 // Options groups all the optional plugins into pprof.
55
 type Options struct {
66
 type Options struct {
56
-	Writer  Writer
57
-	Flagset FlagSet
58
-	Fetch   Fetcher
59
-	Sym     Symbolizer
60
-	Obj     ObjTool
61
-	UI      UI
67
+	Writer     Writer
68
+	Flagset    FlagSet
69
+	Fetch      Fetcher
70
+	Sym        Symbolizer
71
+	Obj        ObjTool
72
+	UI         UI
73
+	HTTPServer func(*HTTPServerArgs) error
62
 }
74
 }
63
 
75
 
64
 // Writer provides a mechanism to write data under a certain name,
76
 // Writer provides a mechanism to write data under a certain name,

+ 1
- 1
internal/driver/driver.go Parādīt failu

54
 	}
54
 	}
55
 
55
 
56
 	if src.HTTPHostport != "" {
56
 	if src.HTTPHostport != "" {
57
-		return serveWebInterface(src.HTTPHostport, p, o, true)
57
+		return serveWebInterface(src.HTTPHostport, p, o)
58
 	}
58
 	}
59
 	return interactive(p, o)
59
 	return interactive(p, o)
60
 }
60
 }

+ 2
- 2
internal/driver/webui.go Parādīt failu

82
 	FlameGraph template.JS
82
 	FlameGraph template.JS
83
 }
83
 }
84
 
84
 
85
-func serveWebInterface(hostport string, p *profile.Profile, o *plugin.Options, wantBrowser bool) error {
85
+func serveWebInterface(hostport string, p *profile.Profile, o *plugin.Options) error {
86
 	host, port, err := getHostAndPort(hostport)
86
 	host, port, err := getHostAndPort(hostport)
87
 	if err != nil {
87
 	if err != nil {
88
 		return err
88
 		return err
117
 		},
117
 		},
118
 	}
118
 	}
119
 
119
 
120
-	if wantBrowser {
120
+	if o.UI.IsTerminal() {
121
 		go openBrowser("http://"+args.Hostport, o)
121
 		go openBrowser("http://"+args.Hostport, o)
122
 	}
122
 	}
123
 	return server(args)
123
 	return server(args)

+ 8
- 2
internal/driver/webui_test.go Parādīt failu

31
 	"github.com/google/pprof/profile"
31
 	"github.com/google/pprof/profile"
32
 )
32
 )
33
 
33
 
34
+type nonTerminalUI struct {
35
+	plugin.UI
36
+}
37
+
38
+func (*nonTerminalUI) IsTerminal() bool { return false }
39
+
34
 func TestWebInterface(t *testing.T) {
40
 func TestWebInterface(t *testing.T) {
35
 	if runtime.GOOS == "nacl" {
41
 	if runtime.GOOS == "nacl" {
36
 		t.Skip("test assumes tcp available")
42
 		t.Skip("test assumes tcp available")
55
 	// Start server and wait for it to be initialized
61
 	// Start server and wait for it to be initialized
56
 	go serveWebInterface("unused:1234", prof, &plugin.Options{
62
 	go serveWebInterface("unused:1234", prof, &plugin.Options{
57
 		Obj:        fakeObjTool{},
63
 		Obj:        fakeObjTool{},
58
-		UI:         &stdUI{},
64
+		UI:         &nonTerminalUI{&stdUI{}}, // don't open browser
59
 		HTTPServer: creator,
65
 		HTTPServer: creator,
60
-	}, false)
66
+	})
61
 	<-serverCreated
67
 	<-serverCreated
62
 	defer server.Close()
68
 	defer server.Close()
63
 
69