forked from go-kratos/kratos
-
Notifications
You must be signed in to change notification settings - Fork 0
/
ratelimit.go
62 lines (54 loc) · 1.37 KB
/
ratelimit.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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
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()
}
}