llink/go/internal/pusher/redis_bridge.go

package pusher

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/flowy-live/llink/internal/utils/flog"

	"github.com/redis/go-redis/v9"
)

const (
	podHeartbeatInterval = 30 * time.Second
	podHeartbeatTTL      = 60 * time.Second
	cleanupInterval      = 60 * time.Second

	// Redis key prefixes
	channelConnsPrefix = "pusher:ch:"
	channelConnsSuffix = ":conns"
	podKeyPrefix       = "pusher:pod:"
	pubsubPrefix       = "pusher:events:"
)

// Wire format for cross-pod Pub/Sub.
type redisEvent struct {
	Type    string          `json:"type"`              // "join", "leave", "message"
	HumanID string          `json:"humanId,omitempty"` // who triggered the event
	PodID   string          `json:"podId,omitempty"`   // originating pod
	Payload json.RawMessage `json:"payload,omitempty"` // message events only
}

// RedisBridge handles cross-pod coordination via Redis Pub/Sub and presence
// tracking.
type RedisBridge struct {
	client *redis.Client
	podID  string
	hub    *Hub // wired post-construction; see SetHub
}

func NewRedisBridge(client *redis.Client, podID string) *RedisBridge {
	return &RedisBridge{
		client: client,
		podID:  podID,
	}
}

// SetHub resolves the circular dependency between Hub and RedisBridge.
func (rb *RedisBridge) SetHub(hub *Hub) {
	rb.hub = hub
}

// --- Presence management ---

// Subscribe records the connection in Redis and returns the channel's
// current deduplicated presence set.
func (rb *RedisBridge) Subscribe(ctx context.Context, channelID, connID, humanID string) ([]string, error) {
	key := channelConnsKey(channelID)
	field := rb.connField(connID)

	// Snapshot before the insert so multi-tab joins don't double-emit.
	existingMembers, err := rb.client.HVals(ctx, key).Result()
	if err != nil && err != redis.Nil {
		return nil, fmt.Errorf("failed to get channel members: %w", err)
	}

	wasPresent := containsString(existingMembers, humanID)

	if err := rb.client.HSet(ctx, key, field, humanID).Err(); err != nil {
		return nil, fmt.Errorf("failed to add connection to channel: %w", err)
	}

	if !wasPresent {
		rb.publishEvent(ctx, channelID, redisEvent{
			Type:    TypeJoin,
			HumanID: humanID,
			PodID:   rb.podID,
		})
	}

	allMembers, err := rb.client.HVals(ctx, key).Result()
	if err != nil {
		return nil, fmt.Errorf("failed to get channel members: %w", err)
	}
	return deduplicateStrings(allMembers), nil
}

func (rb *RedisBridge) Unsubscribe(ctx context.Context, channelID, connID, humanID string) error {
	key := channelConnsKey(channelID)
	field := rb.connField(connID)

	if err := rb.client.HDel(ctx, key, field).Err(); err != nil {
		return fmt.Errorf("failed to remove connection from channel: %w", err)
	}

	// Only emit leave once this humanID has no tabs left in the channel.
	remainingMembers, err := rb.client.HVals(ctx, key).Result()
	if err != nil && err != redis.Nil {
		return fmt.Errorf("failed to get remaining members: %w", err)
	}

	if !containsString(remainingMembers, humanID) {
		rb.publishEvent(ctx, channelID, redisEvent{
			Type:    TypeLeave,
			HumanID: humanID,
			PodID:   rb.podID,
		})
	}

	if len(remainingMembers) == 0 {
		rb.client.Del(ctx, key)
	}

	return nil
}

func (rb *RedisBridge) Broadcast(ctx context.Context, channelID, humanID string, payload json.RawMessage) {
	rb.publishEvent(ctx, channelID, redisEvent{
		Type:    TypeMessage,
		HumanID: humanID,
		PodID:   rb.podID,
		Payload: payload,
	})
}

func (rb *RedisBridge) GetPresence(ctx context.Context, channelIDs []string) (map[string][]string, error) {
	result := make(map[string][]string, len(channelIDs))
	for _, chID := range channelIDs {
		members, err := rb.client.HVals(ctx, channelConnsKey(chID)).Result()
		if err != nil && err != redis.Nil {
			return nil, fmt.Errorf("failed to get presence for %s: %w", chID, err)
		}
		result[chID] = deduplicateStrings(members)
	}
	return result, nil
}

// Returns every humanID with at least one active connection cluster-wide.
func (rb *RedisBridge) GetAllConnectedHumanIDs(ctx context.Context) ([]string, error) {
	allHumanIDs := make(map[string]bool)
	var cursor uint64

	for {
		keys, nextCursor, err := rb.client.Scan(ctx, cursor, channelConnsPrefix+"*"+channelConnsSuffix, 100).Result()
		if err != nil {
			return nil, fmt.Errorf("failed to scan channel keys: %w", err)
		}

		for _, key := range keys {
			members, err := rb.client.HVals(ctx, key).Result()
			if err != nil && err != redis.Nil {
				continue
			}
			for _, humanID := range members {
				allHumanIDs[humanID] = true
			}
		}

		cursor = nextCursor
		if cursor == 0 {
			break
		}
	}

	result := make([]string, 0, len(allHumanIDs))
	for id := range allHumanIDs {
		result = append(result, id)
	}
	return result, nil
}

// --- Pub/Sub listener ---

// Listen forwards Redis Pub/Sub events to the local hub; blocks until ctx is cancelled.
func (rb *RedisBridge) Listen(ctx context.Context) {
	pubsub := rb.client.PSubscribe(ctx, pubsubPrefix+"*")
	defer pubsub.Close()

	ch := pubsub.Channel()
	for {
		select {
		case <-ctx.Done():
			return
		case msg, ok := <-ch:
			if !ok {
				return
			}
			rb.handlePubSubMessage(msg)
		}
	}
}

func (rb *RedisBridge) handlePubSubMessage(msg *redis.Message) {
	// Topic: "pusher:events:{channelID}".
	channelID := strings.TrimPrefix(msg.Channel, pubsubPrefix)
	if channelID == "" {
		return
	}

	var event redisEvent
	if err := json.Unmarshal([]byte(msg.Payload), &event); err != nil {
		flog.Error("failed to parse pub/sub event", "error", err)
		return
	}

	// Same-pod events were already handled by the local hub.
	if event.PodID == rb.podID {
		return
	}

	if rb.hub == nil {
		return
	}

	rb.hub.remoteEventCh <- &remoteEvent{
		channelID: channelID,
		event:     event,
	}
}

// --- Heartbeat + cleanup ---

// Heartbeat refreshes this pod's liveness key and reaps stale pods on a tick.
func (rb *RedisBridge) Heartbeat(ctx context.Context) {
	podKey := podKeyPrefix + rb.podID

	rb.client.Set(ctx, podKey, "alive", podHeartbeatTTL)

	heartbeatTicker := time.NewTicker(podHeartbeatInterval)
	cleanupTicker := time.NewTicker(cleanupInterval)
	defer heartbeatTicker.Stop()
	defer cleanupTicker.Stop()

	for {
		select {
		case <-ctx.Done():
			// On shutdown, drop our pod key and reclaim our connection slots.
			rb.client.Del(context.Background(), podKey)
			rb.cleanupPod(context.Background(), rb.podID)
			return
		case <-heartbeatTicker.C:
			rb.client.Set(ctx, podKey, "alive", podHeartbeatTTL)
		case <-cleanupTicker.C:
			rb.cleanupStalePods(ctx)
		}
	}
}

func (rb *RedisBridge) cleanupStalePods(ctx context.Context) {
	// Collect every pod referenced in channel-conn hashes, then drop those
	// whose liveness key has expired.
	var cursor uint64
	knownPods := make(map[string]bool)
	alivePods := make(map[string]bool)

	for {
		keys, nextCursor, err := rb.client.Scan(ctx, cursor, channelConnsPrefix+"*"+channelConnsSuffix, 100).Result()
		if err != nil {
			flog.Error("failed to scan channel keys", "error", err)
			return
		}

		for _, key := range keys {
			fields, err := rb.client.HKeys(ctx, key).Result()
			if err != nil {
				continue
			}
			for _, field := range fields {
				podID := extractPodID(field)
				if podID != "" {
					knownPods[podID] = true
				}
			}
		}

		cursor = nextCursor
		if cursor == 0 {
			break
		}
	}

	for podID := range knownPods {
		exists, err := rb.client.Exists(ctx, podKeyPrefix+podID).Result()
		if err != nil {
			continue
		}
		if exists > 0 {
			alivePods[podID] = true
		}
	}

	for podID := range knownPods {
		if !alivePods[podID] {
			flog.Info("cleaning up stale pod", "podId", podID)
			rb.cleanupPod(ctx, podID)
		}
	}
}

func (rb *RedisBridge) cleanupPod(ctx context.Context, podID string) {
	var cursor uint64
	for {
		keys, nextCursor, err := rb.client.Scan(ctx, cursor, channelConnsPrefix+"*"+channelConnsSuffix, 100).Result()
		if err != nil {
			return
		}

		for _, key := range keys {
			fields, err := rb.client.HGetAll(ctx, key).Result()
			if err != nil {
				continue
			}

			channelID := extractChannelID(key)
			for field, humanID := range fields {
				if extractPodID(field) == podID {
					rb.client.HDel(ctx, key, field)
					remaining, _ := rb.client.HVals(ctx, key).Result()
					if !containsString(remaining, humanID) {
						rb.publishEvent(ctx, channelID, redisEvent{
							Type:    TypeLeave,
							HumanID: humanID,
							PodID:   rb.podID,
						})
					}
				}
			}
		}

		cursor = nextCursor
		if cursor == 0 {
			break
		}
	}
}

// --- Helpers ---

func (rb *RedisBridge) connField(connID string) string {
	return rb.podID + ":" + connID
}

func (rb *RedisBridge) publishEvent(ctx context.Context, channelID string, event redisEvent) {
	data, err := json.Marshal(event)
	if err != nil {
		flog.Error("failed to marshal event", "error", err)
		return
	}
	if err := rb.client.Publish(ctx, pubsubPrefix+channelID, data).Err(); err != nil {
		flog.Error("failed to publish event", "channelId", channelID, "error", err)
	}
}

func channelConnsKey(channelID string) string {
	return channelConnsPrefix + channelID + channelConnsSuffix
}

// "pusher:ch:{channelID}:conns" → channelID
func extractChannelID(redisKey string) string {
	s := strings.TrimPrefix(redisKey, channelConnsPrefix)
	s = strings.TrimSuffix(s, channelConnsSuffix)
	return s
}

// "{podID}:{connID}" → podID
func extractPodID(field string) string {
	parts := strings.SplitN(field, ":", 2)
	if len(parts) == 2 {
		return parts[0]
	}
	return ""
}

func containsString(slice []string, s string) bool {
	for _, v := range slice {
		if v == s {
			return true
		}
	}
	return false
}

func deduplicateStrings(slice []string) []string {
	seen := make(map[string]bool, len(slice))
	result := make([]string, 0, len(slice))
	for _, s := range slice {
		if !seen[s] {
			seen[s] = true
			result = append(result, s)
		}
	}
	return result
}