From 93c1d39df3b00740d5adc4cda508fc1d43973a93 Mon Sep 17 00:00:00 2001 From: Autumn-27 <53844290+Autumn-27@users.noreply.github.com> Date: Thu, 1 Oct 2026 06:43:29 -0400 Subject: [PATCH] =?UTF-8?q?fix(notify):=20=E6=B1=87=E6=80=BB=E6=A8=A1?= =?UTF-8?q?=E5=BC=8F=E6=8C=89=E6=B6=88=E6=81=AF=E8=AE=A1=E9=87=8F=E9=99=90?= =?UTF-8?q?=E6=B5=81=EF=BC=8C=E6=89=B9=E6=AC=A1=E5=A4=A7=E5=B0=8F=E4=B8=8D?= =?UTF-8?q?=E5=86=8D=E8=A2=AB=E6=AF=8F=E8=BD=AE=E8=AF=B7=E6=B1=82=E9=A2=84?= =?UTF-8?q?=E7=AE=97=E6=8E=90=E6=AD=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 令牌桶的额度此前被同时当成两个量纲用：对实时模式它是「发几条消息」， 对汇总模式却被直接当成「一批装几条漏洞」传给 ClaimDigestBatch。 notifyMaxSendsPerChannelPerTick（=5）是由租约时长倒推的 HTTP 请求预算， 而一个汇总批次只发一次请求，两者本不该相乘。稳态下更糟：rate_per_min=20 的渠道在 3 秒的 tick 里只补得到 1 个令牌，于是每条汇总消息只装 1 个漏洞， 头部还写着「近 30 分钟新增 1 个漏洞」——digest 退化成带汇总文案的实时推送， MaxDigestBatchSize(500) 永不可达。 改为统一按**消息条数**计量：一批漏洞合成一条消息、消耗一个令牌， rate_per_min 对 digest 依然生效（每分钟最多这么多条汇总消息）， 而批次大小只受 MaxDigestBatchSize 这个内存上界约束。 ClaimDigestBatch 原先的注释声称「单批同时受限流额度与内存上界约束」， 那正是本问题的来源——限流额度本不该充当批次大小。注释一并改掉， 免得下次又按「rate_per_min 对 digest 不生效」把它改回去。 顺带让 takeTokens 接受 want 上限：原先整桶抽空后才由调用方截断， 多出来的令牌（飞书默认 100/min 时是 95 个）直接丢弃，渠道这一轮 没有待发投递时也照扣，注释声称的突发容量在任何情况下都做不到。 决策收在 digestTickPlan 里由新增用例直接钉住。现有端到端用例都手动给 stepDigest 传一个够大的 limit，绕过了 step 里的额度计算，所以覆盖不到这里。 --- db/notification_delivery.go | 12 +++++-- server/notifier.go | 69 ++++++++++++++++++++++++++++--------- server/notify_api_test.go | 68 ++++++++++++++++++++++++++++++++---- 3 files changed, 123 insertions(+), 26 deletions(-) diff --git a/db/notification_delivery.go b/db/notification_delivery.go index e0d417f..36a7188 100644 --- a/db/notification_delivery.go +++ b/db/notification_delivery.go @@ -171,9 +171,15 @@ func (d *DB) ClaimDigestBatch(ctx context.Context, channelID int64, limit int, l if limit <= 0 { return nil, nil } - // 单批同时受两个上界约束：调用方的限流额度（本轮还能发几条）与内存上界 - // MaxDigestBatchSize。此前只用了后者，于是 digest 渠道的 rate_per_min - // 完全不生效——令牌被 takeTokens 扣掉却没人用，等于白扣。 + // limit 是**内存上界**，调用方传 MaxDigestBatchSize；这里再夹一道， + // 防止调用方传进一个更大的值。 + // + // 刻意不接受「限流额度」充当批次大小：限流的单位是消息条数——一个批次只发 + // 一条消息、消耗一个令牌，由 server 层的 takeTokens 扣除——与「一批装几条 + // 漏洞」是两个不同的量纲。曾经为了让 rate_per_min 对 digest 生效而把每轮 + // 请求预算传进来当批次大小，结果 rate=20/min 的渠道每批只装 1 条漏洞， + // digest 退化成带汇总文案的实时推送。要改限流请改 takeTokens 的 want， + // 不要动这里。 if limit > MaxDigestBatchSize { limit = MaxDigestBatchSize } diff --git a/server/notifier.go b/server/notifier.go index 0b4fcae..2e7d43b 100644 --- a/server/notifier.go +++ b/server/notifier.go @@ -136,24 +136,50 @@ func (n *Notifier) step(ctx context.Context) { if !ch.IsEnabled() { continue } - // 先问令牌桶这一轮还能发几条，再按这个数量去领——顺序不能反， - // 否则被限流挡下的投递已经消耗过重试次数。 - allow := n.takeTokens(ch.ID, ch.RatePerMin, time.Now()) - if allow > notifyMaxSendsPerChannelPerTick { - allow = notifyMaxSendsPerChannelPerTick - } - if allow <= 0 { + // 令牌桶的计量单位是**消息条数**（等价于 HTTP 请求数），不是漏洞条数。 + // 实时模式下两者相同（一条漏洞一条消息）；汇总模式下一整批漏洞合成 + // 一条消息，所以只消耗一个令牌。 + // + // 两种模式都先问令牌桶、再按额度去领——顺序不能反，否则被限流挡下的 + // 投递已经消耗过重试次数。 + now := time.Now() + if ch.Mode == db.NotifyModeDigest { + tokens, claimLimit := digestTickPlan() + if n.takeTokens(ch.ID, ch.RatePerMin, tokens, now) <= 0 { + continue + } + n.stepDigest(ctx, ch, claimLimit, baseURL) continue } - if ch.Mode == db.NotifyModeDigest { - // 汇总也受同一份额度约束：领到的条数不能超过本轮令牌数。 - n.stepDigest(ctx, ch, allow, baseURL) + allow := n.takeTokens(ch.ID, ch.RatePerMin, notifyMaxSendsPerChannelPerTick, now) + if allow <= 0 { continue } n.stepRealtime(ctx, ch, allow, baseURL) } } +// digestTickPlan 返回汇总渠道本轮的令牌消耗与批次大小上界。 +// +// 两个返回值是**两个不同的量纲**，这正是独立成函数的理由： +// +// - tokens 是消息条数。一批漏洞合成一条消息、发一次 HTTP 请求，所以恒为 1。 +// rate_per_min 因此仍然对 digest 生效（每分钟最多这么多条汇总消息）。 +// - claimLimit 是这一批最多装几条漏洞。它只受内存上界约束，与请求预算无关。 +// +// 曾经为了让 rate_per_min 对 digest 生效，把每轮请求预算 +// （notifyMaxSendsPerChannelPerTick，由租约倒推而来）直接当批次大小传下去。 +// 后果是 rate_per_min=20 的渠道在 3 秒的 tick 里只补到 1 个令牌，于是每条汇总 +// 消息只装 1 个漏洞——digest 退化成「带汇总文案的实时推送」，读者收到的是一串 +// 「近 30 分钟新增 1 个漏洞」，而 db.MaxDigestBatchSize 永不可达。 +// +// 这个症状在端到端测试里不容易发现（现有用例都手动传一个够大的 limit 给 +// stepDigest，绕过了 step 里的额度计算），所以把决策收在这里由 +// TestDigestTickPlanDecouplesBatchSizeFromSendBudget 直接钉住。 +func digestTickPlan() (tokens, claimLimit int) { + return 1, db.MaxDigestBatchSize +} + // stepRealtime 领取并投递某渠道的实时任务，一条漏洞一条消息。 func (n *Notifier) stepRealtime(ctx context.Context, ch *db.NotificationChannel, allow int, baseURL string) { deliveries, err := n.pg.ClaimRealtimeDeliveries(ctx, ch.ID, allow, notifyLease) @@ -422,14 +448,23 @@ func (n *Notifier) itemFor(ctx context.Context, snap notify.Snapshot, baseURL st return item, nil } -// takeTokens 从渠道令牌桶里取出本轮可投递的条数。 +// takeTokens 从渠道令牌桶里取走**最多 want 个**令牌，返回实际取到的数量。 // -// 桶容量为该渠道每分钟上限：积压时最多一次性冲这么多（平台规则允许）， -// 之后按恒定速率补充。ratePerMin<=0 表示不限流，返回一个有限但足够大的值， -// 防止单轮循环被无限积压拖住。 -func (n *Notifier) takeTokens(channelID int64, ratePerMin int, now time.Time) int { +// 一个令牌 = 一条消息（一次 HTTP 请求）。实时模式下调用方要几条就传几条； +// 汇总模式下一整批漏洞只发一条消息，传 1。 +// +// 桶容量为该渠道每分钟上限，按恒定速率补充。ratePerMin<=0 表示不限流， +// 返回一个有限但足够大的值，防止单轮循环被无限积压拖住。 +// +// want 这个上限是必需的：没有它就只能把桶整个抽空，而调用方自己还有每轮上限， +// 多取的令牌既用不上、又在下次补充前凭空消失——攒下来的突发容量永远不可达， +// 连「这一轮没有任何待发投递」都会照扣一笔。 +func (n *Notifier) takeTokens(channelID int64, ratePerMin, want int, now time.Time) int { + if want <= 0 { + return 0 + } if ratePerMin <= 0 { - return notifyUnlimitedBurstPerTick + return min(want, notifyUnlimitedBurstPerTick) } n.mu.Lock() defer n.mu.Unlock() @@ -446,7 +481,7 @@ func (n *Notifier) takeTokens(channelID int64, ratePerMin int, now time.Time) in // 加一个极小 epsilon 再取整：令牌数是浮点累加出来的，分两次补满时 // 0.5 + 0.5 可能得到 0.9999999999，直接 int() 会被截成 0—— // 数学上已满的桶却取不出令牌。1e-9 远小于一个令牌，不会放过真正的欠额。 - take := int(b.tokens + 1e-9) + take := min(int(b.tokens+1e-9), want) if take <= 0 { return 0 } diff --git a/server/notify_api_test.go b/server/notify_api_test.go index d389f01..fe3be74 100644 --- a/server/notify_api_test.go +++ b/server/notify_api_test.go @@ -856,28 +856,84 @@ func TestNotifyRateLimitDoesNotConsumeRetryBudget(t *testing.T) { n := &Notifier{buckets: map[int64]*notifyBucket{}} now := time.Now() // 每分钟 1 条：满桶时最多 1 条。 - if got := n.takeTokens(1, 1, now); got != 1 { + if got := n.takeTokens(1, 1, notifyMaxSendsPerChannelPerTick, now); got != 1 { t.Fatalf("满桶时每分钟 1 条应取 1 个令牌，得到 %d", got) } - if got := n.takeTokens(1, 1, now.Add(time.Millisecond)); got != 0 { + if got := n.takeTokens(1, 1, notifyMaxSendsPerChannelPerTick, now.Add(time.Millisecond)); got != 0 { t.Fatalf("令牌耗尽后应立即返回 0，得到 %d", got) } - if got := n.takeTokens(1, 1, now.Add(30*time.Second)); got != 0 { + if got := n.takeTokens(1, 1, notifyMaxSendsPerChannelPerTick, now.Add(30*time.Second)); got != 0 { t.Fatalf("半程不应补满一个令牌，得到 %d", got) } - if got := n.takeTokens(1, 1, now.Add(time.Minute)); got != 1 { + if got := n.takeTokens(1, 1, notifyMaxSendsPerChannelPerTick, now.Add(time.Minute)); got != 1 { t.Fatalf("满一个周期应补回 1 个令牌，得到 %d", got) } // 不限流渠道走有限上限，避免单轮被无限积压拖住。 - if got := n.takeTokens(2, 0, now); got != notifyUnlimitedBurstPerTick { + if got := n.takeTokens(2, 0, notifyUnlimitedBurstPerTick+10, now); got != notifyUnlimitedBurstPerTick { t.Fatalf("不限流应返回每轮上限 %d，得到 %d", notifyUnlimitedBurstPerTick, got) } // 渠道之间的令牌桶互相独立。 - if got := n.takeTokens(1, 1, now.Add(time.Millisecond)); got != 0 { + if got := n.takeTokens(1, 1, notifyMaxSendsPerChannelPerTick, now.Add(time.Millisecond)); got != 0 { t.Fatalf("渠道 1 的桶应仍然为空，得到 %d", got) } } +// TestNotifyTakeTokensKeepsUnusedTokens 锁住「只取 want 个」的语义。 +// +// 曾经的实现把桶整个抽空后才由调用方截断，于是 rate=100/min 的渠道攒满桶、 +// 一轮只用 5 条，剩下 95 个令牌直接丢弃；渠道这一轮没有待发投递时同样照扣。 +// 结果是注释声称的「积压时可以一次性冲 rate_per_min 条」在任何情况下都做不到。 +func TestNotifyTakeTokensKeepsUnusedTokens(t *testing.T) { + n := &Notifier{buckets: map[int64]*notifyBucket{}} + now := time.Now() + // 桶初始为满（100），本轮只要 5 个。 + if got := n.takeTokens(1, 100, 5, now); got != 5 { + t.Fatalf("want=5 时应恰好取 5 个令牌，得到 %d", got) + } + // 关键断言：余下的 95 个必须还在桶里，而不是被抽空丢弃。 + // 不推进时间，确保取到的只可能来自存量而非补充。 + if got := n.takeTokens(1, 100, 95, now); got != 95 { + t.Fatalf("剩余令牌应仍可取用（期望 95），得到 %d——桶被整轮抽空了", got) + } + if got := n.takeTokens(1, 100, 1, now); got != 0 { + t.Fatalf("桶已取尽，应返回 0，得到 %d", got) + } + // want<=0 不应扣减任何令牌（空轮不收费）。 + n2 := &Notifier{buckets: map[int64]*notifyBucket{}} + if got := n2.takeTokens(1, 20, 0, now); got != 0 { + t.Fatalf("want=0 应返回 0，得到 %d", got) + } + if got := n2.takeTokens(1, 20, 20, now); got != 20 { + t.Fatalf("want=0 的那次不该消耗令牌，应仍可取满 20，得到 %d", got) + } +} + +// TestDigestTickPlanDecouplesBatchSizeFromSendBudget 钉住汇总模式的两个量纲。 +// +// 汇总批次的大小一旦跟每轮请求预算挂上，rate_per_min=20 的渠道就只能在每个 +// 3 秒 tick 里补到 1 个令牌，于是每条汇总消息只装 1 个漏洞——功能上等于没有 +// 汇总，而消息头部还写着「近 30 分钟新增 1 个漏洞」。这个退化不会报错， +// 现有的端到端用例也看不出来（它们手动给 stepDigest 传一个够大的 limit， +// 绕过了 step 里的额度计算），所以在这里直接断言决策本身。 +func TestDigestTickPlanDecouplesBatchSizeFromSendBudget(t *testing.T) { + tokens, claimLimit := digestTickPlan() + // 一批 = 一条消息 = 一次请求 = 一个令牌。令牌的单位是消息，不是漏洞。 + if tokens != 1 { + t.Fatalf("汇总一批只发一条消息，应恰好消耗 1 个令牌，得到 %d", tokens) + } + if claimLimit != db.MaxDigestBatchSize { + t.Fatalf("汇总批次大小应为内存上界 db.MaxDigestBatchSize=%d，得到 %d", + db.MaxDigestBatchSize, claimLimit) + } + // 关键关系：批次大小必须远大于每轮请求预算。两者一旦同量级， + // 说明又把「发几条消息」和「一批装几条漏洞」混成了一个数。 + if claimLimit <= notifyMaxSendsPerChannelPerTick { + t.Fatalf("汇总批次大小 %d 不应受每轮请求预算 %d 约束——"+ + "请求预算是由租约倒推的「发几次请求」，与「一批装几条漏洞」是两个量纲", + claimLimit, notifyMaxSendsPerChannelPerTick) + } +} + // TestNotifyTickBudgetFitsWithinLease 是又一条防漂移断言。 // // 单渠道每轮的投递条数上限（notifyMaxSendsPerChannelPerTick）是从租约时长倒推的：