1
0
mirror of https://github.com/golang/go synced 2024-11-19 16:34:49 -07:00

net/http: Set a timeout on Request.Context when using TimeoutHandler

In TimeoutHandler, use a request whose context has been configured with
the handler's timeout

Fixes #20712

Change-Id: Ie670148f85fdad46841ff29232042309e15665ae
Reviewed-on: https://go-review.googlesource.com/46412
Run-TryBot: Tom Bergan <tombergan@google.com>
Reviewed-by: Tom Bergan <tombergan@google.com>
This commit is contained in:
Michael Fraenkel 2017-06-21 21:58:42 -04:00 committed by Tom Bergan
parent 2b079c3c04
commit 9e7b30b463
2 changed files with 15 additions and 12 deletions

View File

@ -63,9 +63,14 @@ func SetPendingDialHooks(before, after func()) {
func SetTestHookServerServe(fn func(*Server, net.Listener)) { testHookServerServe = fn } func SetTestHookServerServe(fn func(*Server, net.Listener)) { testHookServerServe = fn }
func NewTestTimeoutHandler(handler Handler, ch <-chan time.Time) Handler { func NewTestTimeoutHandler(handler Handler, ch <-chan time.Time) Handler {
ctx, cancel := context.WithCancel(context.Background())
go func() {
<-ch
cancel()
}()
return &timeoutHandler{ return &timeoutHandler{
handler: handler, handler: handler,
testTimeout: ch, testContext: ctx,
// (no body) // (no body)
} }
} }

View File

@ -3032,9 +3032,9 @@ type timeoutHandler struct {
body string body string
dt time.Duration dt time.Duration
// When set, no timer will be created and this channel will // When set, no context will be created and this context will
// be used instead. // be used instead.
testTimeout <-chan time.Time testContext context.Context
} }
func (h *timeoutHandler) errorBody() string { func (h *timeoutHandler) errorBody() string {
@ -3045,12 +3045,13 @@ func (h *timeoutHandler) errorBody() string {
} }
func (h *timeoutHandler) ServeHTTP(w ResponseWriter, r *Request) { func (h *timeoutHandler) ServeHTTP(w ResponseWriter, r *Request) {
var t *time.Timer ctx := h.testContext
timeout := h.testTimeout if ctx == nil {
if timeout == nil { var cancelCtx context.CancelFunc
t = time.NewTimer(h.dt) ctx, cancelCtx = context.WithTimeout(r.Context(), h.dt)
timeout = t.C defer cancelCtx()
} }
r = r.WithContext(ctx)
done := make(chan struct{}) done := make(chan struct{})
tw := &timeoutWriter{ tw := &timeoutWriter{
w: w, w: w,
@ -3073,10 +3074,7 @@ func (h *timeoutHandler) ServeHTTP(w ResponseWriter, r *Request) {
} }
w.WriteHeader(tw.code) w.WriteHeader(tw.code)
w.Write(tw.wbuf.Bytes()) w.Write(tw.wbuf.Bytes())
if t != nil { case <-ctx.Done():
t.Stop()
}
case <-timeout:
tw.mu.Lock() tw.mu.Lock()
defer tw.mu.Unlock() defer tw.mu.Unlock()
w.WriteHeader(StatusServiceUnavailable) w.WriteHeader(StatusServiceUnavailable)