diff --git a/model/user.go b/model/user.go index cee4b023..1b2ec7e6 100644 --- a/model/user.go +++ b/model/user.go @@ -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 diff --git a/model/utils.go b/model/utils.go index 61734332..1c28340b 100644 --- a/model/utils.go +++ b/model/utils.go @@ -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) }