fix: optimize batch update process

This commit is contained in:
CaIon
2026-05-27 13:01:13 +08:00
parent 74985fa877
commit 5b86ce0d70
2 changed files with 43 additions and 11 deletions
+17
View File
@@ -984,6 +984,23 @@ func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) {
//} //}
} }
func updateUserQuotaUsedQuotaAndRequestCount(id int, quota int, usedQuota int, requestCount int) {
if quota == 0 && usedQuota == 0 && requestCount == 0 {
return
}
err := DB.Model(&User{}).Where("id = ?", id).Updates(
map[string]interface{}{
"quota": gorm.Expr("quota + ?", quota),
"used_quota": gorm.Expr("used_quota + ?", usedQuota),
"request_count": gorm.Expr("request_count + ?", requestCount),
},
).Error
if err != nil {
common.SysLog("failed to batch update user quota, used quota and request count: " + err.Error())
}
}
func updateUserUsedQuota(id int, quota int) { func updateUserUsedQuota(id int, quota int) {
err := DB.Model(&User{}).Where("id = ?", id).Updates( err := DB.Model(&User{}).Where("id = ?", id).Updates(
map[string]interface{}{ map[string]interface{}{
+26 -11
View File
@@ -67,33 +67,48 @@ func batchUpdate() {
} }
common.SysLog("batch update started") common.SysLog("batch update started")
stores := make([]map[int]int, BatchUpdateTypeCount)
for i := 0; i < BatchUpdateTypeCount; i++ { for i := 0; i < BatchUpdateTypeCount; i++ {
batchUpdateLocks[i].Lock() batchUpdateLocks[i].Lock()
store := batchUpdateStores[i] stores[i] = 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 i, store := range stores {
if i == BatchUpdateTypeUserQuota || i == BatchUpdateTypeUsedQuota || i == BatchUpdateTypeRequestCount {
continue
}
for key, value := range store { for key, value := range store {
switch i { switch i {
case BatchUpdateTypeUserQuota:
err := increaseUserQuota(key, value)
if err != nil {
common.SysLog("failed to batch update user quota: " + err.Error())
}
case BatchUpdateTypeTokenQuota: case BatchUpdateTypeTokenQuota:
err := increaseTokenQuota(key, value) err := increaseTokenQuota(key, value)
if err != nil { if err != nil {
common.SysLog("failed to batch update token quota: " + err.Error()) common.SysLog("failed to batch update token quota: " + err.Error())
} }
case BatchUpdateTypeUsedQuota:
updateUserUsedQuota(key, value)
case BatchUpdateTypeRequestCount:
updateUserRequestCount(key, value)
case BatchUpdateTypeChannelUsedQuota: case BatchUpdateTypeChannelUsedQuota:
updateChannelUsedQuota(key, value) updateChannelUsedQuota(key, value)
} }
} }
} }
userQuotaStore := stores[BatchUpdateTypeUserQuota]
usedQuotaStore := stores[BatchUpdateTypeUsedQuota]
requestCountStore := stores[BatchUpdateTypeRequestCount]
userIDs := make(map[int]struct{}, len(userQuotaStore)+len(usedQuotaStore)+len(requestCountStore))
for key := range userQuotaStore {
userIDs[key] = struct{}{}
}
for key := range usedQuotaStore {
userIDs[key] = struct{}{}
}
for key := range requestCountStore {
userIDs[key] = struct{}{}
}
for key := range userIDs {
updateUserQuotaUsedQuotaAndRequestCount(key, userQuotaStore[key], usedQuotaStore[key], requestCountStore[key])
}
common.SysLog("batch update finished") common.SysLog("batch update finished")
} }