service backend

api/pkg/messaging/internal/natsb/subscription.go

@@ -0,0 +1,78 @@
+package natsb
+
+import (
+	"sync"
+
+	"github.com/nats-io/nats.go"
+	me "github.com/tech/sendico/pkg/messaging/envelope"
+	"github.com/tech/sendico/pkg/mlogger"
+	"go.uber.org/zap"
+)
+
+type TopicSubscription struct {
+	sub       *nats.Subscription
+	consumers map[<-chan me.Envelope]chan me.Envelope
+	mu        sync.Mutex
+	logger    mlogger.Logger
+}
+
+func NewTopicSubscription(logger mlogger.Logger, nc *nats.Conn, subject string) (*TopicSubscription, error) {
+	ts := &TopicSubscription{
+		consumers: make(map[<-chan me.Envelope]chan me.Envelope),
+		logger:    logger.Named(subject),
+	}
+
+	sub, err := nc.Subscribe(subject, ts.handleMessage)
+	if err != nil {
+		logger.Error("Error subscribing to subject", zap.String("subject", subject), zap.Error(err))
+		return nil, err
+	}
+	ts.sub = sub
+
+	return ts, nil
+}
+
+func (ts *TopicSubscription) handleMessage(m *nats.Msg) {
+	ts.logger.Debug("Received message", zap.String("subject", m.Subject))
+
+	envelope, err := me.Deserialize(m.Data)
+	if err != nil {
+		ts.logger.Warn("Failed to deserialize message", zap.String("subject", m.Subject), zap.Error(err))
+		return // Do not push invalid data to the channels
+	}
+
+	ts.mu.Lock()
+	defer ts.mu.Unlock()
+	for _, c := range ts.consumers {
+		select {
+		case c <- envelope:
+		default:
+			ts.logger.Warn("Consumer is slow or not receiving messages", zap.String("subject", m.Subject))
+		}
+	}
+}
+
+func (ts *TopicSubscription) AddConsumer(messageChan chan me.Envelope) {
+	ts.mu.Lock()
+	ts.consumers[messageChan] = messageChan
+	ts.mu.Unlock()
+}
+
+func (ts *TopicSubscription) RemoveConsumer(messageChan <-chan me.Envelope) {
+	ts.mu.Lock()
+	if c, ok := ts.consumers[messageChan]; ok {
+		delete(ts.consumers, messageChan)
+		close(c)
+	}
+	ts.mu.Unlock()
+}
+
+func (ts *TopicSubscription) HasConsumers() bool {
+	ts.mu.Lock()
+	defer ts.mu.Unlock()
+	return len(ts.consumers) > 0
+}
+
+func (ts *TopicSubscription) Unsubscribe() error {
+	return ts.sub.Drain()
+}