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"

	"iop/packages/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
)

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

	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. Start must be called separately.
func NewConnector(edgeInfo config.EdgeInfo, cpConf config.EdgeControlPlaneConf, version string, logger *zap.Logger) *Connector {
	c := &Connector{
		edgeInfo: edgeInfo,
		cpConf:   cpConf,
		version:  version,
		logger:   logger,
	}
	c.state.Store(int32(StateDisconnected))
	return c
}

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

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

	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
}

// cpParserMap returns the ParserMap for Control Plane messages.
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)
		},
	}
}