mirror of
https://github.com/go-kratos/kratos.git
synced 2025-03-17 21:07:54 +02:00
add warden rate limit middleware
This commit is contained in:
parent
b6a75550a7
commit
9917255196
63
pkg/net/rpc/warden/internal/ratelimiter/ratelimiter.go
Normal file
63
pkg/net/rpc/warden/internal/ratelimiter/ratelimiter.go
Normal file
@ -0,0 +1,63 @@
|
||||
package ratelimiter
|
||||
|
||||
import (
|
||||
"context"
|
||||
"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"
|
||||
|
||||
"google.golang.org/grpc"
|
||||
)
|
||||
|
||||
const (
|
||||
_statName = "go_grpc_bbr"
|
||||
)
|
||||
|
||||
var (
|
||||
stats = prom.New().WithState("go_grpc_bbr", []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 {
|
||||
stats.Incr(_statName, uri)
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
done(limit.DoneInfo{Op: limit.Success})
|
||||
b.printStats(uri, limiter)
|
||||
}()
|
||||
resp, err = handler(ctx, req)
|
||||
return
|
||||
}
|
||||
}
|
Loading…
x
Reference in New Issue
Block a user