core/server/cluster/registry.go

package cluster

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

	"armature/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

	// Optional callbacks — set after construction via setters.
	sandboxStats   func() (exec, errors uint64, avgMs float64)
	triggerFires   func() int64
	extensionCount func() int
}

// SetSandboxStats registers a callback for sandbox execution counters.
func (r *Registry) SetSandboxStats(fn func() (uint64, uint64, float64)) {
	r.sandboxStats = fn
}

// SetTriggerFireCount registers a callback for trigger fire count.
func (r *Registry) SetTriggerFireCount(fn func() int64) {
	r.triggerFires = fn
}

// SetExtensionCount registers a callback for active extension count.
func (r *Registry) SetExtensionCount(fn func() int) {
	r.extensionCount = fn
}

// 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,
		"stack_in_use":  m.StackInuse,
		"gc_cycles":     m.NumGC,
		"gc_pause_ns":   m.PauseNs[(m.NumGC+255)%256],
		"gc_cpu_pct":    m.GCCPUFraction * 100,
		"uptime_sec":    time.Since(r.startTime).Seconds(),
		"ws_clients":    r.hub.ConnCount(),
	}

	if r.extensionCount != nil {
		stats["extensions_loaded"] = r.extensionCount()
	}
	if r.sandboxStats != nil {
		exec, errors, avgMs := r.sandboxStats()
		stats["starlark_exec_total"] = exec
		stats["starlark_errors_total"] = errors
		stats["starlark_avg_duration_ms"] = avgMs
	}
	if r.triggerFires != nil {
		stats["trigger_fires_total"] = r.triggerFires()
	}

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