fix: 修复高并发下,高额度用户使用低额度令牌没有预扣费的问题

This commit is contained in:
CaIon 2024-01-10 14:23:23 +08:00
parent 042d55cfd3
commit a3b726dd82
2 changed files with 20 additions and 4 deletions

View File

@ -52,6 +52,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
tokenId := c.GetInt("token_id") tokenId := c.GetInt("token_id")
userId := c.GetInt("id") userId := c.GetInt("id")
group := c.GetString("group") group := c.GetString("group")
tokenQuota := c.GetInt("token_quota")
startTime := time.Now() startTime := time.Now()
var textRequest GeneralOpenAIRequest var textRequest GeneralOpenAIRequest
@ -261,10 +262,20 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
return errorWrapper(err, "decrease_user_quota_failed", http.StatusInternalServerError) return errorWrapper(err, "decrease_user_quota_failed", http.StatusInternalServerError)
} }
if userQuota > 100*preConsumedQuota { if userQuota > 100*preConsumedQuota {
// in this case, we do not pre-consume quota // 用户额度充足,判断令牌额度是否充足
// because the user has enough quota if tokenQuota != -1 {
preConsumedQuota = 0 // 非无限令牌,判断令牌额度是否充足
common.LogInfo(c.Request.Context(), fmt.Sprintf("user %d has enough quota %d, trusted and no need to pre-consume", userId, userQuota)) if tokenQuota > 100*preConsumedQuota {
// 令牌额度充足,信任令牌
preConsumedQuota = 0
common.LogInfo(c.Request.Context(), fmt.Sprintf("user %d quota %d and token %d quota %d are enough, trusted and no need to pre-consume", userId, userQuota, tokenId, tokenQuota))
}
} else {
// in this case, we do not pre-consume quota
// because the user has enough quota
preConsumedQuota = 0
common.LogInfo(c.Request.Context(), fmt.Sprintf("user %d with unlimited token has enough quota %d, trusted and no need to pre-consume", userId, userQuota))
}
} }
if preConsumedQuota > 0 { if preConsumedQuota > 0 {
userQuota, err = model.PreConsumeTokenQuota(tokenId, preConsumedQuota) userQuota, err = model.PreConsumeTokenQuota(tokenId, preConsumedQuota)

View File

@ -115,6 +115,11 @@ func TokenAuth() func(c *gin.Context) {
c.Set("id", token.UserId) c.Set("id", token.UserId)
c.Set("token_id", token.Id) c.Set("token_id", token.Id)
c.Set("token_name", token.Name) c.Set("token_name", token.Name)
if !token.UnlimitedQuota {
c.Set("token_quota", token.RemainQuota)
} else {
c.Set("token_quota", -1)
}
if token.ModelLimitsEnabled { if token.ModelLimitsEnabled {
c.Set("token_model_limit_enabled", true) c.Set("token_model_limit_enabled", true)
c.Set("token_model_limit", token.GetModelLimitsMap()) c.Set("token_model_limit", token.GetModelLimitsMap())