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

package wire

import (
	"context"
	"fmt"
	"net"
	"strconv"
	"sync"
	"time"

	"go.uber.org/zap"

	proto_socket "git.toki-labs.com/toki/proto-socket/go"
	iop "iop/proto/gen/iop"
)

// EdgeServer hosts the Control Plane-Edge native TCP wire endpoint. It accepts
// Edge connections over proto-socket, handles the hello handshake, and keeps an
// internal connection registry. It does not connect to or schedule Node
// directly; the Control Plane observes and controls the system through Edge.
type EdgeServer struct {
	server   *proto_socket.TcpServer
	registry *EdgeRegistry
	logger   *zap.Logger
	host     string
	port     int

	// activeMu guards activeClients, the live proto-socket client per edge_id
	// the Control Plane uses to send typed requests (e.g. status) to a connected
	// Edge. token mirrors the registry connection token so a stale connection's
	// late disconnect does not evict a newer reconnection's client.
	activeMu      sync.Mutex
	activeClients map[string]activeEdgeClient
}

type activeEdgeClient struct {
	client *proto_socket.TcpClient
	token  uint64
}

// NewEdgeServer creates an EdgeServer bound to the given TCP listen address.
func NewEdgeServer(listenAddr string, logger *zap.Logger) (*EdgeServer, error) {
	host, portStr, err := net.SplitHostPort(listenAddr)
	if err != nil {
		return nil, fmt.Errorf("invalid edge listen address %q: %w", listenAddr, err)
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, fmt.Errorf("invalid port in edge listen address %q: %w", listenAddr, err)
	}

	registry := NewEdgeRegistry()
	es := &EdgeServer{
		registry:      registry,
		logger:        logger,
		host:          host,
		port:          port,
		activeClients: make(map[string]activeEdgeClient),
	}

	es.server = proto_socket.NewTcpServer(host, port, func(conn net.Conn) *proto_socket.TcpClient {
		client := proto_socket.NewTcpClient(conn, EdgeHeartbeatIntervalSec, EdgeHeartbeatWaitSec, EdgeParserMap())

		// edgeID and token are set once the hello handshake identifies the
		// connection so the disconnect listener can flip the right registry entry.
		// The token lets a stale connection's late disconnect be ignored after the
		// same edge_id reconnected. The hello and disconnect callbacks may run on
		// different goroutines, so guard them.
		var (
			idMu   sync.Mutex
			edgeID string
			token  uint64
		)

		proto_socket.AddListenerTyped[*iop.EdgeNodeEvent](&client.Communicator, func(event *iop.EdgeNodeEvent) {
			idMu.Lock()
			id := edgeID
			tok := token
			idMu.Unlock()
			if id == "" {
				return
			}
			registry.RecordNodeEventFromConnection(id, tok, event, time.Now())
		})

		proto_socket.AddRequestListenerTyped(&client.Communicator, func(req *iop.EdgeHelloRequest) (*iop.EdgeHelloResponse, error) {
			now := time.Now()
			if req.GetEdgeId() == "" {
				logger.Warn("edge hello rejected: missing edge_id")
				return &iop.EdgeHelloResponse{
					Accepted:           false,
					Protocol:           Protocol,
					ServerTimeUnixNano: now.UnixNano(),
					Reason:             "edge_id is required",
				}, nil
			}

			connToken := registry.MarkConnected(req, now)
			idMu.Lock()
			edgeID = req.GetEdgeId()
			token = connToken
			idMu.Unlock()
			es.setActiveClient(req.GetEdgeId(), connToken, client)

			logger.Info("edge hello accepted",
				zap.String("edge_id", req.GetEdgeId()),
				zap.String("edge_name", req.GetEdgeName()),
				zap.String("version", req.GetVersion()),
			)
			return &iop.EdgeHelloResponse{
				Accepted:           true,
				Protocol:           Protocol,
				ServerTimeUnixNano: now.UnixNano(),
				Message:            "Edge enrolled on IOP Control Plane",
			}, nil
		})

		client.AddDisconnectListener(func(c *proto_socket.TcpClient) {
			idMu.Lock()
			id := edgeID
			tok := token
			idMu.Unlock()
			if id == "" {
				return
			}
			info := c.DisconnectInfo()
			if registry.MarkDisconnected(id, tok, time.Now(), info) {
				es.clearActiveClient(id, tok)
			}
			logger.Info("edge disconnected",
				zap.String("edge_id", id),
				zap.String("reason", info.Reason),
			)
		})

		return client
	})

	return es, nil
}

// Start begins accepting Edge connections on the configured TCP endpoint.
func (s *EdgeServer) Start(ctx context.Context) error {
	s.logger.Info("starting edge wire TCP server",
		zap.String("host", s.host),
		zap.Int("port", s.port),
		zap.String("transport", EdgeTransport),
	)
	return s.server.Start(ctx)
}

// Stop stops the Edge wire TCP server and closes active connections.
func (s *EdgeServer) Stop() error {
	s.logger.Info("stopping edge wire TCP server")
	return s.server.Stop()
}

// Registry exposes the internal Edge connection registry.
func (s *EdgeServer) Registry() *EdgeRegistry { return s.registry }

// setActiveClient records the live proto-socket client for a connected Edge so
// the Control Plane can send typed requests back to it. A newer connection
// always supersedes an older one for the same edge_id.
func (s *EdgeServer) setActiveClient(edgeID string, token uint64, client *proto_socket.TcpClient) {
	s.activeMu.Lock()
	defer s.activeMu.Unlock()
	s.activeClients[edgeID] = activeEdgeClient{client: client, token: token}
}

// clearActiveClient removes the active client for an Edge, but only when token
// still matches the tracked connection so a stale disconnect cannot evict a
// newer reconnection.
func (s *EdgeServer) clearActiveClient(edgeID string, token uint64) {
	s.activeMu.Lock()
	defer s.activeMu.Unlock()
	if entry, ok := s.activeClients[edgeID]; ok && entry.token == token {
		delete(s.activeClients, edgeID)
	}
}

// activeClient returns the live proto-socket client for a connected Edge.
func (s *EdgeServer) activeClient(edgeID string) (*proto_socket.TcpClient, bool) {
	s.activeMu.Lock()
	defer s.activeMu.Unlock()
	entry, ok := s.activeClients[edgeID]
	if !ok {
		return nil, false
	}
	return entry.client, true
}

// RequestStatus asks a connected Edge for its Edge-owned node snapshot over the
// wire and returns the typed response. It observes only what the Edge reports;
// it does not connect to or schedule Node directly. It fails when the Edge is
// not currently connected.
func (s *EdgeServer) RequestStatus(edgeID string, timeout time.Duration) (*iop.EdgeStatusResponse, error) {
	client, ok := s.activeClient(edgeID)
	if !ok {
		return nil, fmt.Errorf("edge %q is not connected", edgeID)
	}
	return proto_socket.SendRequestTyped[*iop.EdgeStatusRequest, *iop.EdgeStatusResponse](
		&client.Communicator,
		&iop.EdgeStatusRequest{RequestId: fmt.Sprintf("edge-status-%d", time.Now().UnixNano())},
		timeout,
	)
}

// Host returns the host the server is listening on.
func (s *EdgeServer) Host() string { return s.host }

// Port returns the port the server is listening on.
func (s *EdgeServer) Port() int { return s.port }