Implement TimeoutHandler

gemini.go

@@ -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")
 )

timeout.go

@@ -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
+}