From b6a75550a72e639f1e3ba25f24d2f393103c0799 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=A8=8B=E9=A3=9E?= <fay.cheng.cn@gmail.com>
Date: Thu, 16 May 2019 15:04:43 +0800
Subject: [PATCH] add bm rate limit middleware

---
 pkg/net/http/blademaster/ratelimit.go | 62 +++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)
 create mode 100644 pkg/net/http/blademaster/ratelimit.go

diff --git a/pkg/net/http/blademaster/ratelimit.go b/pkg/net/http/blademaster/ratelimit.go
new file mode 100644
index 000000000..10fd6abc6
--- /dev/null
+++ b/pkg/net/http/blademaster/ratelimit.go
@@ -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()
+	}
+}