Commit 4787e70b authored by Brad Fitzpatrick's avatar Brad Fitzpatrick

http: handler timeout support

Fixes #213

R=r, rsc
CC=golang-dev
https://golang.org/cl/4432043
parent 84c7e83b
......@@ -32,3 +32,10 @@ func (t *Transport) IdleConnCountForTesting(cacheKey string) int {
}
return len(conns)
}
func NewTestTimeoutHandler(handler Handler, ch <-chan int64) Handler {
f := func() <-chan int64 {
return ch
}
return &timeoutHandler{handler, f, ""}
}
......@@ -662,3 +662,54 @@ func TestServerConsumesRequestBody(t *testing.T) {
}
}
}
func TestTimeoutHandler(t *testing.T) {
sendHi := make(chan bool, 1)
writeErrors := make(chan os.Error, 1)
sayHi := HandlerFunc(func(w ResponseWriter, r *Request) {
<-sendHi
_, werr := w.Write([]byte("hi"))
writeErrors <- werr
})
timeout := make(chan int64, 1) // write to this to force timeouts
ts := httptest.NewServer(NewTestTimeoutHandler(sayHi, timeout))
defer ts.Close()
// Succeed without timing out:
sendHi <- true
res, _, err := Get(ts.URL)
if err != nil {
t.Error(err)
}
if g, e := res.StatusCode, StatusOK; g != e {
t.Errorf("got res.StatusCode %d; expected %d", g, e)
}
body, _ := ioutil.ReadAll(res.Body)
if g, e := string(body), "hi"; g != e {
t.Errorf("got body %q; expected %q", g, e)
}
if g := <-writeErrors; g != nil {
t.Errorf("got unexpected Write error on first request: %v", g)
}
// Times out:
timeout <- 1
res, _, err = Get(ts.URL)
if err != nil {
t.Error(err)
}
if g, e := res.StatusCode, StatusServiceUnavailable; g != e {
t.Errorf("got res.StatusCode %d; expected %d", g, e)
}
body, _ = ioutil.ReadAll(res.Body)
if !strings.Contains(string(body), "<title>Timeout</title>") {
t.Errorf("expected timeout body; got %q", string(body))
}
// Now make the previously-timed out handler speak again,
// which verifies the panic is handled:
sendHi <- true
if g, e := <-writeErrors, ErrHandlerTimeout; g != e {
t.Errorf("expected Write error of %v; got %v", e, g)
}
}
......@@ -22,6 +22,7 @@ import (
"path"
"strconv"
"strings"
"sync"
"time"
)
......@@ -898,3 +899,89 @@ func ListenAndServeTLS(addr string, certFile string, keyFile string, handler Han
tlsListener := tls.NewListener(conn, config)
return Serve(tlsListener, handler)
}
// TimeoutHandler returns a Handler that runs h with the given time limit.
//
// The new Handler calls h.ServeHTTP to handle each request, but if a
// call runs for more than ns nanoseconds, the handler responds with
// a 503 Service Unavailable error and the given message in its body.
// (If msg is empty, a suitable default message will be sent.)
// After such a timeout, writes by h to its ResponseWriter will return
// ErrHandlerTimeout.
func TimeoutHandler(h Handler, ns int64, msg string) Handler {
f := func() <-chan int64 {
return time.After(ns)
}
return &timeoutHandler{h, f, msg}
}
// ErrHandlerTimeout is returned on ResponseWriter Write calls
// in handlers which have timed out.
var ErrHandlerTimeout = os.NewError("http: Handler timeout")
type timeoutHandler struct {
handler Handler
timeout func() <-chan int64 // returns channel producing a timeout
body string
}
func (h *timeoutHandler) errorBody() string {
if h.body != "" {
return h.body
}
return "<html><head><title>Timeout</title></head><body><h1>Timeout</h1></body></html>"
}
func (h *timeoutHandler) ServeHTTP(w ResponseWriter, r *Request) {
done := make(chan bool)
tw := &timeoutWriter{w: w}
go func() {
h.handler.ServeHTTP(tw, r)
done <- true
}()
select {
case <-done:
return
case <-h.timeout():
tw.mu.Lock()
defer tw.mu.Unlock()
if !tw.wroteHeader {
tw.w.WriteHeader(StatusServiceUnavailable)
tw.w.Write([]byte(h.errorBody()))
}
tw.timedOut = true
}
}
type timeoutWriter struct {
w ResponseWriter
mu sync.Mutex
timedOut bool
wroteHeader bool
}
func (tw *timeoutWriter) Header() Header {
return tw.w.Header()
}
func (tw *timeoutWriter) Write(p []byte) (int, os.Error) {
tw.mu.Lock()
timedOut := tw.timedOut
tw.mu.Unlock()
if timedOut {
return 0, ErrHandlerTimeout
}
return tw.w.Write(p)
}
func (tw *timeoutWriter) WriteHeader(code int) {
tw.mu.Lock()
if tw.timedOut || tw.wroteHeader {
tw.mu.Unlock()
return
}
tw.wroteHeader = true
tw.mu.Unlock()
tw.w.WriteHeader(code)
}
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment