63 lines
1.5 KiB
Go
63 lines
1.5 KiB
Go
package platformapi
|
|
|
|
import (
|
|
"math"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/juju/ratelimit"
|
|
)
|
|
|
|
type LimiterConfig struct {
|
|
MaxAccessCount int
|
|
TimeGapMilliSecond int
|
|
}
|
|
|
|
type Limiter struct {
|
|
limitBucketMap map[string]*ratelimit.Bucket
|
|
limitConfigMap map[string]*LimiterConfig
|
|
defaultConfig *LimiterConfig
|
|
sync.RWMutex
|
|
}
|
|
|
|
func New(config map[string]*LimiterConfig, defaultConfig ...*LimiterConfig) *Limiter {
|
|
limiter := &Limiter{
|
|
limitConfigMap: config,
|
|
limitBucketMap: make(map[string]*ratelimit.Bucket),
|
|
}
|
|
// for k, v := range config {
|
|
// limiter.limitBucketMap[k] = ratelimit.NewBucketWithQuantum(time.Duration(v.TimeGapMilliSecond)*time.Millisecond, int64(v.MaxAccessCount), int64(v.MaxAccessCount))
|
|
// }
|
|
if len(defaultConfig) > 0 {
|
|
limiter.defaultConfig = defaultConfig[0]
|
|
}
|
|
return limiter
|
|
}
|
|
|
|
func (l *Limiter) AccessAPI(apiName, shopID string) {
|
|
bucketKey := shopID + "." + apiName
|
|
l.RLock()
|
|
bucket := l.limitBucketMap[bucketKey]
|
|
l.RUnlock()
|
|
|
|
if bucket == nil {
|
|
l.Lock()
|
|
if bucket = l.limitBucketMap[bucketKey]; bucket == nil {
|
|
v := l.limitConfigMap[apiName]
|
|
if v == nil {
|
|
v = l.defaultConfig
|
|
}
|
|
if v != nil {
|
|
bucket = ratelimit.NewBucketWithQuantum(time.Duration(v.TimeGapMilliSecond)*time.Millisecond, int64(v.MaxAccessCount), int64(v.MaxAccessCount))
|
|
} else {
|
|
bucket = ratelimit.NewBucket(time.Second, math.MaxInt64) // 占位用
|
|
}
|
|
l.limitBucketMap[bucketKey] = bucket
|
|
}
|
|
l.Unlock()
|
|
}
|
|
if bucket != nil {
|
|
bucket.Wait(1)
|
|
}
|
|
}
|