6748b006b7
This commit refactors the logging mechanism across the application by replacing direct logger calls with a centralized logging approach using the `common` package. Key changes include: - Replaced instances of `logger.SysLog` and `logger.FatalLog` with `common.SysLog` and `common.FatalLog` for consistent logging practices. - Updated resource initialization error handling to utilize the new logging structure, enhancing maintainability and readability. - Minor adjustments to improve code clarity and organization throughout various modules. This change aims to streamline logging and improve the overall architecture of the codebase.
74 lines
3.0 KiB
Go
74 lines
3.0 KiB
Go
package service
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"github.com/bytedance/gopkg/util/gopool"
|
|
"github.com/gin-gonic/gin"
|
|
"net/http"
|
|
"one-api/common"
|
|
"one-api/logger"
|
|
"one-api/model"
|
|
relaycommon "one-api/relay/common"
|
|
"one-api/types"
|
|
)
|
|
|
|
func ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, preConsumedQuota int) {
|
|
if preConsumedQuota != 0 {
|
|
gopool.Go(func() {
|
|
relayInfoCopy := *relayInfo
|
|
|
|
err := PostConsumeQuota(&relayInfoCopy, -preConsumedQuota, 0, false)
|
|
if err != nil {
|
|
common.SysLog("error return pre-consumed quota: " + err.Error())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// PreConsumeQuota checks if the user has enough quota to pre-consume.
|
|
// It returns the pre-consumed quota if successful, or an error if not.
|
|
func PreConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) (int, *types.NewAPIError) {
|
|
userQuota, err := model.GetUserQuota(relayInfo.UserId, false)
|
|
if err != nil {
|
|
return 0, types.NewError(err, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry())
|
|
}
|
|
if userQuota <= 0 {
|
|
return 0, types.NewErrorWithStatusCode(errors.New("user quota is not enough"), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
|
|
}
|
|
if userQuota-preConsumedQuota < 0 {
|
|
return 0, types.NewErrorWithStatusCode(fmt.Errorf("pre-consume quota failed, user quota: %s, need quota: %s", logger.FormatQuota(userQuota), logger.FormatQuota(preConsumedQuota)), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
|
|
}
|
|
relayInfo.UserQuota = userQuota
|
|
if userQuota > 100*preConsumedQuota {
|
|
// 用户额度充足，判断令牌额度是否充足
|
|
if !relayInfo.TokenUnlimited {
|
|
// 非无限令牌，判断令牌额度是否充足
|
|
tokenQuota := c.GetInt("token_quota")
|
|
if tokenQuota > 100*preConsumedQuota {
|
|
// 令牌额度充足，信任令牌
|
|
preConsumedQuota = 0
|
|
logger.LogInfo(c, fmt.Sprintf("user %d quota %s and token %d quota %d are enough, trusted and no need to pre-consume", relayInfo.UserId, logger.FormatQuota(userQuota), relayInfo.TokenId, tokenQuota))
|
|
}
|
|
} else {
|
|
// in this case, we do not pre-consume quota
|
|
// because the user has enough quota
|
|
preConsumedQuota = 0
|
|
logger.LogInfo(c, fmt.Sprintf("user %d with unlimited token has enough quota %s, trusted and no need to pre-consume", relayInfo.UserId, logger.FormatQuota(userQuota)))
|
|
}
|
|
}
|
|
|
|
if preConsumedQuota > 0 {
|
|
err := PreConsumeTokenQuota(relayInfo, preConsumedQuota)
|
|
if err != nil {
|
|
return 0, types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
|
|
}
|
|
err = model.DecreaseUserQuota(relayInfo.UserId, preConsumedQuota)
|
|
if err != nil {
|
|
return 0, types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
|
|
}
|
|
}
|
|
relayInfo.FinalPreConsumedQuota = preConsumedQuota
|
|
return preConsumedQuota, nil
|
|
}
|