github.com/kyma-incubator/compass/components/director@v0.0.0-20230623144113-d764f56ff805/pkg/handler/timeout_handler.go (about)

     1  package handler
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"encoding/json"
     7  	"html/template"
     8  	"net/http"
     9  	"strings"
    10  	"time"
    11  
    12  	"github.com/kyma-incubator/compass/components/director/pkg/apperrors"
    13  	"github.com/kyma-incubator/compass/components/director/pkg/log"
    14  )
    15  
    16  const (
    17  	// HeaderContentTypeKey missing godoc
    18  	HeaderContentTypeKey = "Content-Type"
    19  	// HeaderContentTypeValue missing godoc
    20  	HeaderContentTypeValue = "application/json;charset=UTF-8"
    21  )
    22  
    23  // WithTimeout returns timeout http.Handler
    24  func WithTimeout(h http.Handler, timeout time.Duration) (http.Handler, error) {
    25  	msg, err := json.Marshal(apperrors.NewOperationTimeoutError())
    26  	if err != nil {
    27  		return nil, err
    28  	}
    29  
    30  	return WithTimeoutWithErrorMessage(h, timeout, msg)
    31  }
    32  
    33  // WithTimeoutWithErrorMessage returns timeout http.Handler with provided error message
    34  func WithTimeoutWithErrorMessage(h http.Handler, timeout time.Duration, msg []byte) (http.Handler, error) {
    35  	preTimoutLoggingHandler := newTimeoutLoggingHandler(h, timeout, msg)
    36  	timeoutHandler := http.TimeoutHandler(preTimoutLoggingHandler, timeout, string(msg))
    37  	postTimeoutLoggingHandler := newTimeoutLoggingHandler(timeoutHandler, timeout, msg)
    38  
    39  	return newContentTypeHandler(postTimeoutLoggingHandler), nil
    40  }
    41  
    42  func newTimeoutLoggingHandler(h http.Handler, timeout time.Duration, msg []byte) http.Handler {
    43  	return &timeoutLoggingHandler{
    44  		h:       h,
    45  		timeout: timeout,
    46  		msg:     msg,
    47  	}
    48  }
    49  
    50  type timeoutLoggingHandler struct {
    51  	h       http.Handler
    52  	timeout time.Duration
    53  	msg     []byte
    54  }
    55  
    56  // ServeHTTP missing godoc
    57  func (h *timeoutLoggingHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    58  	ctx := log.ContextWithLogger(r.Context(), log.LoggerWithCorrelationID(r))
    59  	r = r.WithContext(ctx)
    60  
    61  	timoutRW := &timoutLoggingResponseWriter{
    62  		ResponseWriter: w,
    63  		method:         template.HTMLEscapeString(r.Method),
    64  		url:            template.HTMLEscapeString(r.URL.String()),
    65  		timeout:        h.timeout,
    66  		msg:            h.msg,
    67  		requestStart:   time.Now(),
    68  		ctx:            ctx,
    69  	}
    70  
    71  	h.h.ServeHTTP(timoutRW, r)
    72  }
    73  
    74  type timoutLoggingResponseWriter struct {
    75  	http.ResponseWriter
    76  	method       string
    77  	url          string
    78  	timeout      time.Duration
    79  	msg          []byte
    80  	requestStart time.Time
    81  	ctx          context.Context
    82  }
    83  
    84  // Write missing godoc
    85  func (lrw *timoutLoggingResponseWriter) Write(b []byte) (int, error) {
    86  	if bytes.Equal(lrw.msg, b) && time.Since(lrw.requestStart) > lrw.timeout {
    87  		log.C(lrw.ctx).Warnf("%s request to %s timed out after %s", lrw.method, lrw.url, lrw.timeout)
    88  	}
    89  
    90  	i, err := lrw.ResponseWriter.Write(b)
    91  
    92  	if err != nil && strings.Contains(err.Error(), http.ErrHandlerTimeout.Error()) && time.Since(lrw.requestStart) > lrw.timeout {
    93  		log.C(lrw.ctx).Warnf("Finished processing %s request to %s due to exceeded timeout of %s. Request processing terminated %s after timeout.",
    94  			lrw.method, lrw.url, lrw.timeout, time.Since(lrw.requestStart)-lrw.timeout)
    95  	}
    96  	return i, err
    97  }
    98  
    99  func newContentTypeHandler(h http.Handler) http.Handler {
   100  	return &contentTypeHandler{
   101  		h: h,
   102  	}
   103  }
   104  
   105  type contentTypeHandler struct {
   106  	h http.Handler
   107  }
   108  
   109  // ServeHTTP missing godoc
   110  func (h *contentTypeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
   111  	w.Header().Set(HeaderContentTypeKey, HeaderContentTypeValue)
   112  	h.h.ServeHTTP(w, r)
   113  }