forked from zeromicro/go-zero
-
Notifications
You must be signed in to change notification settings - Fork 0
/
breakerhandler.go
41 lines (36 loc) · 1.13 KB
/
breakerhandler.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
package handler
import (
"fmt"
"net/http"
"strings"
"github.com/tal-tech/go-zero/core/breaker"
"github.com/tal-tech/go-zero/core/logx"
"github.com/tal-tech/go-zero/core/stat"
"github.com/tal-tech/go-zero/rest/httpx"
"github.com/tal-tech/go-zero/rest/internal/security"
)
const breakerSeparator = "://"
func BreakerHandler(method, path string, metrics *stat.Metrics) func(http.Handler) http.Handler {
brk := breaker.NewBreaker(breaker.WithName(strings.Join([]string{method, path}, breakerSeparator)))
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
promise, err := brk.Allow()
if err != nil {
metrics.AddDrop()
logx.Errorf("[http] dropped, %s - %s - %s",
r.RequestURI, httpx.GetRemoteAddr(r), r.UserAgent())
w.WriteHeader(http.StatusServiceUnavailable)
return
}
cw := &security.WithCodeResponseWriter{Writer: w}
defer func() {
if cw.Code < http.StatusInternalServerError {
promise.Accept()
} else {
promise.Reject(fmt.Sprintf("%d %s", cw.Code, http.StatusText(cw.Code)))
}
}()
next.ServeHTTP(cw, r)
})
}
}