!1 Merge back to master

Merge pull request !1 from dev branch

server/libs/limiter/limiter.go

@@ -0,0 +1,5 @@
+package limiter
+
+type Limiter interface {
+	Access(string, int16) bool
+}

server/libs/limiter/rate_limiter.go

@@ -0,0 +1,220 @@
+package limiter
+
+import (
+	"sync"
+	"time"
+)
+
+func now() int32 {
+	return int32(time.Now().Unix())
+}
+
+func afterCyc(cyc int32) int32 {
+	return int32(time.Now().Unix()) + cyc
+}
+
+func afterTtl(ttl int32) int32 {
+	return int32(time.Now().Unix()) + ttl
+}
+
+type Bucket struct {
+	Refresh int32
+	Tokens  int16
+}
+
+func NewBucket(cyc int32, cap int16) *Bucket {
+	return &Bucket{
+		Refresh: afterCyc(cyc),
+		Tokens:  cap,
+	}
+}
+
+type Item struct {
+	Expired int32
+	Buckets map[int16]*Bucket
+}
+
+func NewItem(ttl int32) *Item {
+	return &Item{
+		Expired: afterTtl(ttl),
+		Buckets: make(map[int16]*Bucket),
+	}
+}
+
+type RateLimiter struct {
+	items      map[string]*Item
+	bucketCap  int16
+	customCaps map[int16]int16
+	cap        int64
+	cyc        int32 // how much time, item autoclean will be executed, bucket will be refreshed
+	ttl        int32 // how much time, item will be expired(but not cleaned)
+	mux        sync.RWMutex
+	snapshot   map[string]map[int16]*Bucket
+}
+
+func NewRateLimiter(cap int64, ttl int32, cyc int32, bucketCap int16, customCaps map[int16]int16) Limiter {
+	if cap < 1 || ttl < 1 || cyc < 1 || bucketCap < 1 {
+		panic("cap | bucketCap | ttl | cycle cant be less than 1")
+	}
+
+	limiter := &RateLimiter{
+		items:      make(map[string]*Item, cap),
+		bucketCap:  bucketCap,
+		customCaps: customCaps,
+		cap:        cap,
+		ttl:        ttl,
+		cyc:        cyc,
+	}
+
+	go limiter.autoClean()
+
+	return limiter
+}
+
+func (limiter *RateLimiter) getBucketCap(opId int16) int16 {
+	bucketCap, existed := limiter.customCaps[opId]
+	if !existed {
+		return limiter.bucketCap
+	}
+	return bucketCap
+}
+
+func (limiter *RateLimiter) Access(itemId string, opId int16) bool {
+	limiter.mux.Lock()
+	defer limiter.mux.Unlock()
+
+	item, itemExisted := limiter.items[itemId]
+	if !itemExisted {
+		if int64(len(limiter.items)) >= limiter.cap {
+			return false
+		}
+
+		limiter.items[itemId] = NewItem(limiter.ttl)
+		limiter.items[itemId].Buckets[opId] = NewBucket(limiter.cyc, limiter.getBucketCap(opId)-1)
+		return true
+	}
+
+	bucket, bucketExisted := item.Buckets[opId]
+	if !bucketExisted {
+		item.Buckets[opId] = NewBucket(limiter.cyc, limiter.getBucketCap(opId)-1)
+		return true
+	}
+
+	if bucket.Refresh > now() {
+		if bucket.Tokens > 0 {
+			bucket.Tokens--
+			return true
+		}
+		return false
+	}
+
+	bucket.Refresh = afterCyc(limiter.cyc)
+	bucket.Tokens = limiter.getBucketCap(opId) - 1
+	return true
+}
+
+func (limiter *RateLimiter) GetCap() int64 {
+	return limiter.cap
+}
+
+func (limiter *RateLimiter) GetSize() int64 {
+	limiter.mux.RLock()
+	defer limiter.mux.RUnlock()
+	return int64(len(limiter.items))
+}
+
+func (limiter *RateLimiter) ExpandCap(cap int64) bool {
+	limiter.mux.RLock()
+	defer limiter.mux.RUnlock()
+
+	if cap <= int64(len(limiter.items)) {
+		return false
+	}
+
+	limiter.cap = cap
+	return true
+}
+
+func (limiter *RateLimiter) GetTTL() int32 {
+	return limiter.ttl
+}
+
+func (limiter *RateLimiter) UpdateTTL(ttl int32) bool {
+	if ttl < 1 {
+		return false
+	}
+
+	limiter.ttl = ttl
+	return true
+}
+
+func (limiter *RateLimiter) GetCyc() int32 {
+	return limiter.cyc
+}
+
+func (limiter *RateLimiter) UpdateCyc(cyc int32) bool {
+	if limiter.cyc < 1 {
+		return false
+	}
+
+	limiter.cyc = cyc
+	return true
+}
+
+func (limiter *RateLimiter) Snapshot() map[string]map[int16]*Bucket {
+	return limiter.snapshot
+}
+
+func (limiter *RateLimiter) autoClean() {
+	for {
+		if limiter.cyc == 0 {
+			break
+		}
+		time.Sleep(time.Duration(int64(limiter.cyc) * 1000000000))
+		limiter.clean()
+	}
+}
+
+// clean may add affect other operations, do frequently?
+func (limiter *RateLimiter) clean() {
+	limiter.snapshot = make(map[string]map[int16]*Bucket)
+	now := now()
+
+	limiter.mux.RLock()
+	defer limiter.mux.RUnlock()
+	for key, item := range limiter.items {
+		if item.Expired <= now {
+			delete(limiter.items, key)
+		} else {
+			limiter.snapshot[key] = item.Buckets
+		}
+	}
+}
+
+// Only for test
+func (limiter *RateLimiter) exist(id string) bool {
+	limiter.mux.RLock()
+	defer limiter.mux.RUnlock()
+
+	_, existed := limiter.items[id]
+	return existed
+}
+
+// Only for test
+func (limiter *RateLimiter) truncate() {
+	limiter.mux.RLock()
+	defer limiter.mux.RUnlock()
+
+	for key, _ := range limiter.items {
+		delete(limiter.items, key)
+	}
+}
+
+// Only for test
+func (limiter *RateLimiter) get(id string) (*Item, bool) {
+	limiter.mux.RLock()
+	defer limiter.mux.RUnlock()
+
+	item, existed := limiter.items[id]
+	return item, existed
+}

server/libs/limiter/rate_limiter_test.go

@@ -0,0 +1,161 @@
+package limiter
+
+import (
+	"fmt"
+	"math/rand"
+	"testing"
+	"time"
+)
+
+var rnd = rand.New(rand.NewSource(time.Now().UnixNano()))
+
+const rndCap = 10000
+const addCap = 1
+
+// how to set time
+// extend: wait can be greater than ttl/2
+// cyc is smaller than ttl and wait, then it can be clean in time
+const cap = 40
+const ttl = 3
+const cyc = 1
+const bucketCap = 2
+const id1 = "id1"
+const id2 = "id2"
+const op1 int16 = 0
+const op2 int16 = 1
+
+var customCaps = map[int16]int16{
+	op2: 1000,
+}
+
+const wait = 1
+
+var limiter = NewRateLimiter(cap, ttl, cyc, bucketCap, customCaps).(*RateLimiter)
+
+func printItem(id string) {
+	item, existed := limiter.get(id1)
+	if existed {
+		fmt.Println("expired, now, existed", item.Expired, now(), existed)
+		for id, bucket := range item.Buckets {
+			fmt.Println("\tid, bucket", id, bucket)
+		}
+	} else {
+		fmt.Println("not existed")
+	}
+}
+
+var idSeed = 0
+
+func randId() string {
+	idSeed++
+	return fmt.Sprintf("%d", idSeed)
+}
+
+func TestAccess(t *testing.T) {
+	func(t *testing.T) {
+		canAccess := limiter.Access(id1, op1)
+		if !canAccess {
+			t.Fatal("access: fail")
+		}
+
+		for i := 0; i < bucketCap; i++ {
+			canAccess = limiter.Access(id1, op1)
+		}
+
+		if canAccess {
+			t.Fatal("access: fail to deny access")
+		}
+
+		time.Sleep(time.Duration(limiter.GetCyc()) * time.Second)
+
+		canAccess = limiter.Access(id1, op1)
+		if !canAccess {
+			t.Fatal("access: fail to refresh tokens")
+		}
+	}(t)
+}
+
+func TestCap(t *testing.T) {
+	originalCap := limiter.GetCap()
+	fmt.Printf("cap:info: %d\n", originalCap)
+
+	ok := limiter.ExpandCap(originalCap + addCap)
+
+	if !ok || limiter.GetCap() != originalCap+addCap {
+		t.Fatal("cap: fail to expand")
+	}
+
+	ok = limiter.ExpandCap(limiter.GetSize() - addCap)
+	if ok {
+		t.Fatal("cap: shrink cap")
+	}
+
+	ids := []string{}
+	for limiter.GetSize() < limiter.GetCap() {
+		id := randId()
+		ids = append(ids, id)
+
+		ok := limiter.Access(id, 0)
+		if !ok {
+			t.Fatal("cap: not full")
+		}
+	}
+
+	if limiter.GetSize() != limiter.GetCap() {
+		t.Fatal("cap: incorrect size")
+	}
+
+	if limiter.Access(randId(), 0) {
+		t.Fatal("cap: more than cap")
+	}
+
+	limiter.truncate()
+}
+
+func TestTtl(t *testing.T) {
+	var addTtl int32 = 1
+	originalTTL := limiter.GetTTL()
+	fmt.Printf("ttl:info: %d\n", originalTTL)
+
+	limiter.UpdateTTL(originalTTL + addTtl)
+	if limiter.GetTTL() != originalTTL+addTtl {
+		t.Fatal("ttl: update fail")
+	}
+}
+
+func cycTest(t *testing.T) {
+	var addCyc int32 = 1
+	originalCyc := limiter.GetCyc()
+	fmt.Printf("cyc:info: %d\n", originalCyc)
+
+	limiter.UpdateCyc(originalCyc + addCyc)
+	if limiter.GetCyc() != originalCyc+addCyc {
+		t.Fatal("cyc: update fail")
+	}
+}
+
+func autoCleanTest(t *testing.T) {
+	ids := []string{
+		randId(),
+		randId(),
+	}
+
+	for _, id := range ids {
+		ok := limiter.Access(id, 0)
+		if ok {
+			t.Fatal("autoClean: warning: add fail")
+		}
+	}
+
+	time.Sleep(time.Duration(limiter.GetTTL()+wait) * time.Second)
+
+	for _, id := range ids {
+		_, exist := limiter.get(id)
+		if exist {
+			t.Fatal("autoClean: item still exist")
+		}
+	}
+}
+
+// func snapshotTest(t *testing.T) {
+// }