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) {
 | 
					func UpdateUserUsedQuotaAndRequestCount(id int, quota int) {
 | 
				
			||||||
	if common.BatchUpdateEnabled {
 | 
						if common.BatchUpdateEnabled {
 | 
				
			||||||
		addNewRecord(BatchUpdateTypeUsedQuotaAndRequestCount, id, quota)
 | 
							addNewRecord(BatchUpdateTypeUsedQuota, id, quota)
 | 
				
			||||||
 | 
							addNewRecord(BatchUpdateTypeRequestCount, id, 1)
 | 
				
			||||||
		return
 | 
							return
 | 
				
			||||||
	}
 | 
						}
 | 
				
			||||||
	updateUserUsedQuotaAndRequestCount(id, quota, 1)
 | 
						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) {
 | 
					func GetUsernameById(id int) (username string) {
 | 
				
			||||||
	DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username)
 | 
						DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username)
 | 
				
			||||||
	return username
 | 
						return username
 | 
				
			||||||
 
 | 
				
			|||||||
@@ -6,13 +6,13 @@ import (
 | 
				
			|||||||
	"time"
 | 
						"time"
 | 
				
			||||||
)
 | 
					)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
const BatchUpdateTypeCount = 4 // if you add a new type, you need to add a new map and a new lock
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
const (
 | 
					const (
 | 
				
			||||||
	BatchUpdateTypeUserQuota = iota
 | 
						BatchUpdateTypeUserQuota = iota
 | 
				
			||||||
	BatchUpdateTypeTokenQuota
 | 
						BatchUpdateTypeTokenQuota
 | 
				
			||||||
	BatchUpdateTypeUsedQuotaAndRequestCount
 | 
						BatchUpdateTypeUsedQuota
 | 
				
			||||||
	BatchUpdateTypeChannelUsedQuota
 | 
						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
 | 
					var batchUpdateStores []map[int]int
 | 
				
			||||||
@@ -51,7 +51,7 @@ func batchUpdate() {
 | 
				
			|||||||
		store := batchUpdateStores[i]
 | 
							store := 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 key, value := range store {
 | 
							for key, value := range store {
 | 
				
			||||||
			switch i {
 | 
								switch i {
 | 
				
			||||||
			case BatchUpdateTypeUserQuota:
 | 
								case BatchUpdateTypeUserQuota:
 | 
				
			||||||
@@ -64,8 +64,10 @@ func batchUpdate() {
 | 
				
			|||||||
				if err != nil {
 | 
									if err != nil {
 | 
				
			||||||
					common.SysError("failed to batch update token quota: " + err.Error())
 | 
										common.SysError("failed to batch update token quota: " + err.Error())
 | 
				
			||||||
				}
 | 
									}
 | 
				
			||||||
			case BatchUpdateTypeUsedQuotaAndRequestCount:
 | 
								case BatchUpdateTypeUsedQuota:
 | 
				
			||||||
				updateUserUsedQuotaAndRequestCount(key, value, 1) // TODO: count is incorrect
 | 
									updateUserUsedQuota(key, value)
 | 
				
			||||||
 | 
								case BatchUpdateTypeRequestCount:
 | 
				
			||||||
 | 
									updateUserRequestCount(key, value)
 | 
				
			||||||
			case BatchUpdateTypeChannelUsedQuota:
 | 
								case BatchUpdateTypeChannelUsedQuota:
 | 
				
			||||||
				updateChannelUsedQuota(key, value)
 | 
									updateChannelUsedQuota(key, value)
 | 
				
			||||||
			}
 | 
								}
 | 
				
			||||||
 
 | 
				
			|||||||
		Reference in New Issue
	
	Block a user