iop/apps/control-plane/internal/wire/edge_registry.go

package wire

import (
	"sort"
	"sync"
	"time"

	proto_socket "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"
	iop "iop/proto/gen/iop"
)

const recentNodeEventCapacity = 128
const recentCommandCapacity = 128
const commandEventCapacity = 64

// EdgeConnectionState captures the Control Plane's internal view of a single
// Edge wire connection.
type EdgeConnectionState struct {
	EdgeID         string
	EdgeName       string
	Version        string
	Capabilities   []string
	Protocol       string
	Connected      bool
	LastSeen       time.Time
	DisconnectInfo proto_socket.DisconnectInfo

	// token identifies the specific connection that produced this state. It is an
	// internal registry detail (not an external API contract) used to ignore a
	// stale connection's late disconnect after the same edge_id reconnected.
	token uint64
}

// EdgeNodeEventRecord is the Control Plane's copied recent view of one node
// lifecycle event as received from a specific Edge connection.
type EdgeNodeEventRecord struct {
	EdgeID     string
	Event      *iop.EdgeNodeEvent
	ReceivedAt time.Time
}

// EdgeCommandRecord is the Control Plane's bounded audit view of one command
// dispatched to an Edge: the request summary, the latest result, and a bounded
// list of relayed lifecycle events. It is a control/audit view; durable command
// history is intentionally outside this in-memory baseline.
type EdgeCommandRecord struct {
	EdgeID         string
	CommandID      string
	Operation      string
	TargetSelector string
	Status         string
	Summary        string
	Error          string
	CreatedAt      time.Time
	UpdatedAt      time.Time
	Events         []EdgeCommandEventRecord
}

// EdgeCommandEventRecord is the Control Plane's copied view of one command
// lifecycle event relayed from an Edge during command execution.
type EdgeCommandEventRecord struct {
	Phase      string
	Summary    string
	OccurredAt time.Time
	ReceivedAt time.Time
}

// EdgeRegistry stores Edge connection state keyed by edge_id behind a mutex.
type EdgeRegistry struct {
	mu               sync.Mutex
	edges            map[string]*EdgeConnectionState
	nextToken        uint64
	recentNodeEvents []EdgeNodeEventRecord
	commands         []*EdgeCommandRecord
}

// NewEdgeRegistry returns an empty Edge connection registry.
func NewEdgeRegistry() *EdgeRegistry {
	return &EdgeRegistry{edges: make(map[string]*EdgeConnectionState)}
}

// MarkConnected records or refreshes an Edge as connected from its hello and
// returns the connection token to hand back on disconnect. A reconnection
// supersedes any earlier connection for the same edge_id.
func (r *EdgeRegistry) MarkConnected(req *iop.EdgeHelloRequest, at time.Time) uint64 {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.nextToken++
	token := r.nextToken
	r.edges[req.GetEdgeId()] = &EdgeConnectionState{
		EdgeID:       req.GetEdgeId(),
		EdgeName:     req.GetEdgeName(),
		Version:      req.GetVersion(),
		Capabilities: append([]string(nil), req.GetCapabilities()...),
		Protocol:     Protocol,
		Connected:    true,
		LastSeen:     at,
		token:        token,
	}
	return token
}

// MarkDisconnected flips a tracked Edge to disconnected and records why, but
// only when token still matches the current connection. A stale token from a
// superseded connection is ignored so a reconnection is not cleared. It returns
// true only when the current connection was flipped, so callers (e.g. the
// EdgeServer active client map) can apply the same stale-safe cleanup.
func (r *EdgeRegistry) MarkDisconnected(edgeID string, token uint64, at time.Time, info proto_socket.DisconnectInfo) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	state, ok := r.edges[edgeID]
	if !ok || state.token != token {
		return false
	}
	state.Connected = false
	state.LastSeen = at
	state.DisconnectInfo = info
	return true
}

// Snapshot returns a copy of the tracked state for an Edge.
func (r *EdgeRegistry) Snapshot(edgeID string) (EdgeConnectionState, bool) {
	r.mu.Lock()
	defer r.mu.Unlock()
	state, ok := r.edges[edgeID]
	if !ok {
		return EdgeConnectionState{}, false
	}
	snapshot := *state
	snapshot.Capabilities = append([]string(nil), state.Capabilities...)
	return snapshot, true
}

// Snapshots returns a stable, copied view of all tracked Edge states.
func (r *EdgeRegistry) Snapshots() []EdgeConnectionState {
	r.mu.Lock()
	defer r.mu.Unlock()
	ids := make([]string, 0, len(r.edges))
	for id := range r.edges {
		ids = append(ids, id)
	}
	sort.Strings(ids)
	snapshots := make([]EdgeConnectionState, 0, len(ids))
	for _, id := range ids {
		state := r.edges[id]
		snapshot := *state
		snapshot.Capabilities = append([]string(nil), state.Capabilities...)
		snapshots = append(snapshots, snapshot)
	}
	return snapshots
}

// RecordNodeEvent stores a copied EdgeNodeEvent under the identified Edge. It
// is intended for tests and non-connection seed paths that do not have an
// active connection token.
func (r *EdgeRegistry) RecordNodeEvent(edgeID string, event *iop.EdgeNodeEvent, receivedAt time.Time) {
	r.recordNodeEvent(edgeID, 0, false, event, receivedAt)
}

// RecordNodeEventFromConnection stores a copied EdgeNodeEvent only when the
// supplied token still identifies the current connection for edgeID. It returns
// false for stale superseded connections.
func (r *EdgeRegistry) RecordNodeEventFromConnection(edgeID string, token uint64, event *iop.EdgeNodeEvent, receivedAt time.Time) bool {
	return r.recordNodeEvent(edgeID, token, true, event, receivedAt)
}

// recordNodeEvent appends to the bounded recent buffer. The recent buffer is
// global to the in-memory registry; durable event history is intentionally
// outside this baseline.
func (r *EdgeRegistry) recordNodeEvent(edgeID string, token uint64, requireToken bool, event *iop.EdgeNodeEvent, receivedAt time.Time) bool {
	if edgeID == "" || event == nil {
		return false
	}
	if receivedAt.IsZero() {
		receivedAt = time.Now()
	}
	copied := cloneEdgeNodeEvent(event)
	if copied.GetTimestamp() == 0 {
		copied.Timestamp = receivedAt.UnixNano()
	}

	r.mu.Lock()
	defer r.mu.Unlock()
	if requireToken {
		state, ok := r.edges[edgeID]
		if !ok || state.token != token {
			return false
		}
	}
	if len(r.recentNodeEvents) >= recentNodeEventCapacity {
		r.recentNodeEvents = append(r.recentNodeEvents[:0], r.recentNodeEvents[1:]...)
	}
	r.recentNodeEvents = append(r.recentNodeEvents, EdgeNodeEventRecord{
		EdgeID:     edgeID,
		Event:      copied,
		ReceivedAt: receivedAt,
	})
	return true
}

// NodeEvents returns copied recent node lifecycle events for an Edge. When
// nodeID is non-empty, only events for that node are returned.
func (r *EdgeRegistry) NodeEvents(edgeID, nodeID string) []EdgeNodeEventRecord {
	if edgeID == "" {
		return nil
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	out := make([]EdgeNodeEventRecord, 0, len(r.recentNodeEvents))
	for _, record := range r.recentNodeEvents {
		if record.EdgeID != edgeID {
			continue
		}
		if nodeID != "" && record.Event.GetNodeId() != nodeID {
			continue
		}
		out = append(out, cloneEdgeNodeEventRecord(record))
	}
	if len(out) == 0 {
		return nil
	}
	return out
}

// RecordCommandRequest stores or refreshes a pending command audit record for a
// command the Control Plane is dispatching to an Edge. It is bounded; the oldest
// record is evicted past capacity.
func (r *EdgeRegistry) RecordCommandRequest(edgeID string, req *iop.EdgeCommandRequest, at time.Time) {
	if edgeID == "" || req == nil || req.GetCommandId() == "" {
		return
	}
	if at.IsZero() {
		at = time.Now()
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if rec := r.findCommandLocked(edgeID, req.GetCommandId()); rec != nil {
		rec.Operation = req.GetOperation()
		rec.TargetSelector = req.GetTargetSelector()
		rec.UpdatedAt = at
		return
	}
	rec := r.appendCommandLocked(edgeID, req.GetCommandId(), at)
	rec.Operation = req.GetOperation()
	rec.TargetSelector = req.GetTargetSelector()
	rec.Status = "pending"
}

// RecordCommandResult updates the audit record for a command with the Edge's
// response status/summary. A missing record is created so a result is never
// silently dropped.
func (r *EdgeRegistry) RecordCommandResult(edgeID string, resp *iop.EdgeCommandResponse, at time.Time) {
	if edgeID == "" || resp == nil || resp.GetCommandId() == "" {
		return
	}
	if at.IsZero() {
		at = time.Now()
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	rec := r.findCommandLocked(edgeID, resp.GetCommandId())
	if rec == nil {
		rec = r.appendCommandLocked(edgeID, resp.GetCommandId(), at)
	}
	status := resp.GetStatus()
	if status == "" {
		status = "ok"
	}
	rec.Status = status
	rec.Summary = resp.GetSummary()
	rec.Error = resp.GetError()
	rec.UpdatedAt = at
}

// RecordCommandEventFromConnection appends a bounded command lifecycle event to
// the matching command audit record, but only when token still identifies the
// current connection for edgeID. It returns false for stale superseded
// connections.
func (r *EdgeRegistry) RecordCommandEventFromConnection(edgeID string, token uint64, event *iop.EdgeCommandEvent, receivedAt time.Time) bool {
	return r.recordCommandEvent(edgeID, token, true, event, receivedAt)
}

// RecordCommandEvent stores a command lifecycle event without a connection
// token. It is intended for tests and non-connection seed paths.
func (r *EdgeRegistry) RecordCommandEvent(edgeID string, event *iop.EdgeCommandEvent, receivedAt time.Time) {
	r.recordCommandEvent(edgeID, 0, false, event, receivedAt)
}

func (r *EdgeRegistry) recordCommandEvent(edgeID string, token uint64, requireToken bool, event *iop.EdgeCommandEvent, receivedAt time.Time) bool {
	if edgeID == "" || event == nil || event.GetCommandId() == "" {
		return false
	}
	if receivedAt.IsZero() {
		receivedAt = time.Now()
	}
	occurredAt := receivedAt
	if event.GetOccurredAt() != 0 {
		occurredAt = time.Unix(0, event.GetOccurredAt())
	}

	r.mu.Lock()
	defer r.mu.Unlock()
	if requireToken {
		state, ok := r.edges[edgeID]
		if !ok || state.token != token {
			return false
		}
	}
	rec := r.findCommandLocked(edgeID, event.GetCommandId())
	if rec == nil {
		rec = r.appendCommandLocked(edgeID, event.GetCommandId(), receivedAt)
		rec.Status = "running"
	}
	if len(rec.Events) >= commandEventCapacity {
		rec.Events = append(rec.Events[:0], rec.Events[1:]...)
	}
	rec.Events = append(rec.Events, EdgeCommandEventRecord{
		Phase:      event.GetPhase(),
		Summary:    event.GetSummary(),
		OccurredAt: occurredAt,
		ReceivedAt: receivedAt,
	})
	rec.UpdatedAt = receivedAt
	return true
}

// Commands returns copied command audit records for an Edge, oldest first.
func (r *EdgeRegistry) Commands(edgeID string) []EdgeCommandRecord {
	if edgeID == "" {
		return nil
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	out := make([]EdgeCommandRecord, 0, len(r.commands))
	for _, rec := range r.commands {
		if rec.EdgeID != edgeID {
			continue
		}
		out = append(out, cloneEdgeCommandRecord(rec))
	}
	if len(out) == 0 {
		return nil
	}
	return out
}

// Command returns a single copied command audit record.
func (r *EdgeRegistry) Command(edgeID, commandID string) (EdgeCommandRecord, bool) {
	r.mu.Lock()
	defer r.mu.Unlock()
	rec := r.findCommandLocked(edgeID, commandID)
	if rec == nil {
		return EdgeCommandRecord{}, false
	}
	return cloneEdgeCommandRecord(rec), true
}

// findCommandLocked returns the stored command record for an edge_id/command_id
// pair. Callers must hold r.mu.
func (r *EdgeRegistry) findCommandLocked(edgeID, commandID string) *EdgeCommandRecord {
	for _, rec := range r.commands {
		if rec.EdgeID == edgeID && rec.CommandID == commandID {
			return rec
		}
	}
	return nil
}

// appendCommandLocked appends a new bounded command record. Callers must hold
// r.mu.
func (r *EdgeRegistry) appendCommandLocked(edgeID, commandID string, at time.Time) *EdgeCommandRecord {
	if len(r.commands) >= recentCommandCapacity {
		r.commands = append(r.commands[:0], r.commands[1:]...)
	}
	rec := &EdgeCommandRecord{
		EdgeID:    edgeID,
		CommandID: commandID,
		CreatedAt: at,
		UpdatedAt: at,
	}
	r.commands = append(r.commands, rec)
	return rec
}

func cloneEdgeCommandRecord(rec *EdgeCommandRecord) EdgeCommandRecord {
	clone := *rec
	clone.Events = append([]EdgeCommandEventRecord(nil), rec.Events...)
	return clone
}

// Len returns the number of tracked Edge connections.
func (r *EdgeRegistry) Len() int {
	r.mu.Lock()
	defer r.mu.Unlock()
	return len(r.edges)
}

func cloneEdgeNodeEventRecord(record EdgeNodeEventRecord) EdgeNodeEventRecord {
	return EdgeNodeEventRecord{
		EdgeID:     record.EdgeID,
		Event:      cloneEdgeNodeEvent(record.Event),
		ReceivedAt: record.ReceivedAt,
	}
}

func cloneEdgeNodeEvent(event *iop.EdgeNodeEvent) *iop.EdgeNodeEvent {
	if event == nil {
		return nil
	}
	copied, ok := proto.Clone(event).(*iop.EdgeNodeEvent)
	if !ok {
		return &iop.EdgeNodeEvent{}
	}
	return copied
}