diff --git a/model/user.go b/model/user.go index 0eb4f53a..e40ac3d2 100644 --- a/model/user.go +++ b/model/user.go @@ -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) { err := DB.Model(&User{}).Where("id = ?", id).Updates( map[string]interface{}{ diff --git a/model/utils.go b/model/utils.go index adfd8e13..b1793706 100644 --- a/model/utils.go +++ b/model/utils.go @@ -67,33 +67,48 @@ func batchUpdate() { } common.SysLog("batch update started") + stores := make([]map[int]int, BatchUpdateTypeCount) for i := 0; i < BatchUpdateTypeCount; i++ { batchUpdateLocks[i].Lock() - store := batchUpdateStores[i] + stores[i] = batchUpdateStores[i] batchUpdateStores[i] = make(map[int]int) 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 { switch i { - case BatchUpdateTypeUserQuota: - err := increaseUserQuota(key, value) - if err != nil { - common.SysLog("failed to batch update user quota: " + err.Error()) - } case BatchUpdateTypeTokenQuota: err := increaseTokenQuota(key, value) if err != nil { common.SysLog("failed to batch update token quota: " + err.Error()) } - case BatchUpdateTypeUsedQuota: - updateUserUsedQuota(key, value) - case BatchUpdateTypeRequestCount: - updateUserRequestCount(key, value) case BatchUpdateTypeChannelUsedQuota: 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") }