package rate import ( "sync" "time" ctime "github.com/cherry-game/cherry/extend/time" clog "github.com/cherry-game/cherry/logger" ) type Buckets struct { bucketMap map[string]*Bucket // key:name, value:*Bucket lock *sync.RWMutex // lock } func NewBuckets() *Buckets { return &Buckets{ bucketMap: map[string]*Bucket{}, lock: &sync.RWMutex{}, } } func (p *Buckets) Add(name string, interval time.Duration, capacity, quantum int64) (*Bucket, bool) { if name == "" { return nil, false } p.lock.Lock() defer p.lock.Unlock() bucket, found := p.bucketMap[name] if found { return bucket, found } bucket = NewBucketWithQuantum( interval*time.Second, capacity, quantum, ) p.bucketMap[name] = bucket return bucket, true } func (p *Buckets) Get(name string) (*Bucket, bool) { p.lock.RLock() defer p.lock.RUnlock() bucket, found := p.bucketMap[name] return bucket, found } func (p *Buckets) Remove(name string) { p.lock.Lock() defer p.lock.Unlock() delete(p.bucketMap, name) } func (p *Buckets) Clean(beforeMinute int64) { p.lock.Lock() defer p.lock.Unlock() for name, bucket := range p.bucketMap { now := ctime.NewTime(time.Now(), false) lastTime := ctime.NewTime(bucket.lastTime, false) minute := lastTime.DiffInMinutes(&now) if minute >= beforeMinute { delete(p.bucketMap, name) clog.Infof("clean bucket. [name = %s]", name) } } }