aiChat/internal/model/limit.go

package model

import (
	"context"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/juju/ratelimit"
	"sync"
	"time"
)

var (
	ChatLimit = NewUserLimiter()
)

type UserLimiter struct {
	sync.Mutex
	limiters map[int64]*ratelimit.Bucket
}

func NewUserLimiter() *UserLimiter {
	return &UserLimiter{
		limiters: make(map[int64]*ratelimit.Bucket),
	}
}

func (l *UserLimiter) GetBucket(ctx context.Context, userID int64) *ratelimit.Bucket {
	l.Lock()
	defer l.Unlock()
	bucket, ok := l.limiters[userID]
	if !ok {
		bucket = ratelimit.NewBucket(time.Minute, g.Cfg().MustGet(ctx, "limit.rateMinute").Int64())
		l.limiters[userID] = bucket
	}
	return bucket
}