fancy-dns/internal/cluster/cluster.go

// Package cluster keeps several fancy-dns instances in agreement using
// Valkey (Redis protocol): configuration changes are broadcast, blocklist
// digests are compared so a lagging instance can be spotted, and one instance
// is elected to fetch lists so N instances do not hammer the same mirror N
// times.
package cluster

import (
	"context"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

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

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/config"
)

// EventKind names a broadcast message.
type EventKind string

const (
	// EventConfig carries a full configuration document.
	EventConfig EventKind = "config"
	// EventReload asks peers to re-read the config file from disk (the
	// normal case when the file is managed by git or an operator).
	EventReload EventKind = "reload"
	// EventBlocklists asks the leader to refresh list sources.
	EventBlocklists EventKind = "blocklists"
	// EventHeartbeat publishes this instance's state.
	EventHeartbeat EventKind = "heartbeat"
	// EventBye announces a clean shutdown.
	EventBye EventKind = "bye"
)

// Event is a cluster message.
type Event struct {
	Kind    EventKind      `json:"kind"`
	NodeID  string         `json:"node_id"`
	At      time.Time      `json:"at"`
	Version int64          `json:"version,omitempty"`
	Config  string         `json:"config,omitempty"`
	State   *NodeState     `json:"state,omitempty"`
	Sources []string       `json:"sources,omitempty"`
	Extra   map[string]any `json:"extra,omitempty"`
}

// NodeState is what an instance publishes about itself.
type NodeState struct {
	NodeID     string    `json:"node_id"`
	Addr       string    `json:"addr"`
	Version    string    `json:"version"`
	StartedAt  time.Time `json:"started_at"`
	Generation uint64    `json:"generation"`
	// BlocklistDigest lets peers detect that they hold different rules.
	BlocklistDigest string    `json:"blocklist_digest"`
	BlocklistRules  int       `json:"blocklist_rules"`
	BlocklistBuilt  time.Time `json:"blocklist_built_at"`
	ConfigDigest    string    `json:"config_digest"`
	Listeners       []string  `json:"listeners"`
	// Leader is true when this node holds the list-refresh lease.
	Leader bool `json:"leader"`
}

// Handler receives cluster events from peers.
type Handler func(context.Context, Event)

// Options configures a Cluster.
type Options struct {
	// Config is the hot-reloadable config store.
	Config *config.Store
	// Logger is the structured logger.
	Logger *slog.Logger
	// Version is the running build's version string.
	Version string
	// StartedAt is this process's start time.
	StartedAt time.Time
	// Addr is this instance's advertised admin address.
	Addr string
	// Handler receives events from peers (and from this instance's own
	// publishes, so a caller can hook both without two code paths).
	Handler Handler
	// StateFunc supplies the live state published with heartbeats.
	StateFunc func() *NodeState
	// Digests is called on peer state to compare rules.
	Digests func() (blocklist string, rules int, built time.Time, cfgDigest string)
}

// Cluster is the Valkey-backed coordination layer. A disabled Cluster is a
// valid, fully-functional no-op: every method is safe to call.
type Cluster struct {
	cfg  *config.Store
	log  *slog.Logger
	opt  Options
	node string

	client atomic.Pointer[redis.Client]
	// connected is true while the pubsub loop is healthy.
	connected atomic.Bool
	// leader is true while this instance holds the refresh lease.
	leader atomic.Bool
	// lastEvent is when a peer message last arrived, for the health check.
	lastEvent atomic.Int64
	// seen is the last state per node id.
	seenMu sync.RWMutex
	seen   map[string]NodeState

	stop chan struct{}
	wg   sync.WaitGroup
}

// New builds a Cluster. Call Start to connect.
func New(opt Options) *Cluster {
	log := opt.Logger
	if log == nil {
		log = slog.Default()
	}
	c := &Cluster{
		cfg:  opt.Config,
		log:  log.With("component", "cluster"),
		opt:  opt,
		node: opt.Version,
		seen: make(map[string]NodeState),
		stop: make(chan struct{}),
	}
	if cfg := opt.Config.Get(); cfg != nil && cfg.Cluster.NodeID != "" {
		c.node = cfg.Cluster.NodeID
	}
	if c.node == "" {
		host, err := os.Hostname()
		if err != nil || host == "" {
			host = "fancy-dns"
		}
		c.node = host
	}
	c.node = c.node + "-" + shortID()
	return c
}

func shortID() string {
	b := make([]byte, 4)
	if _, err := cryptoRead(b); err != nil {
		return "0000"
	}
	return hex.EncodeToString(b)
}

// NodeID is this instance's cluster identity.
func (c *Cluster) NodeID() string { return c.node }

// Enabled reports whether clustering is configured on.
func (c *Cluster) Enabled() bool {
	if c == nil || c.cfg == nil {
		return false
	}
	return c.cfg.Get().Cluster.Enabled
}

// Connected reports whether the pubsub connection is up.
func (c *Cluster) Connected() bool { return c.connected.Load() }

// IsLeader reports whether this instance holds the refresh lease.
func (c *Cluster) IsLeader() bool { return c.leader.Load() }

// Start connects and subscribes, then keeps the connection alive until ctx is
// cancelled. A Valkey outage is not fatal: the server runs standalone and
// reconnects in the background.
func (c *Cluster) Start(ctx context.Context) {
	if !c.Enabled() {
		c.log.Info("cluster disabled")
		return
	}
	cfg := c.cfg.Get()
	cl, err := c.dial(cfg)
	if err != nil {
		c.log.Warn("cluster: cannot connect, running standalone", "error", err)
		return
	}
	c.client.Store(cl)
	c.connected.Store(true)
	c.log.Info("cluster connected", "addrs", cfg.Cluster.Valkey.Addrs, "node", c.node)

	c.publishState(ctx)

	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		c.maintain(ctx)
	}()
}

// dial builds a Valkey client from config.
func (c *Cluster) dial(cfg *config.Config) (*redis.Client, error) {
	v := cfg.Cluster.Valkey
	cl := redis.NewClient(&redis.Options{
		Addr:         strings.Join(v.Addrs, ","),
		Username:     v.Username,
		Password:     v.Password,
		DB:           v.DB,
		PoolSize:     max(1, v.PoolSize),
		DialTimeout:  v.DialTimeout.D(),
		ReadTimeout:  v.ReadTimeout.D(),
		WriteTimeout: v.ReadTimeout.D(),
	})
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := cl.Ping(ctx).Err(); err != nil {
		_ = cl.Close()
		return nil, err
	}
	return cl, nil
}

// maintain runs the pubsub receive loop and the periodic heartbeat, reconnecting
// as needed.
func (c *Cluster) maintain(ctx context.Context) {
	for {
		if ctx.Err() != nil || c.stopping() {
			return
		}
		if err := c.runSession(ctx); err != nil {
			c.log.Warn("cluster session ended, retrying", "error", err)
			c.connected.Store(false)
			select {
			case <-ctx.Done():
				return
			case <-c.stop:
				return
			case <-time.After(3 * time.Second):
			}
			continue
		}
		return
	}
}

func (c *Cluster) stopping() bool {
	select {
	case <-c.stop:
		return true
	default:
		return false
	}
}

// runSession subscribes and dispatches until the connection drops.
func (c *Cluster) runSession(ctx context.Context) error {
	cl := c.client.Load()
	if cl == nil {
		return errors.New("cluster: no client")
	}
	sub := cl.Subscribe(ctx, c.channel(EventHeartbeat), c.channel(EventConfig), c.channel(EventReload), c.channel(EventBlocklists), c.channel(EventBye))
	defer func() { _ = sub.Close() }()

	// Confirm the subscription is live before declaring the session healthy.
	if _, err := sub.Receive(ctx); err != nil {
		return fmt.Errorf("cluster: subscribe: %w", err)
	}
	c.connected.Store(true)
	c.log.Info("cluster subscribed", "channel", c.channel(EventHeartbeat))

	hb := time.NewTicker(c.cfg.Get().Cluster.PeerTTL.D() / 2)
	defer hb.Stop()

	ch := sub.Channel()
	for {
		select {
		case <-ctx.Done():
			return nil
		case <-c.stop:
			return nil
		case <-hb.C:
			c.publishState(ctx)
			c.renewLease(ctx)
		case msg, ok := <-ch:
			if !ok {
				return errors.New("cluster: subscription closed")
			}
			var ev Event
			if err := json.Unmarshal([]byte(msg.Payload), &ev); err != nil {
				continue // not ours / malformed; ignore rather than crash
			}
			if ev.NodeID == c.node {
				continue // our own echo
			}
			c.lastEvent.Store(time.Now().Unix())
			if ev.State != nil {
				c.seenMu.Lock()
				c.seen[ev.NodeID] = *ev.State
				c.seenMu.Unlock()
			}
			if ev.Kind == EventBye {
				c.seenMu.Lock()
				delete(c.seen, ev.NodeID)
				c.seenMu.Unlock()
				continue
			}
			if c.opt.Handler != nil {
				c.opt.Handler(ctx, ev)
			}
		}
	}
}

func (c *Cluster) channel(k EventKind) string {
	prefix := c.cfg.Get().Cluster.Valkey.Prefix
	if prefix == "" {
		prefix = "fancydns"
	}
	return prefix + ":events:" + string(k)
}

func (c *Cluster) heartbeatKey() string {
	prefix := c.cfg.Get().Cluster.Valkey.Prefix
	if prefix == "" {
		prefix = "fancydns"
	}
	return prefix + ":nodes"
}

// State returns this instance's current state.
func (c *Cluster) State() *NodeState {
	st := &NodeState{
		NodeID:     c.node,
		Addr:       c.opt.Addr,
		Version:    c.opt.Version,
		StartedAt:  c.opt.StartedAt,
		Generation: c.cfg.Generation(),
		Leader:     c.leader.Load(),
	}
	if f := c.opt.StateFunc; f != nil {
		if live := f(); live != nil {
			live.NodeID = c.node
			live.Leader = c.leader.Load()
			st = live
		}
	}
	if d := c.opt.Digests; d != nil {
		bl, rules, built, cfgDigest := d()
		st.BlocklistDigest = bl
		st.BlocklistRules = rules
		st.BlocklistBuilt = built
		st.ConfigDigest = cfgDigest
	}
	return st
}

// publishState broadcasts this instance's heartbeat.
func (c *Cluster) publishState(ctx context.Context) {
	cl := c.client.Load()
	if cl == nil {
		return
	}
	st := c.State()
	ev := Event{Kind: EventHeartbeat, NodeID: c.node, At: time.Now(), State: st}
	data, err := json.Marshal(ev)
	if err != nil {
		return
	}
	pipe := cl.Pipeline()
	pipe.Publish(ctx, c.channel(EventHeartbeat), data)
	// Also keep a keyed record so a late-joining instance can discover the
	// existing set without waiting a full heartbeat interval.
	pipe.HSet(ctx, c.heartbeatKey(), c.node, mustJSON(st))
	pipe.Expire(ctx, c.heartbeatKey(), c.cfg.Get().Cluster.PeerTTL.D()*4)
	_, _ = pipe.Exec(ctx)
}

func mustJSON(v any) string {
	b, err := json.Marshal(v)
	if err != nil {
		return "{}"
	}
	return string(b)
}

// Peers returns the last known state of other instances.
func (c *Cluster) Peers() []NodeState {
	c.seenMu.RLock()
	defer c.seenMu.RUnlock()
	out := make([]NodeState, 0, len(c.seen))
	for _, s := range c.seen {
		out = append(out, s)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].NodeID < out[j].NodeID })
	return out
}

// BroadcastConfig publishes a full configuration document to peers.
func (c *Cluster) BroadcastConfig(ctx context.Context, cfg *config.Config) error {
	cl := c.client.Load()
	if !c.Enabled() || cl == nil {
		return nil
	}
	data, err := config.Marshal(cfg)
	if err != nil {
		return err
	}
	ev := Event{Kind: EventConfig, NodeID: c.node, At: time.Now(), Version: cfg.Version, Config: string(data)}
	payload, err := json.Marshal(ev)
	if err != nil {
		return err
	}
	return cl.Publish(ctx, c.channel(EventConfig), payload).Err()
}

// BroadcastReload asks peers to re-read their config file from disk.
func (c *Cluster) BroadcastReload(ctx context.Context) error {
	cl := c.client.Load()
	if !c.Enabled() || cl == nil {
		return nil
	}
	ev := Event{Kind: EventReload, NodeID: c.node, At: time.Now()}
	payload, err := json.Marshal(ev)
	if err != nil {
		return err
	}
	return cl.Publish(ctx, c.channel(EventReload), payload).Err()
}

// BroadcastRefresh asks the leader to refresh list sources. The names slice is
// advisory: an empty slice means "everything".
func (c *Cluster) BroadcastRefresh(ctx context.Context, names []string) error {
	cl := c.client.Load()
	if !c.Enabled() || cl == nil {
		return nil
	}
	ev := Event{Kind: EventBlocklists, NodeID: c.node, At: time.Now(), Sources: names}
	payload, err := json.Marshal(ev)
	if err != nil {
		return err
	}
	return cl.Publish(ctx, c.channel(EventBlocklists), payload).Err()
}

// Digest computes a stable digest of a configuration document.
func Digest(cfg *config.Config) string {
	data, err := config.Marshal(cfg)
	if err != nil {
		return ""
	}
	sum := sha1.Sum(data)
	return hex.EncodeToString(sum[:8])
}

// leaseKey is the Valkey key guarding the list-refresh leader.
func (c *Cluster) leaseKey() string {
	prefix := c.cfg.Get().Cluster.Valkey.Prefix
	if prefix == "" {
		prefix = "fancydns"
	}
	return prefix + ":leader"
}

// renewLease acquires or extends the refresh lease. Only the leader fetches
// lists; followers serve their compiled copy and keep their cache warm, so a
// fleet of N nodes does not multiply the load on a blocklist mirror.
func (c *Cluster) renewLease(ctx context.Context) {
	cl := c.client.Load()
	if cl == nil || !c.Enabled() {
		c.leader.Store(false)
		return
	}
	ttl := c.cfg.Get().Cluster.Valkey.LeaderTTL.D()
	if ttl <= 0 {
		ttl = 30 * time.Second
	}
	// SET NX with a TTL: only the first caller wins the lease. Everyone else
	// just extends it when they already hold it.
	ok, err := cl.SetNX(ctx, c.leaseKey(), c.node, ttl).Result()
	if err != nil {
		c.leader.Store(false)
		return
	}
	if ok {
		c.leader.Store(true)
		return
	}
	holder, err := c.leaseHolder(ctx)
	if err == nil && holder == c.node {
		// We already hold it: extend.
		_ = cl.Expire(ctx, c.leaseKey(), ttl).Err()
		c.leader.Store(true)
		return
	}
	c.leader.Store(false)
}

// leaseHolder reads who holds the refresh lease.
func (c *Cluster) leaseHolder(ctx context.Context) (string, error) {
	cl := c.client.Load()
	if cl == nil {
		return "", errors.New("no client")
	}
	v, err := cl.Get(ctx, c.leaseKey()).Result()
	if err != nil {
		return "", err
	}
	return v, nil
}

// IsLeader reports the current lease holder, for the UI.
func (c *Cluster) LeaseHolder(ctx context.Context) string {
	if !c.Enabled() {
		return ""
	}
	h, err := c.leaseHolder(ctx)
	if err != nil {
		return ""
	}
	return h
}

// Status is the cluster view for the API.
type Status struct {
	Enabled    bool        `json:"enabled"`
	Connected  bool        `json:"connected"`
	NodeID     string      `json:"node_id"`
	Leader     bool        `json:"leader"`
	LeaseOwner string      `json:"lease_owner,omitempty"`
	Peers      []NodeState `json:"peers"`
	// Divergent lists peers whose blocklist digest differs from ours, which
	// is the one thing an operator really needs to see.
	Divergent []string  `json:"divergent"`
	Channel   string    `json:"channel"`
	LastEvent time.Time `json:"last_event,omitempty"`
}

// Status snapshots the cluster state.
func (c *Cluster) Status(ctx context.Context) Status {
	st := Status{
		Enabled:   c.Enabled(),
		Connected: c.Connected(),
		NodeID:    c.node,
		Leader:    c.leader.Load(),
		Channel:   c.channel(EventHeartbeat),
	}
	if c.Enabled() {
		st.LeaseOwner = c.LeaseHolder(ctx)
	}
	st.Peers = c.Peers()
	own := c.State()
	for _, p := range st.Peers {
		if p.BlocklistDigest != "" && own.BlocklistDigest != "" && p.BlocklistDigest != own.BlocklistDigest {
			st.Divergent = append(st.Divergent, p.NodeID)
		}
	}
	sort.Strings(st.Divergent)
	if v := c.lastEvent.Load(); v > 0 {
		st.LastEvent = time.Unix(v, 0)
	}
	return st
}

// Discover lists the node ids currently registered, for a fresh instance that
// has not heard a heartbeat yet.
func (c *Cluster) Discover(ctx context.Context) []NodeState {
	if !c.Enabled() {
		return nil
	}
	cl := c.client.Load()
	if cl == nil {
		return nil
	}
	fields, err := cl.HGetAll(ctx, c.heartbeatKey()).Result()
	if err != nil {
		return nil
	}
	out := make([]NodeState, 0, len(fields))
	ids := make([]string, 0, len(fields))
	for id := range fields {
		ids = append(ids, id)
	}
	sort.Strings(ids)
	for _, id := range ids {
		if id == c.node {
			continue
		}
		var st NodeState
		if err := json.Unmarshal([]byte(fields[id]), &st); err == nil {
			out = append(out, st)
		}
	}
	return out
}

// Stop closes the connection and announces the departure.
func (c *Cluster) Stop() {
	select {
	case <-c.stop:
	default:
		close(c.stop)
	}
	cl := c.client.Load()
	if cl == nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	ev := Event{Kind: EventBye, NodeID: c.node, At: time.Now()}
	if payload, err := json.Marshal(ev); err == nil {
		_ = cl.Publish(ctx, c.channel(EventBye), payload).Err()
	}
	// Only release the lease if we hold it, so a shutdown does not hand the
	// job to a node that is also going away.
	if c.leader.Load() {
		cl.Eval(ctx, `if redis.call("GET", KEYS[1]) == ARGV[1] then return redis.call("DEL", KEYS[1]) else return 0 end`, []string{c.leaseKey()}, c.node)
	}
	_ = cl.Close()
	c.connected.Store(false)
	c.leader.Store(false)
}

// String renders the node id for logs.
func (c *Cluster) String() string { return strings.TrimSpace(c.node) }

func cryptoRead(b []byte) (int, error) {
	h := sha1.Sum([]byte(time.Now().String() + string(b)))
	copy(b, h[:4])
	return len(b), nil
}