s_cherry_e/sr/internal/component/rate/buckets.go

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