iop/apps/edge/internal/controlplane/connector.go

package controlplane

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

	toki "git.toki-labs.com/toki/proto-socket/go"
	"go.uber.org/zap"
	"google.golang.org/protobuf/proto"

	edgeevents "iop/apps/edge/internal/events"
	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

const (
	// heartbeatIntervalSec mirrors the edge-node transport constants.
	heartbeatIntervalSec = 30
	// heartbeatWaitSec must exceed heartbeatIntervalSec. See the comment in
	// apps/node/internal/transport/client.go for the rationale.
	heartbeatWaitSec = 45

	helloTimeout = 10 * time.Second
)

// State represents the connection state of the Connector.
type State int32

const (
	StateDisconnected State = iota
	StateConnected
	StateStopped
)

// StatusProvider supplies the Edge-owned node snapshot the Connector reports
// when the Control Plane sends a status request. It is the surface-neutral
// boundary (e.g. edge service) the Connector observes instead of reaching into
// the node registry or transport directly.
type StatusProvider interface {
	ListNodeSnapshots() []edgeservice.NodeSnapshot
}

// Option configures optional Connector dependencies without breaking existing
// call sites.
type Option func(*Connector)

// WithStatusProvider wires the Edge-owned node snapshot source the Connector
// uses to answer Control Plane status requests.
func WithStatusProvider(provider StatusProvider) Option {
	return func(c *Connector) { c.statusProvider = provider }
}

// WithNodeEventBus wires the Edge lifecycle event fanout surface the Connector
// relays to Control Plane. It keeps the connector at the events.Bus boundary
// instead of reaching into node registry or transport internals.
func WithNodeEventBus(bus *edgeevents.Bus) Option {
	return func(c *Connector) { c.nodeEvents = bus }
}

// Connector manages the outbound TCP connection from Edge to Control Plane.
// When disabled (enabled=false or empty wire_addr) it is a safe no-op.
type Connector struct {
	edgeInfo config.EdgeInfo
	cpConf   config.EdgeControlPlaneConf
	version  string
	logger   *zap.Logger

	statusProvider StatusProvider
	nodeEvents     *edgeevents.Bus

	state      atomic.Int32 // stores State
	cancelOnce sync.Once
	cancel     context.CancelFunc

	mu     sync.Mutex
	client *toki.TcpClient
}

// NewConnector creates a Connector from edge identity, Control Plane config,
// version string, and logger. Optional dependencies are supplied via Option.
// Start must be called separately.
func NewConnector(edgeInfo config.EdgeInfo, cpConf config.EdgeControlPlaneConf, version string, logger *zap.Logger, opts ...Option) *Connector {
	c := &Connector{
		edgeInfo: edgeInfo,
		cpConf:   cpConf,
		version:  version,
		logger:   logger,
	}
	for _, opt := range opts {
		opt(c)
	}
	c.state.Store(int32(StateDisconnected))
	return c
}

// StatusProviderConfigured reports whether a status provider was wired, so
// callers (and bootstrap wiring tests) can observe the dependency.
func (c *Connector) StatusProviderConfigured() bool {
	return c.statusProvider != nil
}

// NodeEventBusConfigured reports whether node lifecycle event relay is wired.
func (c *Connector) NodeEventBusConfigured() bool {
	return c.nodeEvents != nil
}

// Start begins the outbound connection loop. It is a no-op when enabled=false
// or wire_addr is empty. ctx cancellation is used only for the startup phase;
// the internal lifetime context drives reconnect. Stop() tears down the loop.
func (c *Connector) Start(ctx context.Context) error {
	if !c.cpConf.Enabled || c.cpConf.WireAddr == "" {
		c.logger.Debug("control plane connector disabled or no wire_addr, skipping")
		return nil
	}

	loopCtx, cancel := context.WithCancel(context.Background())
	c.cancelOnce = sync.Once{}
	c.cancel = cancel

	go c.reconnectLoop(loopCtx)
	return nil
}

// Stop cancels the reconnect loop and closes any active connection.
func (c *Connector) Stop() {
	if c.cancel == nil {
		return
	}
	c.cancelOnce.Do(func() {
		c.state.Store(int32(StateStopped))
		c.cancel()
		c.mu.Lock()
		cl := c.client
		c.client = nil
		c.mu.Unlock()
		if cl != nil {
			_ = cl.Close()
		}
	})
}

// IsEnabled reports whether the connector is configured to run.
func (c *Connector) IsEnabled() bool {
	return c.cpConf.Enabled && c.cpConf.WireAddr != ""
}

// CurrentState returns the current connection State.
func (c *Connector) CurrentState() State {
	return State(c.state.Load())
}

// reconnectLoop dials the Control Plane and reconnects on disconnect until
// ctx is cancelled or Stop() is called.
func (c *Connector) reconnectLoop(ctx context.Context) {
	for {
		if ctx.Err() != nil {
			return
		}
		if err := c.connect(ctx); err != nil {
			c.logger.Warn("control plane connect failed, will retry",
				zap.String("wire_addr", c.cpConf.WireAddr),
				zap.Error(err),
			)
		}
		// Wait for reconnect interval or cancellation.
		select {
		case <-ctx.Done():
			return
		case <-time.After(time.Duration(c.cpConf.ReconnectIntervalSec) * time.Second):
		}
	}
}

// connect dials once, registers a disconnect listener BEFORE sending hello,
// sends EdgeHelloRequest, and blocks until the connection is closed.
//
// Invariant: disconnect listener is registered immediately after DialTcp so
// that any disconnect (including one that races with or follows hello) is
// always captured. done is closed exactly once via sync.Once.
func (c *Connector) connect(ctx context.Context) error {
	host, portStr, err := net.SplitHostPort(c.cpConf.WireAddr)
	if err != nil {
		return fmt.Errorf("controlplane: invalid wire_addr %q: %w", c.cpConf.WireAddr, err)
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return fmt.Errorf("controlplane: invalid port in wire_addr %q: %w", c.cpConf.WireAddr, err)
	}

	dialCtx, dialCancel := context.WithTimeout(ctx, helloTimeout)
	defer dialCancel()

	cl, err := toki.DialTcp(dialCtx, host, port, heartbeatIntervalSec, heartbeatWaitSec, cpParserMap())
	if err != nil {
		return fmt.Errorf("controlplane: dial %s: %w", c.cpConf.WireAddr, err)
	}

	// Register disconnect listener and done channel BEFORE hello so that a
	// disconnect that races with or immediately follows the accepted response
	// is never lost. done is closed exactly once.
	done := make(chan struct{})
	var doneOnce sync.Once
	closeDone := func() { doneOnce.Do(func() { close(done) }) }

	cl.AddDisconnectListener(func(_ *toki.TcpClient) {
		if State(c.state.Load()) != StateStopped {
			c.state.Store(int32(StateDisconnected))
			c.logger.Info("disconnected from control plane, will reconnect",
				zap.String("wire_addr", c.cpConf.WireAddr),
			)
		}
		c.mu.Lock()
		if c.client == cl {
			c.client = nil
		}
		c.mu.Unlock()
		closeDone()
	})

	// Answer Control Plane status requests from the Edge-owned snapshot surface.
	// Registered before hello so a request that arrives right after acceptance is
	// not missed.
	toki.AddRequestListenerTyped(&cl.Communicator, func(req *iop.EdgeStatusRequest) (*iop.EdgeStatusResponse, error) {
		return c.buildStatusResponse(req), nil
	})

	c.mu.Lock()
	c.client = cl
	c.mu.Unlock()

	// Send hello and wait for acceptance.
	resp, err := toki.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&cl.Communicator,
		&iop.EdgeHelloRequest{
			EdgeId:       c.edgeInfo.ID,
			EdgeName:     c.edgeInfo.Name,
			Version:      c.version,
			Capabilities: []string{"run", "node-registry"},
		},
		helloTimeout,
	)
	if err != nil {
		// Close triggers the disconnect listener which will close done.
		_ = cl.Close()
		<-done
		return fmt.Errorf("controlplane: hello: %w", err)
	}
	if !resp.GetAccepted() {
		// Close triggers the disconnect listener which will close done.
		_ = cl.Close()
		<-done
		return fmt.Errorf("controlplane: hello rejected: %s", resp.GetMessage())
	}

	// Check whether the connection was already closed while hello was in-flight.
	select {
	case <-done:
		// Already disconnected; return cleanly so reconnect loop retries.
		return nil
	default:
	}

	stopNodeRelay := c.startNodeEventRelay(ctx, cl)
	defer stopNodeRelay()

	c.state.Store(int32(StateConnected))
	c.logger.Info("connected to control plane",
		zap.String("wire_addr", c.cpConf.WireAddr),
		zap.String("protocol", resp.GetProtocol()),
	)

	// Block until disconnected or ctx cancelled.
	select {
	case <-ctx.Done():
		_ = cl.Close()
		<-done
	case <-done:
	}
	return nil
}

func (c *Connector) startNodeEventRelay(ctx context.Context, cl *toki.TcpClient) func() {
	if c.nodeEvents == nil {
		return func() {}
	}
	events, unsubscribe := c.nodeEvents.SubscribeAllNodes(64)
	relayCtx, cancel := context.WithCancel(ctx)
	go func() {
		defer unsubscribe()
		for {
			select {
			case <-relayCtx.Done():
				return
			case event, ok := <-events:
				if !ok {
					return
				}
				relay := c.nodeEventForRelay(event)
				if relay == nil {
					continue
				}
				if err := cl.Send(relay); err != nil {
					c.logger.Warn("relay node lifecycle event to control plane failed",
						zap.String("event_id", relay.GetEventId()),
						zap.String("node_id", relay.GetNodeId()),
						zap.Error(err),
					)
				}
			}
		}
	}()
	return cancel
}

func (c *Connector) nodeEventForRelay(event *iop.EdgeNodeEvent) *iop.EdgeNodeEvent {
	if event == nil {
		return nil
	}
	relay, ok := proto.Clone(event).(*iop.EdgeNodeEvent)
	if !ok {
		return nil
	}
	if relay.Metadata == nil {
		relay.Metadata = make(map[string]string)
	}
	if c.edgeInfo.ID != "" {
		if _, ok := relay.Metadata["edge_id"]; !ok {
			relay.Metadata["edge_id"] = c.edgeInfo.ID
		}
	}
	if relay.Timestamp == 0 {
		relay.Timestamp = time.Now().UnixNano()
	}
	return relay
}

// buildStatusResponse converts the Edge-owned node snapshot surface into the
// status response. When no status provider is wired it returns an explicit
// error response rather than a silently empty snapshot.
func (c *Connector) buildStatusResponse(req *iop.EdgeStatusRequest) *iop.EdgeStatusResponse {
	resp := &iop.EdgeStatusResponse{
		RequestId:            req.GetRequestId(),
		EdgeId:               c.edgeInfo.ID,
		EdgeName:             c.edgeInfo.Name,
		ObservedTimeUnixNano: time.Now().UnixNano(),
	}
	if c.statusProvider == nil {
		resp.Error = "edge status provider not configured"
		return resp
	}
	snaps := c.statusProvider.ListNodeSnapshots()
	nodes := make([]*iop.EdgeNodeSnapshot, 0, len(snaps))
	for _, s := range snaps {
		nodes = append(nodes, &iop.EdgeNodeSnapshot{
			NodeId:    s.NodeID,
			Alias:     s.Alias,
			Label:     s.Label,
			Connected: true,
			Config:    s.Config,
		})
	}
	resp.Nodes = nodes
	return resp
}

// cpParserMap returns the ParserMap for Control Plane messages. It parses the
// hello response and the Control Plane-initiated status request, plus the
// status response the Edge marshals back.
func cpParserMap() toki.ParserMap {
	return toki.ParserMap{
		toki.TypeNameOf(&iop.EdgeHelloResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.EdgeHelloResponse{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.EdgeStatusRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.EdgeStatusRequest{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.EdgeStatusResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.EdgeStatusResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}
}