go-zero/ngin/handler/maxconnshandler.go

38 lines
758 B
Go
Raw Normal View History

2020-07-29 18:00:04 +08:00
package handler
2020-07-26 17:09:05 +08:00
import (
"net/http"
"zero/core/logx"
"zero/core/syncx"
2020-07-29 18:00:04 +08:00
"zero/ngin/internal"
2020-07-26 17:09:05 +08:00
)
func MaxConns(n int) func(http.Handler) http.Handler {
if n <= 0 {
return func(next http.Handler) http.Handler {
return next
}
}
return func(next http.Handler) http.Handler {
latchLimiter := syncx.NewLimit(n)
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if latchLimiter.TryBorrow() {
defer func() {
if err := latchLimiter.Return(); err != nil {
logx.Error(err)
}
}()
next.ServeHTTP(w, r)
} else {
2020-07-29 18:00:04 +08:00
internal.Errorf(r, "Concurrent connections over %d, rejected with code %d",
2020-07-26 17:09:05 +08:00
n, http.StatusServiceUnavailable)
w.WriteHeader(http.StatusServiceUnavailable)
}
})
}
}