diff --git a/server/server_test.go b/server/server_test.go index 62a28869..8997f6e5 100644 --- a/server/server_test.go +++ b/server/server_test.go @@ -6026,3 +6026,34 @@ func TestServer_Account_RemainingMatchesAdmissions(t *testing.T) { s.resetStats() } } + +func TestServer_SubscriptionLimit_ReleasedWhenStreamCloses(t *testing.T) { + c := newTestConfig(t, "") + c.VisitorSubscriptionLimit = 1 + s := newTestServer(t, c) + + // The open stream takes the only slot + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + request(t, s, "GET", "/mytopic/json", "", nil, func(r *http.Request) { + *r = *r.WithContext(ctx) + }) + close(done) + }() + waitFor(t, func() bool { + s.visitors.subscriptions.mu.Lock() + defer s.visitors.subscriptions.mu.Unlock() + return len(s.visitors.subscriptions.counts) == 1 + }) + rr := request(t, s, "GET", "/mytopic/json?poll=1", "", nil) + require.Equal(t, 429, rr.Code) + require.Equal(t, 42903, toHTTPError(t, rr.Body.String()).Code) + + // Closing the stream gives it back + cancel() + <-done + rr = request(t, s, "GET", "/mytopic/json?poll=1", "", nil) + require.Equal(t, 200, rr.Code) + require.Empty(t, s.visitors.subscriptions.counts) +} diff --git a/server/visitor.go b/server/visitor.go index 7487c550..8b9c1f6b 100644 --- a/server/visitor.go +++ b/server/visitor.go @@ -237,6 +237,12 @@ func (v *visitor) enqueueUserStats(delta *user.Stats) { v.userManager.EnqueueUserStats(v.user.ID, delta) } +// SubscriptionAllowed takes one of the visitor key's subscription slots on this node, if one is free. +// Every successful call must be paired with RemoveSubscription. +func (v *visitor) SubscriptionAllowed() bool { + return v.subscriptions.Acquire(v.key, int64(v.config.VisitorSubscriptionLimit)) +} + // TopicCreationAllowed returns true if the visitor is allowed to cause a new topic to be // inserted into the server's in-memory topic map. Returns true if no limiter is configured. func (v *visitor) TopicCreationAllowed() bool { @@ -275,6 +281,11 @@ func (v *visitor) BandwidthAllowed(bytes int64) bool { return v.limiter.Allow(v.key, v.bandwidthRule(), bytes).OK } +// RemoveSubscription gives back a slot taken by SubscriptionAllowed +func (v *visitor) RemoveSubscription() { + v.subscriptions.Release(v.key) +} + // Keepalive marks the visitor and its key as seen. The subscribe handlers call it on the visitor a // topic keeps as its rate visitor; nothing else does, so Stale is the subscriber's liveness. func (v *visitor) Keepalive() { diff --git a/server/visitor_subscriptions.go b/server/visitor_subscriptions.go index 0f4bf81b..09b801f2 100644 --- a/server/visitor_subscriptions.go +++ b/server/visitor_subscriptions.go @@ -37,14 +37,3 @@ func (g *subscriptionGauge) Release(key string) { delete(g.counts, key) } } - -// SubscriptionAllowed takes one of the visitor key's subscription slots on this node, if one is free. -// Every successful call must be paired with RemoveSubscription. -func (v *visitor) SubscriptionAllowed() bool { - return v.subscriptions.Acquire(v.key, int64(v.config.VisitorSubscriptionLimit)) -} - -// RemoveSubscription gives back a slot taken by SubscriptionAllowed -func (v *visitor) RemoveSubscription() { - v.subscriptions.Release(v.key) -} diff --git a/server/visitor_subscriptions_test.go b/server/visitor_subscriptions_test.go index 2f65b190..cc2bfefd 100644 --- a/server/visitor_subscriptions_test.go +++ b/server/visitor_subscriptions_test.go @@ -1,9 +1,6 @@ package server import ( - "context" - "net/http" - "net/netip" "sync" "sync/atomic" "testing" @@ -52,49 +49,3 @@ func TestSubscriptionGauge_ConcurrentNeverExceedsLimit(t *testing.T) { } require.Empty(t, g.counts) } - -func TestVisitor_SubscriptionLimitSharedPerKey(t *testing.T) { - c := newTestConfig(t, "") - c.VisitorSubscriptionLimit = 2 - s := newTestServer(t, c) - ip := netip.MustParseAddr("9.9.9.9") - v1 := s.visitor(ip, nil) - v2 := s.visitor(ip, nil) // A second request from the same IP shares the slots - require.True(t, v1.SubscriptionAllowed()) - require.True(t, v2.SubscriptionAllowed()) - require.False(t, v1.SubscriptionAllowed()) - require.True(t, s.visitor(netip.MustParseAddr("1.2.3.4"), nil).SubscriptionAllowed()) - v1.RemoveSubscription() - require.True(t, v2.SubscriptionAllowed()) -} - -func TestServer_SubscriptionLimit_ReleasedWhenStreamCloses(t *testing.T) { - c := newTestConfig(t, "") - c.VisitorSubscriptionLimit = 1 - s := newTestServer(t, c) - - // The open stream takes the only slot - ctx, cancel := context.WithCancel(context.Background()) - done := make(chan struct{}) - go func() { - request(t, s, "GET", "/mytopic/json", "", nil, func(r *http.Request) { - *r = *r.WithContext(ctx) - }) - close(done) - }() - waitFor(t, func() bool { - s.visitors.subscriptions.mu.Lock() - defer s.visitors.subscriptions.mu.Unlock() - return len(s.visitors.subscriptions.counts) == 1 - }) - rr := request(t, s, "GET", "/mytopic/json?poll=1", "", nil) - require.Equal(t, 429, rr.Code) - require.Equal(t, 42903, toHTTPError(t, rr.Body.String()).Code) - - // Closing the stream gives it back - cancel() - <-done - rr = request(t, s, "GET", "/mytopic/json?poll=1", "", nil) - require.Equal(t, 200, rr.Code) - require.Empty(t, s.visitors.subscriptions.counts) -} diff --git a/server/visitor_test.go b/server/visitor_test.go index 00673c13..559d888d 100644 --- a/server/visitor_test.go +++ b/server/visitor_test.go @@ -265,6 +265,21 @@ func TestVisitor_UserIsPerRequest(t *testing.T) { require.True(t, anon.SubscriptionAllowed()) } +func TestVisitor_SubscriptionLimitSharedPerKey(t *testing.T) { + c := newTestConfig(t, "") + c.VisitorSubscriptionLimit = 2 + s := newTestServer(t, c) + ip := netip.MustParseAddr("9.9.9.9") + v1 := s.visitor(ip, nil) + v2 := s.visitor(ip, nil) // A second request from the same IP shares the slots + require.True(t, v1.SubscriptionAllowed()) + require.True(t, v2.SubscriptionAllowed()) + require.False(t, v1.SubscriptionAllowed()) + require.True(t, s.visitor(netip.MustParseAddr("1.2.3.4"), nil).SubscriptionAllowed()) + v1.RemoveSubscription() + require.True(t, v2.SubscriptionAllowed()) +} + func TestVisitor_FirebasePenaltyIsSharedPerKey(t *testing.T) { s := newTestServer(t, newTestConfig(t, "")) ip := netip.MustParseAddr("9.9.9.9")