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

// 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
}

// 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
}

// 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
}

// 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
}