fix: 上行分发前统一限速,send 不再单独扣桶

This commit is contained in:
Nixevol
2026-09-30 16:21:05 +08:00
parent f139b9ed9b
commit aa24e974af
7 changed files with 155 additions and 13 deletions
+8
View File
@@ -57,3 +57,11 @@ func (r *rateLimiter) allow(endpointID string, now time.Time) bool {
b.tokens--
return true
}
// AllowRequest 消耗该端 1 个请求令牌;允许则 true。rps<=0 时不限速。
func (a *App) AllowRequest(endpointID string) bool {
if a == nil {
return true
}
return a.rates.allow(endpointID, a.now())
}
+22
View File
@@ -0,0 +1,22 @@
package message
import (
"testing"
)
func TestAllowRequestBurstAndAckExemptBucket(t *testing.T) {
t.Parallel()
lim := defaultTestLimits()
lim.RequestsPerSecond = 50
lim.RequestBurst = 100
app, _ := openTestApp(t, lim)
allowed := 0
for i := 0; i < 150; i++ {
if app.AllowRequest("alice") {
allowed++
}
}
if allowed != 100 {
t.Fatalf("allowed=%d want 100 (burst)", allowed)
}
}
-3
View File
@@ -21,9 +21,6 @@ func (a *App) Submit(ctx context.Context, senderID string, conn port.ConnInfo, r
return SubmitResult{}, errCode(protocol.CodeBadRequest, "invalid sender")
}
now := a.now()
if !a.rates.allow(senderID, now) {
return SubmitResult{}, errCode(protocol.CodeRateLimited, "request rate exceeded")
}
if err := req.Validate(a.protocolLimits()); err != nil {
return SubmitResult{}, err
+6 -7
View File
@@ -371,16 +371,15 @@ SELECT kind FROM talk_grants WHERE sender_id=? AND target_id=?`, "bob", "alice")
app, db := openTestApp(t, lim)
insertEndpoint(t, db, "alice", "", 1, 0)
insertEndpoint(t, db, "bob", "", 1, 0)
if !app.AllowRequest("alice") || !app.AllowRequest("alice") {
t.Fatal("burst should allow first two")
}
if app.AllowRequest("alice") {
t.Fatal("third request should be rate limited")
}
ctx := context.Background()
if _, err := app.Submit(ctx, "alice", port.ConnInfo{}, baseSend("r1", "bob")); err != nil {
t.Fatal(err)
}
if _, err := app.Submit(ctx, "alice", port.ConnInfo{}, baseSend("r2", "bob")); err != nil {
t.Fatal(err)
}
_, err := app.Submit(ctx, "alice", port.ConnInfo{}, baseSend("r3", "bob"))
if protoCode(err) != protocol.CodeRateLimited {
t.Fatalf("want rate_limited got %v", err)
}
})
}