3a25007df2
Add user RPM and concurrency middleware, enforce channel concurrency for relay selection and retry paths, and wire checks into relay routes.
49 lines
1.0 KiB
Go
49 lines
1.0 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/service"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
var (
|
|
getUserByIdForConcurrency = model.GetUserById
|
|
checkUserRPMForConcurrency = service.CheckUserRPM
|
|
acquireUserConcurrencyForCheck = service.AcquireUserConcurrency
|
|
)
|
|
|
|
func ConcurrencyCheck() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
userId := c.GetInt("id")
|
|
if userId <= 0 {
|
|
c.Next()
|
|
return
|
|
}
|
|
|
|
user, err := getUserByIdForConcurrency(userId, false)
|
|
if err != nil || user == nil {
|
|
c.Next()
|
|
return
|
|
}
|
|
|
|
if user.RpmLimit > 0 && !checkUserRPMForConcurrency(userId, user.RpmLimit) {
|
|
abortWithOpenAiMessage(c, http.StatusTooManyRequests, "请求频率超过限制")
|
|
return
|
|
}
|
|
|
|
if user.ConcurrencyLimit > 0 {
|
|
allowed, release := acquireUserConcurrencyForCheck(userId, user.ConcurrencyLimit)
|
|
if !allowed {
|
|
abortWithOpenAiMessage(c, http.StatusTooManyRequests, "用户并发请求数超过限制")
|
|
return
|
|
}
|
|
defer release()
|
|
}
|
|
|
|
c.Next()
|
|
}
|
|
}
|