fix: Change pubsub interface to use callbacks.

generics/setttl.go

@@ -50,6 +50,7 @@ func (s *SetWithTTL[T]) Remove(es ...T) {
 }
 
 // Contains returns true if the SetWithTTL contains `e`.
+// We don't have to clean up first because the test checks the TTL.
 func (s *SetWithTTL[T]) Contains(e T) bool {
 	s.mut.RLock()
 	item, ok := s.Items[e]

pubsub/pubsub.go

@@ -11,20 +11,22 @@ import (
 // pubsub.Start()
 // defer pubsub.Stop()
 // ctx := context.Background()
-// pubsub.Publish(ctx, "topic", "message")
-// sub := pubsub.Subscribe(ctx, "topic")
-// for msg := range sub.Channel() {
+// sub := pubsub.Subscribe(ctx, "topic", func(msg string) {
 // 	fmt.Println(msg)
 // }
+// pubsub.Publish(ctx, "topic", "message")
 // sub.Close()  // optional
 // pubsub.Close()
 
 type PubSub interface {
 	// Publish sends a message to all subscribers of the specified topic.
 	Publish(ctx context.Context, topic, message string) error
-	// Subscribe returns a Subscription that will receive all messages published to the specified topic.
+	// Subscribe returns a Subscription to the specified topic.
+	// The callback will be called for each message published to the topic.
 	// There is no unsubscribe method; close the subscription to stop receiving messages.
-	Subscribe(ctx context.Context, topic string) Subscription
+	// The subscription only exists to provide a way to stop receiving messages; if you don't need to stop,
+	// you can ignore the return value.
+	Subscribe(ctx context.Context, topic string, callback func(msg string)) Subscription
 	// Close shuts down all topics and the pubsub connection.
 	Close()
 
@@ -35,9 +37,7 @@ type PubSub interface {
 }
 
 type Subscription interface {
-	// Channel returns the channel that will receive all messages published to the topic.
-	Channel() <-chan string
-	// Close stops the subscription and closes the channel. Calling this is optional;
-	// the topic will be closed when the pubsub connection is closed.
+	// Close stops the subscription which means the callback will no longer be called.
+	// Optional; the topic will be closed when the pubsub connection is closed.
 	Close()
 }

pubsub/pubsub_goredis.go

@@ -34,7 +34,7 @@ var _ PubSub = (*GoRedisPubSub)(nil)
 type GoRedisSubscription struct {
 	topic  string
 	pubsub *redis.PubSub
-	ch     chan string
+	cb     func(msg string)
 	done   chan struct{}
 	once   sync.Once
 }
@@ -109,11 +109,15 @@ func (ps *GoRedisPubSub) Publish(ctx context.Context, topic, message string) err
 	return ps.client.Publish(ctx, topic, message).Err()
 }
 
-func (ps *GoRedisPubSub) Subscribe(ctx context.Context, topic string) Subscription {
+// Subscribe creates a new Subscription to the given topic, and calls the provided callback
+// whenever a message is received on that topic.
+// Note that the same topic is Subscribed to multiple times, this will incur a separate
+// connection to Redis for each Subscription.
+func (ps *GoRedisPubSub) Subscribe(ctx context.Context, topic string, callback func(string)) Subscription {
 	sub := &GoRedisSubscription{
 		topic:  topic,
 		pubsub: ps.client.Subscribe(ctx, topic),
-		ch:     make(chan string, 100),
+		cb:     callback,
 		done:   make(chan struct{}),
 	}
 	ps.mut.Lock()
@@ -124,30 +128,20 @@ func (ps *GoRedisPubSub) Subscribe(ctx context.Context, topic string) Subscripti
 		for {
 			select {
 			case <-sub.done:
-				close(sub.ch)
 				return
 			case msg := <-redisch:
 				if msg == nil {
 					continue
 				}
-				select {
-				case sub.ch <- msg.Payload:
-				default:
-					ps.Logger.Warn().WithField("topic", topic).Logf("Dropping subscription message because channel is full")
-				}
+				go sub.cb(msg.Payload)
 			}
 		}
 	}()
 	return sub
 }
 
-func (s *GoRedisSubscription) Channel() <-chan string {
-	return s.ch
-}
-
 func (s *GoRedisSubscription) Close() {
 	s.once.Do(func() {
-		s.pubsub.Close()
 		close(s.done)
 	})
 }

pubsub/pubsub_local.go

@@ -9,29 +9,29 @@ import (
 
 // LocalPubSub is a PubSub implementation that uses local channels to send messages; it does
 // not communicate with any external processes.
+// subs are individual channels for each subscription
 type LocalPubSub struct {
 	Config *config.Config `inject:""`
-	subs   []*LocalSubscription
-	topics map[string]chan string
+	topics map[string][]*LocalSubscription
 	mut    sync.RWMutex
 }
 
 // Ensure that LocalPubSub implements PubSub
 var _ PubSub = (*LocalPubSub)(nil)
 
 type LocalSubscription struct {
+	ps    *LocalPubSub
 	topic string
-	ch    chan string
-	done  chan struct{}
+	cb    func(string)
+	mut   sync.RWMutex
 }
 
 // Ensure that LocalSubscription implements Subscription
 var _ Subscription = (*LocalSubscription)(nil)
 
 // Start initializes the LocalPubSub
 func (ps *LocalPubSub) Start() error {
-	ps.subs = make([]*LocalSubscription, 0)
-	ps.topics = make(map[string]chan string)
+	ps.topics = make(map[string][]*LocalSubscription)
 	return nil
 }
 
@@ -44,59 +44,51 @@ func (ps *LocalPubSub) Stop() error {
 func (ps *LocalPubSub) Close() {
 	ps.mut.Lock()
 	defer ps.mut.Unlock()
-	for _, sub := range ps.subs {
-		sub.Close()
+	for _, subs := range ps.topics {
+		for i := range subs {
+			subs[i].cb = nil
+		}
 	}
-	ps.subs = nil
+	ps.topics = make(map[string][]*LocalSubscription, 0)
 }
 
-func (ps *LocalPubSub) ensureTopic(topic string) chan string {
+func (ps *LocalPubSub) ensureTopic(topic string) {
 	if _, ok := ps.topics[topic]; !ok {
-		ps.topics[topic] = make(chan string, 100)
+		ps.topics[topic] = make([]*LocalSubscription, 0)
 	}
-	return ps.topics[topic]
 }
 
 func (ps *LocalPubSub) Publish(ctx context.Context, topic, message string) error {
-	ps.mut.RLock()
-	ch := ps.ensureTopic(topic)
-	ps.mut.RUnlock()
-	select {
-	case ch <- message:
-	case <-ctx.Done():
-		return ctx.Err()
+	ps.mut.Lock()
+	defer ps.mut.Unlock()
+	ps.ensureTopic(topic)
+	for _, sub := range ps.topics[topic] {
+		// don't wait around for slow consumers
+		if sub.cb != nil {
+			go sub.cb(message)
+		}
 	}
 	return nil
 }
 
-func (ps *LocalPubSub) Subscribe(ctx context.Context, topic string) Subscription {
+func (ps *LocalPubSub) Subscribe(ctx context.Context, topic string, callback func(msg string)) Subscription {
 	ps.mut.Lock()
-	defer ps.mut.Unlock()
-	ch := ps.ensureTopic(topic)
-	sub := &LocalSubscription{
-		topic: topic,
-		ch:    ch,
-		done:  make(chan struct{}),
-	}
-	ps.subs = append(ps.subs, sub)
-	go func() {
-		for {
-			select {
-			case <-sub.done:
-				close(ch)
-				return
-			case msg := <-ch:
-				sub.ch <- msg
-			}
-		}
-	}()
+	ps.ensureTopic(topic)
+	sub := &LocalSubscription{ps: ps, topic: topic, cb: callback}
+	ps.topics[topic] = append(ps.topics[topic], sub)
+	ps.mut.Unlock()
 	return sub
 }
 
-func (s *LocalSubscription) Channel() <-chan string {
-	return s.ch
-}
-
 func (s *LocalSubscription) Close() {
-	close(s.done)
+	s.ps.mut.RLock()
+	for _, sub := range s.ps.topics[s.topic] {
+		if sub == s {
+			sub.mut.Lock()
+			sub.cb = nil
+			sub.mut.Unlock()
+			return
+		}
+	}
+	s.ps.mut.RUnlock()
 }