eta_api/utils/ws/limiter.go

package ws

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

var (
	limiterManagers map[string]*LimiterManger
	limiterOnce     sync.Once
	limters         = map[string]string{
		CONNECT_LIMITER: LIMITER_KEY,
		QA_LIMITER:      CONNECT_LIMITER_KEY,
	}
)

const (
	CONNECT_LIMITER     = "connetLimiter"
	QA_LIMITER          = "qaLimiter"
	LIMITER_KEY         = "llm_chat_key_user_%d"
	CONNECT_LIMITER_KEY = "llm_chat_connect_key_user_%d"
)

type RateLimiter struct {
	LastRequest time.Time
	*rate.Limiter
}
type LimiterManger struct {
	sync.RWMutex
	limiterMap map[string]*RateLimiter
}

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

// GetLimiter 获取或创建用户的限流器
func (qalm *LimiterManger) GetLimiter(token string) *RateLimiter {
	qalm.Lock()
	defer qalm.Unlock()
	if limiter, exists := qalm.limiterMap[token]; exists {
		return limiter
	}

	// 创建一个新的限流器，例如每10秒1个请求
	limiter := &RateLimiter{
		Limiter: rate.NewLimiter(rate.Every(10*time.Second), 1),
	}
	qalm.limiterMap[token] = limiter
	return limiter
}
func (qalm *LimiterManger) 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(key string) *LimiterManger {
	limiterOnce.Do(func() {
		if limiterManagers == nil {
			limiterManagers = make(map[string]*LimiterManger, len(limters))
		}
		for key, _ := range limters {
			limiterManagers[key] = &LimiterManger{
				limiterMap: make(map[string]*RateLimiter),
			}
		}
	})
	return limiterManagers[key]
}

func Allow(userId int, limiter string) bool {
	tokenKey := limters[limiter]
	if tokenKey == "" {
		return false
	}
	token := fmt.Sprintf(tokenKey, userId)
	handler := getInstance(limiter)
	if handler == nil {
		return false
	}
	return handler.Allow(token)
}