52 lines
1.2 KiB
Go
52 lines
1.2 KiB
Go
package platformapi
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/juju/ratelimit"
|
|
)
|
|
|
|
type LimiterConfig struct {
|
|
MaxAccessCount int
|
|
TimeGapMilliSecond int
|
|
}
|
|
|
|
type Limiter struct {
|
|
limitConfig map[string]*ratelimit.Bucket
|
|
defaultConfig *LimiterConfig
|
|
sync.RWMutex
|
|
}
|
|
|
|
func New(config map[string]*LimiterConfig, defaultConfig ...*LimiterConfig) *Limiter {
|
|
limiter := &Limiter{
|
|
limitConfig: make(map[string]*ratelimit.Bucket),
|
|
}
|
|
for k, v := range config {
|
|
limiter.limitConfig[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 string) {
|
|
l.RLock()
|
|
bucket := l.limitConfig[apiName]
|
|
l.RUnlock()
|
|
|
|
if bucket == nil && l.defaultConfig != nil {
|
|
l.Lock()
|
|
if bucket = l.limitConfig[apiName]; bucket == nil {
|
|
v := l.defaultConfig
|
|
bucket = ratelimit.NewBucketWithQuantum(time.Duration(v.TimeGapMilliSecond)*time.Millisecond, int64(v.MaxAccessCount), int64(v.MaxAccessCount))
|
|
l.limitConfig[apiName] = bucket
|
|
}
|
|
l.Unlock()
|
|
}
|
|
if bucket != nil {
|
|
bucket.Wait(1)
|
|
}
|
|
}
|