middleware/timeout/timeout.go

package timeout

import (
	"context"
	"net/http"
	"time"

	"go.pkg.cx/middleware"
)

// DefaultOptions represents default timeout middleware options
var DefaultOptions = Options(
	SetTimeout(time.Second*30),
	SetResponseHandler(RespondWithTimeout),
)

// Options turns a list of option instances into an option
func Options(opts ...Option) Option {
	return func(t *timeout) {
		for _, opt := range opts {
			opt(t)
		}
	}
}

// Option configures timeout middleware
type Option func(t *timeout)

// SetTimeout sets request timeout
func SetTimeout(limit time.Duration) Option {
	return func(t *timeout) {
		t.timeout = limit
	}
}

// SetResponseHandler sets response handler
func SetResponseHandler(fn middleware.ResponseHandle) Option {
	return func(t *timeout) {
		t.responseHandler = fn
	}
}

type timeout struct {
	timeout         time.Duration
	responseHandler middleware.ResponseHandle
}

// Middleware is a timeout middleware that cancels ctx after a given timeout
// and return a 504 Gateway Timeout error to the client.
//
// It's required that you select the ctx.Done() channel to check for the signal
// if the context has reached its deadline and return, otherwise the timeout
// signal will be just ignored
func Middleware(opts ...Option) func(next http.Handler) http.Handler {
	t := &timeout{}
	opts = append([]Option{DefaultOptions}, opts...)

	for _, opt := range opts {
		opt(t)
	}

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ctx, cancel := context.WithTimeout(r.Context(), t.timeout)
			defer func() {
				cancel()
				if ctx.Err() == context.DeadlineExceeded {
					t.responseHandler(w, r, ctx.Err())
					return
				}
			}()

			next.ServeHTTP(w, r.WithContext(ctx))
		})
	}
}

// RespondWithTimeout is a default response handler
func RespondWithTimeout(w http.ResponseWriter, _ *http.Request, _ error) {
	http.Error(w, http.StatusText(http.StatusGatewayTimeout), http.StatusGatewayTimeout)
}