github.com/lingyao2333/mo-zero@v1.4.1/rest/handler/breakerhandler.go (about)

     1  package handler
     2  
     3  import (
     4  	"fmt"
     5  	"net/http"
     6  	"strings"
     7  
     8  	"github.com/lingyao2333/mo-zero/core/breaker"
     9  	"github.com/lingyao2333/mo-zero/core/logx"
    10  	"github.com/lingyao2333/mo-zero/core/stat"
    11  	"github.com/lingyao2333/mo-zero/rest/httpx"
    12  	"github.com/lingyao2333/mo-zero/rest/internal/response"
    13  )
    14  
    15  const breakerSeparator = "://"
    16  
    17  // BreakerHandler returns a break circuit middleware.
    18  func BreakerHandler(method, path string, metrics *stat.Metrics) func(http.Handler) http.Handler {
    19  	brk := breaker.NewBreaker(breaker.WithName(strings.Join([]string{method, path}, breakerSeparator)))
    20  	return func(next http.Handler) http.Handler {
    21  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    22  			promise, err := brk.Allow()
    23  			if err != nil {
    24  				metrics.AddDrop()
    25  				logx.Errorf("[http] dropped, %s - %s - %s",
    26  					r.RequestURI, httpx.GetRemoteAddr(r), r.UserAgent())
    27  				w.WriteHeader(http.StatusServiceUnavailable)
    28  				return
    29  			}
    30  
    31  			cw := &response.WithCodeResponseWriter{Writer: w}
    32  			defer func() {
    33  				if cw.Code < http.StatusInternalServerError {
    34  					promise.Accept()
    35  				} else {
    36  					promise.Reject(fmt.Sprintf("%d %s", cw.Code, http.StatusText(cw.Code)))
    37  				}
    38  			}()
    39  			next.ServeHTTP(cw, r)
    40  		})
    41  	}
    42  }