ntfy/server/topic.go
2023-02-21 20:16:03 -06:00

158 lines
4.4 KiB
Go

package server
import (
"math/rand"
"sync"
"time"
"heckel.io/ntfy/log"
)
// topic represents a channel to which subscribers can subscribe, and publishers
// can publish a message
type topic struct {
ID string
subscribers map[int]*topicSubscriber
lastVisitor *visitor
lastVisitorExpires time.Time
mu sync.Mutex
}
type topicSubscriber struct {
subscriber subscriber
visitor *visitor // User ID associated with this subscription, may be empty
cancel func()
subscriberRateLimit bool
}
// subscriber is a function that is called for every new message on a topic
type subscriber func(v *visitor, msg *message) error
// newTopic creates a new topic
func newTopic(id string) *topic {
return &topic{
ID: id,
subscribers: make(map[int]*topicSubscriber),
}
}
// Subscribe subscribes to this topic
func (t *topic) Subscribe(s subscriber, visitor *visitor, cancel func(), subscriberRateLimit bool) int {
t.mu.Lock()
defer t.mu.Unlock()
subscriberID := rand.Int()
t.subscribers[subscriberID] = &topicSubscriber{
visitor: visitor, // May be empty
subscriber: s,
cancel: cancel,
subscriberRateLimit: subscriberRateLimit,
}
// if no subscriber is already handling the rate limit
if t.lastVisitor == nil && subscriberRateLimit {
t.lastVisitor = visitor
t.lastVisitorExpires = time.Time{}
}
return subscriberID
}
func (t *topic) Stale() bool {
t.mu.Lock()
defer t.mu.Unlock()
// if Time is initialized (not the zero value) and the expiry time has passed
if !t.lastVisitorExpires.IsZero() && t.lastVisitorExpires.Before(time.Now()) {
t.lastVisitor = nil
}
return len(t.subscribers) == 0 && t.lastVisitor == nil
}
func (t *topic) Billee() *visitor {
t.mu.Lock()
defer t.mu.Unlock()
return t.lastVisitor
}
// Unsubscribe removes the subscription from the list of subscribers
func (t *topic) Unsubscribe(id int) {
t.mu.Lock()
defer t.mu.Unlock()
deletingSub := t.subscribers[id]
delete(t.subscribers, id)
// look for an active subscriber (in random order) that wants to handle the rate limit
for _, v := range t.subscribers {
if v.subscriberRateLimit {
t.lastVisitor = v.visitor
t.lastVisitorExpires = time.Time{}
return
}
}
// if no active subscriber is found, count it towards the leaving subscriber
if deletingSub.subscriberRateLimit {
t.lastVisitor = deletingSub.visitor
t.lastVisitorExpires = time.Now().Add(subscriberBilledValidity)
}
}
// Publish asynchronously publishes to all subscribers
func (t *topic) Publish(v *visitor, m *message) error {
go func() {
// We want to lock the topic as short as possible, so we make a shallow copy of the
// subscribers map here. Actually sending out the messages then doesn't have to lock.
subscribers := t.subscribersCopy()
if len(subscribers) > 0 {
logvm(v, m).Tag(tagPublish).Debug("Forwarding to %d subscriber(s)", len(subscribers))
for _, s := range subscribers {
// We call the subscriber functions in their own Go routines because they are blocking, and
// we don't want individual slow subscribers to be able to block others.
go func(s subscriber) {
if err := s(v, m); err != nil {
logvm(v, m).Tag(tagPublish).Err(err).Warn("Error forwarding to subscriber")
}
}(s.subscriber)
}
} else {
logvm(v, m).Tag(tagPublish).Trace("No stream or WebSocket subscribers, not forwarding")
}
}()
return nil
}
// SubscribersCount returns the number of subscribers to this topic
func (t *topic) SubscribersCount() int {
t.mu.Lock()
defer t.mu.Unlock()
return len(t.subscribers)
}
// CancelSubscribers calls the cancel function for all subscribers, forcing
func (t *topic) CancelSubscribers(exceptUserID string) {
t.mu.Lock()
defer t.mu.Unlock()
for _, s := range t.subscribers {
if s.visitor.MaybeUserID() != exceptUserID {
// TODO: Shouldn't this log the IP for anonymous visitors? It was s.userID before my change.
log.Tag(tagSubscribe).Field("topic", t.ID).Debug("Canceling subscriber %s", s.visitor.MaybeUserID())
s.cancel()
}
}
}
// subscribersCopy returns a shallow copy of the subscribers map
func (t *topic) subscribersCopy() map[int]*topicSubscriber {
t.mu.Lock()
defer t.mu.Unlock()
subscribers := make(map[int]*topicSubscriber)
for k, sub := range t.subscribers {
subscribers[k] = &topicSubscriber{
visitor: sub.visitor,
subscriber: sub.subscriber,
cancel: sub.cancel,
}
}
return subscribers
}