Commit 70ea899e by CaIon

fix(model): centralize row locking in transactional flows

parent bae799cc
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"})
}
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")
}
...@@ -149,7 +149,7 @@ func Redeem(key string, userId int) (quota int, err error) { ...@@ -149,7 +149,7 @@ func Redeem(key string, userId int) (quota int, err error) {
} }
common.RandomSleep() common.RandomSleep()
err = DB.Transaction(func(tx *gorm.DB) error { 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 { if err != nil {
return errors.New("无效的兑换码") return errors.New("无效的兑换码")
} }
...@@ -159,15 +159,23 @@ func Redeem(key string, userId int) (quota int, err error) { ...@@ -159,15 +159,23 @@ func Redeem(key string, userId int) (quota int, err error) {
if redemption.ExpiredTime != 0 && redemption.ExpiredTime < common.GetTimestamp() { if redemption.ExpiredTime != 0 && redemption.ExpiredTime < common.GetTimestamp() {
return errors.New("该兑换码已过期") return errors.New("该兑换码已过期")
} }
err = tx.Model(&User{}).Where("id = ?", userId).Update("quota", gorm.Expr("quota + ?", redemption.Quota)).Error // Compare-and-swap on status: only the transaction that flips
if err != nil { // enabled -> used may credit quota, so a concurrent redeem of the
return err // 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() if result.RowsAffected == 0 {
redemption.Status = common.RedemptionCodeStatusUsed return errors.New("该兑换码已被使用")
redemption.UsedUserId = userId }
err = tx.Save(redemption).Error return tx.Model(&User{}).Where("id = ?", userId).Update("quota", gorm.Expr("quota + ?", redemption.Quota)).Error
return err
}) })
if err != nil { if err != nil {
common.SysError("redemption failed: " + err.Error()) common.SysError("redemption failed: " + err.Error())
......
package model package model
import ( import (
"sync"
"testing" "testing"
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
...@@ -98,3 +99,83 @@ func TestSearchRedemptionsFiltersAndPaginates(t *testing.T) { ...@@ -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")
}
...@@ -565,7 +565,7 @@ func CompleteSubscriptionOrder(tradeNo string, providerPayload string, expectedP ...@@ -565,7 +565,7 @@ func CompleteSubscriptionOrder(tradeNo string, providerPayload string, expectedP
var upgradeGroup string var upgradeGroup string
err := DB.Transaction(func(tx *gorm.DB) error { err := DB.Transaction(func(tx *gorm.DB) error {
var order SubscriptionOrder 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 return ErrSubscriptionOrderNotFound
} }
if expectedPaymentProvider != "" && order.PaymentProvider != expectedPaymentProvider { if expectedPaymentProvider != "" && order.PaymentProvider != expectedPaymentProvider {
...@@ -668,7 +668,7 @@ func ExpireSubscriptionOrder(tradeNo string, expectedPaymentProvider string) err ...@@ -668,7 +668,7 @@ func ExpireSubscriptionOrder(tradeNo string, expectedPaymentProvider string) err
} }
return DB.Transaction(func(tx *gorm.DB) error { return DB.Transaction(func(tx *gorm.DB) error {
var order SubscriptionOrder 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 return ErrSubscriptionOrderNotFound
} }
if expectedPaymentProvider != "" && order.PaymentProvider != expectedPaymentProvider { if expectedPaymentProvider != "" && order.PaymentProvider != expectedPaymentProvider {
...@@ -751,7 +751,7 @@ func PurchaseSubscriptionWithBalance(userId int, planId int) error { ...@@ -751,7 +751,7 @@ func PurchaseSubscriptionWithBalance(userId int, planId int) error {
} }
var user User 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 return err
} }
if requiredQuota > 0 && user.Quota < requiredQuota { if requiredQuota > 0 && user.Quota < requiredQuota {
...@@ -899,7 +899,7 @@ func AdminInvalidateUserSubscription(userSubscriptionId int) (string, error) { ...@@ -899,7 +899,7 @@ func AdminInvalidateUserSubscription(userSubscriptionId int) (string, error) {
var userId int var userId int
err := DB.Transaction(func(tx *gorm.DB) error { err := DB.Transaction(func(tx *gorm.DB) error {
var sub UserSubscription var sub UserSubscription
if err := tx.Set("gorm:query_option", "FOR UPDATE"). if err := lockForUpdate(tx).
Where("id = ?", userSubscriptionId).First(&sub).Error; err != nil { Where("id = ?", userSubscriptionId).First(&sub).Error; err != nil {
return err return err
} }
...@@ -944,7 +944,7 @@ func AdminDeleteUserSubscription(userSubscriptionId int) (string, error) { ...@@ -944,7 +944,7 @@ func AdminDeleteUserSubscription(userSubscriptionId int) (string, error) {
var userId int var userId int
err := DB.Transaction(func(tx *gorm.DB) error { err := DB.Transaction(func(tx *gorm.DB) error {
var sub UserSubscription var sub UserSubscription
if err := tx.Set("gorm:query_option", "FOR UPDATE"). if err := lockForUpdate(tx).
Where("id = ?", userSubscriptionId).First(&sub).Error; err != nil { Where("id = ?", userSubscriptionId).First(&sub).Error; err != nil {
return err return err
} }
...@@ -1178,7 +1178,7 @@ func PreConsumeUserSubscription(requestId string, userId int, modelName string, ...@@ -1178,7 +1178,7 @@ func PreConsumeUserSubscription(requestId string, userId int, modelName string,
} }
var subs []UserSubscription 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). Where("user_id = ? AND status = ? AND end_time > ?", userId, "active", now).
Order("end_time asc, id asc"). Order("end_time asc, id asc").
Find(&subs).Error; err != nil { Find(&subs).Error; err != nil {
...@@ -1251,7 +1251,7 @@ func RefundSubscriptionPreConsume(requestId string) error { ...@@ -1251,7 +1251,7 @@ func RefundSubscriptionPreConsume(requestId string) error {
} }
return DB.Transaction(func(tx *gorm.DB) error { return DB.Transaction(func(tx *gorm.DB) error {
var record SubscriptionPreConsumeRecord 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 { Where("request_id = ?", requestId).First(&record).Error; err != nil {
return err return err
} }
...@@ -1295,7 +1295,7 @@ func ResetDueSubscriptions(limit int) (int, error) { ...@@ -1295,7 +1295,7 @@ func ResetDueSubscriptions(limit int) (int, error) {
} }
err = DB.Transaction(func(tx *gorm.DB) error { err = DB.Transaction(func(tx *gorm.DB) error {
var locked UserSubscription 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). Where("id = ? AND next_reset_time > 0 AND next_reset_time <= ?", subCopy.Id, now).
First(&locked).Error; err != nil { First(&locked).Error; err != nil {
return nil return nil
...@@ -1362,7 +1362,7 @@ func PostConsumeUserSubscriptionDelta(userSubscriptionId int, delta int64) error ...@@ -1362,7 +1362,7 @@ func PostConsumeUserSubscriptionDelta(userSubscriptionId int, delta int64) error
} }
return DB.Transaction(func(tx *gorm.DB) error { return DB.Transaction(func(tx *gorm.DB) error {
var sub UserSubscription var sub UserSubscription
if err := tx.Set("gorm:query_option", "FOR UPDATE"). if err := lockForUpdate(tx).
Where("id = ?", userSubscriptionId). Where("id = ?", userSubscriptionId).
First(&sub).Error; err != nil { First(&sub).Error; err != nil {
return err return err
......
...@@ -91,7 +91,7 @@ func UpdatePendingTopUpStatus(tradeNo string, expectedPaymentProvider string, ta ...@@ -91,7 +91,7 @@ func UpdatePendingTopUpStatus(tradeNo string, expectedPaymentProvider string, ta
return DB.Transaction(func(tx *gorm.DB) error { return DB.Transaction(func(tx *gorm.DB) error {
topUp := &TopUp{} 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 return ErrTopUpNotFound
} }
if expectedPaymentProvider != "" && topUp.PaymentProvider != expectedPaymentProvider { if expectedPaymentProvider != "" && topUp.PaymentProvider != expectedPaymentProvider {
...@@ -120,7 +120,7 @@ func Recharge(referenceId string, customerId string, callerIp string) (err error ...@@ -120,7 +120,7 @@ func Recharge(referenceId string, customerId string, callerIp string) (err error
} }
err = DB.Transaction(func(tx *gorm.DB) 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 { if err != nil {
return errors.New("充值订单不存在") return errors.New("充值订单不存在")
} }
...@@ -335,7 +335,7 @@ func ManualCompleteTopUp(tradeNo string, callerIp string) error { ...@@ -335,7 +335,7 @@ func ManualCompleteTopUp(tradeNo string, callerIp string) error {
err := DB.Transaction(func(tx *gorm.DB) error { err := DB.Transaction(func(tx *gorm.DB) error {
topUp := &TopUp{} 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("充值订单不存在") return errors.New("充值订单不存在")
} }
...@@ -403,7 +403,7 @@ func RechargeCreem(referenceId string, customerEmail string, customerName string ...@@ -403,7 +403,7 @@ func RechargeCreem(referenceId string, customerEmail string, customerName string
} }
err = DB.Transaction(func(tx *gorm.DB) 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 { if err != nil {
return errors.New("充值订单不存在") return errors.New("充值订单不存在")
} }
...@@ -478,7 +478,7 @@ func RechargeWaffo(tradeNo string, callerIp string) (err error) { ...@@ -478,7 +478,7 @@ func RechargeWaffo(tradeNo string, callerIp string) (err error) {
} }
err = DB.Transaction(func(tx *gorm.DB) 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 { if err != nil {
return errors.New("充值订单不存在") return errors.New("充值订单不存在")
} }
...@@ -541,7 +541,7 @@ func RechargeWaffoPancake(tradeNo string) (err error) { ...@@ -541,7 +541,7 @@ func RechargeWaffoPancake(tradeNo string) (err error) {
} }
err = DB.Transaction(func(tx *gorm.DB) 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 { if err != nil {
return errors.New("充值订单不存在") return errors.New("充值订单不存在")
} }
......
...@@ -456,7 +456,7 @@ func (user *User) TransferAffQuotaToQuota(quota int) error { ...@@ -456,7 +456,7 @@ func (user *User) TransferAffQuotaToQuota(quota int) error {
defer tx.Rollback() // 确保在函数退出时事务能回滚 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 { if err != nil {
return err return err
} }
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment