forked from all_user/stat_center
111 lines
2.4 KiB
Go
111 lines
2.4 KiB
Go
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)
|
|
}
|
|
})
|
|
}
|