34 lines
667 B
Go
34 lines
667 B
Go
package platformapi
|
|
|
|
import (
|
|
"time"
|
|
|
|
"github.com/juju/ratelimit"
|
|
)
|
|
|
|
type LimiterConfig struct {
|
|
MaxAccessCount int
|
|
TimeGapInSecond int
|
|
}
|
|
|
|
type Limiter struct {
|
|
limitConfig map[string]*ratelimit.Bucket
|
|
}
|
|
|
|
func New(config map[string]*LimiterConfig) *Limiter {
|
|
limiter := &Limiter{
|
|
limitConfig: make(map[string]*ratelimit.Bucket),
|
|
}
|
|
for k, v := range config {
|
|
limiter.limitConfig[k] = ratelimit.NewBucketWithQuantum(time.Duration(v.TimeGapInSecond)*time.Second, int64(v.MaxAccessCount), int64(v.MaxAccessCount))
|
|
}
|
|
return limiter
|
|
}
|
|
|
|
func (l *Limiter) AccessAPI(apiName string) {
|
|
bucket := l.limitConfig[apiName]
|
|
if bucket != nil {
|
|
bucket.Wait(1)
|
|
}
|
|
}
|