gitdeploy/vendor/github.com/go-chi/chi/middleware/timeout.go

50 lines
1.2 KiB
Go

package middleware
import (
"context"
"net/http"
"time"
)
// Timeout is a 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.
//
// ie. a route/handler may look like:
//
// r.Get("/long", func(w http.ResponseWriter, r *http.Request) {
// ctx := r.Context()
// processTime := time.Duration(rand.Intn(4)+1) * time.Second
//
// select {
// case <-ctx.Done():
// return
//
// case <-time.After(processTime):
// // The above channel simulates some hard work.
// }
//
// w.Write([]byte("done"))
// })
//
func Timeout(timeout time.Duration) func(next http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
fn := func(w http.ResponseWriter, r *http.Request) {
ctx, cancel := context.WithTimeout(r.Context(), timeout)
defer func() {
cancel()
if ctx.Err() == context.DeadlineExceeded {
w.WriteHeader(http.StatusGatewayTimeout)
}
}()
r = r.WithContext(ctx)
next.ServeHTTP(w, r)
}
return http.HandlerFunc(fn)
}
}