add warden rate limit middleware

pkg/net/rpc/warden/internal/ratelimiter/ratelimiter.go

@@ -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
+	}
+}