Implement TimeoutHandler
This commit is contained in:
parent
64dbb3eecb
commit
5922cff2e5
@ -24,5 +24,9 @@ var (
|
||||
// While any panic from ServeGemini aborts the response to the client,
|
||||
// panicking with ErrAbortHandler also suppresses logging of a stack
|
||||
// trace to the server's error log.
|
||||
ErrAbortHandler = errors.New("net/http: abort Handler")
|
||||
ErrAbortHandler = errors.New("gemini: abort Handler")
|
||||
|
||||
// ErrHandlerTimeout is returned on ResponseWriter Write calls
|
||||
// in handlers which have timed out.
|
||||
ErrHandlerTimeout = errors.New("gemini: Handler timeout")
|
||||
)
|
||||
|
115
timeout.go
Normal file
115
timeout.go
Normal file
@ -0,0 +1,115 @@
|
||||
package gemini
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TimeoutHandler returns a Handler that runs h with the given time limit.
|
||||
//
|
||||
// The new Handler calls h.ServeGemini to handle each request, but
|
||||
// if a call runs for longer than its time limit, the handler responds with a
|
||||
// 40 Temporary Failure error. After such a timeout, writes by h to its
|
||||
// ResponseWriter will return ErrHandlerTimeout.
|
||||
//
|
||||
// TimeoutHandler does not support the Hijacker or Flusher interfaces.
|
||||
func TimeoutHandler(h Handler, dt time.Duration) Handler {
|
||||
return &timeoutHandler{
|
||||
h: h,
|
||||
dt: dt,
|
||||
}
|
||||
}
|
||||
|
||||
type timeoutHandler struct {
|
||||
h Handler
|
||||
dt time.Duration
|
||||
}
|
||||
|
||||
func (t *timeoutHandler) ServeGemini(w ResponseWriter, r *Request) {
|
||||
ctx := r.Context
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
var cancelCtx func()
|
||||
ctx, cancelCtx = context.WithTimeout(ctx, t.dt)
|
||||
defer cancelCtx()
|
||||
|
||||
done := make(chan struct{})
|
||||
tw := &timeoutWriter{}
|
||||
panicChan := make(chan interface{}, 1)
|
||||
go func() {
|
||||
defer func() {
|
||||
if p := recover(); p != nil {
|
||||
panicChan <- p
|
||||
}
|
||||
}()
|
||||
t.h.ServeGemini(tw, r)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case p := <-panicChan:
|
||||
panic(p)
|
||||
case <-done:
|
||||
tw.mu.Lock()
|
||||
defer tw.mu.Unlock()
|
||||
if !tw.wroteHeader {
|
||||
tw.status = StatusSuccess
|
||||
}
|
||||
w.WriteHeader(tw.status, tw.meta)
|
||||
w.Write(tw.b.Bytes())
|
||||
case <-ctx.Done():
|
||||
tw.mu.Lock()
|
||||
defer tw.mu.Unlock()
|
||||
w.WriteHeader(StatusTemporaryFailure, "Timeout")
|
||||
tw.timedOut = true
|
||||
}
|
||||
}
|
||||
|
||||
type timeoutWriter struct {
|
||||
mu sync.Mutex
|
||||
b bytes.Buffer
|
||||
status int
|
||||
meta string
|
||||
mediatype string
|
||||
wroteHeader bool
|
||||
timedOut bool
|
||||
}
|
||||
|
||||
func (w *timeoutWriter) MediaType(mediatype string) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
w.mediatype = mediatype
|
||||
}
|
||||
|
||||
func (w *timeoutWriter) Write(b []byte) (int, error) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
if w.timedOut {
|
||||
return 0, ErrHandlerTimeout
|
||||
}
|
||||
if !w.wroteHeader {
|
||||
w.writeHeaderLocked(StatusSuccess, w.mediatype)
|
||||
}
|
||||
return w.b.Write(b)
|
||||
}
|
||||
|
||||
func (w *timeoutWriter) WriteHeader(status int, meta string) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
if w.timedOut {
|
||||
return
|
||||
}
|
||||
w.writeHeaderLocked(status, meta)
|
||||
}
|
||||
|
||||
func (w *timeoutWriter) writeHeaderLocked(status int, meta string) {
|
||||
if w.wroteHeader {
|
||||
return
|
||||
}
|
||||
w.status = status
|
||||
w.meta = meta
|
||||
w.wroteHeader = true
|
||||
}
|
Loading…
Reference in New Issue
Block a user