diff --git a/model/locking.go b/model/locking.go new file mode 100644 index 00000000..8c6313b8 --- /dev/null +++ b/model/locking.go @@ -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"}) +} diff --git a/model/locking_test.go b/model/locking_test.go new file mode 100644 index 00000000..2443490e --- /dev/null +++ b/model/locking_test.go @@ -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") +} diff --git a/model/redemption.go b/model/redemption.go index f29f867c..23985ef4 100644 --- a/model/redemption.go +++ b/model/redemption.go @@ -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()) diff --git a/model/redemption_test.go b/model/redemption_test.go index c9da024b..0ba2e8e8 100644 --- a/model/redemption_test.go +++ b/model/redemption_test.go @@ -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") +} diff --git a/model/subscription.go b/model/subscription.go index 18c822c4..e5f51823 100644 --- a/model/subscription.go +++ b/model/subscription.go @@ -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 diff --git a/model/topup.go b/model/topup.go index 8d722508..92cb276b 100644 --- a/model/topup.go +++ b/model/topup.go @@ -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("充值订单不存在") } diff --git a/model/user.go b/model/user.go index c0239317..08ffabcc 100644 --- a/model/user.go +++ b/model/user.go @@ -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 }