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) {
|
func UpdateUserUsedQuotaAndRequestCount(id int, quota int) {
|
||||||
if common.BatchUpdateEnabled {
|
if common.BatchUpdateEnabled {
|
||||||
addNewRecord(BatchUpdateTypeUsedQuotaAndRequestCount, id, quota)
|
addNewRecord(BatchUpdateTypeUsedQuota, id, quota)
|
||||||
|
addNewRecord(BatchUpdateTypeRequestCount, id, 1)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
updateUserUsedQuotaAndRequestCount(id, quota, 1)
|
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) {
|
func GetUsernameById(id int) (username string) {
|
||||||
DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username)
|
DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username)
|
||||||
return username
|
return username
|
||||||
|
@ -6,13 +6,13 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
const BatchUpdateTypeCount = 4 // if you add a new type, you need to add a new map and a new lock
|
|
||||||
|
|
||||||
const (
|
const (
|
||||||
BatchUpdateTypeUserQuota = iota
|
BatchUpdateTypeUserQuota = iota
|
||||||
BatchUpdateTypeTokenQuota
|
BatchUpdateTypeTokenQuota
|
||||||
BatchUpdateTypeUsedQuotaAndRequestCount
|
BatchUpdateTypeUsedQuota
|
||||||
BatchUpdateTypeChannelUsedQuota
|
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
|
var batchUpdateStores []map[int]int
|
||||||
@ -51,7 +51,7 @@ func batchUpdate() {
|
|||||||
store := batchUpdateStores[i]
|
store := batchUpdateStores[i]
|
||||||
batchUpdateStores[i] = make(map[int]int)
|
batchUpdateStores[i] = make(map[int]int)
|
||||||
batchUpdateLocks[i].Unlock()
|
batchUpdateLocks[i].Unlock()
|
||||||
|
// TODO: maybe we can combine updates with same key?
|
||||||
for key, value := range store {
|
for key, value := range store {
|
||||||
switch i {
|
switch i {
|
||||||
case BatchUpdateTypeUserQuota:
|
case BatchUpdateTypeUserQuota:
|
||||||
@ -64,8 +64,10 @@ func batchUpdate() {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
common.SysError("failed to batch update token quota: " + err.Error())
|
common.SysError("failed to batch update token quota: " + err.Error())
|
||||||
}
|
}
|
||||||
case BatchUpdateTypeUsedQuotaAndRequestCount:
|
case BatchUpdateTypeUsedQuota:
|
||||||
updateUserUsedQuotaAndRequestCount(key, value, 1) // TODO: count is incorrect
|
updateUserUsedQuota(key, value)
|
||||||
|
case BatchUpdateTypeRequestCount:
|
||||||
|
updateUserRequestCount(key, value)
|
||||||
case BatchUpdateTypeChannelUsedQuota:
|
case BatchUpdateTypeChannelUsedQuota:
|
||||||
updateChannelUsedQuota(key, value)
|
updateChannelUsedQuota(key, value)
|
||||||
}
|
}
|
||||||
|
Loading…
Reference in New Issue
Block a user