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) } }