Keep a topic's rate visitor alive by the topic, for any stream with its key

A stream's keepalives refreshed only its own visitor, so a topic that held
the visitor of a newer stream or a poll from the same key lost it after a
day although a stream was still open. The topic now keeps when a
subscriber with its rate visitor's key was last seen, and a stream takes a
stale topic back.

Claude-Session: https://claude.ai/code/session_01DeXuhSjq5Gta5qPJgSceFe

server/server.go

@@ -1635,7 +1635,8 @@ func (s *Server) handleSubscribeHTTP(w http.ResponseWriter, r *http.Request, v *
 		}
 		return nil
 	}
-	if err := s.maybeSetRateVisitors(r, v, topics, poll); err != nil {
+	rateTopics, err := s.maybeSetRateVisitors(r, v, topics)
+	if err != nil {
 		return err
 	}
 	w.Header().Set("Access-Control-Allow-Origin", s.config.AccessControlAllowOrigin) // CORS, allow cross-origin requests
@@ -1690,6 +1691,9 @@ func (s *Server) handleSubscribeHTTP(w http.ResponseWriter, r *http.Request, v *
 			for _, t := range topics {
 				t.Keepalive()
 			}
+			for _, t := range rateTopics {
+				t.KeepRateVisitor(v)
+			}
 			if err := sub(v, model.NewKeepaliveMessage(topicsStr)); err != nil { // Send keepalive message
 				return err
 			}
@@ -1729,6 +1733,10 @@ func (s *Server) handleSubscribeWS(w http.ResponseWriter, r *http.Request, v *vi
 		return err
 	}
 	defer conn.Close()
+	rateTopics, err := s.maybeSetRateVisitors(r, v, topics) // Before the keepalive loop below reads them
+	if err != nil {
+		return err
+	}
 
 	// Subscription connections can be canceled externally, see topic.CancelSubscribersExceptUser
 	cancelCtx, cancel := context.WithCancel(context.Background())
@@ -1782,6 +1790,9 @@ func (s *Server) handleSubscribeWS(w http.ResponseWriter, r *http.Request, v *vi
 				for _, t := range topics {
 					t.Keepalive()
 				}
+				for _, t := range rateTopics {
+					t.KeepRateVisitor(v)
+				}
 				if err := ping(); err != nil {
 					return err
 				}
@@ -1799,9 +1810,6 @@ func (s *Server) handleSubscribeWS(w http.ResponseWriter, r *http.Request, v *vi
 		}
 		return conn.WriteJSON(msg)
 	}
-	if err := s.maybeSetRateVisitors(r, v, topics, poll); err != nil {
-		return err
-	}
 	w.Header().Set("Access-Control-Allow-Origin", s.config.AccessControlAllowOrigin) // CORS, allow cross-origin requests
 	if poll {
 		for _, t := range topics {
@@ -1865,11 +1873,12 @@ func parseSubscribeParams(r *http.Request) (poll bool, since model.SinceMarker,
 // - or the topic is reserved, and v.user is the owner
 // - or the topic is not reserved, and v.user has write access
 //
-// This only applies to UnifiedPush topics ("up...").
-func (s *Server) maybeSetRateVisitors(r *http.Request, v *visitor, topics []*topic, poll bool) error {
+// This only applies to UnifiedPush topics ("up..."). It returns the topics it set the rate visitor on, for
+// the stream's keepalives (see topic.KeepRateVisitor).
+func (s *Server) maybeSetRateVisitors(r *http.Request, v *visitor, topics []*topic) ([]*topic, error) {
 	// Bail out if not enabled
 	if !s.config.VisitorSubscriberRateLimiting {
-		return nil
+		return nil, nil
 	}
 
 	// Make a list of topics that we'll actually set the RateVisitor on
@@ -1880,12 +1889,13 @@ func (s *Server) maybeSetRateVisitors(r *http.Request, v *visitor, topics []*top
 		}
 	}
 	if len(eligibleRateTopics) == 0 {
-		return nil
+		return nil, nil
 	}
 
 	// If access controls are turned off, v has access to everything, and we can set the rate visitor
 	if s.userManager == nil {
-		return s.setRateVisitors(r, v, eligibleRateTopics, poll)
+		s.setRateVisitors(r, v, eligibleRateTopics)
+		return eligibleRateTopics, nil
 	}
 
 	// If access controls are enabled, only set rate visitor if
@@ -1898,7 +1908,7 @@ func (s *Server) maybeSetRateVisitors(r *http.Request, v *visitor, topics []*top
 		}
 		ownerUserID, err := s.userManager.ReservationOwner(t.ID)
 		if err != nil {
-			return err
+			return nil, err
 		}
 		if ownerUserID == "" {
 			if err := s.userManager.Authorize(v.User(), t.ID, user.PermissionWrite); err == nil {
@@ -1908,18 +1918,18 @@ func (s *Server) maybeSetRateVisitors(r *http.Request, v *visitor, topics []*top
 			writableRateTopics = append(writableRateTopics, t)
 		}
 	}
-	return s.setRateVisitors(r, v, writableRateTopics, poll)
+	s.setRateVisitors(r, v, writableRateTopics)
+	return writableRateTopics, nil
 }
 
-func (s *Server) setRateVisitors(r *http.Request, v *visitor, rateTopics []*topic, poll bool) error {
+func (s *Server) setRateVisitors(r *http.Request, v *visitor, rateTopics []*topic) {
 	for _, t := range rateTopics {
 		logvr(v, r).
 			Tag(tagSubscribe).
 			With(t).
 			Debug("Setting visitor as rate visitor for topic %s", t.ID)
-		t.SetRateVisitor(v, poll)
+		t.SetRateVisitor(v)
 	}
-	return nil
 }
 
 // sendOldMessages selects old messages from the messageCache and calls sub for each of them. It uses since as the

server/server_test.go

@@ -3637,9 +3637,9 @@ func TestServer_SubscriberRateLimiting_VisitorExpiration(t *testing.T) {
 		require.Equal(t, int64(1), s.topics["upAAAAAAAAAAAA"].rateVisitor.Stats().Messages)
 
 		// Expire visitor: its subscriber has been gone for a day, and its key is pruned after a day without use
-		s.topics["upAAAAAAAAAAAA"].rateVisitor.seen.Store(time.Now().Add(-25 * time.Hour).UnixNano())
+		ageRateVisitor(s.topics["upAAAAAAAAAAAA"], 25*time.Hour)
 		s.limiter.Prune(-time.Hour) // Everything is idle for longer than -1h
-		require.True(t, s.topics["upAAAAAAAAAAAA"].rateVisitor.Stale())
+		require.True(t, s.topics["upAAAAAAAAAAAA"].rateVisitorStaleNoLock())
 
 		// Publish message again, observe that rateVisitor is not used anymore and is reset
 		response = request(t, s, "POST", "/upAAAAAAAAAAAA", "some message", nil)
@@ -3651,9 +3651,9 @@ func TestServer_SubscriberRateLimiting_VisitorExpiration(t *testing.T) {
 }
 
 func TestServer_SubscriberRateLimiting_PollKeepsStreamAsRateVisitor(t *testing.T) {
-	// A poll from the subscriber of an open stream must not take the stream's place as the topic's
-	// rate visitor: only the stream's keepalives keep it alive, so UnifiedPush publishes would fail
-	// a day after the poll while the stream is still open
+	// A poll from the subscriber of an open stream sets its own visitor on the topic; the stream's
+	// keepalives must keep that one alive too, or UnifiedPush publishes fail a day after the poll
+	// while the stream is still open
 	c := newTestConfig(t, "")
 	c.VisitorSubscriberRateLimiting = true
 	c.KeepaliveInterval = 100 * time.Millisecond
@@ -3701,16 +3701,92 @@ func TestServer_SubscriberRateLimiting_PollKeepsPollingSubscriberAlive(t *testin
 	rr = request(t, s, "GET", "/upAAAAAAAAAAAA/json?poll=1", "", nil, subscriberFn)
 	require.Equal(t, 200, rr.Code)
 	to.mu.Lock()
-	seen := time.Unix(0, to.rateVisitor.seen.Load())
+	seen := to.rateSeen
 	to.mu.Unlock()
 	require.WithinDuration(t, time.Now(), seen, time.Minute)
 }
 
+func TestServer_SubscriberRateLimiting_OverlappingStreamsKeepTopicAlive(t *testing.T) {
+	// Two streams from the same subscriber overlap (e.g. a reconnect), and the newer one closes: the
+	// older one's keepalives must keep the topic's rate visitor alive, though the topic holds the newer one
+	c := newTestConfig(t, "")
+	c.VisitorSubscriberRateLimiting = true
+	c.KeepaliveInterval = 100 * time.Millisecond
+	s := newTestServer(t, c)
+	subscriberFn := func(r *http.Request) {
+		r.RemoteAddr = "1.2.3.4:1234"
+	}
+
+	// Stream A, then stream B, then B closes
+	ctxA, cancelA := context.WithCancel(context.Background())
+	defer cancelA()
+	go func() {
+		request(t, s, "GET", "/upAAAAAAAAAAAA/json", "", nil, subscriberFn, func(r *http.Request) {
+			*r = *r.WithContext(ctxA)
+		})
+	}()
+	waitFor(t, func() bool {
+		to := lockedTopic(s, "upAAAAAAAAAAAA")
+		return to != nil && to.SubscribersCount() == 1
+	})
+	ctxB, cancelB := context.WithCancel(context.Background())
+	doneB := make(chan struct{})
+	go func() {
+		defer close(doneB)
+		request(t, s, "GET", "/upAAAAAAAAAAAA/json", "", nil, subscriberFn, func(r *http.Request) {
+			*r = *r.WithContext(ctxB)
+		})
+	}()
+	waitFor(t, func() bool {
+		return lockedTopic(s, "upAAAAAAAAAAAA").SubscribersCount() == 2
+	})
+	cancelB()
+	<-doneB
+
+	// A day later, stream A is still there
+	ageRateVisitor(lockedTopic(s, "upAAAAAAAAAAAA"), 25*time.Hour)
+	time.Sleep(3 * c.KeepaliveInterval)
+	response := request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil)
+	require.Equal(t, 200, response.Code)
+}
+
+func TestServer_SubscriberRateLimiting_StreamReclaimsTopicFromGonePoller(t *testing.T) {
+	// A poll from another key (e.g. the same device over IPv4) takes the topic from an open stream
+	// (IPv6). Once the poller is gone for a day, the stream's keepalives must take the topic back.
+	c := newTestConfig(t, "")
+	c.VisitorSubscriberRateLimiting = true
+	c.KeepaliveInterval = 100 * time.Millisecond
+	s := newTestServer(t, c)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go func() {
+		request(t, s, "GET", "/upAAAAAAAAAAAA/json", "", nil, func(r *http.Request) {
+			r.RemoteAddr = "[2001:db8::1]:1234"
+			*r = *r.WithContext(ctx)
+		})
+	}()
+	waitFor(t, func() bool {
+		to := lockedTopic(s, "upAAAAAAAAAAAA")
+		return to != nil && to.SubscribersCount() == 1
+	})
+	rr := request(t, s, "GET", "/upAAAAAAAAAAAA/json?poll=1", "", nil, func(r *http.Request) {
+		r.RemoteAddr = "1.2.3.4:1234"
+	})
+	require.Equal(t, 200, rr.Code)
+
+	// A day later, only the stream is still there; the publish is billed to the stream's key
+	ageRateVisitor(lockedTopic(s, "upAAAAAAAAAAAA"), 25*time.Hour)
+	time.Sleep(3 * c.KeepaliveInterval)
+	response := request(t, s, "PUT", "/upAAAAAAAAAAAA?up=1", "some message", nil)
+	require.Equal(t, 200, response.Code)
+	require.Equal(t, int64(1), s.visitor(netip.MustParseAddr("2001:db8::1"), nil).Stats().Messages)
+}
+
 // ageRateVisitor makes the topic's rate visitor look as if it was last seen d ago
 func ageRateVisitor(t *topic, d time.Duration) {
 	t.mu.Lock()
 	defer t.mu.Unlock()
-	t.rateVisitor.seen.Store(time.Now().Add(-d).UnixNano())
+	t.rateSeen = time.Now().Add(-d)
 }
 
 // lockedTopic returns the topic with the given ID, or nil, reading the topics map under the server's lock

server/topic.go

@@ -30,6 +30,7 @@ type topic struct {
 	ID          string
 	subscribers map[int]*topicSubscriber
 	rateVisitor *visitor
+	rateSeen    time.Time // When a subscriber with the rate visitor's key was last seen (see KeepRateVisitor)
 	lastAccess  time.Time
 	delivered   map[string]time.Time // When recent messages went out to the subscribers, by message ID
 	mu          sync.RWMutex
@@ -109,7 +110,7 @@ func (t *topic) ExpireDelivered() {
 func (t *topic) Stale() bool {
 	t.mu.Lock()
 	defer t.mu.Unlock()
-	if t.rateVisitor != nil && !t.rateVisitor.Stale() {
+	if !t.rateVisitorStaleNoLock() {
 		return false
 	}
 	return len(t.subscribers) == 0 && time.Since(t.lastAccess) > topicExpungeAfter
@@ -121,26 +122,45 @@ func (t *topic) LastAccess() time.Time {
 	return t.lastAccess
 }
 
-func (t *topic) SetRateVisitor(v *visitor, poll bool) {
+// SetRateVisitor makes the subscriber's visitor the topic's rate visitor; the latest subscriber wins
+func (t *topic) SetRateVisitor(v *visitor) {
 	t.mu.Lock()
 	defer t.mu.Unlock()
-	if poll && t.rateVisitor != nil && t.rateVisitor.key == v.key {
-		t.rateVisitor.Keepalive() // Replacing it would leave an open stream keeping alive a visitor the topic no longer has
-	} else {
-		t.rateVisitor = v
-	}
+	t.rateVisitor = v
+	t.rateSeen = time.Now()
 	t.lastAccess = time.Now()
 }
 
+// KeepRateVisitor is called on the keepalives of a stream whose visitor may be the topic's rate visitor.
+// A stream with the rate visitor's key keeps it alive, whichever of its subscribers set it. A stream
+// with another key takes the topic back once the rate visitor is stale.
+func (t *topic) KeepRateVisitor(v *visitor) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	if t.rateVisitor != nil && t.rateVisitor.key == v.key {
+		t.rateSeen = time.Now()
+	} else if t.rateVisitorStaleNoLock() {
+		t.rateVisitor = v
+		t.rateSeen = time.Now()
+	}
+}
+
+// RateVisitor returns the topic's rate visitor, or nil if there is none or it is stale
 func (t *topic) RateVisitor() *visitor {
 	t.mu.Lock()
 	defer t.mu.Unlock()
-	if t.rateVisitor != nil && t.rateVisitor.Stale() {
+	if t.rateVisitorStaleNoLock() {
 		t.rateVisitor = nil
 	}
 	return t.rateVisitor
 }
 
+// rateVisitorStaleNoLock is true if there is no rate visitor, or no subscriber with its key has been
+// seen for visitorExpungeAfter
+func (t *topic) rateVisitorStaleNoLock() bool {
+	return t.rateVisitor == nil || time.Since(t.rateSeen) > visitorExpungeAfter
+}
+
 // Unsubscribe removes the subscription from the list of subscribers
 func (t *topic) Unsubscribe(id int) {
 	t.mu.Lock()

server/visitor.go

@@ -2,7 +2,6 @@ package server
 
 import (
 	"net/netip"
-	"sync/atomic"
 	"time"
 
 	"heckel.io/ntfy/v2/limit"
@@ -30,11 +29,10 @@ const (
 // key, so two visitors with the same key share one set of buckets and counts.
 type visitor struct {
 	env    *visitorEnv
-	key    string       // The limiter key: "ip:<prefix>" or "user:<id>", see vis.Key
-	ip     netip.Addr   // Visitor IP address
-	user   *user.User   // Only set if authenticated user, otherwise nil
-	limits *vis.Limits  // Resolved from the config or the user's tier
-	seen   atomic.Int64 // Unix nanos of the last keepalive: the liveness of a subscriber kept as a topic's rate visitor (see Stale)
+	key    string      // The limiter key: "ip:<prefix>" or "user:<id>", see vis.Key
+	ip     netip.Addr  // Visitor IP address
+	user   *user.User  // Only set if authenticated user, otherwise nil
+	limits *vis.Limits // Resolved from the config or the user's tier
 }
 
 // visitorEnv is what every visitor of a server shares; it is built once, so building a visitor per
@@ -103,7 +101,6 @@ func newVisitor(env *visitorEnv, ip netip.Addr, u *user.User) *visitor {
 		env.limiter.Seed(v.key, usage)
 	}
 	env.limiter.Touch(v.key) // The key is seen now, as the visitor's request is (after Seed, which only fills unknown keys)
-	v.seen.Store(time.Now().UnixNano())
 	return v
 }
 
@@ -112,7 +109,6 @@ func (v *visitor) Context() log.Context {
 	fields := log.Context{
 		"visitor_id":                 v.key,
 		"visitor_ip":                 v.ip.String(),
-		"visitor_seen":               util.FormatTime(time.Unix(0, v.seen.Load())),
 		"visitor_messages":           info.Stats.Messages,
 		"visitor_messages_limit":     info.Limits.MessageLimit,
 		"visitor_messages_remaining": info.Stats.MessagesRemaining,
@@ -288,10 +284,8 @@ func (v *visitor) RemoveSubscription() {
 	v.env.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.
+// Keepalive marks the visitor's key as seen, so the limiter keeps the key of a long-lived subscription
 func (v *visitor) Keepalive() {
-	v.seen.Store(time.Now().UnixNano())
 	v.env.limiter.Touch(v.key)
 }
 
@@ -304,13 +298,6 @@ func (v *visitor) BandwidthLimiter() util.Limiter {
 	return v.env.limiter.Limiter(v.key, v.bandwidthRule())
 }
 
-// Stale is true when the visitor has not been kept alive for visitorExpungeAfter. It reads the
-// visitor's own keepalive, not the limiter key's last use: publishers charging a topic's rate
-// visitor must not keep a vanished subscriber's limits on the topic.
-func (v *visitor) Stale() bool {
-	return time.Since(time.Unix(0, v.seen.Load())) > visitorExpungeAfter
-}
-
 func (v *visitor) Stats() *user.Stats {
 	usage := v.env.limiter.Usage(v.key)
 	return &user.Stats{

server/visitor_test.go

@@ -345,16 +345,9 @@ func TestVisitor_AccountLimiter_NotDeniedByConcurrentLogin(t *testing.T) {
 	require.True(t, vip.AccountActionAllowed())
 }
 
-func TestVisitor_NewVisitorIsNotStale(t *testing.T) {
-	// A visitor is seen when it is built for a request, like it was when visitors were shared
-	s := newTestServer(t, newTestConfig(t, ""))
-	v := s.visitor(netip.MustParseAddr("9.9.9.9"), nil)
-	require.False(t, v.Stale())
-}
-
 func TestVisitor_RateVisitorKeptForExemptSubscriber(t *testing.T) {
 	// A subscriber from a request-limit-exempt prefix spends nothing before it becomes a topic's
-	// rate visitor; the topic must still keep it. Fails if a new visitor counts as stale.
+	// rate visitor; the topic must still keep it. Fails if a new rate visitor counts as stale.
 	c := newTestConfig(t, "")
 	c.VisitorSubscriberRateLimiting = true
 	c.VisitorRequestExemptPrefixes = []netip.Prefix{netip.MustParsePrefix("9.9.9.0/24")}
@@ -364,7 +357,7 @@ func TestVisitor_RateVisitorKeptForExemptSubscriber(t *testing.T) {
 	v := s.visitor(netip.MustParseAddr("9.9.9.9"), nil)
 	topics, err := s.topicsFromIDs(v, "mytopic")
 	require.Nil(t, err)
-	require.Nil(t, s.setRateVisitors(httptest.NewRequest("GET", "/mytopic/json", nil), v, topics, false))
+	s.setRateVisitors(httptest.NewRequest("GET", "/mytopic/json", nil), v, topics)
 	require.NotNil(t, topics[0].RateVisitor())
 }
 
@@ -738,11 +731,11 @@ func TestVisitor_RateVisitorLivenessIsTheSubscribers(t *testing.T) {
 	v := s.visitor(netip.MustParseAddr("9.9.9.9"), nil)
 	topics, err := s.topicsFromIDs(v, "mytopic")
 	require.Nil(t, err)
-	require.Nil(t, s.setRateVisitors(httptest.NewRequest("GET", "/mytopic/json", nil), v, topics, false))
+	s.setRateVisitors(httptest.NewRequest("GET", "/mytopic/json", nil), v, topics)
 	s.limiter.Prune(0) // The key is gone, the subscriber is not
 	require.NotNil(t, topics[0].RateVisitor())
-	v.seen.Store(time.Now().Add(-2 * visitorExpungeAfter).UnixNano()) // The subscriber is gone
-	require.True(t, s.limiter.Allow(v.key, v.messageRule(), 1).OK)    // A publisher charges the key
+	ageRateVisitor(topics[0], 2*visitorExpungeAfter)               // The subscriber is gone
+	require.True(t, s.limiter.Allow(v.key, v.messageRule(), 1).OK) // A publisher charges the key
 	require.Nil(t, topics[0].RateVisitor())
 }