fix: when cached quota is too low, force refresh it

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
+	if err != nil {
-		return fetchAndUpdateUserQuota(id)
+		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
 	}

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)
 	}

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())
 	}

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())
 		}

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())
 	}