Create & Init Project...

library/net/http/blademaster/middleware/rate/BUILD

@@ -0,0 +1,53 @@
+package(default_visibility = ["//visibility:public"])
+
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_test",
+    "go_library",
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = ["limit_test.go"],
+    embed = [":go_default_library"],
+    rundir = ".",
+    tags = ["automanaged"],
+    deps = ["//library/net/http/blademaster:go_default_library"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["limit.go"],
+    importpath = "go-common/library/net/http/blademaster/middleware/rate",
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//library/log:go_default_library",
+        "//library/net/http/blademaster:go_default_library",
+        "//vendor/golang.org/x/time/rate:go_default_library",
+    ],
+)
+
+go_test(
+    name = "go_default_xtest",
+    srcs = ["example_test.go"],
+    tags = ["automanaged"],
+    deps = [
+        "//library/net/http/blademaster:go_default_library",
+        "//library/net/http/blademaster/middleware/rate:go_default_library",
+    ],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [":package-srcs"],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)

library/net/http/blademaster/middleware/rate/CHANGELOG.md

@@ -0,0 +1,5 @@
+### business/blademaster/rate
+
+##### Version 1.0.0
+
+1. 完成基本功能与测试

library/net/http/blademaster/middleware/rate/CONTRIBUTORS.md

@@ -0,0 +1,6 @@
+# Author
+lintnaghui
+caoguoliang
+ 
+# Reviewer
+maojian

library/net/http/blademaster/middleware/rate/OWNERS

@@ -0,0 +1,9 @@
+# See the OWNERS docs at https://go.k8s.io/owners
+
+approvers:
+- caoguoliang
+- lintnaghui
+reviewers:
+- caoguoliang
+- lintnaghui
+- maojian

library/net/http/blademaster/middleware/rate/README.md

@@ -0,0 +1,13 @@
+#### business/blademaster/rate
+
+##### 项目简介
+
+blademaster 的 rate middleware，主要用于限制内部调用的频率
+
+##### 编译环境
+
+- **请只用 Golang v1.8.x 以上版本编译执行**
+
+##### 依赖包
+
+- No other dependency

library/net/http/blademaster/middleware/rate/example_test.go

@@ -0,0 +1,28 @@
+package rate_test
+
+import (
+	"go-common/library/net/http/blademaster"
+	"go-common/library/net/http/blademaster/middleware/rate"
+)
+
+// This example create a rate middleware instance and attach to a blademaster engine,
+// it will protect '/ping' API frequency with specified policy.
+// If any internal service who requests this API more frequently than 1 req/second,
+// a StatusTooManyRequests error will be raised.
+func Example() {
+	lim := rate.New(&rate.Config{
+		URLs: map[string]*rate.Limit{
+			"/ping": &rate.Limit{Limit: 1, Burst: 2},
+		},
+		Apps: map[string]*rate.Limit{
+			"a-secret-app-key": &rate.Limit{Limit: 1, Burst: 2},
+		},
+	})
+
+	engine := blademaster.Default()
+	engine.Use(lim)
+	engine.GET("/ping", func(c *blademaster.Context) {
+		c.String(200, "%s", "pong")
+	})
+	engine.Run(":18080")
+}

library/net/http/blademaster/middleware/rate/limit.go

@@ -0,0 +1,137 @@
+package rate
+
+import (
+	"net/http"
+	"sync/atomic"
+
+	"go-common/library/log"
+	bm "go-common/library/net/http/blademaster"
+
+	"golang.org/x/time/rate"
+)
+
+const (
+	_defBurst = 100
+)
+
+// Limiter controls how frequently events are allowed to happen.
+type Limiter struct {
+	apps atomic.Value
+	urls atomic.Value
+}
+
+// Config limitter  conf.
+type Config struct {
+	Apps map[string]*Limit
+	URLs map[string]*Limit
+}
+
+// Limit limit conf.
+type Limit struct {
+	Limit rate.Limit
+	Burst int
+}
+
+// New return Limiter.
+func New(conf *Config) (l *Limiter) {
+	l = &Limiter{}
+	l.apps.Store(make(map[string]*rate.Limiter))
+	l.urls.Store(make(map[string]*rate.Limiter))
+	if conf != nil {
+		l.Reload(conf)
+	}
+	return
+}
+
+// Reload reload limit conf.
+func (l *Limiter) Reload(c *Config) {
+	if c == nil {
+		return
+	}
+	var (
+		ok  bool
+		al  *rate.Limiter
+		ul  *rate.Limiter
+		as  map[string]*rate.Limiter
+		nas map[string]*rate.Limiter
+		us  map[string]*rate.Limiter
+		nus map[string]*rate.Limiter
+	)
+	if as, ok = l.apps.Load().(map[string]*rate.Limiter); !ok {
+		log.Error("apps limiter load map hava no data ")
+		return
+	}
+	nas = make(map[string]*rate.Limiter, len(as))
+	for k, v := range as {
+		nas[k] = v
+	}
+	for k, v := range c.Apps {
+		if al, ok = nas[k]; !ok || (al.Burst() != v.Burst || al.Limit() != v.Limit) {
+			nas[k] = rate.NewLimiter(v.fix())
+		}
+	}
+	l.apps.Store(nas)
+
+	if us, ok = l.urls.Load().(map[string]*rate.Limiter); !ok {
+		log.Error("urls limiter load map hava no data ")
+		return
+	}
+	nus = make(map[string]*rate.Limiter, len(us))
+	for k, v := range us {
+		nus[k] = v
+	}
+	for k, v := range c.URLs {
+		if ul, ok = nus[k]; !ok || (ul.Burst() != v.Burst || ul.Limit() != v.Limit) {
+			nus[k] = rate.NewLimiter(v.fix())
+		}
+	}
+	l.urls.Store(nus)
+}
+
+func (l *Limit) fix() (lim rate.Limit, b int) {
+	lim = rate.Inf
+	b = _defBurst
+	if l.Limit <= 0 {
+		lim = rate.Inf
+	} else {
+		lim = l.Limit
+	}
+	if l.Burst > 0 {
+		b = l.Burst
+	}
+	return
+}
+
+// Allow reports whether event may happen at time now.
+func (l *Limiter) Allow(appKey, path string) bool {
+	if as, ok := l.apps.Load().(map[string]*rate.Limiter); ok {
+		if lim, ok := as[appKey]; ok {
+			if !lim.Allow() {
+				return false
+			}
+		}
+	}
+	if us, ok := l.urls.Load().(map[string]*rate.Limiter); ok {
+		if lim, ok := us[path]; ok {
+			if !lim.Allow() {
+				return false
+			}
+		}
+	}
+	return true
+}
+
+func (l *Limiter) ServeHTTP(c *bm.Context) {
+	req := c.Request
+	appkey := req.Form.Get("appkey")
+	path := req.URL.Path
+	if !l.Allow(appkey, path) {
+		c.AbortWithStatus(http.StatusTooManyRequests)
+		return
+	}
+}
+
+// Handler is router allow handle.
+func (l *Limiter) Handler() bm.HandlerFunc {
+	return l.ServeHTTP
+}

library/net/http/blademaster/middleware/rate/limit_test.go

@@ -0,0 +1,131 @@
+package rate
+
+import (
+	"context"
+	"io/ioutil"
+	"net/http"
+	"testing"
+	"time"
+
+	bm "go-common/library/net/http/blademaster"
+)
+
+func TestLimiterUrl(t *testing.T) {
+	l := New(&Config{
+		URLs: map[string]*Limit{"/limit/test": &Limit{Limit: 1, Burst: 2}},
+	})
+	if !l.Allow("testApp", "/limit/test") {
+		t.Logf("request should pass,but blocked")
+		t.FailNow()
+	}
+	if !l.Allow("testApp", "/limit/test") {
+		t.Logf("request should pass,but blocked")
+		t.FailNow()
+	}
+	if l.Allow("testApp", "/limit/test") {
+		t.Logf("request should block,but passed")
+		t.FailNow()
+	}
+}
+
+func TestLimiterApp(t *testing.T) {
+	l := New(&Config{
+		Apps: map[string]*Limit{"testApp": &Limit{Limit: 1, Burst: 2}},
+	})
+	if !l.Allow("testApp", "/limit/test") {
+		t.Logf("request should pass,but blocked")
+		t.FailNow()
+	}
+	if !l.Allow("testApp", "/limit/test") {
+		t.Logf("request should pass,but blocked")
+		t.FailNow()
+	}
+	if l.Allow("testApp", "/limit/test") {
+		t.Logf("request should block,but passed")
+		t.FailNow()
+	}
+}
+
+func TestLimiterUrlApp(t *testing.T) {
+	l := New(&Config{
+		Apps: map[string]*Limit{"testApp": &Limit{Limit: 2, Burst: 1}},
+		URLs: map[string]*Limit{"/limit/test": &Limit{Limit: 2, Burst: 1}},
+	})
+	if !l.Allow("testApp", "/limit/test") {
+		t.Logf("request should pass,but blocked")
+		t.FailNow()
+	}
+	if l.Allow("testApp", "/limit/test") {
+		t.Logf("request should block,but passed")
+		t.FailNow()
+	}
+	l.Reload(&Config{
+		Apps: map[string]*Limit{"testApp": &Limit{Limit: 1, Burst: 2}},
+		URLs: map[string]*Limit{"/limit/test": &Limit{Limit: 1, Burst: 2}},
+	})
+	if !l.Allow("testApp", "/limit/test") {
+		t.Logf("request should pass,but blocked")
+		t.FailNow()
+	}
+	if !l.Allow("testApp", "/limit/test") {
+		t.Logf("request should pass,but blocked")
+		t.FailNow()
+	}
+	if l.Allow("testApp", "/limit/test") {
+		t.Logf("request should block,but passed")
+		t.FailNow()
+	}
+}
+
+func TestLimiterHandler(t *testing.T) {
+	l := New(&Config{
+		Apps: map[string]*Limit{"testApp": &Limit{Limit: 1, Burst: 1}},
+		URLs: map[string]*Limit{"/limit/test": &Limit{Limit: 2, Burst: 4}},
+	})
+	engine := bm.New()
+	engine.Use(l.Handler())
+	engine.GET("/limit/test", func(c *bm.Context) {
+		c.String(200, "pass")
+	})
+	go engine.Run(":18080")
+	defer func() {
+		engine.Server().Shutdown(context.TODO())
+	}()
+
+	time.Sleep(time.Millisecond * 20)
+	code, content, err := httpGet("http://127.0.0.1:18080/limit/test?appkey=testApp")
+	if err != nil {
+		t.Logf("http get failed,err:=%v", err)
+		t.FailNow()
+	}
+	if code != 200 || string(content) != "pass" {
+		t.Logf("request should pass by limiter,but blocked")
+		t.FailNow()
+	}
+
+	_, content, err = httpGet("http://127.0.0.1:18080/limit/test?appkey=testApp")
+	if err != nil {
+		t.Logf("http get failed,err:=%v", err)
+		t.FailNow()
+	}
+	if string(content) == "pass" {
+		t.Logf("request should block by limiter,but passed")
+		t.FailNow()
+	}
+}
+
+func httpGet(url string) (code int, content []byte, err error) {
+	resp, err := http.Get(url)
+	if err != nil {
+
+		return
+	}
+	defer resp.Body.Close()
+
+	content, err = ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return
+	}
+	code = resp.StatusCode
+	return
+}