fix: fix request count not updated correctly when using batch update
This commit is contained in:
parent
fa41ca9805
commit
f073592d39
@ -309,7 +309,8 @@ func GetRootUserEmail() (email string) {
|
||||
|
||||
func UpdateUserUsedQuotaAndRequestCount(id int, quota int) {
|
||||
if common.BatchUpdateEnabled {
|
||||
addNewRecord(BatchUpdateTypeUsedQuotaAndRequestCount, id, quota)
|
||||
addNewRecord(BatchUpdateTypeUsedQuota, id, quota)
|
||||
addNewRecord(BatchUpdateTypeRequestCount, id, 1)
|
||||
return
|
||||
}
|
||||
updateUserUsedQuotaAndRequestCount(id, quota, 1)
|
||||
@ -327,6 +328,24 @@ func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) {
|
||||
}
|
||||
}
|
||||
|
||||
func updateUserUsedQuota(id int, quota int) {
|
||||
err := DB.Model(&User{}).Where("id = ?", id).Updates(
|
||||
map[string]interface{}{
|
||||
"used_quota": gorm.Expr("used_quota + ?", quota),
|
||||
},
|
||||
).Error
|
||||
if err != nil {
|
||||
common.SysError("failed to update user used quota: " + err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func updateUserRequestCount(id int, count int) {
|
||||
err := DB.Model(&User{}).Where("id = ?", id).Update("request_count", gorm.Expr("request_count + ?", count)).Error
|
||||
if err != nil {
|
||||
common.SysError("failed to update user request count: " + err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func GetUsernameById(id int) (username string) {
|
||||
DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username)
|
||||
return username
|
||||
|
@ -6,13 +6,13 @@ import (
|
||||
"time"
|
||||
)
|
||||
|
||||
const BatchUpdateTypeCount = 4 // if you add a new type, you need to add a new map and a new lock
|
||||
|
||||
const (
|
||||
BatchUpdateTypeUserQuota = iota
|
||||
BatchUpdateTypeTokenQuota
|
||||
BatchUpdateTypeUsedQuotaAndRequestCount
|
||||
BatchUpdateTypeUsedQuota
|
||||
BatchUpdateTypeChannelUsedQuota
|
||||
BatchUpdateTypeRequestCount
|
||||
BatchUpdateTypeCount // if you add a new type, you need to add a new map and a new lock
|
||||
)
|
||||
|
||||
var batchUpdateStores []map[int]int
|
||||
@ -51,7 +51,7 @@ func batchUpdate() {
|
||||
store := batchUpdateStores[i]
|
||||
batchUpdateStores[i] = make(map[int]int)
|
||||
batchUpdateLocks[i].Unlock()
|
||||
|
||||
// TODO: maybe we can combine updates with same key?
|
||||
for key, value := range store {
|
||||
switch i {
|
||||
case BatchUpdateTypeUserQuota:
|
||||
@ -64,8 +64,10 @@ func batchUpdate() {
|
||||
if err != nil {
|
||||
common.SysError("failed to batch update token quota: " + err.Error())
|
||||
}
|
||||
case BatchUpdateTypeUsedQuotaAndRequestCount:
|
||||
updateUserUsedQuotaAndRequestCount(key, value, 1) // TODO: count is incorrect
|
||||
case BatchUpdateTypeUsedQuota:
|
||||
updateUserUsedQuota(key, value)
|
||||
case BatchUpdateTypeRequestCount:
|
||||
updateUserRequestCount(key, value)
|
||||
case BatchUpdateTypeChannelUsedQuota:
|
||||
updateChannelUsedQuota(key, value)
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user