Jaymr/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(model): centralize row locking in transactional flows

Browse Source
This commit is contained in:
CaIon
2026-07-07 12:40:09 +08:00
parent bae799ccb1
commit 70ea899e37
7 changed files with 177 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/locking.go
+25
View File
@@ -0,0 +1,25 @@
package model
import (
"github.com/QuantumNous/new-api/common"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
// lockForUpdate makes the next query emit SELECT ... FOR UPDATE so the matched
// rows stay locked until the surrounding transaction ends.
//
// GORM v2 silently ignores the legacy `Set("gorm:query_option", "FOR UPDATE")`
// from GORM v1, so that form does not lock anything. Always use this helper
// instead.
//
// SQLite has no FOR UPDATE syntax (the clause would be a syntax error), so it
// is skipped there; SQLite's single-writer model makes one of two conflicting
// transactions fail instead of both committing.
func lockForUpdate(tx *gorm.DB) *gorm.DB {
if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
return tx
}
return tx.Clauses(clause.Locking{Strength: "UPDATE"})
}
model/locking_test.go
+38
View File
@@ -0,0 +1,38 @@
package model
import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
"gorm.io/gorm/utils/tests"
)
// lockForUpdate must emit FOR UPDATE on databases that support it and skip
// it on SQLite, where the syntax does not exist.
//
// The dummy dialector is used because SQLite drivers strip locking clauses
// from the generated SQL, which would mask what the helper itself does.
func TestLockForUpdateEmitsRowLock(t *testing.T) {
dummyDB, err := gorm.Open(tests.DummyDialector{}, &gorm.Config{DryRun: true})
require.NoError(t, err)
buildSQL := func() string {
var rows []Redemption
return lockForUpdate(dummyDB).Where("id = ?", 1).Find(&rows).Statement.SQL.String()
}
t.Cleanup(func() {
common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
})
common.SetDatabaseTypes(common.DatabaseTypeMySQL, common.DatabaseTypeSQLite)
assert.Contains(t, buildSQL(), "FOR UPDATE")
common.SetDatabaseTypes(common.DatabaseTypePostgreSQL, common.DatabaseTypeSQLite)
assert.Contains(t, buildSQL(), "FOR UPDATE")
common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
assert.NotContains(t, buildSQL(), "FOR UPDATE")
}
model/redemption.go
+17 -9
View File
@@ -149,7 +149,7 @@ func Redeem(key string, userId int) (quota int, err error) {
}
common.RandomSleep()
err = DB.Transaction(func(tx *gorm.DB) error {
err := tx.Set("gorm:query_option", "FOR UPDATE").Where(keyCol+" = ?", key).First(redemption).Error
err := lockForUpdate(tx).Where(keyCol+" = ?", key).First(redemption).Error
if err != nil {
return errors.New("无效的兑换码")
}
@@ -159,15 +159,23 @@ func Redeem(key string, userId int) (quota int, err error) {
if redemption.ExpiredTime != 0 && redemption.ExpiredTime < common.GetTimestamp() {
return errors.New("该兑换码已过期")
}
err = tx.Model(&User{}).Where("id = ?", userId).Update("quota", gorm.Expr("quota + ?", redemption.Quota)).Error
if err != nil {
return err
// Compare-and-swap on status: only the transaction that flips
// enabled -> used may credit quota, so a concurrent redeem of the
// same code loses here even without a row lock (e.g. on SQLite).
result := tx.Model(&Redemption{}).
Where("id = ? AND status = ?", redemption.Id, common.RedemptionCodeStatusEnabled).
Updates(map[string]interface{}{
"redeemed_time": common.GetTimestamp(),
"status": common.RedemptionCodeStatusUsed,
"used_user_id": userId,
})
if result.Error != nil {
return result.Error
}
redemption.RedeemedTime = common.GetTimestamp()
redemption.Status = common.RedemptionCodeStatusUsed
redemption.UsedUserId = userId
err = tx.Save(redemption).Error
return err
if result.RowsAffected == 0 {
return errors.New("该兑换码已被使用")
}
return tx.Model(&User{}).Where("id = ?", userId).Update("quota", gorm.Expr("quota + ?", redemption.Quota)).Error
})
if err != nil {
common.SysError("redemption failed: " + err.Error())
model/redemption_test.go
+81
View File
@@ -1,6 +1,7 @@
package model
import (
"sync"
"testing"
"github.com/QuantumNous/new-api/common"
@@ -98,3 +99,83 @@ func TestSearchRedemptionsFiltersAndPaginates(t *testing.T) {
})
}
}
func setupRedeemFixture(t *testing.T, quota int) (userId int, key string) {
t.Helper()
require.NoError(t, DB.AutoMigrate(&Redemption{}))
require.NoError(t, DB.Session(&gorm.Session{AllowGlobalUpdate: true}).Unscoped().Delete(&Redemption{}).Error)
t.Cleanup(func() {
require.NoError(t, DB.Session(&gorm.Session{AllowGlobalUpdate: true}).Unscoped().Delete(&Redemption{}).Error)
DB.Exec("DELETE FROM users")
DB.Exec("DELETE FROM logs")
})
user := &User{Username: "redeem-user", Password: "password", Status: common.UserStatusEnabled, Quota: 0}
require.NoError(t, DB.Create(user).Error)
key = "10000000000000000000000000000001"
redemption := &Redemption{
Name: "redeem-test",
Key: key,
Status: common.RedemptionCodeStatusEnabled,
Quota: quota,
CreatedTime: common.GetTimestamp(),
}
require.NoError(t, DB.Create(redemption).Error)
return user.Id, key
}
func TestRedeemCreditsQuotaExactlyOnce(t *testing.T) {
userId, key := setupRedeemFixture(t, 500)
quota, err := Redeem(key, userId)
require.NoError(t, err)
assert.Equal(t, 500, quota)
var user User
require.NoError(t, DB.First(&user, "id = ?", userId).Error)
assert.Equal(t, 500, user.Quota)
var redemption Redemption
require.NoError(t, DB.First(&redemption, "name = ?", "redeem-test").Error)
assert.Equal(t, common.RedemptionCodeStatusUsed, redemption.Status)
assert.Equal(t, userId, redemption.UsedUserId)
// Redeeming the same code again must fail and must not credit quota.
_, err = Redeem(key, userId)
require.Error(t, err)
require.NoError(t, DB.First(&user, "id = ?", userId).Error)
assert.Equal(t, 500, user.Quota)
}
// Exactly one of several concurrent redeems of the same code may win, and
// quota must be credited exactly once.
func TestRedeemConcurrentSingleSuccess(t *testing.T) {
userId, key := setupRedeemFixture(t, 300)
const goroutines = 5
successes := make([]bool, goroutines)
var wg sync.WaitGroup
wg.Add(goroutines)
for i := 0; i < goroutines; i++ {
go func(idx int) {
defer wg.Done()
if _, err := Redeem(key, userId); err == nil {
successes[idx] = true
}
}(i)
}
wg.Wait()
successCount := 0
for _, ok := range successes {
if ok {
successCount++
}
}
assert.Equal(t, 1, successCount, "exactly one concurrent redeem should succeed")
var user User
require.NoError(t, DB.First(&user, "id = ?", userId).Error)
assert.Equal(t, 300, user.Quota, "quota must be credited exactly once")
}
model/subscription.go
+9 -9
View File
@@ -565,7 +565,7 @@ func CompleteSubscriptionOrder(tradeNo string, providerPayload string, expectedP
var upgradeGroup string
err := DB.Transaction(func(tx *gorm.DB) error {
var order SubscriptionOrder
if err := tx.Set("gorm:query_option", "FOR UPDATE").Where(refCol+" = ?", tradeNo).First(&order).Error; err != nil {
if err := lockForUpdate(tx).Where(refCol+" = ?", tradeNo).First(&order).Error; err != nil {
return ErrSubscriptionOrderNotFound
}
if expectedPaymentProvider != "" && order.PaymentProvider != expectedPaymentProvider {
@@ -668,7 +668,7 @@ func ExpireSubscriptionOrder(tradeNo string, expectedPaymentProvider string) err
}
return DB.Transaction(func(tx *gorm.DB) error {
var order SubscriptionOrder
if err := tx.Set("gorm:query_option", "FOR UPDATE").Where(refCol+" = ?", tradeNo).First(&order).Error; err != nil {
if err := lockForUpdate(tx).Where(refCol+" = ?", tradeNo).First(&order).Error; err != nil {
return ErrSubscriptionOrderNotFound
}
if expectedPaymentProvider != "" && order.PaymentProvider != expectedPaymentProvider {
@@ -751,7 +751,7 @@ func PurchaseSubscriptionWithBalance(userId int, planId int) error {
}
var user User
if err := tx.Set("gorm:query_option", "FOR UPDATE").Where("id = ?", userId).First(&user).Error; err != nil {
if err := lockForUpdate(tx).Where("id = ?", userId).First(&user).Error; err != nil {
return err
}
if requiredQuota > 0 && user.Quota < requiredQuota {
@@ -899,7 +899,7 @@ func AdminInvalidateUserSubscription(userSubscriptionId int) (string, error) {
var userId int
err := DB.Transaction(func(tx *gorm.DB) error {
var sub UserSubscription
if err := tx.Set("gorm:query_option", "FOR UPDATE").
if err := lockForUpdate(tx).
Where("id = ?", userSubscriptionId).First(&sub).Error; err != nil {
return err
}
@@ -944,7 +944,7 @@ func AdminDeleteUserSubscription(userSubscriptionId int) (string, error) {
var userId int
err := DB.Transaction(func(tx *gorm.DB) error {
var sub UserSubscription
if err := tx.Set("gorm:query_option", "FOR UPDATE").
if err := lockForUpdate(tx).
Where("id = ?", userSubscriptionId).First(&sub).Error; err != nil {
return err
}
@@ -1178,7 +1178,7 @@ func PreConsumeUserSubscription(requestId string, userId int, modelName string,
}
var subs []UserSubscription
if err := tx.Set("gorm:query_option", "FOR UPDATE").
if err := lockForUpdate(tx).
Where("user_id = ? AND status = ? AND end_time > ?", userId, "active", now).
Order("end_time asc, id asc").
Find(&subs).Error; err != nil {
@@ -1251,7 +1251,7 @@ func RefundSubscriptionPreConsume(requestId string) error {
}
return DB.Transaction(func(tx *gorm.DB) error {
var record SubscriptionPreConsumeRecord
if err := tx.Set("gorm:query_option", "FOR UPDATE").
if err := lockForUpdate(tx).
Where("request_id = ?", requestId).First(&record).Error; err != nil {
return err
}
@@ -1295,7 +1295,7 @@ func ResetDueSubscriptions(limit int) (int, error) {
}
err = DB.Transaction(func(tx *gorm.DB) error {
var locked UserSubscription
if err := tx.Set("gorm:query_option", "FOR UPDATE").
if err := lockForUpdate(tx).
Where("id = ? AND next_reset_time > 0 AND next_reset_time <= ?", subCopy.Id, now).
First(&locked).Error; err != nil {
return nil
@@ -1362,7 +1362,7 @@ func PostConsumeUserSubscriptionDelta(userSubscriptionId int, delta int64) error
}
return DB.Transaction(func(tx *gorm.DB) error {
var sub UserSubscription
if err := tx.Set("gorm:query_option", "FOR UPDATE").
if err := lockForUpdate(tx).
Where("id = ?", userSubscriptionId).
First(&sub).Error; err != nil {
return err
model/topup.go
+6 -6
View File
@@ -91,7 +91,7 @@ func UpdatePendingTopUpStatus(tradeNo string, expectedPaymentProvider string, ta
return DB.Transaction(func(tx *gorm.DB) error {
topUp := &TopUp{}
if err := tx.Set("gorm:query_option", "FOR UPDATE").Where(refCol+" = ?", tradeNo).First(topUp).Error; err != nil {
if err := lockForUpdate(tx).Where(refCol+" = ?", tradeNo).First(topUp).Error; err != nil {
return ErrTopUpNotFound
}
if expectedPaymentProvider != "" && topUp.PaymentProvider != expectedPaymentProvider {
@@ -120,7 +120,7 @@ func Recharge(referenceId string, customerId string, callerIp string) (err error
}
err = DB.Transaction(func(tx *gorm.DB) error {
err := tx.Set("gorm:query_option", "FOR UPDATE").Where(refCol+" = ?", referenceId).First(topUp).Error
err := lockForUpdate(tx).Where(refCol+" = ?", referenceId).First(topUp).Error
if err != nil {
return errors.New("充值订单不存在")
}
@@ -335,7 +335,7 @@ func ManualCompleteTopUp(tradeNo string, callerIp string) error {
err := DB.Transaction(func(tx *gorm.DB) error {
topUp := &TopUp{}
// 行级锁，避免并发补单
if err := tx.Set("gorm:query_option", "FOR UPDATE").Where(refCol+" = ?", tradeNo).First(topUp).Error; err != nil {
if err := lockForUpdate(tx).Where(refCol+" = ?", tradeNo).First(topUp).Error; err != nil {
return errors.New("充值订单不存在")
}
@@ -403,7 +403,7 @@ func RechargeCreem(referenceId string, customerEmail string, customerName string
}
err = DB.Transaction(func(tx *gorm.DB) error {
err := tx.Set("gorm:query_option", "FOR UPDATE").Where(refCol+" = ?", referenceId).First(topUp).Error
err := lockForUpdate(tx).Where(refCol+" = ?", referenceId).First(topUp).Error
if err != nil {
return errors.New("充值订单不存在")
}
@@ -478,7 +478,7 @@ func RechargeWaffo(tradeNo string, callerIp string) (err error) {
}
err = DB.Transaction(func(tx *gorm.DB) error {
err := tx.Set("gorm:query_option", "FOR UPDATE").Where(refCol+" = ?", tradeNo).First(topUp).Error
err := lockForUpdate(tx).Where(refCol+" = ?", tradeNo).First(topUp).Error
if err != nil {
return errors.New("充值订单不存在")
}
@@ -541,7 +541,7 @@ func RechargeWaffoPancake(tradeNo string) (err error) {
}
err = DB.Transaction(func(tx *gorm.DB) error {
err := tx.Set("gorm:query_option", "FOR UPDATE").Where(refCol+" = ?", tradeNo).First(topUp).Error
err := lockForUpdate(tx).Where(refCol+" = ?", tradeNo).First(topUp).Error
if err != nil {
return errors.New("充值订单不存在")
}
model/user.go
+1 -1
View File
@@ -456,7 +456,7 @@ func (user *User) TransferAffQuotaToQuota(quota int) error {
defer tx.Rollback() // 确保在函数退出时事务能回滚
// 加锁查询用户以确保数据一致性
err := tx.Set("gorm:query_option", "FOR UPDATE").First(&user, user.Id).Error
err := lockForUpdate(tx).First(&user, user.Id).Error
if err != nil {
return err
}