timeout.go


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

package gemini

import (
	"bytes"
	"context"
	"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.
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(ctx context.Context, w *ResponseWriter, r *Request) {
	ctx, cancel := context.WithTimeout(ctx, t.dt)
	defer cancel()

	conn := w.Hijack()

	var b bytes.Buffer
	w.reset(nopCloser{&b})

	done := make(chan struct{})
	go func() {
		t.h.ServeGemini(ctx, w, r)
		close(done)
	}()

	select {
	case <-done:
		conn.Write(b.Bytes())
	case <-ctx.Done():
		w.reset(conn)
		w.WriteHeader(StatusTemporaryFailure, "Timeout")
	}
}