refactor: user cache logic
This commit is contained in:
231
model/cache.go
231
model/cache.go
@@ -6,20 +6,13 @@ import (
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"one-api/common"
|
||||
"one-api/constant"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
var (
|
||||
TokenCacheSeconds = common.SyncFrequency
|
||||
UserId2GroupCacheSeconds = common.SyncFrequency
|
||||
UserId2QuotaCacheSeconds = common.SyncFrequency
|
||||
UserId2StatusCacheSeconds = common.SyncFrequency
|
||||
)
|
||||
|
||||
// 仅用于定时同步缓存
|
||||
var token2UserId = make(map[string]int)
|
||||
var token2UserIdLock sync.RWMutex
|
||||
@@ -29,7 +22,7 @@ func cacheSetToken(token *Token) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
err = common.RedisSet(fmt.Sprintf("token:%s", token.Key), string(jsonBytes), time.Duration(TokenCacheSeconds)*time.Second)
|
||||
err = common.RedisSet(fmt.Sprintf("token:%s", token.Key), string(jsonBytes), time.Duration(constant.TokenCacheSeconds)*time.Second)
|
||||
if err != nil {
|
||||
common.SysError(fmt.Sprintf("failed to set token %s to redis: %s", token.Key, err.Error()))
|
||||
return err
|
||||
@@ -57,7 +50,7 @@ func CacheGetTokenByKey(key string) (*Token, error) {
|
||||
return token, nil
|
||||
}
|
||||
// 如果缓存中存在,则续期时间
|
||||
err = common.RedisExpire(fmt.Sprintf("token:%s", key), time.Duration(TokenCacheSeconds)*time.Second)
|
||||
err = common.RedisExpire(fmt.Sprintf("token:%s", key), time.Duration(constant.TokenCacheSeconds)*time.Second)
|
||||
err = json.Unmarshal([]byte(tokenObjectString), &token)
|
||||
return token, err
|
||||
}
|
||||
@@ -101,109 +94,105 @@ func SyncTokenCache(frequency int) {
|
||||
}
|
||||
}
|
||||
|
||||
func CacheGetUserGroup(id int) (group string, err error) {
|
||||
if !common.RedisEnabled {
|
||||
return GetUserGroup(id)
|
||||
}
|
||||
group, err = common.RedisGet(fmt.Sprintf("user_group:%d", id))
|
||||
if err != nil {
|
||||
group, err = GetUserGroup(id)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
err = common.RedisSet(fmt.Sprintf("user_group:%d", id), group, time.Duration(UserId2GroupCacheSeconds)*time.Second)
|
||||
if err != nil {
|
||||
common.SysError("Redis set user group error: " + err.Error())
|
||||
}
|
||||
}
|
||||
return group, err
|
||||
}
|
||||
|
||||
func CacheGetUsername(id int) (username string, err error) {
|
||||
if !common.RedisEnabled {
|
||||
return GetUsernameById(id)
|
||||
}
|
||||
username, err = common.RedisGet(fmt.Sprintf("user_name:%d", id))
|
||||
if err != nil {
|
||||
username, err = GetUsernameById(id)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
err = common.RedisSet(fmt.Sprintf("user_name:%d", id), username, time.Duration(UserId2GroupCacheSeconds)*time.Second)
|
||||
if err != nil {
|
||||
common.SysError("Redis set user group error: " + err.Error())
|
||||
}
|
||||
}
|
||||
return username, err
|
||||
}
|
||||
|
||||
func CacheGetUserQuota(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 {
|
||||
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 {
|
||||
common.SysError("Redis set user quota error: " + err.Error())
|
||||
}
|
||||
return quota, err
|
||||
}
|
||||
quota, err = strconv.Atoi(quotaString)
|
||||
return quota, err
|
||||
}
|
||||
|
||||
func CacheUpdateUserQuota(id int) error {
|
||||
if !common.RedisEnabled {
|
||||
return nil
|
||||
}
|
||||
quota, err := GetUserQuota(id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return cacheSetUserQuota(id, quota)
|
||||
}
|
||||
|
||||
func cacheSetUserQuota(id int, quota int) error {
|
||||
err := common.RedisSet(fmt.Sprintf("user_quota:%d", id), fmt.Sprintf("%d", quota), time.Duration(UserId2QuotaCacheSeconds)*time.Second)
|
||||
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, error) {
|
||||
if !common.RedisEnabled {
|
||||
return IsUserEnabled(userId)
|
||||
}
|
||||
enabled, err := common.RedisGet(fmt.Sprintf("user_enabled:%d", userId))
|
||||
if err == nil {
|
||||
return enabled == "1", nil
|
||||
}
|
||||
|
||||
userEnabled, err := IsUserEnabled(userId)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
enabled = "0"
|
||||
if userEnabled {
|
||||
enabled = "1"
|
||||
}
|
||||
err = common.RedisSet(fmt.Sprintf("user_enabled:%d", userId), enabled, time.Duration(UserId2StatusCacheSeconds)*time.Second)
|
||||
if err != nil {
|
||||
common.SysError("Redis set user enabled error: " + err.Error())
|
||||
}
|
||||
return userEnabled, err
|
||||
}
|
||||
//func CacheGetUserGroup(id int) (group string, err error) {
|
||||
// if !common.RedisEnabled {
|
||||
// return GetUserGroup(id)
|
||||
// }
|
||||
// group, err = common.RedisGet(fmt.Sprintf("user_group:%d", id))
|
||||
// if err != nil {
|
||||
// group, err = GetUserGroup(id)
|
||||
// if err != nil {
|
||||
// return "", err
|
||||
// }
|
||||
// err = common.RedisSet(fmt.Sprintf("user_group:%d", id), group, time.Duration(constant.UserId2GroupCacheSeconds)*time.Second)
|
||||
// if err != nil {
|
||||
// common.SysError("Redis set user group error: " + err.Error())
|
||||
// }
|
||||
// }
|
||||
// return group, err
|
||||
//}
|
||||
//
|
||||
//func CacheGetUsername(id int) (username string, err error) {
|
||||
// if !common.RedisEnabled {
|
||||
// return GetUsernameById(id)
|
||||
// }
|
||||
// username, err = common.RedisGet(fmt.Sprintf("user_name:%d", id))
|
||||
// if err != nil {
|
||||
// username, err = GetUsernameById(id)
|
||||
// if err != nil {
|
||||
// return "", err
|
||||
// }
|
||||
// err = common.RedisSet(fmt.Sprintf("user_name:%d", id), username, time.Duration(constant.UserId2GroupCacheSeconds)*time.Second)
|
||||
// if err != nil {
|
||||
// common.SysError("Redis set user group error: " + err.Error())
|
||||
// }
|
||||
// }
|
||||
// return username, err
|
||||
//}
|
||||
//
|
||||
//func CacheGetUserQuota(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 {
|
||||
// quota, err = GetUserQuota(id)
|
||||
// if err != nil {
|
||||
// return 0, err
|
||||
// }
|
||||
// return quota, nil
|
||||
// }
|
||||
// quota, err = strconv.Atoi(quotaString)
|
||||
// return quota, nil
|
||||
//}
|
||||
//
|
||||
//func CacheUpdateUserQuota(id int) error {
|
||||
// if !common.RedisEnabled {
|
||||
// return nil
|
||||
// }
|
||||
// quota, err := GetUserQuota(id)
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
// return cacheSetUserQuota(id, quota)
|
||||
//}
|
||||
//
|
||||
//func cacheSetUserQuota(id int, quota int) error {
|
||||
// err := common.RedisSet(fmt.Sprintf("user_quota:%d", id), fmt.Sprintf("%d", quota), time.Duration(constant.UserId2QuotaCacheSeconds)*time.Second)
|
||||
// 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, error) {
|
||||
// if !common.RedisEnabled {
|
||||
// return IsUserEnabled(userId)
|
||||
// }
|
||||
// enabled, err := common.RedisGet(fmt.Sprintf("user_enabled:%d", userId))
|
||||
// if err == nil {
|
||||
// return enabled == "1", nil
|
||||
// }
|
||||
//
|
||||
// userEnabled, err := IsUserEnabled(userId)
|
||||
// if err != nil {
|
||||
// return false, err
|
||||
// }
|
||||
// enabled = "0"
|
||||
// if userEnabled {
|
||||
// enabled = "1"
|
||||
// }
|
||||
// err = common.RedisSet(fmt.Sprintf("user_enabled:%d", userId), enabled, time.Duration(constant.UserId2StatusCacheSeconds)*time.Second)
|
||||
// if err != nil {
|
||||
// common.SysError("Redis set user enabled error: " + err.Error())
|
||||
// }
|
||||
// return userEnabled, err
|
||||
//}
|
||||
|
||||
var group2model2channels map[string]map[string][]*Channel
|
||||
var channelsIDM map[int]*Channel
|
||||
@@ -344,12 +333,12 @@ func CacheGetChannel(id int) (*Channel, error) {
|
||||
}
|
||||
|
||||
func CacheUpdateChannelStatus(id int, status int) {
|
||||
if (!common.MemoryCacheEnabled) {
|
||||
return
|
||||
}
|
||||
channelSyncLock.Lock()
|
||||
defer channelSyncLock.Unlock()
|
||||
if channel, ok := channelsIDM[id]; ok {
|
||||
channel.Status = status
|
||||
}
|
||||
if !common.MemoryCacheEnabled {
|
||||
return
|
||||
}
|
||||
channelSyncLock.Lock()
|
||||
defer channelSyncLock.Unlock()
|
||||
if channel, ok := channelsIDM[id]; ok {
|
||||
channel.Status = status
|
||||
}
|
||||
}
|
||||
|
||||
@@ -81,7 +81,7 @@ func RecordLog(userId int, logType int, content string) {
|
||||
if logType == LogTypeConsume && !common.LogConsumeEnabled {
|
||||
return
|
||||
}
|
||||
username, _ := CacheGetUsername(userId)
|
||||
username, _ := GetUsernameById(userId, false)
|
||||
log := &Log{
|
||||
UserId: userId,
|
||||
Username: username,
|
||||
@@ -102,7 +102,7 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
|
||||
if !common.LogConsumeEnabled {
|
||||
return
|
||||
}
|
||||
username, _ := CacheGetUsername(userId)
|
||||
username, _ := GetUsernameById(userId, false)
|
||||
otherStr := common.MapToJsonStr(other)
|
||||
log := &Log{
|
||||
UserId: userId,
|
||||
|
||||
@@ -258,37 +258,29 @@ func decreaseTokenQuota(id int, quota int) (err error) {
|
||||
return err
|
||||
}
|
||||
|
||||
func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) (userQuota int, err error) {
|
||||
func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) error {
|
||||
if quota < 0 {
|
||||
return 0, errors.New("quota 不能为负数!")
|
||||
return errors.New("quota 不能为负数!")
|
||||
}
|
||||
if !relayInfo.IsPlayground {
|
||||
token, err := GetTokenById(relayInfo.TokenId)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
return err
|
||||
}
|
||||
if !token.UnlimitedQuota && token.RemainQuota < quota {
|
||||
return 0, errors.New("令牌额度不足")
|
||||
return errors.New("令牌额度不足")
|
||||
}
|
||||
}
|
||||
userQuota, err = GetUserQuota(relayInfo.UserId)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if userQuota < quota {
|
||||
return 0, errors.New(fmt.Sprintf("用户额度不足,剩余额度为 %d", userQuota))
|
||||
}
|
||||
if !relayInfo.IsPlayground {
|
||||
err = DecreaseTokenQuota(relayInfo.TokenId, quota)
|
||||
err := DecreaseTokenQuota(relayInfo.TokenId, quota)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
return err
|
||||
}
|
||||
}
|
||||
err = DecreaseUserQuota(relayInfo.UserId, quota)
|
||||
return userQuota - quota, err
|
||||
return nil
|
||||
}
|
||||
|
||||
func PostConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, userQuota int, quota int, preConsumedQuota int, sendEmail bool) (err error) {
|
||||
func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, userQuota int, quota int, preConsumedQuota int, sendEmail bool) (err error) {
|
||||
|
||||
if quota > 0 {
|
||||
err = DecreaseUserQuota(relayInfo.UserId, quota)
|
||||
|
||||
183
model/user.go
183
model/user.go
@@ -6,7 +6,8 @@ import (
|
||||
"one-api/common"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/bytedance/gopkg/util/gopool"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
@@ -107,7 +108,7 @@ func SearchUsers(keyword string, group string) ([]*User, error) {
|
||||
return users, err
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
err = nil
|
||||
|
||||
query := DB.Unscoped().Omit("password")
|
||||
@@ -251,14 +252,12 @@ func (user *User) Update(updatePassword bool) error {
|
||||
}
|
||||
newUser := *user
|
||||
DB.First(&user, user.Id)
|
||||
err = DB.Model(user).Updates(newUser).Error
|
||||
if err == nil {
|
||||
if common.RedisEnabled {
|
||||
_ = common.RedisSet(fmt.Sprintf("user_group:%d", user.Id), user.Group, time.Duration(UserId2GroupCacheSeconds)*time.Second)
|
||||
_ = common.RedisSet(fmt.Sprintf("user_quota:%d", user.Id), strconv.Itoa(user.Quota), time.Duration(UserId2QuotaCacheSeconds)*time.Second)
|
||||
}
|
||||
if err = DB.Model(user).Updates(newUser).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return err
|
||||
|
||||
// 更新缓存
|
||||
return updateUserCache(user)
|
||||
}
|
||||
|
||||
func (user *User) Edit(updatePassword bool) error {
|
||||
@@ -269,6 +268,7 @@ func (user *User) Edit(updatePassword bool) error {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
newUser := *user
|
||||
updates := map[string]interface{}{
|
||||
"username": newUser.Username,
|
||||
@@ -279,23 +279,26 @@ func (user *User) Edit(updatePassword bool) error {
|
||||
if updatePassword {
|
||||
updates["password"] = newUser.Password
|
||||
}
|
||||
|
||||
DB.First(&user, user.Id)
|
||||
err = DB.Model(user).Updates(updates).Error
|
||||
if err == nil {
|
||||
if common.RedisEnabled {
|
||||
_ = common.RedisSet(fmt.Sprintf("user_group:%d", user.Id), user.Group, time.Duration(UserId2GroupCacheSeconds)*time.Second)
|
||||
_ = common.RedisSet(fmt.Sprintf("user_quota:%d", user.Id), strconv.Itoa(user.Quota), time.Duration(UserId2QuotaCacheSeconds)*time.Second)
|
||||
}
|
||||
if err = DB.Model(user).Updates(updates).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return err
|
||||
|
||||
// 更新缓存
|
||||
return updateUserCache(user)
|
||||
}
|
||||
|
||||
func (user *User) Delete() error {
|
||||
if user.Id == 0 {
|
||||
return errors.New("id 为空!")
|
||||
}
|
||||
err := DB.Delete(user).Error
|
||||
return err
|
||||
if err := DB.Delete(user).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 清除缓存
|
||||
return invalidateUserCache(user.Id)
|
||||
}
|
||||
|
||||
func (user *User) HardDelete() error {
|
||||
@@ -409,15 +412,33 @@ func IsAdmin(userId int) bool {
|
||||
return user.Role >= common.RoleAdminUser
|
||||
}
|
||||
|
||||
func IsUserEnabled(userId int) (bool, error) {
|
||||
if userId == 0 {
|
||||
return false, errors.New("user id is empty")
|
||||
// IsUserEnabled checks user status from Redis first, falls back to DB if needed
|
||||
func IsUserEnabled(id int, fromDB bool) (status bool, err error) {
|
||||
defer func() {
|
||||
// Update Redis cache asynchronously on successful DB read
|
||||
if common.RedisEnabled {
|
||||
gopool.Go(func() {
|
||||
if err := updateUserStatusCache(id, status); err != nil {
|
||||
common.SysError("failed to update user status cache: " + err.Error())
|
||||
}
|
||||
})
|
||||
}
|
||||
}()
|
||||
if !fromDB && common.RedisEnabled {
|
||||
// Try Redis first
|
||||
status, err := getUserStatusCache(id)
|
||||
if err == nil {
|
||||
return status == common.UserStatusEnabled, nil
|
||||
}
|
||||
// Don't return error - fall through to DB
|
||||
}
|
||||
|
||||
var user User
|
||||
err := DB.Where("id = ?", userId).Select("status").Find(&user).Error
|
||||
err = DB.Where("id = ?", id).Select("status").Find(&user).Error
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
return user.Status == common.UserStatusEnabled, nil
|
||||
}
|
||||
|
||||
@@ -433,14 +454,33 @@ func ValidateAccessToken(token string) (user *User) {
|
||||
return nil
|
||||
}
|
||||
|
||||
func GetUserQuota(id int) (quota int, err error) {
|
||||
// GetUserQuota gets quota from Redis first, falls back to DB if needed
|
||||
func GetUserQuota(id int, fromDB bool) (quota int, err error) {
|
||||
defer func() {
|
||||
// Update Redis cache asynchronously on successful DB read
|
||||
if common.RedisEnabled && err == nil {
|
||||
gopool.Go(func() {
|
||||
if err := updateUserQuotaCache(id, quota); err != nil {
|
||||
common.SysError("failed to update user quota cache: " + err.Error())
|
||||
}
|
||||
})
|
||||
}
|
||||
}()
|
||||
if !fromDB && common.RedisEnabled {
|
||||
quota, err := getUserQuotaCache(id)
|
||||
if err == nil {
|
||||
return quota, nil
|
||||
}
|
||||
// Don't return error - fall through to DB
|
||||
//common.SysError("failed to get user quota from cache: " + err.Error())
|
||||
}
|
||||
|
||||
err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find("a).Error
|
||||
if err != nil {
|
||||
if common.RedisEnabled {
|
||||
go cacheSetUserQuota(id, quota)
|
||||
}
|
||||
return 0, err
|
||||
}
|
||||
return quota, err
|
||||
|
||||
return quota, nil
|
||||
}
|
||||
|
||||
func GetUserUsedQuota(id int) (quota int, err error) {
|
||||
@@ -453,20 +493,49 @@ func GetUserEmail(id int) (email string, err error) {
|
||||
return email, err
|
||||
}
|
||||
|
||||
func GetUserGroup(id int) (group string, err error) {
|
||||
// GetUserGroup gets group from Redis first, falls back to DB if needed
|
||||
func GetUserGroup(id int, fromDB bool) (group string, err error) {
|
||||
defer func() {
|
||||
// Update Redis cache asynchronously on successful DB read
|
||||
if common.RedisEnabled && err == nil {
|
||||
gopool.Go(func() {
|
||||
if err := updateUserGroupCache(id, group); err != nil {
|
||||
common.SysError("failed to update user group cache: " + err.Error())
|
||||
}
|
||||
})
|
||||
}
|
||||
}()
|
||||
if !fromDB && common.RedisEnabled {
|
||||
group, err := getUserGroupCache(id)
|
||||
if err == nil {
|
||||
return group, nil
|
||||
}
|
||||
// Don't return error - fall through to DB
|
||||
}
|
||||
|
||||
groupCol := "`group`"
|
||||
if common.UsingPostgreSQL {
|
||||
groupCol = `"group"`
|
||||
}
|
||||
|
||||
err = DB.Model(&User{}).Where("id = ?", id).Select(groupCol).Find(&group).Error
|
||||
return group, err
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return group, nil
|
||||
}
|
||||
|
||||
func IncreaseUserQuota(id int, quota int) (err error) {
|
||||
if quota < 0 {
|
||||
return errors.New("quota 不能为负数!")
|
||||
}
|
||||
gopool.Go(func() {
|
||||
err := cacheIncrUserQuota(id, quota)
|
||||
if err != nil {
|
||||
common.SysError("failed to increase user quota: " + err.Error())
|
||||
}
|
||||
})
|
||||
if common.BatchUpdateEnabled {
|
||||
addNewRecord(BatchUpdateTypeUserQuota, id, quota)
|
||||
return nil
|
||||
@@ -476,6 +545,9 @@ func IncreaseUserQuota(id int, quota int) (err error) {
|
||||
|
||||
func increaseUserQuota(id int, quota int) (err error) {
|
||||
err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -483,6 +555,12 @@ func DecreaseUserQuota(id int, quota int) (err error) {
|
||||
if quota < 0 {
|
||||
return errors.New("quota 不能为负数!")
|
||||
}
|
||||
gopool.Go(func() {
|
||||
err := cacheDecrUserQuota(id, quota)
|
||||
if err != nil {
|
||||
common.SysError("failed to decrease user quota: " + err.Error())
|
||||
}
|
||||
})
|
||||
if common.BatchUpdateEnabled {
|
||||
addNewRecord(BatchUpdateTypeUserQuota, id, -quota)
|
||||
return nil
|
||||
@@ -492,9 +570,23 @@ func DecreaseUserQuota(id int, quota int) (err error) {
|
||||
|
||||
func decreaseUserQuota(id int, quota int) (err error) {
|
||||
err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func DeltaUpdateUserQuota(id int, delta int) (err error) {
|
||||
if delta == 0 {
|
||||
return nil
|
||||
}
|
||||
if delta > 0 {
|
||||
return IncreaseUserQuota(id, delta)
|
||||
} else {
|
||||
return DecreaseUserQuota(id, -delta)
|
||||
}
|
||||
}
|
||||
|
||||
func GetRootUserEmail() (email string) {
|
||||
DB.Model(&User{}).Where("role = ?", common.RoleRootUser).Select("email").Find(&email)
|
||||
return email
|
||||
@@ -518,7 +610,13 @@ func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) {
|
||||
).Error
|
||||
if err != nil {
|
||||
common.SysError("failed to update user used quota and request count: " + err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
//// 更新缓存
|
||||
//if err := invalidateUserCache(id); err != nil {
|
||||
// common.SysError("failed to invalidate user cache: " + err.Error())
|
||||
//}
|
||||
}
|
||||
|
||||
func updateUserUsedQuota(id int, quota int) {
|
||||
@@ -539,9 +637,32 @@ func updateUserRequestCount(id int, count int) {
|
||||
}
|
||||
}
|
||||
|
||||
func GetUsernameById(id int) (username string, err error) {
|
||||
// GetUsernameById gets username from Redis first, falls back to DB if needed
|
||||
func GetUsernameById(id int, fromDB bool) (username string, err error) {
|
||||
defer func() {
|
||||
// Update Redis cache asynchronously on successful DB read
|
||||
if common.RedisEnabled && err == nil {
|
||||
gopool.Go(func() {
|
||||
if err := updateUserNameCache(id, username); err != nil {
|
||||
common.SysError("failed to update user name cache: " + err.Error())
|
||||
}
|
||||
})
|
||||
}
|
||||
}()
|
||||
if !fromDB && common.RedisEnabled {
|
||||
username, err := getUserNameCache(id)
|
||||
if err == nil {
|
||||
return username, nil
|
||||
}
|
||||
// Don't return error - fall through to DB
|
||||
}
|
||||
|
||||
err = DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username).Error
|
||||
return username, err
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return username, nil
|
||||
}
|
||||
|
||||
func IsLinuxDOIdAlreadyTaken(linuxDOId string) bool {
|
||||
|
||||
206
model/user_cache.go
Normal file
206
model/user_cache.go
Normal file
@@ -0,0 +1,206 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"one-api/common"
|
||||
"one-api/constant"
|
||||
"strconv"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Change UserCache struct to userCache
|
||||
type userCache struct {
|
||||
Id int `json:"id"`
|
||||
Group string `json:"group"`
|
||||
Quota int `json:"quota"`
|
||||
Status int `json:"status"`
|
||||
Role int `json:"role"`
|
||||
Username string `json:"username"`
|
||||
}
|
||||
|
||||
// Rename all exported functions to private ones
|
||||
// invalidateUserCache clears all user related cache
|
||||
func invalidateUserCache(userId int) error {
|
||||
if !common.RedisEnabled {
|
||||
return nil
|
||||
}
|
||||
|
||||
keys := []string{
|
||||
fmt.Sprintf(constant.UserGroupKeyFmt, userId),
|
||||
fmt.Sprintf(constant.UserQuotaKeyFmt, userId),
|
||||
fmt.Sprintf(constant.UserEnabledKeyFmt, userId),
|
||||
fmt.Sprintf(constant.UserUsernameKeyFmt, userId),
|
||||
}
|
||||
|
||||
for _, key := range keys {
|
||||
if err := common.RedisDel(key); err != nil {
|
||||
return fmt.Errorf("failed to delete cache key %s: %w", key, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// updateUserGroupCache updates user group cache
|
||||
func updateUserGroupCache(userId int, group string) error {
|
||||
if !common.RedisEnabled {
|
||||
return nil
|
||||
}
|
||||
return common.RedisSet(
|
||||
fmt.Sprintf(constant.UserGroupKeyFmt, userId),
|
||||
group,
|
||||
time.Duration(constant.UserId2QuotaCacheSeconds)*time.Second,
|
||||
)
|
||||
}
|
||||
|
||||
// updateUserQuotaCache updates user quota cache
|
||||
func updateUserQuotaCache(userId int, quota int) error {
|
||||
if !common.RedisEnabled {
|
||||
return nil
|
||||
}
|
||||
return common.RedisSet(
|
||||
fmt.Sprintf(constant.UserQuotaKeyFmt, userId),
|
||||
fmt.Sprintf("%d", quota),
|
||||
time.Duration(constant.UserId2QuotaCacheSeconds)*time.Second,
|
||||
)
|
||||
}
|
||||
|
||||
// updateUserStatusCache updates user status cache
|
||||
func updateUserStatusCache(userId int, userEnabled bool) error {
|
||||
if !common.RedisEnabled {
|
||||
return nil
|
||||
}
|
||||
enabled := "0"
|
||||
if userEnabled {
|
||||
enabled = "1"
|
||||
}
|
||||
return common.RedisSet(
|
||||
fmt.Sprintf(constant.UserEnabledKeyFmt, userId),
|
||||
enabled,
|
||||
time.Duration(constant.UserId2StatusCacheSeconds)*time.Second,
|
||||
)
|
||||
}
|
||||
|
||||
// updateUserNameCache updates username cache
|
||||
func updateUserNameCache(userId int, username string) error {
|
||||
if !common.RedisEnabled {
|
||||
return nil
|
||||
}
|
||||
return common.RedisSet(
|
||||
fmt.Sprintf(constant.UserUsernameKeyFmt, userId),
|
||||
username,
|
||||
time.Duration(constant.UserId2QuotaCacheSeconds)*time.Second,
|
||||
)
|
||||
}
|
||||
|
||||
// updateUserCache updates all user cache fields
|
||||
func updateUserCache(user *User) error {
|
||||
if !common.RedisEnabled {
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := updateUserGroupCache(user.Id, user.Group); err != nil {
|
||||
return fmt.Errorf("update group cache: %w", err)
|
||||
}
|
||||
|
||||
if err := updateUserQuotaCache(user.Id, user.Quota); err != nil {
|
||||
return fmt.Errorf("update quota cache: %w", err)
|
||||
}
|
||||
|
||||
if err := updateUserStatusCache(user.Id, user.Status == common.UserStatusEnabled); err != nil {
|
||||
return fmt.Errorf("update status cache: %w", err)
|
||||
}
|
||||
|
||||
if err := updateUserNameCache(user.Id, user.Username); err != nil {
|
||||
return fmt.Errorf("update username cache: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// getUserGroupCache gets user group from cache
|
||||
func getUserGroupCache(userId int) (string, error) {
|
||||
if !common.RedisEnabled {
|
||||
return "", nil
|
||||
}
|
||||
return common.RedisGet(fmt.Sprintf(constant.UserGroupKeyFmt, userId))
|
||||
}
|
||||
|
||||
// getUserQuotaCache gets user quota from cache
|
||||
func getUserQuotaCache(userId int) (int, error) {
|
||||
if !common.RedisEnabled {
|
||||
return 0, nil
|
||||
}
|
||||
quotaStr, err := common.RedisGet(fmt.Sprintf(constant.UserQuotaKeyFmt, userId))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return strconv.Atoi(quotaStr)
|
||||
}
|
||||
|
||||
// getUserStatusCache gets user status from cache
|
||||
func getUserStatusCache(userId int) (int, error) {
|
||||
if !common.RedisEnabled {
|
||||
return 0, nil
|
||||
}
|
||||
statusStr, err := common.RedisGet(fmt.Sprintf(constant.UserEnabledKeyFmt, userId))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return strconv.Atoi(statusStr)
|
||||
}
|
||||
|
||||
// getUserNameCache gets username from cache
|
||||
func getUserNameCache(userId int) (string, error) {
|
||||
if !common.RedisEnabled {
|
||||
return "", nil
|
||||
}
|
||||
return common.RedisGet(fmt.Sprintf(constant.UserUsernameKeyFmt, userId))
|
||||
}
|
||||
|
||||
// getUserCache gets complete user cache
|
||||
func getUserCache(userId int) (*userCache, error) {
|
||||
if !common.RedisEnabled {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
group, err := getUserGroupCache(userId)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get group cache: %w", err)
|
||||
}
|
||||
|
||||
quota, err := getUserQuotaCache(userId)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get quota cache: %w", err)
|
||||
}
|
||||
|
||||
status, err := getUserStatusCache(userId)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get status cache: %w", err)
|
||||
}
|
||||
|
||||
username, err := getUserNameCache(userId)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get username cache: %w", err)
|
||||
}
|
||||
|
||||
return &userCache{
|
||||
Id: userId,
|
||||
Group: group,
|
||||
Quota: quota,
|
||||
Status: status,
|
||||
Username: username,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Add atomic quota operations
|
||||
func cacheIncrUserQuota(userId int, delta int) error {
|
||||
if !common.RedisEnabled {
|
||||
return nil
|
||||
}
|
||||
key := fmt.Sprintf(constant.UserQuotaKeyFmt, userId)
|
||||
return common.RedisIncr(key, delta)
|
||||
}
|
||||
|
||||
func cacheDecrUserQuota(userId int, delta int) error {
|
||||
return cacheIncrUserQuota(userId, -delta)
|
||||
}
|
||||
Reference in New Issue
Block a user