1 Star 0 Fork 2

QunXiongZhuLu / kratos

Create your Gitee Account
Explore and code with more than 6 million developers,Free private repositories !:)
Sign up
Clone or download
ratelimiter.go 1.60 KB
Copy Edit Raw Blame History
sirius authored 2021-01-04 16:20 . 更新依赖库地址,更新包名称
package ratelimiter
import (
"context"
"sync/atomic"
"time"
"google.golang.org/grpc"
"gitee.com/QunXiongZhuLu/kratos/pkg/log"
limit "gitee.com/QunXiongZhuLu/kratos/pkg/ratelimit"
"gitee.com/QunXiongZhuLu/kratos/pkg/ratelimit/bbr"
"gitee.com/QunXiongZhuLu/kratos/pkg/stat/metric"
)
var (
_metricServerBBR = metric.NewCounterVec(&metric.CounterVecOpts{
Namespace: "grpc_server",
Subsystem: "",
Name: "bbr_total",
Help: "grpc server bbr total.",
Labels: []string{"url"},
})
)
// RateLimiter bbr middleware.
type RateLimiter struct {
group *bbr.Group
logTime int64
}
// New return a ratelimit middleware.
func New(conf *bbr.Config) (s *RateLimiter) {
return &RateLimiter{
group: bbr.NewGroup(conf),
logTime: time.Now().UnixNano(),
}
}
func (b *RateLimiter) printStats(fullMethod 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("grpc.bbr path:%s stat:%+v", fullMethod, limiter.(*bbr.BBR).Stat())
}
}
// Limit is a server interceptor that detects and rejects overloaded traffic.
func (b *RateLimiter) Limit() grpc.UnaryServerInterceptor {
return func(ctx context.Context, req interface{}, args *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
uri := args.FullMethod
limiter := b.group.Get(uri)
done, err := limiter.Allow(ctx)
if err != nil {
_metricServerBBR.Inc(uri)
return
}
defer func() {
done(limit.DoneInfo{Op: limit.Success})
b.printStats(uri, limiter)
}()
resp, err = handler(ctx, req)
return
}
}

Comment ( 0 )

Sign in for post a comment

1
https://gitee.com/QunXiongZhuLu/kratos.git
git@gitee.com:QunXiongZhuLu/kratos.git
QunXiongZhuLu
kratos
kratos
v0.6.2

Search