github.com/xmidt-org/webpa-common@v1.11.9/xhttp/xtimeout/constructor.go (about)

     1  package xtimeout
     2  
     3  import (
     4  	"context"
     5  	"net/http"
     6  	"time"
     7  
     8  	"github.com/xmidt-org/webpa-common/xhttp"
     9  )
    10  
    11  // defaultTimedOut is the default http.Handler used for timeout responses
    12  var defaultTimedOut = xhttp.Constant{Code: http.StatusGatewayTimeout}
    13  
    14  // timeoutHandler is the internal decorator handler that handles timeouts in its ServeHTTP method
    15  type timeoutHandler struct {
    16  	timeout  time.Duration
    17  	timedOut http.Handler
    18  	next     http.Handler
    19  }
    20  
    21  func (th *timeoutHandler) ServeHTTP(response http.ResponseWriter, request *http.Request) {
    22  	var (
    23  		done        = make(chan struct{})
    24  		panics      = make(chan interface{}, 1)
    25  		writer      xhttp.BufferedWriter
    26  		ctx, cancel = context.WithTimeout(request.Context(), th.timeout)
    27  	)
    28  
    29  	defer cancel()
    30  	go func() {
    31  		defer func() {
    32  			if p := recover(); p != nil {
    33  				panics <- p
    34  			}
    35  		}()
    36  
    37  		th.next.ServeHTTP(&writer, request.WithContext(ctx))
    38  		close(done)
    39  	}()
    40  
    41  	select {
    42  	case p := <-panics:
    43  		panic(p) // mimic the behavior of net/http TimeoutHandler
    44  	case <-done:
    45  		writer.WriteTo(response)
    46  	case <-ctx.Done():
    47  		writer.Close()
    48  		th.timedOut.ServeHTTP(response, request)
    49  	}
    50  }
    51  
    52  // Options holds the set of configurable options for a timeout constructor.
    53  type Options struct {
    54  	// Timeout is the time allowed for the decorated handler's ServeHTTP method to run.
    55  	// If unset or nonpositive, no decoration is performed.
    56  	Timeout time.Duration
    57  
    58  	// TimedOut is the optional http.Handler that is executed with the original http.Request
    59  	// whenever a timeout occurs.  If unset, a default handler is used that simply sets the response
    60  	// code to http.StatusGatewayTimeout.
    61  	TimedOut http.Handler
    62  }
    63  
    64  // NewConstructor returns an Alice-style constructor that enforces a timeout for any handler it decorates.
    65  // If timeout is nonpositive, a constructor is returned that does no decoration.
    66  // If timedOut is nil, a default timedOut handler is used that just sets an http.StatusGatewayTimeout response code.
    67  func NewConstructor(o Options) func(http.Handler) http.Handler {
    68  	if o.Timeout <= 0 {
    69  		return xhttp.NilConstructor
    70  	}
    71  
    72  	if o.TimedOut == nil {
    73  		o.TimedOut = defaultTimedOut
    74  	}
    75  
    76  	return func(next http.Handler) http.Handler {
    77  		return &timeoutHandler{
    78  			timeout:  o.Timeout,
    79  			timedOut: o.TimedOut,
    80  			next:     next,
    81  		}
    82  	}
    83  }