net/http/pprof: configure WriteDeadline

Configure write deadline according to requested profiling duration.

Fixes #62358

Change-Id: I2350110ff20a637c7e90bdda57026b0b0d9c87ba
GitHub-Last-Rev: b79ae38defc4d9b58b23a9d13ff86fa863f18b00
GitHub-Pull-Request: golang/go#64360
Reviewed-on: https://go-review.googlesource.com/c/go/+/544756
Reviewed-by: Damien Neil <dneil@google.com>
Reviewed-by: Dmitri Shuralyov <dmitshur@google.com>
Auto-Submit: Damien Neil <dneil@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

1 files changed, 13 insertions, 15 deletions

diff --git a/src/net/http/pprof/pprof.go b/src/net/http/pprof/pprof.go
index bc48f11834..18af18f976 100644
--- a/src/net/http/pprof/pprof.go
+++ b/src/net/http/pprof/pprof.go
@@ -114,9 +114,14 @@ func sleep(r *http.Request, d time.Duration) {
 	}
 }
 
-func durationExceedsWriteTimeout(r *http.Request, seconds float64) bool {
+func configureWriteDeadline(w http.ResponseWriter, r *http.Request, seconds float64) {
 	srv, ok := r.Context().Value(http.ServerContextKey).(*http.Server)
-	return ok && srv.WriteTimeout != 0 && seconds >= srv.WriteTimeout.Seconds()
+	if ok && srv.WriteTimeout > 0 {
+		timeout := srv.WriteTimeout + time.Duration(seconds*float64(time.Second))
+
+		rc := http.NewResponseController(w)
+		rc.SetWriteDeadline(time.Now().Add(timeout))
+	}
 }
 
 func serveError(w http.ResponseWriter, status int, txt string) {
@@ -137,10 +142,7 @@ func Profile(w http.ResponseWriter, r *http.Request) {
 		sec = 30
 	}
 
-	if durationExceedsWriteTimeout(r, float64(sec)) {
-		serveError(w, http.StatusBadRequest, "profile duration exceeds server's WriteTimeout")
-		return
-	}
+	configureWriteDeadline(w, r, float64(sec))
 
 	// Set Content Type assuming StartCPUProfile will work,
 	// because if it does it starts writing.
@@ -166,10 +168,7 @@ func Trace(w http.ResponseWriter, r *http.Request) {
 		sec = 1
 	}
 
-	if durationExceedsWriteTimeout(r, sec) {
-		serveError(w, http.StatusBadRequest, "profile duration exceeds server's WriteTimeout")
-		return
-	}
+	configureWriteDeadline(w, r, sec)
 
 	// Set Content Type assuming trace.Start will work,
 	// because if it does it starts writing.
@@ -273,15 +272,14 @@ func (name handler) serveDeltaProfile(w http.ResponseWriter, r *http.Request, p
 		serveError(w, http.StatusBadRequest, `invalid value for "seconds" - must be a positive integer`)
 		return
 	}
+	// 'name' should be a key in profileSupportsDelta.
 	if !profileSupportsDelta[name] {
 		serveError(w, http.StatusBadRequest, `"seconds" parameter is not supported for this profile type`)
 		return
 	}
-	// 'name' should be a key in profileSupportsDelta.
-	if durationExceedsWriteTimeout(r, float64(sec)) {
-		serveError(w, http.StatusBadRequest, "profile duration exceeds server's WriteTimeout")
-		return
-	}
+
+	configureWriteDeadline(w, r, float64(sec))
+
 	debug, _ := strconv.Atoi(r.FormValue("debug"))
 	if debug != 0 {
 		serveError(w, http.StatusBadRequest, "seconds and debug params are incompatible")