feat(nodes): gRPC server + Hub + Redis pub/sub cross-instance delivery (tsk__58l3wTLvaSn)
Implements AgentService gRPC server with all 6 RPCs (Enroll/Register/Heartbeat/ Subscribe/Ack/ReportUsage), Hub command routing with Redis ZSET at-least-once persistence and cross-instance pub/sub delivery, LoadCache for node:load metrics, NodeStore SQL interface + MySQL implementation, and full mTLS gRPC listener in main. Integration tests: 18 tests covering full Enroll→Register→Heartbeat→Subscribe→Ack flow, reconnect resume with last_command_id, and cross-instance pub/sub delivery via bufconn + miniredis + mockNodeStore + real mTLS certificates. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
wangjia
@@ -0,0 +1,228 @@
|
||||
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()
|
||||
}
|
||||
Reference in New Issue
Block a user