gochat/backend/internal/pubsub/redis_pubsub.go

package pubsub

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"

	"github.com/ThreeDotsLabs/watermill"
	"github.com/ThreeDotsLabs/watermill/message"
	redisstream "github.com/ThreeDotsLabs/watermill-redisstream/pkg/redisstream"
	"github.com/redis/go-redis/v9"

	applogger "github.com/gochat/gochat/pkg/logger"
)

// RedisPubSub implements the PubSub interface using Redis Streams via watermill-redisstream.
// It provides production-grade message delivery with consumer groups for reliable fan-out.
type RedisPubSub struct {
	publisher  *redisstream.Publisher
	subscriber *redisstream.Subscriber
	client     redis.UniversalClient

	mu         sync.Mutex
	subscriptions map[string]context.CancelFunc // track active subscriptions for cleanup
}

// NewRedisPubSub creates a new Redis-backed PubSub instance.
// Falls back to InMemoryPubSub if the Redis client is nil or connection fails.
func NewRedisPubSub(client redis.UniversalClient) (*RedisPubSub, error) {
	if client == nil {
		return nil, fmt.Errorf("redis client is required for RedisPubSub")
	}

	logger := watermill.NopLogger{}

	publisher, err := redisstream.NewPublisher(redisstream.PublisherConfig{
		Client: client,
	}, logger)
	if err != nil {
		return nil, fmt.Errorf("failed to create redis stream publisher: %w", err)
	}

	subscriber, err := redisstream.NewSubscriber(redisstream.SubscriberConfig{
		Client:        client,
		ConsumerGroup: "gochat-pubsub",
	}, logger)
	if err != nil {
		return nil, fmt.Errorf("failed to create redis stream subscriber: %w", err)
	}

	return &RedisPubSub{
		publisher:     publisher,
		subscriber:    subscriber,
		client:        client,
		subscriptions: make(map[string]context.CancelFunc),
	}, nil
}

// Publish sends an event to all subscribers of the given topic via Redis Streams.
func (ps *RedisPubSub) Publish(ctx context.Context, topic string, event Event) error {
	payload, err := json.Marshal(event)
	if err != nil {
		return fmt.Errorf("failed to marshal event: %w", err)
	}

	msg := message.NewMessage(watermill.NewUUID(), payload)
	// Propagate context metadata into message metadata
	msg.Metadata.Set("event_type", event.Type)
	msg.Metadata.Set("account_id", fmt.Sprintf("%d", event.AccountID))

	if err := ps.publisher.Publish(topic, msg); err != nil {
		return fmt.Errorf("failed to publish event to redis stream: %w", err)
	}

	applogger.L().Debugf("Published event %s to topic %s via Redis", event.Type, topic)
	return nil
}

// Subscribe registers a handler for events on the given topic via Redis Streams consumer groups.
func (ps *RedisPubSub) Subscribe(ctx context.Context, topic string, handler EventHandler) error {
	messages, err := ps.subscriber.Subscribe(ctx, topic)
	if err != nil {
		return fmt.Errorf("failed to subscribe to redis stream topic %s: %w", topic, err)
	}

	// Create a cancellable context so we can stop this specific subscription
	subCtx, cancel := context.WithCancel(ctx)

	ps.mu.Lock()
	ps.subscriptions[topic] = cancel
	ps.mu.Unlock()

	go func() {
		for {
			select {
			case <-subCtx.Done():
				applogger.L().Debugf("Subscription to topic %s stopped", topic)
				return
			case msg, ok := <-messages:
				if !ok {
					applogger.L().Debugf("Message channel for topic %s closed", topic)
					ps.mu.Lock()
					delete(ps.subscriptions, topic)
					ps.mu.Unlock()
					return
				}

				var event Event
				if err := json.Unmarshal(msg.Payload, &event); err != nil {
					applogger.L().Errorf("Failed to unmarshal event from redis stream: %v", err)
					msg.Nack()
					continue
				}

				handler(event)
				msg.Ack()
			}
		}
	}()

	applogger.L().Infof("Subscribed to topic %s via Redis Streams (consumer group: gochat-pubsub)", topic)
	return nil
}

// Unsubscribe removes a subscription from the given topic.
func (ps *RedisPubSub) Unsubscribe(ctx context.Context, topic string) error {
	ps.mu.Lock()
	cancel, ok := ps.subscriptions[topic]
	if ok {
		cancel()
		delete(ps.subscriptions, topic)
	}
	ps.mu.Unlock()

	if !ok {
		applogger.L().Debugf("No active subscription found for topic %s", topic)
		return nil
	}

	applogger.L().Infof("Unsubscribed from topic %s", topic)
	return nil
}

// Close shuts down the Redis PubSub publisher and subscriber gracefully.
func (ps *RedisPubSub) Close() error {
	ps.mu.Lock()
	for topic, cancel := range ps.subscriptions {
		cancel()
		delete(ps.subscriptions, topic)
	}
	ps.mu.Unlock()

	if err := ps.subscriber.Close(); err != nil {
		applogger.L().Errorf("Redis PubSub subscriber close error: %v", err)
	}

	if err := ps.publisher.Close(); err != nil {
		applogger.L().Errorf("Redis PubSub publisher close error: %v", err)
		return fmt.Errorf("redis publisher close error: %w", err)
	}

	applogger.L().Info("Redis PubSub closed")
	return nil
}