Reimplement TimeoutHandler
This commit is contained in:
parent
b386a9ba41
commit
9eae88f00c
|
@ -138,6 +138,14 @@ func NewResponseWriter(w io.WriteCloser) *ResponseWriter {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (w *ResponseWriter) reset(wc io.WriteCloser) {
|
||||||
|
w.bw.Reset(wc)
|
||||||
|
*w = ResponseWriter{
|
||||||
|
bw: w.bw,
|
||||||
|
cl: wc,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// SetMediaType sets the media type that will be sent by Write for a
|
// SetMediaType sets the media type that will be sent by Write for a
|
||||||
// successful response. If no media type is set, a default of
|
// successful response. If no media type is set, a default of
|
||||||
// "text/gemini; charset=utf-8" will be used.
|
// "text/gemini; charset=utf-8" will be used.
|
||||||
|
|
79
timeout.go
79
timeout.go
|
@ -1,11 +1,8 @@
|
||||||
// +build ignore
|
|
||||||
|
|
||||||
package gemini
|
package gemini
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"bytes"
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
"sync"
|
|
||||||
"time"
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@ -27,84 +24,26 @@ type timeoutHandler struct {
|
||||||
dt time.Duration
|
dt time.Duration
|
||||||
}
|
}
|
||||||
|
|
||||||
func (t *timeoutHandler) ServeGemini(ctx context.Context, w ResponseWriter, r *Request) {
|
func (t *timeoutHandler) ServeGemini(ctx context.Context, w *ResponseWriter, r *Request) {
|
||||||
ctx, cancel := context.WithTimeout(ctx, t.dt)
|
ctx, cancel := context.WithTimeout(ctx, t.dt)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
|
|
||||||
|
conn := w.Hijack()
|
||||||
|
|
||||||
|
var b bytes.Buffer
|
||||||
|
w.reset(nopCloser{&b})
|
||||||
|
|
||||||
done := make(chan struct{})
|
done := make(chan struct{})
|
||||||
tw := &timeoutWriter{}
|
|
||||||
go func() {
|
go func() {
|
||||||
t.h.ServeGemini(ctx, tw, r)
|
t.h.ServeGemini(ctx, w, r)
|
||||||
close(done)
|
close(done)
|
||||||
}()
|
}()
|
||||||
|
|
||||||
select {
|
select {
|
||||||
case <-done:
|
case <-done:
|
||||||
tw.mu.Lock()
|
conn.Write(b.Bytes())
|
||||||
defer tw.mu.Unlock()
|
|
||||||
if !tw.wroteHeader {
|
|
||||||
tw.status = StatusSuccess
|
|
||||||
}
|
|
||||||
w.WriteHeader(tw.status, tw.meta)
|
|
||||||
w.Write(tw.b.Bytes())
|
|
||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
tw.mu.Lock()
|
w.reset(conn)
|
||||||
defer tw.mu.Unlock()
|
|
||||||
w.WriteHeader(StatusTemporaryFailure, "Timeout")
|
w.WriteHeader(StatusTemporaryFailure, "Timeout")
|
||||||
tw.timedOut = true
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
type timeoutWriter struct {
|
|
||||||
mu sync.Mutex
|
|
||||||
b bytes.Buffer
|
|
||||||
status Status
|
|
||||||
meta string
|
|
||||||
mediatype string
|
|
||||||
wroteHeader bool
|
|
||||||
timedOut bool
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *timeoutWriter) SetMediaType(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 Status, meta string) {
|
|
||||||
w.mu.Lock()
|
|
||||||
defer w.mu.Unlock()
|
|
||||||
if w.timedOut {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
w.writeHeaderLocked(status, meta)
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *timeoutWriter) writeHeaderLocked(status Status, meta string) {
|
|
||||||
if w.wroteHeader {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
w.status = status
|
|
||||||
w.meta = meta
|
|
||||||
w.wroteHeader = true
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *timeoutWriter) Flush() error {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *timeoutWriter) Close() error {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
Loading…
Reference in New Issue
Block a user