mirror of
				https://github.com/linux-do/new-api.git
				synced 2025-11-04 13:23:42 +08:00 
			
		
		
		
	fix: fix request count not updated correctly when using batch update
This commit is contained in:
		@@ -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
 | 
			
		||||
 
 | 
			
		||||
@@ -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)
 | 
			
		||||
			}
 | 
			
		||||
 
 | 
			
		||||
		Reference in New Issue
	
	Block a user