llink/go/internal/pusher/redis_bridge.go

package pusher

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

	"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:"
)

// redisEvent is published/received via Redis Pub/Sub for cross-pod communication.
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"` // for message events
}

// RedisBridge handles cross-pod coordination via Redis Pub/Sub and presence tracking.
type RedisBridge struct {
	client *redis.Client
	podID  string
	hub    *Hub // set after hub is created
}

// NewRedisBridge creates a new Redis bridge for cross-pod coordination.
func NewRedisBridge(client *redis.Client, podID string) *RedisBridge {
	return &RedisBridge{
		client: client,
		podID:  podID,
	}
}

// SetHub sets the hub reference. Called during initialization.
func (rb *RedisBridge) SetHub(hub *Hub) {
	rb.hub = hub
}

// --- Presence management (called by hub goroutine) ---

// Subscribe adds a connection to a channel in Redis.
// Returns the current presence set for the channel.
func (rb *RedisBridge) Subscribe(ctx context.Context, channelID, connID, humanID string) ([]string, error) {
	key := channelConnsKey(channelID)
	field := rb.connField(connID)

	// Check if humanID was already present before adding
	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)

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

	// Publish join event if this is a new humanID in the channel
	if !wasPresent {
		rb.publishEvent(ctx, channelID, redisEvent{
			Type:    TypeJoin,
			HumanID: humanID,
			PodID:   rb.podID,
		})
	}

	// Return deduplicated presence set
	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
}

// Unsubscribe removes a connection from a channel in Redis.
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)
	}

	// Check if this humanID is still present via other connections
	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,
		})
	}

	// Clean up empty channel hash
	if len(remainingMembers) == 0 {
		rb.client.Del(ctx, key)
	}

	return nil
}

// Broadcast publishes a message event to all pods.
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,
	})
}

// GetPresence returns the deduplicated humanIDs for the given channels.
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
}

// GetAllConnectedHumanIDs scans all channel connection hashes in Redis and returns
// the deduplicated set of all humanIDs that have at least one active connection.
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 (runs in its own goroutine) ---

// Listen subscribes to Redis Pub/Sub and forwards events to the local hub.
// Blocks until the context 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) {
	// Extract channel ID from 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 {
		slog.Error("failed to parse pub/sub event", "error", err)
		return
	}

	// Skip events originating from this pod — the local hub already handled them
	if event.PodID == rb.podID {
		return
	}

	if rb.hub == nil {
		return
	}

	// Forward to local hub for delivery to local WebSocket connections
	rb.hub.remoteEventCh <- &remoteEvent{
		channelID: channelID,
		event:     event,
	}
}

// --- Heartbeat + cleanup (runs in its own goroutine) ---

// Heartbeat maintains this pod's liveness key and cleans up stale pods.
func (rb *RedisBridge) Heartbeat(ctx context.Context) {
	podKey := podKeyPrefix + rb.podID

	// Initial heartbeat
	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, remove our pod key and clean up our connections
			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) {
	// Scan all channel conn hashes for pod IDs, then check if those pods are still alive
	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 {
			slog.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
		}
	}

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

	// Clean up dead pods
	for podID := range knownPods {
		if !alivePods[podID] {
			slog.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)
					// Check if this humanID is now gone from the channel
					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 {
		slog.Error("failed to marshal event", "error", err)
		return
	}
	if err := rb.client.Publish(ctx, pubsubPrefix+channelID, data).Err(); err != nil {
		slog.Error("failed to publish event", "channelId", channelID, "error", err)
	}
}

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

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

func extractPodID(field string) string {
	// "{podID}:{connID}" → podID
	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
}