eta_api/utils/ws/limiter.go

package ws

import (
	"fmt"
	"golang.org/x/time/rate"
	"sync"
	"time"
)

var (
	limiterManager *QALimiterManger
	limiterOnce    sync.Once
)

const (
	LIMITER_KEY = "llm_chat_key_user_%d"
)

type QALimiterManger struct {
	sync.RWMutex
	limiterMap map[string]*QALimiter
}

type QALimiter struct {
	LastRequest time.Time
	*rate.Limiter
}

//func (qaLimiter *QALimiter) Allow() bool {
//	return qaLimiter.Limiter.Allow()
//}

// GetLimiter 获取或创建用户的限流器
func (qalm *QALimiterManger) GetLimiter(token string) *QALimiter {
	qalm.Lock()
	defer qalm.Unlock()

	if limiter, exists := qalm.limiterMap[token]; exists {
		return limiter
	}

	// 创建一个新的限流器，例如每10秒1个请求
	limiter := &QALimiter{
		Limiter: rate.NewLimiter(rate.Every(10*time.Second), 1),
	}
	qalm.limiterMap[token] = limiter
	return limiter
}

func (qalm *QALimiterManger) Allow(token string) bool {

	limiter := qalm.GetLimiter(token)
	if limiter.LastRequest.IsZero() {
		limiter.LastRequest = time.Now()
		return limiter.Allow()
	}
	if time.Now().Sub(limiter.LastRequest) < 10*time.Second {
		return false
	}
	limiter.LastRequest = time.Now()
	return limiter.Allow()
}
func getInstance() *QALimiterManger {
	limiterOnce.Do(func() {
		if limiterManager == nil {
			limiterManager = &QALimiterManger{
				limiterMap: make(map[string]*QALimiter),
			}
		}
	})
	return limiterManager
}

func Allow(userId int) bool {
	token := fmt.Sprintf(LIMITER_KEY, userId)
	return getInstance().Allow(token)
}