fix: optimize batch update process
This commit is contained in:
@@ -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
@@ -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")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user