mirror of
				https://github.com/songquanpeng/one-api.git
				synced 2025-11-04 15:53:42 +08:00 
			
		
		
		
	fix: when cached quota is too low, force refresh it
This commit is contained in:
		@@ -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
 | 
			
		||||
	}
 | 
			
		||||
 
 | 
			
		||||
		Reference in New Issue
	
	Block a user