From 03e440b0057faf8599fbcb74780932843a582bc2 Mon Sep 17 00:00:00 2001 From: zizi Date: Wed, 20 May 2026 15:47:00 +0800 Subject: [PATCH] fix: apply affiliate rebates on topups Move rebate processing into model transactions, keep frozen rebates out of transferable quota until release, and trigger top-up rebates idempotently. --- model/payment_method_guard_test.go | 57 +++++++++ model/rebate_process.go | 178 +++++++++++++++++++++++++++++ model/task_cas_test.go | 2 + model/topup.go | 38 ++++-- service/rebate.go | 136 ++-------------------- service/rebate_test.go | 5 +- 6 files changed, 278 insertions(+), 138 deletions(-) create mode 100644 model/rebate_process.go diff --git a/model/payment_method_guard_test.go b/model/payment_method_guard_test.go index 2af8b812..59a8ec25 100644 --- a/model/payment_method_guard_test.go +++ b/model/payment_method_guard_test.go @@ -265,6 +265,63 @@ func TestRechargeWxpay_RejectsMismatchedPaymentProvider(t *testing.T) { assert.Equal(t, 0, getUserQuotaForPaymentGuardTest(t, 183)) } +func TestRechargeWaffo_ProcessesAffiliateRebateIdempotently(t *testing.T) { + truncateTables(t) + + common.OptionMapRWMutex.RLock() + originalOptionMap := common.OptionMap + common.OptionMapRWMutex.RUnlock() + common.OptionMapRWMutex.Lock() + common.OptionMap = map[string]string{ + AffRebateEnabledKey: "true", + AffRebateRatePercentKey: "10", + AffRebateFrozenDaysKey: "0", + } + common.OptionMapRWMutex.Unlock() + t.Cleanup(func() { + common.OptionMapRWMutex.Lock() + common.OptionMap = originalOptionMap + common.OptionMapRWMutex.Unlock() + }) + + inviter := &User{Id: 184, Username: "rebate_inviter", AffCode: "rebate_inviter_code", Status: common.UserStatusEnabled} + invitee := &User{Id: 185, Username: "rebate_invitee", AffCode: "rebate_invitee_code", Status: common.UserStatusEnabled, InviterId: inviter.Id} + require.NoError(t, DB.Create(inviter).Error) + require.NoError(t, DB.Create(invitee).Error) + topUp := &TopUp{ + UserId: invitee.Id, + Amount: 2, + Money: 2, + TradeNo: "waffo-rebate-guard", + PaymentMethod: PaymentMethodWaffo, + PaymentProvider: PaymentProviderWaffo, + Status: common.TopUpStatusPending, + CreateTime: time.Now().Unix(), + } + require.NoError(t, topUp.Insert()) + + require.NoError(t, RechargeWaffo("waffo-rebate-guard", "127.0.0.1")) + require.NoError(t, RechargeWaffo("waffo-rebate-guard", "127.0.0.1")) + + baseQuota := int(decimal.NewFromInt(2).Mul(decimal.NewFromFloat(common.QuotaPerUnit)).IntPart()) + var record RebateRecord + require.NoError(t, DB.First(&record).Error) + assert.Equal(t, topUp.Id, record.OrderId) + assert.Equal(t, RebateOrderTypeTopUp, record.OrderType) + assert.Equal(t, baseQuota, record.OrderAmount) + assert.Equal(t, baseQuota/10, record.RebateAmount) + assert.Equal(t, "released", record.Status) + + var reloadedInviter User + require.NoError(t, DB.First(&reloadedInviter, inviter.Id).Error) + assert.Equal(t, baseQuota/10, reloadedInviter.AffQuota) + assert.Equal(t, baseQuota/10, reloadedInviter.AffHistoryQuota) + + var count int64 + require.NoError(t, DB.Model(&RebateRecord{}).Count(&count).Error) + assert.Equal(t, int64(1), count) +} + func TestCompleteSubscriptionOrder_RejectsMismatchedPaymentProvider(t *testing.T) { truncateTables(t) diff --git a/model/rebate_process.go b/model/rebate_process.go new file mode 100644 index 00000000..54b9a7da --- /dev/null +++ b/model/rebate_process.go @@ -0,0 +1,178 @@ +package model + +import ( + "errors" + "strconv" + "time" + + "github.com/QuantumNous/new-api/common" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +const RebateOrderTypeTopUp = "topup" + +func IsRebateEnabled() bool { + enabled, err := strconv.ParseBool(readRebateOption(AffRebateEnabledKey)) + return err == nil && enabled +} + +func GetGlobalRebateRatePercent() int { + return parseNonNegativeRebateInt(readRebateOption(AffRebateRatePercentKey)) +} + +func GetGlobalRebateFrozenDays() int { + return parseNonNegativeRebateInt(readRebateOption(AffRebateFrozenDaysKey)) +} + +func GetEffectiveRebateRate(inviter *User) int { + if inviter != nil && inviter.AffRebateRatePercent > 0 { + return inviter.AffRebateRatePercent + } + return GetGlobalRebateRatePercent() +} + +func GetEffectiveFrozenDays(inviter *User) int { + if inviter != nil && inviter.AffRebateFrozenDays > 0 { + return inviter.AffRebateFrozenDays + } + return GetGlobalRebateFrozenDays() +} + +func ProcessRebateAfterRecharge(invitee *User, orderId int, orderType string, orderAmount int) error { + return DB.Transaction(func(tx *gorm.DB) error { + return processRebateAfterRechargeTx(tx, invitee, orderId, orderType, orderAmount) + }) +} + +func processTopUpRebateTx(tx *gorm.DB, topUp *TopUp, quotaToAdd int) error { + if topUp == nil || quotaToAdd <= 0 { + return nil + } + var invitee User + if err := tx.First(&invitee, "id = ?", topUp.UserId).Error; err != nil { + return err + } + return processRebateAfterRechargeTx(tx, &invitee, topUp.Id, RebateOrderTypeTopUp, quotaToAdd) +} + +func processRebateAfterRechargeTx(tx *gorm.DB, invitee *User, orderId int, orderType string, orderAmount int) error { + if invitee == nil || invitee.InviterId <= 0 || orderAmount <= 0 { + return nil + } + if !IsRebateEnabled() { + return nil + } + if orderType == "" { + return errors.New("order type is empty") + } + if orderId <= 0 { + return errors.New("order id must be positive") + } + + var inviter User + if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + First(&inviter, "id = ?", invitee.InviterId).Error; err != nil { + return err + } + + var existing RebateRecord + err := tx.Where( + "order_type = ? AND order_id = ? AND inviter_id = ? AND invitee_id = ?", + orderType, orderId, inviter.Id, invitee.Id, + ).First(&existing).Error + if err == nil { + return nil + } + if !errors.Is(err, gorm.ErrRecordNotFound) { + return err + } + + rate := GetEffectiveRebateRate(&inviter) + if rate <= 0 { + return nil + } + rebateAmount := orderAmount * rate / 100 + if rebateAmount <= 0 { + return nil + } + + now := time.Now() + frozenDays := GetEffectiveFrozenDays(&inviter) + record := RebateRecord{ + InviterId: inviter.Id, + InviteeId: invitee.Id, + OrderId: orderId, + OrderType: orderType, + OrderAmount: orderAmount, + RebateAmount: rebateAmount, + RatePercent: rate, + Status: "released", + ReleasedAt: &now, + } + if frozenDays > 0 { + frozenUntil := now.Add(time.Duration(frozenDays) * 24 * time.Hour) + record.Status = "frozen" + record.FrozenUntil = &frozenUntil + record.ReleasedAt = nil + } + + if err := tx.Create(&record).Error; err != nil { + return err + } + + updates := map[string]interface{}{ + "aff_history": gorm.Expr("aff_history + ?", rebateAmount), + } + if frozenDays <= 0 { + updates["aff_quota"] = gorm.Expr("aff_quota + ?", rebateAmount) + } + return tx.Model(&User{}).Where("id = ?", inviter.Id).Updates(updates).Error +} + +func ReleaseExpiredRebates() error { + return DB.Transaction(func(tx *gorm.DB) error { + var records []RebateRecord + if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + Where("status = ? AND frozen_until IS NOT NULL AND frozen_until <= ?", "frozen", time.Now()). + Order("id ASC"). + Find(&records).Error; err != nil { + return err + } + + releasedAt := time.Now() + for _, record := range records { + result := tx.Model(&RebateRecord{}). + Where("id = ? AND status = ?", record.Id, "frozen"). + Updates(map[string]interface{}{ + "status": "released", + "released_at": releasedAt, + }) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + continue + } + if err := tx.Model(&User{}).Where("id = ?", record.InviterId). + Update("aff_quota", gorm.Expr("aff_quota + ?", record.RebateAmount)).Error; err != nil { + return err + } + } + return nil + }) +} + +func readRebateOption(key string) string { + common.OptionMapRWMutex.RLock() + defer common.OptionMapRWMutex.RUnlock() + return common.OptionMap[key] +} + +func parseNonNegativeRebateInt(value string) int { + parsed, err := strconv.Atoi(value) + if err != nil || parsed < 0 { + return 0 + } + return parsed +} diff --git a/model/task_cas_test.go b/model/task_cas_test.go index d9025a17..0be1c146 100644 --- a/model/task_cas_test.go +++ b/model/task_cas_test.go @@ -45,6 +45,7 @@ func TestMain(m *testing.M) { &UserSubscription{}, &PromoCode{}, &PromoCodeUsage{}, + &RebateRecord{}, ); err != nil { panic("failed to migrate: " + err.Error()) } @@ -66,6 +67,7 @@ func truncateTables(t *testing.T) { DB.Exec("DELETE FROM user_subscriptions") DB.Exec("DELETE FROM promo_code_usages") DB.Exec("DELETE FROM promo_codes") + DB.Exec("DELETE FROM rebate_records") }) } diff --git a/model/topup.go b/model/topup.go index ab7cbaf8..ada6985b 100644 --- a/model/topup.go +++ b/model/topup.go @@ -286,8 +286,11 @@ func RechargeEpay(tradeNo string, actualPaymentMethod string, callerIp string) ( } bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp) + if bonusErr != nil { + return bonusErr + } promoBonus = bonus - return bonusErr + return processTopUpRebateTx(tx, topUp, quotaToAdd) }) if err != nil { @@ -345,8 +348,11 @@ func Recharge(referenceId string, customerId string, callerIp string) (err error } bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp) + if bonusErr != nil { + return bonusErr + } promoBonus = bonus - return bonusErr + return processTopUpRebateTx(tx, topUp, int(quota)) }) if err != nil { @@ -584,6 +590,9 @@ func ManualCompleteTopUp(tradeNo string, callerIp string) error { return bonusErr } promoBonus = bonus + if err := processTopUpRebateTx(tx, topUp, quotaToAdd); err != nil { + return err + } userId = topUp.UserId payMoney = topUp.Money @@ -670,8 +679,11 @@ func RechargeCreem(referenceId string, customerEmail string, customerName string } bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp) + if bonusErr != nil { + return bonusErr + } promoBonus = bonus - return bonusErr + return processTopUpRebateTx(tx, topUp, int(quota)) }) if err != nil { @@ -737,8 +749,11 @@ func RechargeWaffo(tradeNo string, callerIp string) (err error) { } bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp) + if bonusErr != nil { + return bonusErr + } promoBonus = bonus - return bonusErr + return processTopUpRebateTx(tx, topUp, quotaToAdd) }) if err != nil { @@ -804,8 +819,11 @@ func RechargeAlipay(tradeNo string, callerIp string) (err error) { } bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp) + if bonusErr != nil { + return bonusErr + } promoBonus = bonus - return bonusErr + return processTopUpRebateTx(tx, topUp, quotaToAdd) }) if err != nil { @@ -871,8 +889,11 @@ func RechargeWxpay(tradeNo string, callerIp string) (err error) { } bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp) + if bonusErr != nil { + return bonusErr + } promoBonus = bonus - return bonusErr + return processTopUpRebateTx(tx, topUp, quotaToAdd) }) if err != nil { @@ -936,8 +957,11 @@ func RechargeWaffoPancake(tradeNo string) (err error) { } bonus, bonusErr := applyPromoCodeBonusTx(tx, topUp) + if bonusErr != nil { + return bonusErr + } promoBonus = bonus - return bonusErr + return processTopUpRebateTx(tx, topUp, quotaToAdd) }) if err != nil { diff --git a/service/rebate.go b/service/rebate.go index d8d72529..298f6dd1 100644 --- a/service/rebate.go +++ b/service/rebate.go @@ -1,151 +1,31 @@ package service -import ( - "errors" - "strconv" - "time" - - "github.com/QuantumNous/new-api/common" - "github.com/QuantumNous/new-api/model" - "gorm.io/gorm" - "gorm.io/gorm/clause" -) +import "github.com/QuantumNous/new-api/model" func IsRebateEnabled() bool { - enabled, err := strconv.ParseBool(readRebateOption(model.AffRebateEnabledKey)) - return err == nil && enabled + return model.IsRebateEnabled() } func GetGlobalRebateRatePercent() int { - return parseNonNegativeRebateInt(readRebateOption(model.AffRebateRatePercentKey)) + return model.GetGlobalRebateRatePercent() } func GetGlobalRebateFrozenDays() int { - return parseNonNegativeRebateInt(readRebateOption(model.AffRebateFrozenDaysKey)) + return model.GetGlobalRebateFrozenDays() } func GetEffectiveRebateRate(inviter *model.User) int { - if inviter != nil && inviter.AffRebateRatePercent > 0 { - return inviter.AffRebateRatePercent - } - return GetGlobalRebateRatePercent() + return model.GetEffectiveRebateRate(inviter) } func GetEffectiveFrozenDays(inviter *model.User) int { - if inviter != nil && inviter.AffRebateFrozenDays > 0 { - return inviter.AffRebateFrozenDays - } - return GetGlobalRebateFrozenDays() + return model.GetEffectiveFrozenDays(inviter) } func ProcessRebateAfterRecharge(invitee *model.User, orderId int, orderType string, orderAmount int) error { - if invitee == nil || invitee.InviterId <= 0 || orderAmount <= 0 { - return nil - } - if !IsRebateEnabled() { - return nil - } - if orderType == "" { - return errors.New("order type is empty") - } - if orderId <= 0 { - return errors.New("order id must be positive") - } - - return model.DB.Transaction(func(tx *gorm.DB) error { - var inviter model.User - if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). - First(&inviter, "id = ?", invitee.InviterId).Error; err != nil { - return err - } - - var existing model.RebateRecord - err := tx.Where( - "order_type = ? AND order_id = ? AND inviter_id = ? AND invitee_id = ?", - orderType, orderId, inviter.Id, invitee.Id, - ).First(&existing).Error - if err == nil { - return nil - } - if !errors.Is(err, gorm.ErrRecordNotFound) { - return err - } - - rate := GetEffectiveRebateRate(&inviter) - if rate <= 0 { - return nil - } - rebateAmount := orderAmount * rate / 100 - if rebateAmount <= 0 { - return nil - } - - now := time.Now() - record := model.RebateRecord{ - InviterId: inviter.Id, - InviteeId: invitee.Id, - OrderId: orderId, - OrderType: orderType, - OrderAmount: orderAmount, - RebateAmount: rebateAmount, - RatePercent: rate, - Status: "released", - ReleasedAt: &now, - } - if frozenDays := GetEffectiveFrozenDays(&inviter); frozenDays > 0 { - frozenUntil := now.Add(time.Duration(frozenDays) * 24 * time.Hour) - record.Status = "frozen" - record.FrozenUntil = &frozenUntil - record.ReleasedAt = nil - } - - if err := tx.Create(&record).Error; err != nil { - return err - } - - return tx.Model(&model.User{}).Where("id = ?", inviter.Id).Updates(map[string]interface{}{ - "aff_quota": gorm.Expr("aff_quota + ?", rebateAmount), - "aff_history": gorm.Expr("aff_history + ?", rebateAmount), - }).Error - }) + return model.ProcessRebateAfterRecharge(invitee, orderId, orderType, orderAmount) } func ReleaseExpiredRebates() error { - return model.DB.Transaction(func(tx *gorm.DB) error { - var records []model.RebateRecord - if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). - Where("status = ? AND frozen_until IS NOT NULL AND frozen_until <= ?", "frozen", time.Now()). - Order("id ASC"). - Find(&records).Error; err != nil { - return err - } - - releasedAt := time.Now() - for _, record := range records { - result := tx.Model(&model.RebateRecord{}). - Where("id = ? AND status = ?", record.Id, "frozen"). - Updates(map[string]interface{}{ - "status": "released", - "released_at": releasedAt, - }) - if result.Error != nil { - return result.Error - } - } - return nil - }) -} - -func readRebateOption(key string) string { - common.OptionMapRWMutex.RLock() - defer common.OptionMapRWMutex.RUnlock() - return common.OptionMap[key] -} - -func parseNonNegativeRebateInt(value string) int { - parsed, err := strconv.Atoi(value) - if err != nil || parsed < 0 { - return 0 - } - return parsed + return model.ReleaseExpiredRebates() } diff --git a/service/rebate_test.go b/service/rebate_test.go index a6b1c0c7..1575daa3 100644 --- a/service/rebate_test.go +++ b/service/rebate_test.go @@ -144,7 +144,7 @@ func TestProcessRebateAfterRechargeUsesGlobalRate(t *testing.T) { require.Nil(t, record.ReleasedAt) inviter := loadRebateUser(t, db, 1) - require.Equal(t, 100, inviter.AffQuota) + require.Equal(t, 0, inviter.AffQuota) require.Equal(t, 100, inviter.AffHistoryQuota) } @@ -167,7 +167,7 @@ func TestProcessRebateAfterRechargeUsesPerUserOverrides(t *testing.T) { require.WithinDuration(t, time.Now().Add(24*time.Hour), *record.FrozenUntil, time.Minute) inviter := loadRebateUser(t, db, 1) - require.Equal(t, 250, inviter.AffQuota) + require.Equal(t, 0, inviter.AffQuota) require.Equal(t, 250, inviter.AffHistoryQuota) } @@ -193,7 +193,6 @@ func TestReleaseExpiredRebatesMarksOnlyDueFrozenRecords(t *testing.T) { past := time.Now().Add(-time.Hour) future := time.Now().Add(time.Hour) require.NoError(t, db.Model(inviter).Updates(map[string]interface{}{ - "aff_quota": 50, "aff_history": 50, }).Error) require.NoError(t, db.Create(&model.RebateRecord{