From a72e5fcc9e9d6fbabd53630c66971ecec8a53f47 Mon Sep 17 00:00:00 2001 From: JustSong Date: Wed, 13 Mar 2024 19:38:44 +0800 Subject: [PATCH] fix: when cached quota is too low, force refresh it --- model/cache.go | 23 ++++++++++++++--------- relay/controller/audio.go | 3 ++- relay/controller/helper.go | 4 ++-- relay/controller/image.go | 4 ++-- relay/util/common.go | 2 +- 5 files changed, 21 insertions(+), 15 deletions(-) diff --git a/model/cache.go b/model/cache.go index 3a9f8023..9fd29e7a 100644 --- a/model/cache.go +++ b/model/cache.go @@ -1,6 +1,7 @@ package model import ( + "context" "encoding/json" "errors" "fmt" @@ -70,38 +71,42 @@ func CacheGetUserGroup(id int) (group string, err error) { return group, err } -func fetchAndUpdateUserQuota(id int) (quota int, err error) { +func fetchAndUpdateUserQuota(ctx context.Context, id int) (quota int, err error) { quota, err = GetUserQuota(id) if err != nil { return 0, err } err = common.RedisSet(fmt.Sprintf("user_quota:%d", id), fmt.Sprintf("%d", quota), time.Duration(UserId2QuotaCacheSeconds)*time.Second) if err != nil { - logger.SysError("Redis set user quota error: " + err.Error()) + logger.Error(ctx, "Redis set user quota error: "+err.Error()) } return } -func CacheGetUserQuota(id int) (quota int, err error) { +func CacheGetUserQuota(ctx context.Context, id int) (quota int, err error) { if !common.RedisEnabled { return GetUserQuota(id) } quotaString, err := common.RedisGet(fmt.Sprintf("user_quota:%d", id)) if err != nil { - return fetchAndUpdateUserQuota(id) + return fetchAndUpdateUserQuota(ctx, id) } quota, err = strconv.Atoi(quotaString) - if quota <= config.PreConsumedQuota { // when user's quota is less than pre-consumed quota, we need to fetch from db - return fetchAndUpdateUserQuota(id) + if err != nil { + return 0, nil } - return quota, err + if quota <= config.PreConsumedQuota { // when user's quota is less than pre-consumed quota, we need to fetch from db + logger.Infof(ctx, "user %d's cached quota is too low: %d, refreshing from db", quota, id) + return fetchAndUpdateUserQuota(ctx, id) + } + return quota, nil } -func CacheUpdateUserQuota(id int) error { +func CacheUpdateUserQuota(ctx context.Context, id int) error { if !common.RedisEnabled { return nil } - quota, err := CacheGetUserQuota(id) + quota, err := CacheGetUserQuota(ctx, id) if err != nil { return err } diff --git a/relay/controller/audio.go b/relay/controller/audio.go index ee8771c9..3d5086e9 100644 --- a/relay/controller/audio.go +++ b/relay/controller/audio.go @@ -22,6 +22,7 @@ import ( ) func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatusCode { + ctx := c.Request.Context() audioModel := "whisper-1" tokenId := c.GetInt("token_id") @@ -58,7 +59,7 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus default: preConsumedQuota = int(float64(config.PreConsumedQuota) * ratio) } - userQuota, err := model.CacheGetUserQuota(userId) + userQuota, err := model.CacheGetUserQuota(ctx, userId) if err != nil { return openai.ErrorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError) } diff --git a/relay/controller/helper.go b/relay/controller/helper.go index 89fc69ce..d0cb424f 100644 --- a/relay/controller/helper.go +++ b/relay/controller/helper.go @@ -118,7 +118,7 @@ func getPreConsumedQuota(textRequest *relaymodel.GeneralOpenAIRequest, promptTok func preConsumeQuota(ctx context.Context, textRequest *relaymodel.GeneralOpenAIRequest, promptTokens int, ratio float64, meta *util.RelayMeta) (int, *relaymodel.ErrorWithStatusCode) { preConsumedQuota := getPreConsumedQuota(textRequest, promptTokens, ratio) - userQuota, err := model.CacheGetUserQuota(meta.UserId) + userQuota, err := model.CacheGetUserQuota(ctx, meta.UserId) if err != nil { return preConsumedQuota, openai.ErrorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError) } @@ -168,7 +168,7 @@ func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *util.R if err != nil { logger.Error(ctx, "error consuming token remain quota: "+err.Error()) } - err = model.CacheUpdateUserQuota(meta.UserId) + err = model.CacheUpdateUserQuota(ctx, meta.UserId) if err != nil { logger.Error(ctx, "error update user quota cache: "+err.Error()) } diff --git a/relay/controller/image.go b/relay/controller/image.go index 3ce3809b..a0428ecb 100644 --- a/relay/controller/image.go +++ b/relay/controller/image.go @@ -79,7 +79,7 @@ func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus modelRatio := common.GetModelRatio(imageRequest.Model) groupRatio := common.GetGroupRatio(meta.Group) ratio := modelRatio * groupRatio - userQuota, err := model.CacheGetUserQuota(meta.UserId) + userQuota, err := model.CacheGetUserQuota(ctx, meta.UserId) quota := int(ratio*imageCostRatio*1000) * imageRequest.N @@ -125,7 +125,7 @@ func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus if err != nil { logger.SysError("error consuming token remain quota: " + err.Error()) } - err = model.CacheUpdateUserQuota(meta.UserId) + err = model.CacheUpdateUserQuota(ctx, meta.UserId) if err != nil { logger.SysError("error update user quota cache: " + err.Error()) } diff --git a/relay/util/common.go b/relay/util/common.go index dbd724b4..ff0e293d 100644 --- a/relay/util/common.go +++ b/relay/util/common.go @@ -161,7 +161,7 @@ func PostConsumeQuota(ctx context.Context, tokenId int, quotaDelta int, totalQuo if err != nil { logger.SysError("error consuming token remain quota: " + err.Error()) } - err = model.CacheUpdateUserQuota(userId) + err = model.CacheUpdateUserQuota(ctx, userId) if err != nil { logger.SysError("error update user quota cache: " + err.Error()) }