core/server/cluster/registry.go

package cluster

import (
	"context"
	"encoding/json"
	"log"
	"os"
	"runtime"
	"sync"
	"time"

	"switchboard-core/store"
)

// ConnCounter provides WebSocket connection count (satisfied by events.Hub).
type ConnCounter interface {
	ConnCount() int
}

// RegistryConfig holds tunable cluster registry parameters.
type RegistryConfig struct {
	HeartbeatInterval time.Duration
	StaleThreshold    time.Duration
}

// Registry manages node self-registration, heartbeat, and stale sweep.
// Follows the workflow.Scanner pattern: stopCh + wg + ticker.
type Registry struct {
	nodeID    string
	endpoint  string
	store     store.ClusterStore
	hub       ConnCounter
	cfg       RegistryConfig
	startTime time.Time
	stopCh    chan struct{}
	wg        sync.WaitGroup
}

// NewRegistry creates a cluster registry instance.
func NewRegistry(nodeID, endpoint string, s store.ClusterStore, hub ConnCounter, cfg RegistryConfig) *Registry {
	return &Registry{
		nodeID:    nodeID,
		endpoint:  endpoint,
		store:     s,
		hub:       hub,
		cfg:       cfg,
		startTime: time.Now(),
		stopCh:    make(chan struct{}),
	}
}

// NodeID returns this node's identifier.
func (r *Registry) NodeID() string { return r.nodeID }

// Start registers the node and launches the heartbeat goroutine.
func (r *Registry) Start() error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := r.store.Register(ctx, r.nodeID, r.endpoint); err != nil {
		return err
	}

	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		ticker := time.NewTicker(r.cfg.HeartbeatInterval)
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
				r.tick()
			case <-r.stopCh:
				return
			}
		}
	}()
	log.Printf("[cluster] started (node=%s, heartbeat=%s, stale=%s)",
		r.nodeID, r.cfg.HeartbeatInterval, r.cfg.StaleThreshold)
	return nil
}

// Stop deregisters the node and waits for the heartbeat goroutine to exit.
func (r *Registry) Stop() {
	close(r.stopCh)
	r.wg.Wait()

	// Best-effort deregister so peers don't have to wait for stale sweep.
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	if err := r.store.Deregister(ctx, r.nodeID); err != nil {
		log.Printf("[cluster] deregister failed: %v", err)
	}
	log.Printf("[cluster] stopped (node=%s)", r.nodeID)
}

// tick runs one heartbeat + sweep cycle.
func (r *Registry) tick() {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	stats := r.collectStats()

	rows, err := r.store.Heartbeat(ctx, r.nodeID, stats)
	if err != nil {
		log.Printf("[cluster] heartbeat error: %v", err)
		return
	}
	if rows == 0 {
		// Self-eviction: this node was swept by a peer.
		// Kubernetes will restart the process → re-register.
		log.Printf("[cluster] FATAL: node %s evicted from registry — shutting down", r.nodeID)
		os.Exit(1)
	}

	swept, err := r.store.SweepStale(ctx, r.cfg.StaleThreshold)
	if err != nil {
		log.Printf("[cluster] sweep error: %v", err)
		return
	}
	if swept > 0 {
		log.Printf("[cluster] swept %d stale node(s)", swept)
	}
}

// collectStats gathers runtime metrics for the stats JSONB column.
func (r *Registry) collectStats() json.RawMessage {
	var m runtime.MemStats
	runtime.ReadMemStats(&m)

	stats := map[string]any{
		"goroutines":  runtime.NumGoroutine(),
		"heap_alloc":  m.HeapAlloc,
		"heap_sys":    m.HeapSys,
		"gc_cycles":   m.NumGC,
		"gc_pause_ns": m.PauseNs[(m.NumGC+255)%256],
		"uptime_sec":  time.Since(r.startTime).Seconds(),
		"ws_clients":  r.hub.ConnCount(),
	}

	data, _ := json.Marshal(stats)
	return data
}