Initial commit

common/redis/limit.go

@@ -0,0 +1,110 @@
+package redis
+
+import (
+	_ "embed"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/go-redis/redis"
+)
+
+//go:embed limit.lua
+var lua string
+
+const (
+	SUCCESS = "success"
+	FAILED  = "failed"
+)
+
+type Bucket struct {
+	Name        string          `json:"na"`
+	Capacity    int64           `json:"ca"`
+	Rate        int64           `json:"ra"`
+	Size        int64           `json:"si"`
+	GenerateAt  int64           `json:"ge_at"`
+	CreatedAt   int64           `json:"cr_at"`
+	UpdatedAt   int64           `json:"up_at"`
+	PassedCount int64           `json:"pa_co"`
+	RejectCount int64           `json:"re_co"`
+	PassedHours map[int64]int64 `json:"pa_ho"`
+	RejectHours map[int64]int64 `json:"re_ho"`
+}
+
+type BucketList []*Bucket
+
+func (a BucketList) Len() int { return len(a) }
+func (a BucketList) Less(i, j int) bool {
+	return a[i].PassedCount > a[j].PassedCount
+}
+func (a BucketList) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
+
+type TokenBucket struct {
+	keyGen   func() string
+	capacity int64
+	rate     int64
+	script   *redis.Script
+	client   *redis.Client
+	expire   time.Duration
+	lock     sync.RWMutex
+	curKey   string
+}
+
+func NewTokenBucket(client *redis.Client, keyGen func() string, capacity, rate int64, expire time.Duration) *TokenBucket {
+	return &TokenBucket{
+		expire:   expire,
+		capacity: capacity,
+		rate:     rate,
+		client:   client,
+		script:   redis.NewScript(lua),
+		keyGen:   keyGen,
+	}
+}
+
+func (s *TokenBucket) TakeTokens(name string, count int64) error {
+	key := s.keyGen()
+	now := time.Now().Local()
+	cmd := []string{
+		key,  //key
+		name, //name
+		fmt.Sprintf("%d", s.capacity),
+		fmt.Sprintf("%d", s.rate),
+		fmt.Sprintf("%d", count),
+		fmt.Sprintf("%d", now.UnixMilli()),
+		fmt.Sprintf("%d", now.Hour()),
+	}
+	ret, err := s.script.Run(s.client, cmd).String()
+	if err != nil {
+		return err
+	}
+	if ret != SUCCESS {
+		return errors.New("sorry, I can't work anymore~")
+	}
+	var same = false
+	if s.lock.TryRLock() {
+		same = key == s.curKey
+		s.lock.RUnlock()
+	} else {
+		return nil
+	}
+	if !same {
+		if s.lock.TryLock() {
+			if ok, err := s.client.Expire(key, s.expire).Result(); err == nil && ok {
+				s.curKey = key
+			}
+			s.lock.Unlock()
+		}
+	}
+	return nil
+}
+
+func (s *TokenBucket) Traverse(callback func(key string, item *Bucket)) error {
+	return NewHash(s.client, s.keyGen()).Scan("", 1000, func(name, value string) {
+		bucket := Bucket{}
+		if err := json.Unmarshal([]byte(value), &bucket); err == nil {
+			callback(name, &bucket)
+		}
+	})
+}