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

700 lines
21 KiB
Go

package controlplane
import (
"context"
"net"
"sync/atomic"
"testing"
"time"
toki "git.toki-labs.com/toki/proto-socket/go"
"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"
)
type fakeStatusProvider struct {
nodes []edgeservice.NodeSnapshot
}
func (f fakeStatusProvider) ListNodeSnapshots() []edgeservice.NodeSnapshot {
return append([]edgeservice.NodeSnapshot(nil), f.nodes...)
}
// fakeCPParserMap returns a parser for EdgeHelloRequest on the server side.
func fakeCPParserMap() toki.ParserMap {
return toki.ParserMap{
toki.TypeNameOf(&iop.EdgeHelloRequest{}): func(b []byte) (proto.Message, error) {
m := &iop.EdgeHelloRequest{}
return m, proto.Unmarshal(b, m)
},
toki.TypeNameOf(&iop.EdgeNodeEvent{}): func(b []byte) (proto.Message, error) {
m := &iop.EdgeNodeEvent{}
return m, proto.Unmarshal(b, m)
},
}
}
// startFakeCPServer starts a fake Control Plane TCP server on a free port and
// calls onHello for each EdgeHelloRequest, returning the provided response.
// It returns the port and a stop function.
func startFakeCPServer(t *testing.T, onHello func(*iop.EdgeHelloRequest) *iop.EdgeHelloResponse) (int, func()) {
t.Helper()
port, err := freeTCPPort()
if err != nil {
t.Fatalf("startFakeCPServer freeTCPPort: %v", err)
}
srv := toki.NewTcpServer("127.0.0.1", port, func(conn net.Conn) *toki.TcpClient {
return toki.NewTcpClient(conn, 0, 0, fakeCPParserMap())
})
srv.OnClientConnected = func(client *toki.TcpClient) {
toki.AddRequestListenerTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
&client.Communicator,
func(req *iop.EdgeHelloRequest) (*iop.EdgeHelloResponse, error) {
return onHello(req), nil
},
)
}
ctx, cancel := context.WithCancel(context.Background())
if err := srv.Start(ctx); err != nil {
cancel()
t.Fatalf("startFakeCPServer Start: %v", err)
}
return port, func() {
cancel()
_ = srv.Stop()
}
}
// startFakeCPServerWithClientAccess starts a fake server that also exposes
// each connected client via a channel, so the test can close clients directly.
func startFakeCPServerWithClientAccess(t *testing.T, onHello func(*iop.EdgeHelloRequest) *iop.EdgeHelloResponse) (int, <-chan *toki.TcpClient, func()) {
t.Helper()
port, err := freeTCPPort()
if err != nil {
t.Fatalf("startFakeCPServerWithClientAccess freeTCPPort: %v", err)
}
clients := make(chan *toki.TcpClient, 16)
srv := toki.NewTcpServer("127.0.0.1", port, func(conn net.Conn) *toki.TcpClient {
return toki.NewTcpClient(conn, 0, 0, fakeCPParserMap())
})
srv.OnClientConnected = func(client *toki.TcpClient) {
clients <- client
toki.AddRequestListenerTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
&client.Communicator,
func(req *iop.EdgeHelloRequest) (*iop.EdgeHelloResponse, error) {
return onHello(req), nil
},
)
}
ctx, cancel := context.WithCancel(context.Background())
if err := srv.Start(ctx); err != nil {
cancel()
t.Fatalf("startFakeCPServerWithClientAccess Start: %v", err)
}
return port, clients, func() {
cancel()
_ = srv.Stop()
}
}
// wireAddr builds a host:port address string from a port int.
func wireAddr(port int) string {
return net.JoinHostPort("127.0.0.1", tcpPortStr(port))
}
// statusCapableCPParserMap parses the Edge hello plus the status request the
// Edge marshals back so the fake Control Plane can both accept hello and read
// the status response.
func statusCapableCPParserMap() toki.ParserMap {
return toki.ParserMap{
toki.TypeNameOf(&iop.EdgeHelloRequest{}): func(b []byte) (proto.Message, error) {
m := &iop.EdgeHelloRequest{}
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)
},
}
}
// startFakeCPServerForStatus starts a fake Control Plane that accepts hello and
// exposes connected clients, so a test can send a status request to the Edge.
func startFakeCPServerForStatus(t *testing.T) (int, <-chan *toki.TcpClient, func()) {
t.Helper()
port, err := freeTCPPort()
if err != nil {
t.Fatalf("startFakeCPServerForStatus freeTCPPort: %v", err)
}
clients := make(chan *toki.TcpClient, 16)
srv := toki.NewTcpServer("127.0.0.1", port, func(conn net.Conn) *toki.TcpClient {
return toki.NewTcpClient(conn, 0, 0, statusCapableCPParserMap())
})
srv.OnClientConnected = func(client *toki.TcpClient) {
clients <- client
toki.AddRequestListenerTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
&client.Communicator,
func(_ *iop.EdgeHelloRequest) (*iop.EdgeHelloResponse, error) {
return &iop.EdgeHelloResponse{Accepted: true, Protocol: "iop/1"}, nil
},
)
}
ctx, cancel := context.WithCancel(context.Background())
if err := srv.Start(ctx); err != nil {
cancel()
t.Fatalf("startFakeCPServerForStatus Start: %v", err)
}
return port, clients, func() {
cancel()
_ = srv.Stop()
}
}
func startFakeCPServerForNodeEvents(t *testing.T) (int, <-chan *iop.EdgeNodeEvent, func()) {
t.Helper()
port, err := freeTCPPort()
if err != nil {
t.Fatalf("startFakeCPServerForNodeEvents freeTCPPort: %v", err)
}
events := make(chan *iop.EdgeNodeEvent, 8)
srv := toki.NewTcpServer("127.0.0.1", port, func(conn net.Conn) *toki.TcpClient {
return toki.NewTcpClient(conn, 0, 0, fakeCPParserMap())
})
srv.OnClientConnected = func(client *toki.TcpClient) {
toki.AddRequestListenerTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
&client.Communicator,
func(_ *iop.EdgeHelloRequest) (*iop.EdgeHelloResponse, error) {
return &iop.EdgeHelloResponse{Accepted: true, Protocol: "iop/1"}, nil
},
)
toki.AddListenerTyped[*iop.EdgeNodeEvent](&client.Communicator, func(event *iop.EdgeNodeEvent) {
events <- event
})
}
ctx, cancel := context.WithCancel(context.Background())
if err := srv.Start(ctx); err != nil {
cancel()
t.Fatalf("startFakeCPServerForNodeEvents Start: %v", err)
}
return port, events, func() {
cancel()
_ = srv.Stop()
}
}
// TestConnectorRespondsToStatusRequestFromProvider verifies that a status
// request from the Control Plane is answered from the configured snapshot
// boundary rather than a direct registry/transport read.
func TestConnectorRespondsToStatusRequestFromProvider(t *testing.T) {
provider := fakeStatusProvider{nodes: []edgeservice.NodeSnapshot{
{NodeID: "node-1", Alias: "alpha", Label: "node0"},
{NodeID: "node-2", Label: "node1"},
}}
port, serverClients, stop := startFakeCPServerForStatus(t)
defer stop()
c := NewConnector(
config.EdgeInfo{ID: "edge-dgx", Name: "DGX Edge"},
config.EdgeControlPlaneConf{
Enabled: true,
WireAddr: wireAddr(port),
ReconnectIntervalSec: 60,
},
"0.2.0",
noopLogger(),
WithStatusProvider(provider),
)
if !c.StatusProviderConfigured() {
t.Fatal("expected status provider to be configured")
}
if err := c.Start(context.Background()); err != nil {
t.Fatalf("Start: %v", err)
}
defer c.Stop()
var client *toki.TcpClient
select {
case client = <-serverClients:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for edge connection")
}
resp, err := toki.SendRequestTyped[*iop.EdgeStatusRequest, *iop.EdgeStatusResponse](
&client.Communicator,
&iop.EdgeStatusRequest{RequestId: "edge-status-test"},
2*time.Second,
)
if err != nil {
t.Fatalf("SendRequestTyped status: %v", err)
}
if resp.GetRequestId() != "edge-status-test" {
t.Errorf("request_id: got %q want %q", resp.GetRequestId(), "edge-status-test")
}
if resp.GetEdgeId() != "edge-dgx" {
t.Errorf("edge_id: got %q want %q", resp.GetEdgeId(), "edge-dgx")
}
if resp.GetError() != "" {
t.Errorf("unexpected error: %q", resp.GetError())
}
if len(resp.GetNodes()) != 2 {
t.Fatalf("nodes: got %d want 2", len(resp.GetNodes()))
}
byID := map[string]*iop.EdgeNodeSnapshot{}
for _, n := range resp.GetNodes() {
byID[n.GetNodeId()] = n
}
if n := byID["node-1"]; n == nil || n.GetAlias() != "alpha" || n.GetLabel() != "node0" {
t.Errorf("node-1 snapshot: %+v", n)
}
if n := byID["node-2"]; n == nil || n.GetLabel() != "node1" {
t.Errorf("node-2 snapshot: %+v", n)
}
}
func TestConnectorRelaysNodeLifecycleEventsToControlPlane(t *testing.T) {
bus := edgeevents.NewBus()
port, events, stop := startFakeCPServerForNodeEvents(t)
defer stop()
c := NewConnector(
config.EdgeInfo{ID: "edge-dgx", Name: "DGX Edge"},
config.EdgeControlPlaneConf{
Enabled: true,
WireAddr: wireAddr(port),
ReconnectIntervalSec: 60,
},
"0.2.0",
noopLogger(),
WithNodeEventBus(bus),
)
if !c.NodeEventBusConfigured() {
t.Fatal("expected node event bus to be configured")
}
if err := c.Start(context.Background()); err != nil {
t.Fatalf("Start: %v", err)
}
defer c.Stop()
deadline := time.Now().Add(3 * time.Second)
for c.CurrentState() != StateConnected {
if time.Now().After(deadline) {
t.Fatalf("connector did not reach connected state, got %v", c.CurrentState())
}
time.Sleep(20 * time.Millisecond)
}
original := &iop.EdgeNodeEvent{
EventId: "evt-node-1",
Type: "node.connected",
Source: "edge",
NodeId: "node-1",
Alias: "alpha",
Reason: "registered",
Metadata: map[string]string{"rack": "r1"},
}
bus.PublishNode(original)
select {
case got := <-events:
if got.GetEventId() != "evt-node-1" || got.GetType() != "node.connected" {
t.Fatalf("unexpected relayed event: %+v", got)
}
if got.GetNodeId() != "node-1" || got.GetAlias() != "alpha" {
t.Fatalf("unexpected relayed node identity: %+v", got)
}
if got.GetMetadata()["rack"] != "r1" {
t.Fatalf("metadata rack: got %q", got.GetMetadata()["rack"])
}
if got.GetMetadata()["edge_id"] != "edge-dgx" {
t.Fatalf("metadata edge_id: got %q want edge-dgx", got.GetMetadata()["edge_id"])
}
if got.GetTimestamp() == 0 {
t.Fatal("expected relay to fill missing timestamp on copied event")
}
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for relayed EdgeNodeEvent")
}
if _, ok := original.GetMetadata()["edge_id"]; ok {
t.Fatal("relay mutated original event metadata")
}
if original.GetTimestamp() != 0 {
t.Fatal("relay mutated original event timestamp")
}
}
// TestConnectorStatusRequestWithoutProviderReturnsError verifies that a status
// request is answered with an explicit error when no provider is wired.
func TestConnectorStatusRequestWithoutProviderReturnsError(t *testing.T) {
port, serverClients, stop := startFakeCPServerForStatus(t)
defer stop()
c := NewConnector(
config.EdgeInfo{ID: "edge-dgx", Name: "DGX Edge"},
config.EdgeControlPlaneConf{
Enabled: true,
WireAddr: wireAddr(port),
ReconnectIntervalSec: 60,
},
"0.2.0",
noopLogger(),
)
if c.StatusProviderConfigured() {
t.Fatal("expected no status provider configured")
}
if err := c.Start(context.Background()); err != nil {
t.Fatalf("Start: %v", err)
}
defer c.Stop()
var client *toki.TcpClient
select {
case client = <-serverClients:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for edge connection")
}
resp, err := toki.SendRequestTyped[*iop.EdgeStatusRequest, *iop.EdgeStatusResponse](
&client.Communicator,
&iop.EdgeStatusRequest{RequestId: "edge-status-test"},
2*time.Second,
)
if err != nil {
t.Fatalf("SendRequestTyped status: %v", err)
}
if resp.GetError() == "" {
t.Error("expected an explicit error when no status provider is configured")
}
if len(resp.GetNodes()) != 0 {
t.Errorf("expected no nodes without provider, got %d", len(resp.GetNodes()))
}
}
// TestConnectorDisabledNoOp verifies that Start returns immediately when
// enabled=false, without touching the network.
func TestConnectorDisabledNoOp(t *testing.T) {
c := NewConnector(
config.EdgeInfo{ID: "edge-test"},
config.EdgeControlPlaneConf{Enabled: false, WireAddr: "127.0.0.1:9999"},
"0.1.0",
noopLogger(),
)
if err := c.Start(context.Background()); err != nil {
t.Fatalf("Start disabled connector: %v", err)
}
if c.IsEnabled() {
t.Fatal("expected IsEnabled()==false")
}
c.Stop() // must not panic
}
// TestConnectorEmptyAddrNoOp verifies that Start returns immediately when
// wire_addr is empty even if enabled=true.
func TestConnectorEmptyAddrNoOp(t *testing.T) {
c := NewConnector(
config.EdgeInfo{ID: "edge-test"},
config.EdgeControlPlaneConf{Enabled: true, WireAddr: ""},
"0.1.0",
noopLogger(),
)
if err := c.Start(context.Background()); err != nil {
t.Fatalf("Start empty-addr connector: %v", err)
}
if c.IsEnabled() {
t.Fatal("expected IsEnabled()==false for empty wire_addr")
}
c.Stop()
}
// TestConnectorSendsEdgeHello verifies that the connector dials and sends the
// correct EdgeHelloRequest fields.
func TestConnectorSendsEdgeHello(t *testing.T) {
helloReceived := make(chan *iop.EdgeHelloRequest, 1)
port, stop := startFakeCPServer(t, func(req *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
select {
case helloReceived <- req:
default:
}
return &iop.EdgeHelloResponse{Accepted: true, Protocol: "iop/1"}
})
defer stop()
c := NewConnector(
config.EdgeInfo{ID: "edge-dgx", Name: "DGX Edge"},
config.EdgeControlPlaneConf{
Enabled: true,
WireAddr: wireAddr(port),
ReconnectIntervalSec: 60,
},
"0.2.0",
noopLogger(),
)
if err := c.Start(context.Background()); err != nil {
t.Fatalf("Start: %v", err)
}
defer c.Stop()
select {
case req := <-helloReceived:
if req.GetEdgeId() != "edge-dgx" {
t.Errorf("edge_id: got %q, want %q", req.GetEdgeId(), "edge-dgx")
}
if req.GetEdgeName() != "DGX Edge" {
t.Errorf("edge_name: got %q, want %q", req.GetEdgeName(), "DGX Edge")
}
if req.GetVersion() != "0.2.0" {
t.Errorf("version: got %q, want %q", req.GetVersion(), "0.2.0")
}
if len(req.GetCapabilities()) == 0 {
t.Error("expected non-empty capabilities")
}
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for EdgeHelloRequest")
}
}
// TestConnectorRejectsRejectedHello verifies that the connector treats a
// rejected hello as a connect failure and does not stay connected.
func TestConnectorRejectsRejectedHello(t *testing.T) {
port, stop := startFakeCPServer(t, func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
return &iop.EdgeHelloResponse{Accepted: false, Message: "not authorized"}
})
defer stop()
c := NewConnector(
config.EdgeInfo{ID: "edge-test"},
config.EdgeControlPlaneConf{
Enabled: true,
WireAddr: wireAddr(port),
ReconnectIntervalSec: 60,
},
"0.1.0",
noopLogger(),
)
if err := c.Start(context.Background()); err != nil {
t.Fatalf("Start: %v", err)
}
defer c.Stop()
// Give connect time to attempt and fail.
time.Sleep(400 * time.Millisecond)
if c.CurrentState() == StateConnected {
t.Error("expected connector not to be in Connected state after rejected hello")
}
}
// TestConnectorStopClosesConnection verifies that Stop() terminates the loop
// and closes an active connection.
func TestConnectorStopClosesConnection(t *testing.T) {
connected := make(chan struct{}, 1)
port, stop := startFakeCPServer(t, func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
select {
case connected <- struct{}{}:
default:
}
return &iop.EdgeHelloResponse{Accepted: true}
})
defer stop()
c := NewConnector(
config.EdgeInfo{ID: "edge-test"},
config.EdgeControlPlaneConf{
Enabled: true,
WireAddr: wireAddr(port),
ReconnectIntervalSec: 60,
},
"0.1.0",
noopLogger(),
)
if err := c.Start(context.Background()); err != nil {
t.Fatalf("Start: %v", err)
}
select {
case <-connected:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for connection")
}
c.Stop()
if c.CurrentState() != StateStopped {
t.Errorf("expected StateStopped after Stop(), got %v", c.CurrentState())
}
}
// TestConnectorReconnectsAfterDisconnect verifies that after the server closes
// the connection, the connector detects the disconnect and re-dials, sending a
// second EdgeHello. This directly exercises the reconnect loop.
func TestConnectorReconnectsAfterDisconnect(t *testing.T) {
var helloCount atomic.Int32
helloSeen := make(chan struct{}, 16)
port, serverClients, stopServer := startFakeCPServerWithClientAccess(t,
func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
helloCount.Add(1)
helloSeen <- struct{}{}
return &iop.EdgeHelloResponse{Accepted: true}
},
)
defer stopServer()
c := NewConnector(
config.EdgeInfo{ID: "edge-test"},
config.EdgeControlPlaneConf{
Enabled: true,
WireAddr: wireAddr(port),
ReconnectIntervalSec: 1, // fast reconnect for test
},
"0.1.0",
noopLogger(),
)
if err := c.Start(context.Background()); err != nil {
t.Fatalf("Start: %v", err)
}
defer c.Stop()
// Wait for first hello / connection.
select {
case <-helloSeen:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for first EdgeHello")
}
// Server closes the first client connection to trigger a disconnect.
select {
case cl := <-serverClients:
_ = cl.Close()
case <-time.After(2 * time.Second):
t.Fatal("timeout getting server-side client to close")
}
// Wait for a second hello; this proves the connector reconnected.
select {
case <-helloSeen:
if helloCount.Load() < 2 {
t.Errorf("expected at least 2 hellos, got %d", helloCount.Load())
}
case <-time.After(5 * time.Second):
t.Fatalf("timeout waiting for reconnect (second EdgeHello); got %d hellos", helloCount.Load())
}
}
// TestConnectorReconnectsAfterAcceptedThenClosed tests the race/regression
// where the Control Plane closes the connection immediately after sending an
// accepted EdgeHelloResponse. The connector must detect the disconnect even
// though it races with the hello response, and must reconnect.
func TestConnectorReconnectsAfterAcceptedThenClosed(t *testing.T) {
var helloCount atomic.Int32
helloSeen := make(chan struct{}, 16)
port, serverClients, stopServer := startFakeCPServerWithClientAccess(t,
func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
helloCount.Add(1)
helloSeen <- struct{}{}
return &iop.EdgeHelloResponse{Accepted: true}
},
)
defer stopServer()
c := NewConnector(
config.EdgeInfo{ID: "edge-test"},
config.EdgeControlPlaneConf{
Enabled: true,
WireAddr: wireAddr(port),
ReconnectIntervalSec: 1,
},
"0.1.0",
noopLogger(),
)
if err := c.Start(context.Background()); err != nil {
t.Fatalf("Start: %v", err)
}
defer c.Stop()
// Wait for first hello, then immediately close the server-side client.
// This simulates Control Plane closing the connection right after accepting.
select {
case <-helloSeen:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for first EdgeHello")
}
select {
case cl := <-serverClients:
_ = cl.Close()
case <-time.After(2 * time.Second):
t.Fatal("timeout getting server-side client to close immediately after hello")
}
// Connector must reconnect and send a second hello. If the disconnect
// listener was not pre-registered, connect() would block forever here.
select {
case <-helloSeen:
if helloCount.Load() < 2 {
t.Errorf("expected at least 2 hellos, got %d", helloCount.Load())
}
case <-time.After(5 * time.Second):
t.Fatalf("connector did not reconnect after accepted-then-closed; got %d hellos", helloCount.Load())
}
}
// TestConnectorStopPreventsReconnect verifies that after Stop() the connector
// does not send additional hellos within a short window.
func TestConnectorStopPreventsReconnect(t *testing.T) {
var helloCount atomic.Int32
helloSeen := make(chan struct{}, 16)
port, serverClients, stopServer := startFakeCPServerWithClientAccess(t,
func(_ *iop.EdgeHelloRequest) *iop.EdgeHelloResponse {
helloCount.Add(1)
helloSeen <- struct{}{}
return &iop.EdgeHelloResponse{Accepted: true}
},
)
defer stopServer()
c := NewConnector(
config.EdgeInfo{ID: "edge-test"},
config.EdgeControlPlaneConf{
Enabled: true,
WireAddr: wireAddr(port),
ReconnectIntervalSec: 1,
},
"0.1.0",
noopLogger(),
)
if err := c.Start(context.Background()); err != nil {
t.Fatalf("Start: %v", err)
}
// Wait for first hello.
select {
case <-helloSeen:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for first EdgeHello")
}
// Stop the connector before closing the server-side client.
c.Stop()
// Now close the server-side client.
select {
case cl := <-serverClients:
_ = cl.Close()
default:
// client may already be gone after Stop
}
// Give a window longer than ReconnectIntervalSec; no new hello should arrive.
countBefore := helloCount.Load()
select {
case <-helloSeen:
t.Errorf("connector sent hello after Stop(); count went from %d to %d", countBefore, helloCount.Load())
case <-time.After(2 * time.Second):
// good; no extra hello
}
}