pangolin/server/internal/nodes/hub.go

package nodes

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"strconv"
	"sync"

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

	agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1"
)

const (
	// cmdQueuePrefix is the Redis ZSET key for each node's pending command queue.
	// Score = command_id (float64 cast of int64); member = JSON-encoded Command.
	// Commands remain in the queue until Ack'd; reconnecting nodes replay from afterID.
	cmdQueuePrefix = "node:cmdq:"

	// cmdPubChannel is the Redis Pub/Sub channel for cross-instance command delivery.
	// Published whenever a command cannot be delivered locally (node connected elsewhere).
	cmdPubChannel = "node:cmd"

	// cmdIDKey is the Redis counter for auto-generating monotonically increasing command IDs.
	cmdIDKey = "pangolin:node:next_cmdid"

	// subChanCap is the size of each local subscriber's delivery channel buffer.
	subChanCap = 256
)

// pubSubPayload is the message body published to cmdPubChannel.
type pubSubPayload struct {
	NodeUUID string           `json:"n"`
	Cmd      *agentv1.Command `json:"c"`
}

// Hub routes commands to connected nodes.
//
// Persistence: every pushed command is stored in a per-node Redis ZSET
// (node:cmdq:{uuid}) so the queue survives instance restarts.
//
// Local delivery: each active Subscribe stream holds a buffered channel.
// When a node is connected locally, commands are sent directly to that channel.
//
// Cross-instance delivery: when the target node is not connected to this instance,
// the command is published on cmdPubChannel so any other instance holding the node's
// stream can pick it up and forward it.
//
// At-least-once semantics: commands are removed from the queue only on Ack.
// A reconnecting node replays from its last acknowledged command_id.
type Hub struct {
	mu   sync.RWMutex
	subs map[string]chan *agentv1.Command

	rdb *redis.Client
}

// NewHub creates a Hub backed by rdb.
// Call Start(ctx) to activate the cross-instance pub/sub goroutine.
func NewHub(rdb *redis.Client) *Hub {
	return &Hub{
		subs: make(map[string]chan *agentv1.Command),
		rdb:  rdb,
	}
}

// Start begins the background pub/sub subscriber goroutine.
// It runs until ctx is cancelled.
func (h *Hub) Start(ctx context.Context) {
	go h.runPubSub(ctx)
}

func (h *Hub) runPubSub(ctx context.Context) {
	sub := h.rdb.Subscribe(ctx, cmdPubChannel)
	defer func() { _ = sub.Close() }()

	ch := sub.Channel()
	for {
		select {
		case <-ctx.Done():
			return
		case msg, ok := <-ch:
			if !ok {
				return
			}
			var p pubSubPayload
			if err := json.Unmarshal([]byte(msg.Payload), &p); err != nil {
				slog.Warn("nodes.hub: bad pub/sub message", "err", err)
				continue
			}
			h.deliverLocal(p.NodeUUID, p.Cmd)
		}
	}
}

// deliverLocal sends cmd to the local subscriber for nodeUUID (non-blocking).
// Returns true if the command was placed on the channel.
func (h *Hub) deliverLocal(nodeUUID string, cmd *agentv1.Command) bool {
	h.mu.RLock()
	ch, ok := h.subs[nodeUUID]
	h.mu.RUnlock()
	if !ok {
		return false
	}
	select {
	case ch <- cmd:
		return true
	default:
		slog.Warn("nodes.hub: subscriber channel full, dropping command",
			"node_uuid", nodeUUID, "command_id", cmd.CommandID)
		return false
	}
}

// Register adds a local subscriber for nodeUUID and returns:
//   - ch: the receive channel on which live commands will arrive.
//   - done: cleanup function that must be called when the stream ends.
//
// Register must be called before Replay to avoid missing commands that arrive
// between the two operations.
func (h *Hub) Register(nodeUUID string) (<-chan *agentv1.Command, func()) {
	ch := make(chan *agentv1.Command, subChanCap)
	h.mu.Lock()
	h.subs[nodeUUID] = ch
	h.mu.Unlock()
	return ch, func() {
		h.mu.Lock()
		// Guard against re-registration: only delete if this is still the active channel.
		if h.subs[nodeUUID] == ch {
			delete(h.subs, nodeUUID)
		}
		h.mu.Unlock()
	}
}

// Replay returns all unacked commands for nodeUUID with command_id > afterID, ordered ascending.
// This is called during Subscribe to re-send commands the agent may have missed.
func (h *Hub) Replay(ctx context.Context, nodeUUID string, afterID int64) ([]*agentv1.Command, error) {
	key := cmdQueuePrefix + nodeUUID
	min := fmt.Sprintf("(%d", afterID) // exclusive: score > afterID
	strs, err := h.rdb.ZRangeByScore(ctx, key, &redis.ZRangeBy{
		Min: min,
		Max: "+inf",
	}).Result()
	if err != nil {
		return nil, fmt.Errorf("nodes.hub.Replay: %w", err)
	}
	cmds := make([]*agentv1.Command, 0, len(strs))
	for _, s := range strs {
		var c agentv1.Command
		if err := json.Unmarshal([]byte(s), &c); err != nil {
			slog.Warn("nodes.hub: corrupt command in queue, skipping",
				"err", err, "node_uuid", nodeUUID)
			continue
		}
		cmds = append(cmds, &c)
	}
	return cmds, nil
}

// Push enqueues cmd for nodeUUID.
// If cmd.CommandID == 0, a new monotonic ID is assigned via Redis INCR.
// The command is persisted in the node's Redis queue then delivered:
//   - directly to the local subscriber channel if this instance holds the stream,
//   - otherwise published on cmdPubChannel for cross-instance delivery.
func (h *Hub) Push(ctx context.Context, nodeUUID string, cmd *agentv1.Command) error {
	// Assign ID if not set.
	if cmd.CommandID == 0 {
		id, err := h.rdb.Incr(ctx, cmdIDKey).Result()
		if err != nil {
			return fmt.Errorf("nodes.hub.Push: generate id: %w", err)
		}
		cmd.CommandID = id
	}

	// Persist to Redis ZSET (durable, survives instance restart).
	data, err := json.Marshal(cmd)
	if err != nil {
		return fmt.Errorf("nodes.hub.Push: marshal: %w", err)
	}
	if err := h.rdb.ZAdd(ctx, cmdQueuePrefix+nodeUUID, redis.Z{
		Score:  float64(cmd.CommandID),
		Member: string(data),
	}).Err(); err != nil {
		return fmt.Errorf("nodes.hub.Push: zadd: %w", err)
	}

	// Try local delivery first (fastest path, no pub/sub latency).
	if h.deliverLocal(nodeUUID, cmd) {
		return nil
	}

	// Publish for cross-instance delivery.  If the node is not connected anywhere,
	// this is a no-op and the node will pick up the command on next reconnect via Replay.
	payload, err := json.Marshal(pubSubPayload{NodeUUID: nodeUUID, Cmd: cmd})
	if err != nil {
		return fmt.Errorf("nodes.hub.Push: marshal pubsub: %w", err)
	}
	if err := h.rdb.Publish(ctx, cmdPubChannel, payload).Err(); err != nil {
		// Non-fatal: command is already durable.
		slog.Warn("nodes.hub: publish failed (command still queued)",
			"node_uuid", nodeUUID, "command_id", cmd.CommandID, "err", err)
	}
	return nil
}

// Broadcast pushes cmd to each nodeUUID with a freshly assigned command_id per node.
// A Broadcast cmd should be passed with CommandID = 0.
func (h *Hub) Broadcast(ctx context.Context, nodeUUIDs []string, cmd *agentv1.Command) error {
	for _, uuid := range nodeUUIDs {
		clone := *cmd
		clone.CommandID = 0 // each push gets an independent ID
		if err := h.Push(ctx, uuid, &clone); err != nil {
			slog.Warn("nodes.hub.Broadcast: push failed",
				"node_uuid", uuid, "err", err)
		}
	}
	return nil
}

// Ack removes commandID from the pending queue for nodeUUID.
// A second Ack for the same commandID is a no-op (idempotent).
func (h *Hub) Ack(ctx context.Context, nodeUUID string, commandID int64) error {
	score := strconv.FormatInt(commandID, 10)
	return h.rdb.ZRemRangeByScore(ctx, cmdQueuePrefix+nodeUUID, score, score).Err()
}