github.com/kaydxh/golang@v0.0.131/pkg/middleware/http-middleware/http/timeout.go (about)

     1  /*
     2   *Copyright (c) 2022, kaydxh
     3   *
     4   *Permission is hereby granted, free of charge, to any person obtaining a copy
     5   *of this software and associated documentation files (the "Software"), to deal
     6   *in the Software without restriction, including without limitation the rights
     7   *to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     8   *copies of the Software, and to permit persons to whom the Software is
     9   *furnished to do so, subject to the following conditions:
    10   *
    11   *The above copyright notice and this permission notice shall be included in all
    12   *copies or substantial portions of the Software.
    13   *
    14   *THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    15   *IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    16   *FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    17   *AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    18   *LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    19   *OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    20   *SOFTWARE.
    21   */
    22  package interceptorhttp
    23  
    24  import (
    25  	"bytes"
    26  	"context"
    27  	"fmt"
    28  	"net/http"
    29  	"time"
    30  
    31  	logs_ "github.com/kaydxh/golang/pkg/logs"
    32  )
    33  
    34  // DefaultTimeoutHandler is a convenient timeout handler which
    35  // simply returns "504 Service timeout".
    36  var DefaultTimeoutHandler = http.HandlerFunc(
    37  	func(w http.ResponseWriter, r *http.Request) {
    38  		w.WriteHeader(http.StatusGatewayTimeout)
    39  		w.Write([]byte("Service timeout"))
    40  	})
    41  
    42  type timeoutWriter struct {
    43  	w http.ResponseWriter
    44  
    45  	status int
    46  	buf    *bytes.Buffer
    47  }
    48  
    49  func (tw timeoutWriter) Header() http.Header {
    50  	return tw.w.Header()
    51  }
    52  
    53  func (tw *timeoutWriter) WriteHeader(status int) {
    54  	tw.status = status
    55  }
    56  
    57  func (tw *timeoutWriter) Write(b []byte) (int, error) {
    58  	if tw.status == 0 {
    59  		tw.status = http.StatusOK
    60  	}
    61  
    62  	return tw.buf.Write(b)
    63  }
    64  
    65  func (tw *timeoutWriter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    66  	http.Error(w, http.StatusText(http.StatusInternalServerError),
    67  		http.StatusInternalServerError)
    68  }
    69  
    70  // Timeout is a middleware that cancels ctx after a given timeout and return
    71  // a 504 Gateway Timeout error to the client.
    72  func Timeout(timeout time.Duration) func(next http.Handler) http.Handler {
    73  	return func(next http.Handler) http.Handler {
    74  		fn := func(w http.ResponseWriter, r *http.Request) {
    75  
    76  			logger := logs_.GetLogger(r.Context())
    77  			background := r.Context()
    78  			tCtx, tCancel := context.WithTimeout(background, timeout)
    79  			cCtx, cCancel := context.WithCancel(background)
    80  			r = r.WithContext(cCtx)
    81  			defer tCancel()
    82  
    83  			tw := &timeoutWriter{w: w, buf: bytes.NewBuffer(nil)}
    84  			go func() {
    85  				//next.ServeHTTP(tw, r)
    86  				next.ServeHTTP(w, r)
    87  				cCancel()
    88  			}()
    89  
    90  			select {
    91  			case <-cCtx.Done():
    92  				fmt.Printf("----normal\n")
    93  			case <-tCtx.Done():
    94  				if err := tCtx.Err(); err == context.DeadlineExceeded {
    95  					logger.WithField("method", fmt.Sprintf("%s %s", r.Method, r.URL.Path)).Errorf("server %v timeout", timeout)
    96  					cCancel()
    97  					//		DefaultTimeoutHandler.ServeHTTP(w, r)
    98  					tw.ServeHTTP(w, r)
    99  				}
   100  			}
   101  		}
   102  		return http.HandlerFunc(fn)
   103  	}
   104  }