net/http/pprof: stop profiling if client's connection closes

Fixes #13833

Change-Id: If0bd5f7dcfc39d34680d11eb998050f0900d5a26
Reviewed-on: https://go-review.googlesource.com/18283
Reviewed-by: Russ Cox <rsc@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>

diff --git a/src/net/http/pprof/pprof.go b/src/net/http/pprof/pprof.go
index fd9154ac2ac36d431276c0e092376f777d7a0c96..63ff22123b7894166b4d78c9bfab0968ade597e2 100644 (file)
--- a/src/net/http/pprof/pprof.go
+++ b/src/net/http/pprof/pprof.go
@@ -80,6 +80,17 @@ func Cmdline(w http.ResponseWriter, r *http.Request) {
        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:
+       }
+}
+
 // 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) {
@@ -100,7 +111,7 @@ func Profile(w http.ResponseWriter, r *http.Request) {
                fmt.Fprintf(w, "Could not enable CPU profiling: %s\n", err)
                return
        }
-       time.Sleep(time.Duration(sec) * time.Second)
+       sleep(w, time.Duration(sec)*time.Second)
        pprof.StopCPUProfile()
 }
 
@@ -124,7 +135,7 @@ func Trace(w http.ResponseWriter, r *http.Request) {
                fmt.Fprintf(w, "Could not enable tracing: %s\n", err)
                return
        }
-       time.Sleep(time.Duration(sec) * time.Second)
+       sleep(w, time.Duration(sec)*time.Second)
        trace.Stop()
 }