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
	limiters        = map[string]LimiterConfig{
		CONNECT_LIMITER: {
			LimiterKey: LIMITER_KEY,
			Duration:   RATE_LIMTER_TIME,
		},
		QA_LIMITER: {
			LimiterKey: CONNECT_LIMITER_KEY,
			Duration:   CONNECT_LIMITER_TIME,
		},
	}
)

type LimiterConfig struct {
	LimiterKey string
	Duration   time.Duration
}

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

	RATE_LIMTER_TIME     = 30 * time.Second
	CONNECT_LIMITER_TIME = 200 * time.Millisecond
)

var ()

type RateLimiter struct {
	LastRequest time.Time
	Duration    time.Duration
	*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, limiterKey string) (limiter *RateLimiter, duration time.Duration) {
	qalm.Lock()
	defer qalm.Unlock()
	if config, ok := limiters[limiterKey]; !ok {
		duration = 0 * time.Second
	} else {
		duration = config.Duration
	}
	if target, exists := qalm.limiterMap[token]; exists {
		limiter = target
		return
	}
	// 创建一个新的限流器，例如每10秒1个请求
	limiter = &RateLimiter{
		Limiter: rate.NewLimiter(rate.Every(duration), 1),
	}
	qalm.limiterMap[token] = limiter
	return
}
func (qalm *LimiterManger) Allow(token string, limiterKey string) bool {
	limiter, duration := qalm.GetLimiter(token, limiterKey)
	if limiter.LastRequest.IsZero() {
		limiter.LastRequest = time.Now()
		return limiter.Allow()
	}
	if time.Now().Sub(limiter.LastRequest) < duration {
		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(limiters))
		}
		for key = range limiters {
			limiterManagers[key] = &LimiterManger{
				limiterMap: make(map[string]*RateLimiter),
			}
		}
	})
	return limiterManagers[key]
}

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