wangjia
605bfa1ec9
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>
84 lines
2.3 KiB
Go
84 lines
2.3 KiB
Go
package nodes
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
const (
|
|
// nodeLoadKeyPrefix is the Redis HASH key prefix for per-node runtime load metrics.
|
|
// Full key: node:load:{node_uuid} (doc/03 §4)
|
|
nodeLoadKeyPrefix = "node:load:"
|
|
|
|
// nodeLoadTTL is the TTL on each node:load hash. A node that stops sending
|
|
// heartbeats has its entry automatically evicted after 90 seconds.
|
|
nodeLoadTTL = 90 * time.Second
|
|
)
|
|
|
|
// NodeLoad holds the runtime metrics reported in each heartbeat.
|
|
type NodeLoad struct {
|
|
OnlineCount int64
|
|
BandwidthUpBps int64
|
|
BandwidthDownBps int64
|
|
CPUPercent float64
|
|
}
|
|
|
|
// LoadCache reads and writes node:load Redis hashes.
|
|
// It accepts redis.Cmdable so it can be used with a *redis.Client or pipeline.
|
|
type LoadCache struct {
|
|
rdb redis.Cmdable
|
|
}
|
|
|
|
// NewLoadCache creates a LoadCache backed by rdb.
|
|
func NewLoadCache(rdb redis.Cmdable) *LoadCache {
|
|
return &LoadCache{rdb: rdb}
|
|
}
|
|
|
|
// Set writes nodeUUID's current load metrics to Redis and refreshes the TTL.
|
|
// The hash fields match the names in doc/03 §4: online_count, bw_up, bw_down, cpu.
|
|
func (c *LoadCache) Set(ctx context.Context, nodeUUID string, load NodeLoad) error {
|
|
key := nodeLoadKeyPrefix + nodeUUID
|
|
pipe := c.rdb.Pipeline()
|
|
pipe.HSet(ctx, key,
|
|
"online_count", load.OnlineCount,
|
|
"bw_up", load.BandwidthUpBps,
|
|
"bw_down", load.BandwidthDownBps,
|
|
"cpu", load.CPUPercent,
|
|
)
|
|
pipe.Expire(ctx, key, nodeLoadTTL)
|
|
if _, err := pipe.Exec(ctx); err != nil {
|
|
return fmt.Errorf("nodes.LoadCache.Set: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Get reads nodeUUID's last-known load. Returns (nil, false, nil) if the key has
|
|
// expired or was never set. Fields missing from the hash default to zero.
|
|
func (c *LoadCache) Get(ctx context.Context, nodeUUID string) (*NodeLoad, bool, error) {
|
|
key := nodeLoadKeyPrefix + nodeUUID
|
|
vals, err := c.rdb.HGetAll(ctx, key).Result()
|
|
if err != nil {
|
|
return nil, false, fmt.Errorf("nodes.LoadCache.Get: %w", err)
|
|
}
|
|
if len(vals) == 0 {
|
|
return nil, false, nil
|
|
}
|
|
var l NodeLoad
|
|
if v, ok := vals["online_count"]; ok {
|
|
_, _ = fmt.Sscan(v, &l.OnlineCount)
|
|
}
|
|
if v, ok := vals["bw_up"]; ok {
|
|
_, _ = fmt.Sscan(v, &l.BandwidthUpBps)
|
|
}
|
|
if v, ok := vals["bw_down"]; ok {
|
|
_, _ = fmt.Sscan(v, &l.BandwidthDownBps)
|
|
}
|
|
if v, ok := vals["cpu"]; ok {
|
|
_, _ = fmt.Sscan(v, &l.CPUPercent)
|
|
}
|
|
return &l, true, nil
|
|
}
|