ntfy/server/topic.go

185 lines
4.7 KiB
Go
Raw Normal View History

2021-10-23 03:26:01 +02:00
package server
import (
2023-02-23 15:38:53 +01:00
"heckel.io/ntfy/log"
2023-03-04 04:22:07 +01:00
"heckel.io/ntfy/util"
2021-10-23 03:26:01 +02:00
"math/rand"
"sync"
2023-03-01 04:25:13 +01:00
"time"
)
const (
2023-03-04 04:22:07 +01:00
// topicExpungeAfter defines how long a topic is active before it is removed from memory.
// This must be larger than matrixRejectPushKeyForUnifiedPushTopicWithoutRateVisitorAfter to give
// time for more requests to come in, so that we can send a {"rejected":["<pushkey>"]} response back.
topicExpungeAfter = 16 * time.Hour
2023-02-23 04:26:43 +01:00
)
2021-10-24 04:49:50 +02:00
// topic represents a channel to which subscribers can subscribe, and publishers
// can publish a message
2021-10-23 03:26:01 +02:00
type topic struct {
2023-02-23 22:03:40 +01:00
ID string
subscribers map[int]*topicSubscriber
rateVisitor *visitor
2023-03-01 04:25:13 +01:00
lastAccess time.Time
2023-02-23 22:03:40 +01:00
mu sync.RWMutex
2021-10-23 03:26:01 +02:00
}
type topicSubscriber struct {
2023-02-24 02:46:53 +01:00
userID string // User ID associated with this subscription, may be empty
2023-02-23 03:33:18 +01:00
subscriber subscriber
cancel func()
}
2021-10-24 04:49:50 +02:00
// subscriber is a function that is called for every new message on a topic
2022-06-01 02:38:56 +02:00
type subscriber func(v *visitor, msg *message) error
2021-10-23 03:26:01 +02:00
// newTopic creates a new topic
2021-12-09 04:57:31 +01:00
func newTopic(id string) *topic {
2021-10-23 03:26:01 +02:00
return &topic{
2021-12-09 04:57:31 +01:00
ID: id,
subscribers: make(map[int]*topicSubscriber),
2023-03-01 04:25:13 +01:00
lastAccess: time.Now(),
2021-10-23 03:26:01 +02:00
}
}
2021-11-09 20:48:25 +01:00
// Subscribe subscribes to this topic
2023-02-24 02:46:53 +01:00
func (t *topic) Subscribe(s subscriber, userID string, cancel func()) int {
2021-10-23 03:26:01 +02:00
t.mu.Lock()
defer t.mu.Unlock()
subscriberID := rand.Int()
t.subscribers[subscriberID] = &topicSubscriber{
2023-02-24 02:46:53 +01:00
userID: userID, // May be empty
2023-02-23 03:33:18 +01:00
subscriber: s,
cancel: cancel,
}
2023-03-01 04:25:13 +01:00
t.lastAccess = time.Now()
2021-10-23 03:26:01 +02:00
return subscriberID
}
2023-02-14 20:58:13 +01:00
func (t *topic) Stale() bool {
2023-02-22 03:16:03 +01:00
t.mu.Lock()
defer t.mu.Unlock()
2023-02-24 02:46:53 +01:00
if t.rateVisitor != nil && !t.rateVisitor.Stale() {
return false
}
2023-03-04 04:22:07 +01:00
return len(t.subscribers) == 0 && time.Since(t.lastAccess) > topicExpungeAfter
}
func (t *topic) LastAccess() time.Time {
t.mu.RLock()
defer t.mu.RUnlock()
return t.lastAccess
2023-02-22 03:04:56 +01:00
}
2023-02-14 20:58:13 +01:00
2023-02-23 03:33:18 +01:00
func (t *topic) SetRateVisitor(v *visitor) {
t.mu.Lock()
defer t.mu.Unlock()
t.rateVisitor = v
2023-03-01 04:25:13 +01:00
t.lastAccess = time.Now()
2023-02-23 03:33:18 +01:00
}
func (t *topic) RateVisitor() *visitor {
t.mu.Lock()
defer t.mu.Unlock()
2023-02-24 02:46:53 +01:00
if t.rateVisitor != nil && t.rateVisitor.Stale() {
t.rateVisitor = nil
}
2023-02-23 03:00:56 +01:00
return t.rateVisitor
2023-02-14 20:58:13 +01:00
}
2021-11-09 20:48:25 +01:00
// Unsubscribe removes the subscription from the list of subscribers
2021-10-29 19:58:14 +02:00
func (t *topic) Unsubscribe(id int) {
2021-10-23 03:26:01 +02:00
t.mu.Lock()
defer t.mu.Unlock()
2023-02-22 03:04:56 +01:00
delete(t.subscribers, id)
2021-10-23 03:26:01 +02:00
}
2021-11-09 20:48:25 +01:00
// Publish asynchronously publishes to all subscribers
2022-06-01 02:38:56 +02:00
func (t *topic) Publish(v *visitor, m *message) error {
2021-11-09 20:48:25 +01:00
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 {
2023-02-04 04:21:50 +01:00
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 {
2023-02-04 04:21:50 +01:00
logvm(v, m).Tag(tagPublish).Err(err).Warn("Error forwarding to subscriber")
}
}(s.subscriber)
2021-11-09 20:48:25 +01:00
}
2022-06-01 22:57:35 +02:00
} else {
2023-02-04 04:21:50 +01:00
logvm(v, m).Tag(tagPublish).Trace("No stream or WebSocket subscribers, not forwarding")
2021-10-23 03:26:01 +02:00
}
2023-03-01 04:25:13 +01:00
t.Keepalive()
2021-11-09 20:48:25 +01:00
}()
2021-10-23 03:26:01 +02:00
return nil
}
2023-03-01 04:25:13 +01:00
// Stats returns the number of subscribers and last access to this topic
func (t *topic) Stats() (int, time.Time) {
2023-02-23 03:00:56 +01:00
t.mu.RLock()
defer t.mu.RUnlock()
2023-03-01 04:25:13 +01:00
return len(t.subscribers), t.lastAccess
}
// Keepalive sets the last access time and ensures that Stale does not return true
func (t *topic) Keepalive() {
t.mu.Lock()
defer t.mu.Unlock()
t.lastAccess = time.Now()
2021-10-23 03:26:01 +02:00
}
// 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 {
2023-02-24 02:46:53 +01:00
if s.userID != exceptUserID {
log.
Tag(tagSubscribe).
2023-02-26 02:23:22 +01:00
With(t).
2023-02-24 02:46:53 +01:00
Fields(log.Context{
2023-02-26 02:23:22 +01:00
"user_id": s.userID,
2023-02-24 02:46:53 +01:00
}).
Debug("Canceling subscriber %s", s.userID)
s.cancel()
}
}
}
2023-02-26 02:23:22 +01:00
func (t *topic) Context() log.Context {
t.mu.RLock()
defer t.mu.RUnlock()
fields := map[string]any{
"topic": t.ID,
"topic_subscribers": len(t.subscribers),
2023-03-04 04:22:07 +01:00
"topic_last_access": util.FormatTime(t.lastAccess),
2023-02-26 02:23:22 +01:00
}
if t.rateVisitor != nil {
for k, v := range t.rateVisitor.Context() {
fields["topic_rate_"+k] = v
}
2023-02-26 02:23:22 +01:00
}
return fields
}
// 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{
2023-02-24 02:46:53 +01:00
userID: sub.userID,
subscriber: sub.subscriber,
cancel: sub.cancel,
}
}
return subscribers
}