ChatGPT/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-10-17 16:43:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
JustSong
2024-03-13 19:38:44 +08:00
parent 0710f8cd66
commit a72e5fcc9e
5 changed files with 21 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
model/cache.go
View File

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