databasus--databasus/backend/internal/util/cache/pubsub.go

package cache_utils

import (
	"context"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/valkey-io/valkey-go"

	"databasus-backend/internal/util/logger"
)

// readyMarkerPrefix tags the self-published probe that Subscribe uses to
// confirm the SUBSCRIBE command has been accepted by Valkey before
// returning. The wrapped handler intercepts these markers and never
// forwards them to the real handler.
const readyMarkerPrefix = "__databasus_subscribe_ready__:"

// subscribeReadyTimeout caps how long Subscribe waits for its own ready
// marker to round-trip. 5s is generous enough to absorb a stressed CI host
// without making subscriber-setup hang dominate startup time.
const subscribeReadyTimeout = 5 * time.Second

// subscribeReadyPollInterval is how often we re-publish the ready marker
// while waiting. Each publish before the subscriber is registered is lost
// (pubsub is fire-and-forget); we keep republishing until the subscriber
// catches one.
const subscribeReadyPollInterval = 25 * time.Millisecond

type PubSubManager struct {
	client        valkey.Client
	subscriptions map[string]context.CancelFunc
	mu            sync.RWMutex
	logger        *slog.Logger
}

func NewPubSubManager() *PubSubManager {
	return &PubSubManager{
		client:        getCache(),
		subscriptions: make(map[string]context.CancelFunc),
		logger:        logger.GetLogger(),
	}
}

func (m *PubSubManager) Subscribe(
	ctx context.Context,
	channel string,
	handler func(message string),
) error {
	m.mu.Lock()
	if _, exists := m.subscriptions[channel]; exists {
		m.mu.Unlock()
		return fmt.Errorf("already subscribed to channel: %s", channel)
	}

	subCtx, cancel := context.WithCancel(ctx)
	m.subscriptions[channel] = cancel
	m.mu.Unlock()

	readyMarker := readyMarkerPrefix + uuid.New().String()
	ready := make(chan struct{})
	readyOnce := sync.Once{}

	wrappedHandler := func(message string) {
		if strings.HasPrefix(message, readyMarkerPrefix) {
			if message == readyMarker {
				readyOnce.Do(func() { close(ready) })
			}
			return
		}

		handler(message)
	}

	go m.subscriptionLoop(subCtx, channel, wrappedHandler)

	if err := m.waitForSubscribeReady(ctx, channel, readyMarker, ready); err != nil {
		cancel()

		m.mu.Lock()
		delete(m.subscriptions, channel)
		m.mu.Unlock()

		return err
	}

	return nil
}

func (m *PubSubManager) Publish(ctx context.Context, channel, message string) error {
	cmd := m.client.B().Publish().Channel(channel).Message(message).Build()
	result := m.client.Do(ctx, cmd)

	if err := result.Error(); err != nil {
		m.logger.Error("Failed to publish message to Redis", "channel", channel, "error", err)
		return fmt.Errorf("failed to publish message: %w", err)
	}

	return nil
}

func (m *PubSubManager) Close() error {
	m.mu.Lock()
	defer m.mu.Unlock()

	for channel, cancel := range m.subscriptions {
		cancel()
		delete(m.subscriptions, channel)
	}

	return nil
}

func (m *PubSubManager) subscriptionLoop(
	ctx context.Context,
	channel string,
	handler func(message string),
) {
	defer func() {
		if r := recover(); r != nil {
			m.logger.Error("Panic in subscription handler", "channel", channel, "panic", r)
		}
	}()

	m.logger.Info("Starting subscription", "channel", channel)

	err := m.client.Receive(
		ctx,
		m.client.B().Subscribe().Channel(channel).Build(),
		func(msg valkey.PubSubMessage) {
			defer func() {
				if r := recover(); r != nil {
					m.logger.Error("Panic in message handler", "channel", channel, "panic", r)
				}
			}()

			handler(msg.Message)
		},
	)

	if err != nil && ctx.Err() == nil {
		m.logger.Error("Subscription error", "channel", channel, "error", err)
	} else if ctx.Err() != nil {
		m.logger.Info("Subscription cancelled", "channel", channel)
	}

	m.mu.Lock()
	delete(m.subscriptions, channel)
	m.mu.Unlock()
}

// waitForSubscribeReady republishes a unique marker on the channel until
// the subscriber's wrapped handler observes it. This proves that the
// SUBSCRIBE command has been accepted by Valkey and our handler is in the
// receive path. Without this handshake, callers race the SUBSCRIBE ack and
// can publish-then-miss the next message; tests papered over it with
// time.Sleep heuristics that broke under load.
func (m *PubSubManager) waitForSubscribeReady(
	ctx context.Context,
	channel string,
	readyMarker string,
	ready <-chan struct{},
) error {
	deadline := time.NewTimer(subscribeReadyTimeout)
	defer deadline.Stop()

	tick := time.NewTicker(subscribeReadyPollInterval)
	defer tick.Stop()

	publishCtx, cancelPublish := context.WithTimeout(ctx, subscribeReadyTimeout)
	defer cancelPublish()

	publish := func() {
		_ = m.client.Do(
			publishCtx,
			m.client.B().Publish().Channel(channel).Message(readyMarker).Build(),
		).Error()
	}

	publish()

	for {
		select {
		case <-ready:
			return nil

		case <-tick.C:
			publish()

		case <-deadline.C:
			return fmt.Errorf("subscribe ready timeout on channel %q after %s", channel, subscribeReadyTimeout)

		case <-ctx.Done():
			return ctx.Err()
		}
	}
}