Skip to content

Commit

Permalink
fix: fix request count not updated correctly when using batch update
Browse files Browse the repository at this point in the history
  • Loading branch information
songquanpeng committed Oct 14, 2023
1 parent fa41ca9 commit f073592
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 7 deletions.
21 changes: 20 additions & 1 deletion model/user.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand Down
14 changes: 8 additions & 6 deletions model/utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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)
}
Expand Down

0 comments on commit f073592

Please sign in to comment.