From 86c2627c24e814dbea9086fbd8e9a904137a9cda Mon Sep 17 00:00:00 2001 From: JustSong Date: Wed, 16 Aug 2023 23:40:24 +0800 Subject: [PATCH] fix: update cache immediately after cache get Co-authored-by: chikasaki <1347283135@qq.com> Co-authored-by: Cruel <157922018@qq.com> --- common/redis.go | 5 +++++ controller/relay-text.go | 6 +++++- model/cache.go | 8 ++++++++ 3 files changed, 18 insertions(+), 1 deletion(-) diff --git a/common/redis.go b/common/redis.go index 1a05721c..12c477b8 100644 --- a/common/redis.go +++ b/common/redis.go @@ -61,3 +61,8 @@ func RedisDel(key string) error { ctx := context.Background() return RDB.Del(ctx, key).Err() } + +func RedisDecrease(key string, value int64) error { + ctx := context.Background() + return RDB.DecrBy(ctx, key, value).Err() +} diff --git a/controller/relay-text.go b/controller/relay-text.go index e8dab514..761ca86f 100644 --- a/controller/relay-text.go +++ b/controller/relay-text.go @@ -194,7 +194,11 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode { if err != nil { return errorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError) } - if userQuota > 10*preConsumedQuota { + err = model.CacheDecreaseUserQuota(userId, preConsumedQuota) + if err != nil { + return errorWrapper(err, "decrease_user_quota_failed", http.StatusInternalServerError) + } + if userQuota > 100*preConsumedQuota { // in this case, we do not pre-consume quota // because the user has enough quota preConsumedQuota = 0 diff --git a/model/cache.go b/model/cache.go index 64666c86..55fbba9b 100644 --- a/model/cache.go +++ b/model/cache.go @@ -95,6 +95,14 @@ func CacheUpdateUserQuota(id int) error { return err } +func CacheDecreaseUserQuota(id int, quota int) error { + if !common.RedisEnabled { + return nil + } + err := common.RedisDecrease(fmt.Sprintf("user_quota:%d", id), int64(quota)) + return err +} + func CacheIsUserEnabled(userId int) bool { if !common.RedisEnabled { return IsUserEnabled(userId)