topic.go 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207
  1. package server
  2. import (
  3. "math/rand"
  4. "sync"
  5. "time"
  6. "heckel.io/ntfy/log"
  7. "heckel.io/ntfy/util"
  8. )
  9. const (
  10. // topicExpungeAfter defines how long a topic is active before it is removed from memory.
  11. // This must be larger than matrixRejectPushKeyForUnifiedPushTopicWithoutRateVisitorAfter to give
  12. // time for more requests to come in, so that we can send a {"rejected":["<pushkey>"]} response back.
  13. topicExpungeAfter = 16 * time.Hour
  14. )
  15. // topic represents a channel to which subscribers can subscribe, and publishers
  16. // can publish a message
  17. type topic struct {
  18. ID string
  19. subscribers map[int]*topicSubscriber
  20. rateVisitor *visitor
  21. lastAccess time.Time
  22. mu sync.RWMutex
  23. }
  24. type topicSubscriber struct {
  25. userID string // User ID associated with this subscription, may be empty
  26. subscriber subscriber
  27. cancel func()
  28. }
  29. // subscriber is a function that is called for every new message on a topic
  30. type subscriber func(v *visitor, msg *message) error
  31. // newTopic creates a new topic
  32. func newTopic(id string) *topic {
  33. return &topic{
  34. ID: id,
  35. subscribers: make(map[int]*topicSubscriber),
  36. lastAccess: time.Now(),
  37. }
  38. }
  39. // Subscribe subscribes to this topic
  40. func (t *topic) Subscribe(s subscriber, userID string, cancel func()) (subscriberID int) {
  41. t.mu.Lock()
  42. defer t.mu.Unlock()
  43. for i := 0; i < 5; i++ { // Best effort retry
  44. subscriberID = rand.Int()
  45. _, exists := t.subscribers[subscriberID]
  46. if !exists {
  47. break
  48. }
  49. }
  50. t.subscribers[subscriberID] = &topicSubscriber{
  51. userID: userID, // May be empty
  52. subscriber: s,
  53. cancel: cancel,
  54. }
  55. t.lastAccess = time.Now()
  56. return subscriberID
  57. }
  58. func (t *topic) Stale() bool {
  59. t.mu.Lock()
  60. defer t.mu.Unlock()
  61. if t.rateVisitor != nil && !t.rateVisitor.Stale() {
  62. return false
  63. }
  64. return len(t.subscribers) == 0 && time.Since(t.lastAccess) > topicExpungeAfter
  65. }
  66. func (t *topic) LastAccess() time.Time {
  67. t.mu.RLock()
  68. defer t.mu.RUnlock()
  69. return t.lastAccess
  70. }
  71. func (t *topic) SetRateVisitor(v *visitor) {
  72. t.mu.Lock()
  73. defer t.mu.Unlock()
  74. t.rateVisitor = v
  75. t.lastAccess = time.Now()
  76. }
  77. func (t *topic) RateVisitor() *visitor {
  78. t.mu.Lock()
  79. defer t.mu.Unlock()
  80. if t.rateVisitor != nil && t.rateVisitor.Stale() {
  81. t.rateVisitor = nil
  82. }
  83. return t.rateVisitor
  84. }
  85. // Unsubscribe removes the subscription from the list of subscribers
  86. func (t *topic) Unsubscribe(id int) {
  87. t.mu.Lock()
  88. defer t.mu.Unlock()
  89. delete(t.subscribers, id)
  90. }
  91. // Publish asynchronously publishes to all subscribers
  92. func (t *topic) Publish(v *visitor, m *message) error {
  93. go func() {
  94. // We want to lock the topic as short as possible, so we make a shallow copy of the
  95. // subscribers map here. Actually sending out the messages then doesn't have to lock.
  96. subscribers := t.subscribersCopy()
  97. if len(subscribers) > 0 {
  98. logvm(v, m).Tag(tagPublish).Debug("Forwarding to %d subscriber(s)", len(subscribers))
  99. for _, s := range subscribers {
  100. // We call the subscriber functions in their own Go routines because they are blocking, and
  101. // we don't want individual slow subscribers to be able to block others.
  102. go func(s subscriber) {
  103. if err := s(v, m); err != nil {
  104. logvm(v, m).Tag(tagPublish).Err(err).Warn("Error forwarding to subscriber")
  105. }
  106. }(s.subscriber)
  107. }
  108. } else {
  109. logvm(v, m).Tag(tagPublish).Trace("No stream or WebSocket subscribers, not forwarding")
  110. }
  111. t.Keepalive()
  112. }()
  113. return nil
  114. }
  115. // Stats returns the number of subscribers and last access to this topic
  116. func (t *topic) Stats() (int, time.Time) {
  117. t.mu.RLock()
  118. defer t.mu.RUnlock()
  119. return len(t.subscribers), t.lastAccess
  120. }
  121. // Keepalive sets the last access time and ensures that Stale does not return true
  122. func (t *topic) Keepalive() {
  123. t.mu.Lock()
  124. defer t.mu.Unlock()
  125. t.lastAccess = time.Now()
  126. }
  127. // CancelSubscribersExceptUser calls the cancel function for all subscribers, forcing
  128. func (t *topic) CancelSubscribersExceptUser(exceptUserID string) {
  129. t.mu.Lock()
  130. defer t.mu.Unlock()
  131. for _, s := range t.subscribers {
  132. if s.userID != exceptUserID {
  133. t.cancelUserSubscriber(s)
  134. }
  135. }
  136. }
  137. // CancelSubscriberUser kills the subscriber with the given user ID
  138. func (t *topic) CancelSubscriberUser(userID string) {
  139. t.mu.RLock()
  140. defer t.mu.RUnlock()
  141. for _, s := range t.subscribers {
  142. if s.userID == userID {
  143. t.cancelUserSubscriber(s)
  144. return
  145. }
  146. }
  147. }
  148. func (t *topic) cancelUserSubscriber(s *topicSubscriber) {
  149. log.
  150. Tag(tagSubscribe).
  151. With(t).
  152. Fields(log.Context{
  153. "user_id": s.userID,
  154. }).
  155. Debug("Canceling subscriber with user ID %s", s.userID)
  156. s.cancel()
  157. }
  158. func (t *topic) Context() log.Context {
  159. t.mu.RLock()
  160. defer t.mu.RUnlock()
  161. fields := map[string]any{
  162. "topic": t.ID,
  163. "topic_subscribers": len(t.subscribers),
  164. "topic_last_access": util.FormatTime(t.lastAccess),
  165. }
  166. if t.rateVisitor != nil {
  167. for k, v := range t.rateVisitor.Context() {
  168. fields["topic_rate_"+k] = v
  169. }
  170. }
  171. return fields
  172. }
  173. // subscribersCopy returns a shallow copy of the subscribers map
  174. func (t *topic) subscribersCopy() map[int]*topicSubscriber {
  175. t.mu.Lock()
  176. defer t.mu.Unlock()
  177. subscribers := make(map[int]*topicSubscriber)
  178. for k, sub := range t.subscribers {
  179. subscribers[k] = &topicSubscriber{
  180. userID: sub.userID,
  181. subscriber: sub.subscriber,
  182. cancel: sub.cancel,
  183. }
  184. }
  185. return subscribers
  186. }