iop/apps/node/internal/transport/client.go

package transport

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

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

	iop "iop/proto/gen/iop"
)

// connectFailureClass classifies a DialEdge failure so the connectivity
// supervisor can distinguish transient failures it should retry from fatal ones
// that must terminate the Node.
type connectFailureClass int

const (
	// connectFailureRetryable marks a transient failure (dial refused, network
	// unreachable, timeout, transient register transport failure) that the
	// supervisor should retry under the reconnect policy.
	connectFailureRetryable connectFailureClass = iota
	// connectFailureFatal marks a non-retryable failure (invalid address or
	// authoritative registration rejection) that retrying cannot recover.
	connectFailureFatal
)

// ConnectError wraps a DialEdge failure with a retry classification while
// preserving the underlying cause for errors.Is/errors.As and any transport
// status it carries. The supervisor consumes the classification through
// IsFatalConnectError rather than comparing error strings.
type ConnectError struct {
	class connectFailureClass
	cause error
}

// Error returns the wrapped cause's message unchanged.
func (e *ConnectError) Error() string {
	if e.cause == nil {
		return "transport: connect error"
	}
	return e.cause.Error()
}

// Unwrap exposes the original cause so errors.Is/errors.As reach it.
func (e *ConnectError) Unwrap() error { return e.cause }

// Retryable reports whether the failure is a transient one the supervisor
// should retry.
func (e *ConnectError) Retryable() bool { return e.class == connectFailureRetryable }

func wrapConnectError(class connectFailureClass, cause error) *ConnectError {
	return &ConnectError{class: class, cause: cause}
}

// IsFatalConnectError reports whether err (or any wrapped cause) is a fatal,
// non-retryable connect failure. Errors that are not classified ConnectErrors
// are treated as retryable transient failures.
func IsFatalConnectError(err error) bool {
	var ce *ConnectError
	if errors.As(err, &ce) {
		return !ce.Retryable()
	}
	return false
}

const (
	heartbeatIntervalSec = 2
	// heartbeatWaitSec is kept above heartbeatIntervalSec as defence in depth:
	// the library wait-timer callback already self-heals stale state (see
	// proto-socket go/base_client.go sendHeartBeat), but a larger wait window
	// gives the peer's next heartbeat an extra chance to overwrite any stray
	// timer on slow or jittery links before it fires.
	heartbeatWaitSec    = 5
	tcpWriteTimeout     = 10 * time.Second
	registerTimeout     = 10 * time.Second
	registerInitialWait = 100 * time.Millisecond
	registerRetryWait   = 250 * time.Millisecond
	registerAttempts    = 3
)

type writeDeadlineConn struct {
	net.Conn
	timeout time.Duration
}

func (c *writeDeadlineConn) Write(b []byte) (int, error) {
	if c.timeout > 0 {
		if err := c.Conn.SetWriteDeadline(time.Now().Add(c.timeout)); err != nil {
			return 0, err
		}
		defer c.Conn.SetWriteDeadline(time.Time{})
	}
	return c.Conn.Write(b)
}

// RegisterResult is returned by DialEdge after successful registration.
type RegisterResult struct {
	Session *Session
	NodeID  string
	Alias   string
	Config  *iop.NodeConfigPayload
}

// DialEdge connects to edge, performs the registration handshake, and returns
// a RegisterResult. Call result.Session.SetHandler after creating node.Node.
func DialEdge(ctx context.Context, addr, token string, logger *zap.Logger) (*RegisterResult, error) {
	host, portStr, err := net.SplitHostPort(addr)
	if err != nil {
		// Local address preparation error: retrying cannot recover it.
		return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: invalid addr %q: %w", addr, err))
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: invalid port %q: %w", portStr, err))
	}
	if port < 1 || port > 65535 {
		// A numeric port outside the TCP range is still a local configuration
		// error. Preserve a standard range cause so callers can inspect it with
		// errors.Is/errors.As instead of relying on the message.
		rangeErr := &strconv.NumError{Func: "Atoi", Num: portStr, Err: strconv.ErrRange}
		return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: invalid port %q: %w", portStr, rangeErr))
	}

	dialer := net.Dialer{KeepAlive: 15 * time.Second}
	conn, err := dialer.DialContext(ctx, "tcp", net.JoinHostPort(host, strconv.Itoa(port)))
	if err != nil {
		// Edge unavailable / network transient failure: retryable.
		return nil, wrapConnectError(connectFailureRetryable, fmt.Errorf("transport: dial edge %s: %w", addr, err))
	}
	client := toki.NewTcpClient(&writeDeadlineConn{Conn: conn, timeout: tcpWriteTimeout}, heartbeatIntervalSec, heartbeatWaitSec, nodeParserMap())

	resp, err := registerWithEdge(ctx, client, token, logger)
	if err != nil {
		_ = client.Close()
		// Register request transport failure or context cancellation: retryable.
		return nil, wrapConnectError(connectFailureRetryable, fmt.Errorf("transport: register: %w", err))
	}
	if !resp.GetAccepted() {
		_ = client.Close()
		// Authoritative registration rejection (bad token/credential): fatal.
		return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: register rejected: %s", resp.GetReason()))
	}

	sess := newSession(client, logger, resp.GetNodeId(), resp.GetAlias())
	logger.Info("registered with edge",
		zap.String("node_id", resp.GetNodeId()),
		zap.String("alias", resp.GetAlias()),
	)
	return &RegisterResult{
		Session: sess,
		NodeID:  resp.GetNodeId(),
		Alias:   resp.GetAlias(),
		Config:  resp.GetConfig(),
	}, nil
}

func registerWithEdge(ctx context.Context, client *toki.TcpClient, token string, logger *zap.Logger) (*iop.RegisterResponse, error) {
	timer := time.NewTimer(registerInitialWait)
	select {
	case <-ctx.Done():
		timer.Stop()
		return nil, ctx.Err()
	case <-timer.C:
	}

	var lastErr error
	for attempt := 1; attempt <= registerAttempts; attempt++ {
		resp, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
			&client.Communicator,
			&iop.RegisterRequest{Token: token},
			registerTimeout,
		)
		if err == nil {
			return resp, nil
		}
		lastErr = err
		if attempt == registerAttempts || !client.IsAlive() {
			break
		}
		if logger != nil {
			logger.Warn("register request failed, retrying",
				zap.Int("attempt", attempt),
				zap.Error(err),
			)
		}
		timer := time.NewTimer(registerRetryWait)
		select {
		case <-ctx.Done():
			timer.Stop()
			return nil, ctx.Err()
		case <-timer.C:
		}
	}
	return nil, lastErr
}