github.com/hellobchain/third_party@v0.0.0-20230331131523-deb0478a2e52/go-chi/chi/middleware/timeout.go (about)

     1  package middleware
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/hellobchain/newcryptosm/http"
     8  )
     9  
    10  // Timeout is a middleware that cancels ctx after a given timeout and return
    11  // a 504 Gateway Timeout error to the client.
    12  //
    13  // It's required that you select the ctx.Done() channel to check for the signal
    14  // if the context has reached its deadline and return, otherwise the timeout
    15  // signal will be just ignored.
    16  //
    17  // ie. a route/handler may look like:
    18  //
    19  //  r.Get("/long", func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
    20  // 	 processTime := time.Duration(rand.Intn(4)+1) * time.Second
    21  //
    22  // 	 select {
    23  // 	 case <-ctx.Done():
    24  // 	 	return
    25  //
    26  // 	 case <-time.After(processTime):
    27  // 	 	 // The above channel simulates some hard work.
    28  // 	 }
    29  //
    30  // 	 w.Write([]byte("done"))
    31  //  })
    32  //
    33  func Timeout(timeout time.Duration) func(next http.Handler) http.Handler {
    34  	return func(next http.Handler) http.Handler {
    35  		fn := func(w http.ResponseWriter, r *http.Request) {
    36  			ctx, cancel := context.WithTimeout(r.Context(), timeout)
    37  			defer func() {
    38  				cancel()
    39  				if ctx.Err() == context.DeadlineExceeded {
    40  					w.WriteHeader(http.StatusGatewayTimeout)
    41  				}
    42  			}()
    43  
    44  			r = r.WithContext(ctx)
    45  			next.ServeHTTP(w, r)
    46  		}
    47  		return http.HandlerFunc(fn)
    48  	}
    49  }