iop/apps/edge/internal/transport/server.go

package transport

import (
	"context"
	"net"
	"strconv"
	"time"

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

	edgenode "iop/apps/edge/internal/node"
	iop "iop/proto/gen/iop"
)

const (
	heartbeatIntervalSec = 30
	heartbeatWaitSec     = 10
	capabilityTimeout    = 5 * time.Second
)

func edgeParserMap() toki.ParserMap {
	return toki.ParserMap{
		toki.TypeNameOf(&iop.RunEvent{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunEvent{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.CapabilityResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.CapabilityResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}
}

// Server wraps proto-socket TcpServer and manages node connections.
type Server struct {
	tcp      *toki.TcpServer
	listen   string
	registry *edgenode.Registry
	logger   *zap.Logger
}

func NewServer(listen string, registry *edgenode.Registry, logger *zap.Logger) (*Server, error) {
	host, portStr, err := net.SplitHostPort(listen)
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, err
	}

	s := &Server{listen: listen, registry: registry, logger: logger}
	s.tcp = toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
		return toki.NewTcpClient(conn, heartbeatIntervalSec, heartbeatWaitSec, edgeParserMap())
	})
	s.tcp.OnClientConnected = s.onNodeConnected
	return s, nil
}

func (s *Server) Start(ctx context.Context) error {
	if err := s.tcp.Start(ctx); err != nil {
		return err
	}
	s.logger.Info("edge listening for nodes", zap.String("addr", s.listen))
	return nil
}

func (s *Server) Stop() error {
	return s.tcp.Stop()
}

func (s *Server) onNodeConnected(client *toki.TcpClient) {
	s.logger.Info("node connection established")

	toki.AddListenerTyped[*iop.RunEvent](&client.Communicator, func(e *iop.RunEvent) {
		s.logger.Debug("run event received",
			zap.String("run_id", e.GetRunId()),
			zap.String("type", e.GetType()),
		)
	})

	go func() {
		resp, err := toki.SendRequestTyped[*iop.CapabilityRequest, *iop.CapabilityResponse](
			&client.Communicator,
			&iop.CapabilityRequest{},
			capabilityTimeout,
		)
		if err != nil {
			s.logger.Warn("capability request failed", zap.Error(err))
			_ = client.Close()
			return
		}

		entry := &edgenode.NodeEntry{
			NodeID:   resp.GetNodeId(),
			Client:   client,
			Adapters: resp.GetAdapters(),
		}
		client.AddDisconnectListener(func(_ *toki.TcpClient) {
			s.registry.Unregister(entry.NodeID)
			s.logger.Info("node unregistered", zap.String("node_id", entry.NodeID))
		})

		s.registry.Register(entry)
		s.logger.Info("node registered",
			zap.String("node_id", entry.NodeID),
			zap.Int("adapters", len(entry.Adapters)),
		)
	}()
}