diff --git a/internal/accounts/weighted.go b/internal/accounts/weighted.go new file mode 100644 index 0000000..f744ace --- /dev/null +++ b/internal/accounts/weighted.go @@ -0,0 +1,48 @@ +package accounts + +import ( + "hash/fnv" + + "github.com/wangjia/pay/config" +) + +// Weighted 按 Weight 加权选账户,决策源为 hash(OutTradeNo)(确定性、可复现、跨节点一致, +// 见计划 D1)。Weight<=0 记 1(对齐 model.Account gorm default:1)。retry 换号由 +// Router 的 ExcludeAccounts 过滤驱动,不依赖随机。 +type Weighted struct{} + +func NewWeighted() *Weighted { return &Weighted{} } + +func (Weighted) Pick(_ string, candidates []config.AccountConfig, hint PickHint) (config.AccountConfig, error) { + if len(candidates) == 0 { + return config.AccountConfig{}, ErrNoAccount + } + if hint.OutTradeNo == "" { + return candidates[0], nil + } + var total uint32 + for _, c := range candidates { + total += weightOf(c) + } + if total == 0 { // 理论不达(weightOf>=1),兜底 + return candidates[0], nil + } + h := fnv.New32a() + _, _ = h.Write([]byte(hint.OutTradeNo)) + target := h.Sum32() % total + var acc uint32 + for _, c := range candidates { + acc += weightOf(c) + if target < acc { + return c, nil + } + } + return candidates[len(candidates)-1], nil // 浮点无关的整数走位,兜底不达 +} + +func weightOf(c config.AccountConfig) uint32 { + if c.Weight <= 0 { + return 1 + } + return uint32(c.Weight) +} diff --git a/internal/accounts/weighted_test.go b/internal/accounts/weighted_test.go new file mode 100644 index 0000000..34a5067 --- /dev/null +++ b/internal/accounts/weighted_test.go @@ -0,0 +1,56 @@ +package accounts_test + +import ( + "testing" + + "github.com/wangjia/pay/config" + "github.com/wangjia/pay/internal/accounts" +) + +func TestWeightedDeterministic(t *testing.T) { + w := accounts.NewWeighted() + cs := cands("a1", "a2", "a3") + // 同一 OutTradeNo 必落同一账户(确定性,可复现)。 + first, _ := w.Pick("fake|global", cs, accounts.PickHint{OutTradeNo: "PAY-XYZ"}) + for i := 0; i < 5; i++ { + got, _ := w.Pick("fake|global", cs, accounts.PickHint{OutTradeNo: "PAY-XYZ"}) + if got.AccountID != first.AccountID { + t.Fatalf("同单号应恒定命中 %s, 第 %d 次得 %s", first.AccountID, i, got.AccountID) + } + } +} + +func TestWeightedDistributesAndRespectsWeight(t *testing.T) { + w := accounts.NewWeighted() + // a1 权重 3,a2 权重 1 → 长跑大致 3:1。 + cs := []config.AccountConfig{ + {AccountID: "a1", Weight: 3}, + {AccountID: "a2", Weight: 1}, + } + counts := map[string]int{} + for i := 0; i < 4000; i++ { + no := "PAY-" + string(rune('A'+i%26)) + string(rune('0'+i/26%10)) + string(rune('a'+i/260%26)) + string(rune('0'+i%7)) + got, _ := w.Pick("fake|global", cs, accounts.PickHint{OutTradeNo: no}) + counts[got.AccountID]++ + } + if counts["a1"] == 0 || counts["a2"] == 0 { + t.Fatalf("两账户都应被选中, got %+v", counts) + } + // a1 应显著多于 a2(3:1 容差:a1 至少是 a2 的 2 倍)。 + if counts["a1"] < counts["a2"]*2 { + t.Fatalf("权重 3:1 未体现, got %+v", counts) + } +} + +func TestWeightedZeroWeightTreatedAsOne(t *testing.T) { + w := accounts.NewWeighted() + cs := []config.AccountConfig{{AccountID: "a1", Weight: 0}, {AccountID: "a2", Weight: 0}} + counts := map[string]int{} + for i := 0; i < 400; i++ { + got, _ := w.Pick("fake|global", cs, accounts.PickHint{OutTradeNo: "N" + string(rune('0'+i%10)) + string(rune('a'+i%26))}) + counts[got.AccountID]++ + } + if counts["a1"] == 0 || counts["a2"] == 0 { + t.Fatalf("weight=0 应当 1 处理、两账户都可选, got %+v", counts) + } +}