mirror of
https://github.com/go-kratos/kratos.git
synced 2025-03-17 21:07:54 +02:00
add bm rate limit middleware
This commit is contained in:
parent
c6550d4439
commit
b6a75550a7
62
pkg/net/http/blademaster/ratelimit.go
Normal file
62
pkg/net/http/blademaster/ratelimit.go
Normal file
@ -0,0 +1,62 @@
|
||||
package blademaster
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/Bilibili/kratos/pkg/log"
|
||||
limit "github.com/Bilibili/kratos/pkg/ratelimit"
|
||||
"github.com/Bilibili/kratos/pkg/ratelimit/bbr"
|
||||
"github.com/Bilibili/kratos/pkg/stat/prom"
|
||||
)
|
||||
|
||||
const (
|
||||
_statName = "go_http_bbr"
|
||||
)
|
||||
|
||||
var (
|
||||
bbrStats = prom.New().WithState("go_http_bbr", []string{"url"})
|
||||
)
|
||||
|
||||
// RateLimiter bbr middleware.
|
||||
type RateLimiter struct {
|
||||
group *bbr.Group
|
||||
logTime int64
|
||||
}
|
||||
|
||||
// New return a ratelimit middleware.
|
||||
func NewRateLimiter(conf *bbr.Config) (s *RateLimiter) {
|
||||
return &RateLimiter{
|
||||
group: bbr.NewGroup(conf),
|
||||
logTime: time.Now().UnixNano(),
|
||||
}
|
||||
}
|
||||
|
||||
func (b *RateLimiter) printStats(routePath string, limiter limit.Limiter) {
|
||||
now := time.Now().UnixNano()
|
||||
if now-atomic.LoadInt64(&b.logTime) > int64(time.Second*3) {
|
||||
atomic.StoreInt64(&b.logTime, now)
|
||||
log.Info("http.bbr path:%s stat:%+v", routePath, limiter.(*bbr.BBR).Stat())
|
||||
}
|
||||
}
|
||||
|
||||
// Limit return a bm handler func.
|
||||
func (b *RateLimiter) Limit() HandlerFunc {
|
||||
return func(c *Context) {
|
||||
uri := fmt.Sprintf("%s://%s%s", c.Request.URL.Scheme, c.Request.Host, c.Request.URL.Path)
|
||||
limiter := b.group.Get(uri)
|
||||
done, err := limiter.Allow(c)
|
||||
if err != nil {
|
||||
bbrStats.Incr(_statName, uri)
|
||||
c.JSON(nil, err)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
done(limit.DoneInfo{Op: limit.Success})
|
||||
b.printStats(uri, limiter)
|
||||
}()
|
||||
c.Next()
|
||||
}
|
||||
}
|
Loading…
x
Reference in New Issue
Block a user