server/pprof/internal/pprof.go - infra/luci/luci-go - Git at Google

 // Copyright 2010 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.

 package internal

 // This code is copied from go/src/net/http/pprof/pprof.go at revision
 // 01c144c410b09d8b56d40e7e9c54fface204aa29 of the official Golang repo.
 //
 // Modifications:
 //  * Different package name, no package doc.
 //  * Removed init() section, since we don't want to autoregister routes in the
 //    default router (that's the main reason we are copy-pasting code instead of
 //    importing the package).
 //  * Modified Index() to put '&tok=...' into URLs.

 import (
 	"bufio"
 	"bytes"
 	"fmt"
 	"html/template"
 	"io"
 	"log"
 	"net/http"
 	"os"
 	"runtime"
 	"runtime/pprof"
 	"runtime/trace"
 	"strconv"
 	"strings"
 	"time"
 )

 // Cmdline responds with the running program's
 // command line, with arguments separated by NUL bytes.
 // The package initialization registers it as /debug/pprof/cmdline.
 func Cmdline(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
 	fmt.Fprintf(w, strings.Join(os.Args, "\x00"))
 }

 func sleep(w http.ResponseWriter, d time.Duration) {
 	var clientGone <-chan bool
 	if cn, ok := w.(http.CloseNotifier); ok {
 		clientGone = cn.CloseNotify()
 	}
 	select {
 	case <-time.After(d):
 	case <-clientGone:
 	}
 }

 func durationExceedsWriteTimeout(r *http.Request, seconds float64) bool {
 	srv, ok := r.Context().Value(http.ServerContextKey).(*http.Server)
 	return ok && srv.WriteTimeout != 0 && seconds >= srv.WriteTimeout.Seconds()
 }

 // Profile responds with the pprof-formatted cpu profile.
 // The package initialization registers it as /debug/pprof/profile.
 func Profile(w http.ResponseWriter, r *http.Request) {
 	sec, _ := strconv.ParseInt(r.FormValue("seconds"), 10, 64)
 	if sec == 0 {
 		sec = 30
 	}

 	if durationExceedsWriteTimeout(r, float64(sec)) {
 		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
 		w.Header().Set("X-Go-Pprof", "1")
 		w.WriteHeader(http.StatusBadRequest)
 		fmt.Fprintln(w, "profile duration exceeds server's WriteTimeout")
 		return
 	}

 	// Set Content Type assuming StartCPUProfile will work,
 	// because if it does it starts writing.
 	w.Header().Set("Content-Type", "application/octet-stream")
 	if err := pprof.StartCPUProfile(w); err != nil {
 		// StartCPUProfile failed, so no writes yet.
 		// Can change header back to text content
 		// and send error code.
 		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
 		w.Header().Set("X-Go-Pprof", "1")
 		w.WriteHeader(http.StatusInternalServerError)
 		fmt.Fprintf(w, "Could not enable CPU profiling: %s\n", err)
 		return
 	}
 	sleep(w, time.Duration(sec)*time.Second)
 	pprof.StopCPUProfile()
 }

 // Trace responds with the execution trace in binary form.
 // Tracing lasts for duration specified in seconds GET parameter, or for 1 second if not specified.
 // The package initialization registers it as /debug/pprof/trace.
 func Trace(w http.ResponseWriter, r *http.Request) {
 	sec, err := strconv.ParseFloat(r.FormValue("seconds"), 64)
 	if sec <= 0 || err != nil {
 		sec = 1
 	}

 	if durationExceedsWriteTimeout(r, sec) {
 		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
 		w.Header().Set("X-Go-Pprof", "1")
 		w.WriteHeader(http.StatusBadRequest)
 		fmt.Fprintln(w, "profile duration exceeds server's WriteTimeout")
 		return
 	}

 	// Set Content Type assuming trace.Start will work,
 	// because if it does it starts writing.
 	w.Header().Set("Content-Type", "application/octet-stream")
 	if err := trace.Start(w); err != nil {
 		// trace.Start failed, so no writes yet.
 		// Can change header back to text content and send error code.
 		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
 		w.Header().Set("X-Go-Pprof", "1")
 		w.WriteHeader(http.StatusInternalServerError)
 		fmt.Fprintf(w, "Could not enable tracing: %s\n", err)
 		return
 	}
 	sleep(w, time.Duration(sec*float64(time.Second)))
 	trace.Stop()
 }

 // Symbol looks up the program counters listed in the request,
 // responding with a table mapping program counters to function names.
 // The package initialization registers it as /debug/pprof/symbol.
 func Symbol(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "text/plain; charset=utf-8")

 	// We have to read the whole POST body before
 	// writing any output. Buffer the output here.
 	var buf bytes.Buffer

 	// We don't know how many symbols we have, but we
 	// do have symbol information. Pprof only cares whether
 	// this number is 0 (no symbols available) or > 0.
 	fmt.Fprintf(&buf, "num_symbols: 1\n")

 	var b *bufio.Reader
 	if r.Method == "POST" {
 		b = bufio.NewReader(r.Body)
 	} else {
 		b = bufio.NewReader(strings.NewReader(r.URL.RawQuery))
 	}

 	for {
 		word, err := b.ReadSlice('+')
 		if err == nil {
 			word = word[0 : len(word)-1] // trim +
 		}
 		pc, _ := strconv.ParseUint(string(word), 0, 64)
 		if pc != 0 {
 			f := runtime.FuncForPC(uintptr(pc))
 			if f != nil {
 				fmt.Fprintf(&buf, "%#x %s\n", pc, f.Name())
 			}
 		}

 		// Wait until here to check for err; the last
 		// symbol will have an err because it doesn't end in +.
 		if err != nil {
 			if err != io.EOF {
 				fmt.Fprintf(&buf, "reading request: %v\n", err)
 			}
 			break
 		}
 	}

 	w.Write(buf.Bytes())
 }

 // Handler returns an HTTP handler that serves the named profile.
 func Handler(name string) http.Handler {
 	return handler(name)
 }

 type handler string

 func (name handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
 	debug, _ := strconv.Atoi(r.FormValue("debug"))
 	p := pprof.Lookup(string(name))
 	if p == nil {
 		w.WriteHeader(404)
 		fmt.Fprintf(w, "Unknown profile: %s\n", name)
 		return
 	}
 	gc, _ := strconv.Atoi(r.FormValue("gc"))
 	if name == "heap" && gc > 0 {
 		runtime.GC()
 	}
 	p.WriteTo(w, debug)
 }

 // Index responds with the pprof-formatted profile named by the request.
 // For example, "/debug/pprof/heap" serves the "heap" profile.
 // Index responds to a request for "/debug/pprof/" with an HTML page
 // listing the available profiles.
 func Index(w http.ResponseWriter, r *http.Request) {
 	if strings.HasPrefix(r.URL.Path, "/debug/pprof/") {
 		name := strings.TrimPrefix(r.URL.Path, "/debug/pprof/")
 		if name != "" {
 			handler(name).ServeHTTP(w, r)
 			return
 		}
 	}

 	// [BEGIN modification]
 	err := indexTmpl.Execute(w, map[string]any{
 		"Tok":      r.FormValue("tok"),
 		"Profiles": pprof.Profiles(),
 	})
 	if err != nil {
 		log.Print(err)
 	}
 	// [END modification]
 }

 var indexTmpl = template.Must(template.New("index").Parse(`<html>
 <head>
 <title>/debug/pprof/</title>
 </head>
 <body>
 /debug/pprof/<br>
 <br>
 profiles:<br>
 <table>
 {{/* [BEGIN modification] */}}
 {{range .Profiles}}
 <tr><td align=right>{{.Count}}<td><a href="{{.Name}}?debug=1&tok={{$.Tok}}">{{.Name}}</a>
 {{end}}
 </table>
 <br>
 <a href="goroutine?debug=2&tok={{.Tok}}">full goroutine stack dump</a><br>
 {{/* [END modification] */}}
 </body>
 </html>
 `))
	// Copyright 2010 The Go Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style
	// license that can be found in the LICENSE file.

	package internal

	// This code is copied from go/src/net/http/pprof/pprof.go at revision
	// 01c144c410b09d8b56d40e7e9c54fface204aa29 of the official Golang repo.
	//
	// Modifications:
	// * Different package name, no package doc.
	// * Removed init() section, since we don't want to autoregister routes in the
	// default router (that's the main reason we are copy-pasting code instead of
	// importing the package).
	// * Modified Index() to put '&tok=...' into URLs.

	import (
	"bufio"
	"bytes"
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"os"
	"runtime"
	"runtime/pprof"
	"runtime/trace"
	"strconv"
	"strings"
	"time"
	)

	// Cmdline responds with the running program's
	// command line, with arguments separated by NUL bytes.
	// The package initialization registers it as /debug/pprof/cmdline.
	func Cmdline(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	fmt.Fprintf(w, strings.Join(os.Args, "\x00"))
	}

	func sleep(w http.ResponseWriter, d time.Duration) {
	var clientGone <-chan bool
	if cn, ok := w.(http.CloseNotifier); ok {
	clientGone = cn.CloseNotify()
	}
	select {
	case <-time.After(d):
	case <-clientGone:
	}
	}

	func durationExceedsWriteTimeout(r *http.Request, seconds float64) bool {
	srv, ok := r.Context().Value(http.ServerContextKey).(*http.Server)
	return ok && srv.WriteTimeout != 0 && seconds >= srv.WriteTimeout.Seconds()
	}

	// Profile responds with the pprof-formatted cpu profile.
	// The package initialization registers it as /debug/pprof/profile.
	func Profile(w http.ResponseWriter, r *http.Request) {
	sec, _ := strconv.ParseInt(r.FormValue("seconds"), 10, 64)
	if sec == 0 {
	sec = 30
	}

	if durationExceedsWriteTimeout(r, float64(sec)) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.Header().Set("X-Go-Pprof", "1")
	w.WriteHeader(http.StatusBadRequest)
	fmt.Fprintln(w, "profile duration exceeds server's WriteTimeout")
	return
	}

	// Set Content Type assuming StartCPUProfile will work,
	// because if it does it starts writing.
	w.Header().Set("Content-Type", "application/octet-stream")
	if err := pprof.StartCPUProfile(w); err != nil {
	// StartCPUProfile failed, so no writes yet.
	// Can change header back to text content
	// and send error code.
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.Header().Set("X-Go-Pprof", "1")
	w.WriteHeader(http.StatusInternalServerError)
	fmt.Fprintf(w, "Could not enable CPU profiling: %s\n", err)
	return
	}
	sleep(w, time.Duration(sec)*time.Second)
	pprof.StopCPUProfile()
	}

	// Trace responds with the execution trace in binary form.
	// Tracing lasts for duration specified in seconds GET parameter, or for 1 second if not specified.
	// The package initialization registers it as /debug/pprof/trace.
	func Trace(w http.ResponseWriter, r *http.Request) {
	sec, err := strconv.ParseFloat(r.FormValue("seconds"), 64)
	if sec <= 0 \|\| err != nil {
	sec = 1
	}

	if durationExceedsWriteTimeout(r, sec) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.Header().Set("X-Go-Pprof", "1")
	w.WriteHeader(http.StatusBadRequest)
	fmt.Fprintln(w, "profile duration exceeds server's WriteTimeout")
	return
	}

	// Set Content Type assuming trace.Start will work,
	// because if it does it starts writing.
	w.Header().Set("Content-Type", "application/octet-stream")
	if err := trace.Start(w); err != nil {
	// trace.Start failed, so no writes yet.
	// Can change header back to text content and send error code.
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.Header().Set("X-Go-Pprof", "1")
	w.WriteHeader(http.StatusInternalServerError)
	fmt.Fprintf(w, "Could not enable tracing: %s\n", err)
	return
	}
	sleep(w, time.Duration(sec*float64(time.Second)))
	trace.Stop()
	}

	// Symbol looks up the program counters listed in the request,
	// responding with a table mapping program counters to function names.
	// The package initialization registers it as /debug/pprof/symbol.
	func Symbol(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")

	// We have to read the whole POST body before
	// writing any output. Buffer the output here.
	var buf bytes.Buffer

	// We don't know how many symbols we have, but we
	// do have symbol information. Pprof only cares whether
	// this number is 0 (no symbols available) or > 0.
	fmt.Fprintf(&buf, "num_symbols: 1\n")

	var b *bufio.Reader
	if r.Method == "POST" {
	b = bufio.NewReader(r.Body)
	} else {
	b = bufio.NewReader(strings.NewReader(r.URL.RawQuery))
	}

	for {
	word, err := b.ReadSlice('+')
	if err == nil {
	word = word[0 : len(word)-1] // trim +
	}
	pc, _ := strconv.ParseUint(string(word), 0, 64)
	if pc != 0 {
	f := runtime.FuncForPC(uintptr(pc))
	if f != nil {
	fmt.Fprintf(&buf, "%#x %s\n", pc, f.Name())
	}
	}

	// Wait until here to check for err; the last
	// symbol will have an err because it doesn't end in +.
	if err != nil {
	if err != io.EOF {
	fmt.Fprintf(&buf, "reading request: %v\n", err)
	}
	break
	}
	}

	w.Write(buf.Bytes())
	}

	// Handler returns an HTTP handler that serves the named profile.
	func Handler(name string) http.Handler {
	return handler(name)
	}

	type handler string

	func (name handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	debug, _ := strconv.Atoi(r.FormValue("debug"))
	p := pprof.Lookup(string(name))
	if p == nil {
	w.WriteHeader(404)
	fmt.Fprintf(w, "Unknown profile: %s\n", name)
	return
	}
	gc, _ := strconv.Atoi(r.FormValue("gc"))
	if name == "heap" && gc > 0 {
	runtime.GC()
	}
	p.WriteTo(w, debug)
	}

	// Index responds with the pprof-formatted profile named by the request.
	// For example, "/debug/pprof/heap" serves the "heap" profile.
	// Index responds to a request for "/debug/pprof/" with an HTML page
	// listing the available profiles.
	func Index(w http.ResponseWriter, r *http.Request) {
	if strings.HasPrefix(r.URL.Path, "/debug/pprof/") {
	name := strings.TrimPrefix(r.URL.Path, "/debug/pprof/")
	if name != "" {
	handler(name).ServeHTTP(w, r)
	return
	}
	}

	// [BEGIN modification]
	err := indexTmpl.Execute(w, map[string]any{
	"Tok": r.FormValue("tok"),
	"Profiles": pprof.Profiles(),
	})
	if err != nil {
	log.Print(err)
	}
	// [END modification]
	}

	var indexTmpl = template.Must(template.New("index").Parse(`<html>
	<head>
	<title>/debug/pprof/</title>
	</head>
	<body>
	/debug/pprof/<br>
	<br>
	profiles:<br>
	<table>
	{{/* [BEGIN modification] */}}
	{{range .Profiles}}
	<tr><td align=right>{{.Count}}<td><a href="{{.Name}}?debug=1&tok={{$.Tok}}">{{.Name}}</a>
	{{end}}
	</table>
	<br>
	<a href="goroutine?debug=2&tok={{.Tok}}">full goroutine stack dump</a><br>
	{{/* [END modification] */}}
	</body>
	</html>
	`))