Edu/services/push-gateway/internal/redisclient/pubsub.go

// Package redisclient manages cross-instance presence and message fanout
// using Redis Pub/Sub plus the online-presence SET described in 02 §3.1 and
// §8.4. It implements the ISSUE-058 startup-rebuild mechanism: when a Hub
// starts it clears stale members belonging to its own instanceID from the
// online SET, then re-adds itself for every locally connected user.
//
// Keys / channels:
//
//	edu:push:online:<userID>           SET (members are instanceID)        TTL 60s
//	edu:push:session:<userID>:<connID> HASH (instanceID, device, lastPing) TTL 60s
//	edu:push:channel:user:<userID>     Pub/Sub channel (directed push)
//	edu:push:channel:broadcast         Pub/Sub channel (broadcast push)
//	edu:push:idempotent:<event_id>     SETNX key (Kafka dedup)            TTL 24h
package redisclient

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

	"github.com/edu-cloud/push-gateway/internal/hub"
	"github.com/edu-cloud/push-gateway/internal/observability"
	"github.com/redis/go-redis/v9"
)

// PresenceTTL is the expiry applied to the online-presence SET and per-session
// HASH. The Hub refreshes it on every heartbeat; if a push-gateway instance
// crashes the SET entries expire naturally within this window.
const PresenceTTL = 60 * time.Second

// IdempotentTTL is the dedup window for Kafka event_id (24h).
const IdempotentTTL = 24 * time.Hour

// CrossInstanceMessage is the payload published to the
// edu:push:channel:user:<userID> and edu:push:channel:broadcast channels.
// Subscribers unmarshal it and dispatch to local Hub connections.
type CrossInstanceMessage struct {
	Kind     string          `json:"kind"`     // "user" | "broadcast"
	UserID   string          `json:"userId"`   // set when Kind == "user"
	Event    string          `json:"event"`
	Data     json.RawMessage `json:"data"`
	Origin   string          `json:"origin"`   // instanceID of publisher (loopback suppression)
	TraceID  string          `json:"traceId,omitempty"`
}

// Client wraps a redis.Client with push-gateway specific presence and
// fanout helpers. The Hub is wired via SetHub so the subscriber can dispatch
// incoming Pub/Sub messages directly to local connections.
type Client struct {
	rdb        *redis.Client
	instanceID string
	hub        *hub.Hub
	metrics    *observability.Metrics
}

// New connects to Redis at url and returns a Client ready for presence and
// Pub/Sub operations. The instanceID identifies this process in the online
// SET (see ISSUE-058).
func New(url, instanceID string, m *observability.Metrics) (*Client, error) {
	if url == "" {
		return nil, errors.New("redisclient: empty URL")
	}
	opts, err := redis.ParseURL(url)
	if err != nil {
		return nil, fmt.Errorf("redisclient: parsing URL %q: %w", url, err)
	}
	rdb := redis.NewClient(opts)
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	if err := rdb.Ping(ctx).Err(); err != nil {
		_ = rdb.Close()
		return nil, fmt.Errorf("redisclient: ping failed: %w", err)
	}
	return &Client{
		rdb:        rdb,
		instanceID: instanceID,
		metrics:    m,
	}, nil
}

// SetHub wires the Hub so the Pub/Sub subscriber can deliver messages locally.
// Also installs presence hooks on the Hub so Register/Unregister keep the
// online SET in sync.
func (c *Client) SetHub(h *hub.Hub) {
	c.hub = h
	h.SetPresenceHooks(c.onRegister, c.onUnregister)
}

// Redis returns the underlying *redis.Client for direct use by callers that
// need bespoke commands (e.g. SETNX for Kafka idempotency).
func (c *Client) Redis() *redis.Client {
	return c.rdb
}

// Close releases the Redis connection.
func (c *Client) Close() error {
	return c.rdb.Close()
}

// Ping is a lightweight liveness probe used by /readyz.
func (c *Client) Ping(ctx context.Context) error {
	return c.rdb.Ping(ctx).Err()
}

// ----- Presence SET (ISSUE-058) -------------------------------------------

// RebuildPresenceOnStartup clears stale SET entries belonging to this
// instanceID and re-adds the instanceID to edu:push:online:<userID> for every
// locally connected user (ISSUE-058 startup rebuild).
//
// This addresses the "ghost member" problem: when a push-gateway instance
// restarts, its previous instanceID entries may still be present in Redis
// (within the 60s TTL window). Re-adding the same instanceID would be a
// no-op, but if the instance previously crashed without removing members we
// must avoid carrying stale entries that no longer correspond to live
// connections. We clear first, then rebuild from the in-memory Hub state.
//
// 60s inconsistency window: between crash and SET expiry, /internal/push may
// report online:true for a user whose connection is gone; msg must tolerate
// delivered:false by retrying or falling back to offline push.
func (c *Client) RebuildPresenceOnStartup(ctx context.Context) error {
	if c.hub == nil {
		return errors.New("redisclient: Hub not set")
	}
	start := time.Now()
	c.metrics.IncRedisSetRebuild()

	// Clear stale members belonging to this instanceID across all online SETs.
	// We scan keys matching edu:push:online:* and SREM our instanceID from each.
	iter := c.rdb.Scan(ctx, 0, "edu:push:online:*", 100).Iterator()
	cleared := 0
	for iter.Next(ctx) {
		key := iter.Val()
		if n, err := c.rdb.SRem(ctx, key, c.instanceID).Result(); err == nil && n > 0 {
			cleared++
		}
	}
	if err := iter.Err(); err != nil {
		return fmt.Errorf("redisclient: scanning online SETs: %w", err)
	}

	// Re-add instanceID for every locally online user.
	readded := 0
	c.hub.ForEachUser(func(userID string) {
		key := onlineKey(userID)
		pipe := c.rdb.Pipeline()
		pipe.SAdd(ctx, key, c.instanceID)
		pipe.Expire(ctx, key, PresenceTTL)
		if _, err := pipe.Exec(ctx); err == nil {
			readded++
		}
	})

	observability.Logger().Info("redis presence rebuilt",
		"cleared", cleared,
		"readded", readded,
		"duration_ms", time.Since(start).Milliseconds(),
	)
	return nil
}

// onRegister is the Hub hook invoked after a new connection is registered.
// Adds the instanceID to the user's online SET and refreshes TTL.
func (c *Client) onRegister(userID, connID string) {
	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
	defer cancel()
	key := onlineKey(userID)
	pipe := c.rdb.Pipeline()
	pipe.SAdd(ctx, key, c.instanceID)
	pipe.Expire(ctx, key, PresenceTTL)
	if _, err := pipe.Exec(ctx); err != nil {
		observability.Logger().Warn("redis SAdd failed",
			"user_id", userID, "err", err)
	}
}

// onUnregister is the Hub hook invoked after a connection is removed.
// Removes the instanceID from the SET and deletes the key when empty.
func (c *Client) onUnregister(userID, connID string) {
	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
	defer cancel()
	key := onlineKey(userID)
	// Only remove if this user has no further local connections (Hub already
	// removed the connection by the time the hook fires).
	if c.hub != nil && !c.hub.HasUser(userID) {
		if err := c.rdb.SRem(ctx, key, c.instanceID).Err(); err != nil {
			observability.Logger().Warn("redis SRem failed",
				"user_id", userID, "err", err)
		}
		// Best-effort DEL when SET becomes empty.
		_ = c.rdb.Del(ctx, key).Err()
	}
}

// RefreshPresence extends the online SET TTL on heartbeat. Called by the
// WebSocket reader goroutine when a Ping control frame arrives.
func (c *Client) RefreshPresence(ctx context.Context, userID string) error {
	return c.rdb.Expire(ctx, onlineKey(userID), PresenceTTL).Err()
}

// IsOnline reports whether userID has any instance in the online SET.
// Returns false when Redis is unreachable; callers should treat the
// /internal/push response as authoritative for delivery.
func (c *Client) IsOnline(ctx context.Context, userID string) (bool, []string, error) {
	members, err := c.rdb.SMembers(ctx, onlineKey(userID)).Result()
	if err != nil {
		return false, nil, err
	}
	return len(members) > 0, members, nil
}

// ----- Pub/Sub fanout ------------------------------------------------------

// PublishUser publishes a directed message to edu:push:channel:user:<userID>.
// All push-gateway instances subscribe to this channel; the one holding the
// user will deliver locally, others drop it.
func (c *Client) PublishUser(ctx context.Context, userID string, msg CrossInstanceMessage) error {
	msg.Kind = "user"
	msg.UserID = userID
	msg.Origin = c.instanceID
	payload, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("redisclient: marshalling user msg: %w", err)
	}
	start := time.Now()
	if err := c.rdb.Publish(ctx, userChannel(userID), payload).Err(); err != nil {
		return fmt.Errorf("redisclient: PUBLISH user: %w", err)
	}
	c.metrics.ObservePubSubLatency("publish", time.Since(start).Seconds())
	return nil
}

// PublishBroadcast publishes a broadcast message to all instances.
func (c *Client) PublishBroadcast(ctx context.Context, msg CrossInstanceMessage) error {
	msg.Kind = "broadcast"
	msg.Origin = c.instanceID
	payload, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("redisclient: marshalling broadcast msg: %w", err)
	}
	start := time.Now()
	if err := c.rdb.Publish(ctx, broadcastChannel(), payload).Err(); err != nil {
		return fmt.Errorf("redisclient: PUBLISH broadcast: %w", err)
	}
	c.metrics.ObservePubSubLatency("publish", time.Since(start).Seconds())
	return nil
}

// SubscribeAll subscribes to all per-user channels (pattern) and the
// broadcast channel, dispatching received messages to the local Hub.
// The returned cancel function stops the subscriber; it blocks until the
// goroutine has fully drained.
func (c *Client) SubscribeAll(ctx context.Context) (cancel func() error, err error) {
	if c.hub == nil {
		return nil, errors.New("redisclient: Hub not set")
	}

	subCtx, subCancel := context.WithCancel(ctx)
	userPubsub := c.rdb.PSubscribe(subCtx, "edu:push:channel:user:*")
	broadcastPubsub := c.rdb.Subscribe(subCtx, broadcastChannel())

	go c.dispatchLoop(subCtx, userPubsub, "user")
	go c.dispatchLoop(subCtx, broadcastPubsub, "broadcast")

	return func() error {
		subCancel()
		if err := userPubsub.Close(); err != nil {
			return err
		}
		return broadcastPubsub.Close()
	}, nil
}

// dispatchLoop consumes Pub/Sub messages and dispatches them to the Hub.
// channelType is "user" or "broadcast" for metric labeling.
func (c *Client) dispatchLoop(ctx context.Context, ps *redis.PubSub, channelType string) {
	ch := ps.Channel()
	for {
		select {
		case <-ctx.Done():
			return
		case msg, ok := <-ch:
			if !ok {
				return
			}
			start := time.Now()
			c.handlePubSubMessage(msg, channelType)
			c.metrics.ObservePubSubLatency("subscribe", time.Since(start).Seconds())
		}
	}
}

// handlePubSubMessage unmarshals the payload and dispatches to the Hub,
// suppressing loopback (messages this instance published itself).
func (c *Client) handlePubSubMessage(msg *redis.Message, channelType string) {
	var crossMsg CrossInstanceMessage
	if err := json.Unmarshal([]byte(msg.Payload), &crossMsg); err != nil {
		observability.Logger().Warn("pubsub: invalid payload",
			"channel", msg.Channel, "err", err)
		return
	}
	// Loopback suppression: do not deliver our own publishes.
	if crossMsg.Origin == c.instanceID {
		return
	}

	payload, err := json.Marshal(map[string]any{
		"type":      "message",
		"event":     crossMsg.Event,
		"data":      crossMsg.Data,
		"timestamp": time.Now().UTC().Format(time.RFC3339),
	})
	if err != nil {
		return
	}

	switch crossMsg.Kind {
	case "user":
		c.hub.SendToUser(crossMsg.UserID, payload)
	case "broadcast":
		c.hub.Broadcast(payload)
	}
}

// ----- Idempotency (Kafka dedup) ------------------------------------------

// DedupEventId performs a SETNX with TTL for Kafka event_id dedup.
// Returns true when the event is fresh (first occurrence), false when it
// has already been consumed.
func (c *Client) DedupEventId(ctx context.Context, eventID string) (bool, error) {
	if eventID == "" {
		return true, nil
	}
	ok, err := c.rdb.SetNX(ctx, idempotentKey(eventID), "1", IdempotentTTL).Result()
	if err != nil {
		return false, err
	}
	return ok, nil
}

// ----- Key helpers ---------------------------------------------------------

func onlineKey(userID string) string {
	return "edu:push:online:" + userID
}

func userChannel(userID string) string {
	return "edu:push:channel:user:" + userID
}

func broadcastChannel() string {
	return "edu:push:channel:broadcast"
}

func idempotentKey(eventID string) string {
	return "edu:push:idempotent:" + eventID
}

// ParseUserChannel extracts the userID from a "edu:push:channel:user:<userID>"
// channel name. Returns empty string when the name does not match.
func ParseUserChannel(channel string) string {
	const prefix = "edu:push:channel:user:"
	if !strings.HasPrefix(channel, prefix) {
		return ""
	}
	return strings.TrimPrefix(channel, prefix)
}