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

fix: apply affiliate rebates on topups

Browse Source 
Move rebate processing into model transactions, keep frozen rebates out of transferable quota until release, and trigger top-up rebates idempotently.
This commit is contained in:
zizi 2026-05-20 15:47:00 +08:00
parent 9264873e2c
commit 03e440b005
6 changed files with 278 additions and 138 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

57
model/payment_method_guard_test.go
View File

@ -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)

178
model/rebate_process.go Normal file
View File

@ -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
}

2
model/task_cas_test.go
View File

@ -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")
})
}

38
model/topup.go
View File

@ -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 {

136
service/rebate.go
View File

@ -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()
}

5
service/rebate_test.go
View File

@ -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{