iop/apps/control-plane/internal/wire/edge_server_test.go

package wire

import (
	"context"
	"fmt"
	"sync"
	"testing"
	"time"

	"go.uber.org/zap/zaptest"

	proto_socket "git.toki-labs.com/toki/proto-socket/go"
	iop "iop/proto/gen/iop"
)

func startEdgeServer(t *testing.T) (*EdgeServer, int) {
	t.Helper()
	logger := zaptest.NewLogger(t)
	port := getFreePort(t)
	listenAddr := fmt.Sprintf("127.0.0.1:%d", port)

	server, err := NewEdgeServer(listenAddr, logger)
	if err != nil {
		t.Fatalf("NewEdgeServer failed: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)

	if err := server.Start(ctx); err != nil {
		t.Fatalf("edge server Start failed: %v", err)
	}
	t.Cleanup(func() { _ = server.Stop() })

	// Give a tiny moment for the listener to bind.
	time.Sleep(50 * time.Millisecond)
	return server, port
}

func dialEdge(t *testing.T, ctx context.Context, port int) *proto_socket.TcpClient {
	t.Helper()
	client, err := proto_socket.DialTcp(
		ctx,
		"127.0.0.1",
		port,
		proto_socket.DefaultHeartbeatIntervalSec,
		proto_socket.DefaultHeartbeatWaitSec,
		EdgeParserMap(),
	)
	if err != nil {
		t.Fatalf("failed to DialTcp: %v", err)
	}
	return client
}

func TestEdgeServerHelloRegistersEdge(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)
	defer client.Close()

	req := &iop.EdgeHelloRequest{
		EdgeId:       "edge-dgx-group",
		EdgeName:     "DGX Group",
		Version:      "1.2.3",
		Capabilities: []string{"node-registry", "run-dispatch"},
	}

	res, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		req,
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}

	if !res.Accepted {
		t.Errorf("expected response.Accepted to be true, got false (reason=%q)", res.Reason)
	}
	if res.Protocol != Protocol {
		t.Errorf("expected response.Protocol to be %q, got %q", Protocol, res.Protocol)
	}
	if res.ServerTimeUnixNano <= 0 {
		t.Errorf("expected server time unix nano to be positive, got %d", res.ServerTimeUnixNano)
	}

	state, ok := server.Registry().Snapshot("edge-dgx-group")
	if !ok {
		t.Fatalf("edge not registered after hello")
	}
	if !state.Connected {
		t.Errorf("expected registry Connected to be true")
	}
	if state.EdgeName != "DGX Group" {
		t.Errorf("edge_name: got %q want %q", state.EdgeName, "DGX Group")
	}
	if len(state.Capabilities) != 2 {
		t.Errorf("capabilities: got %v", state.Capabilities)
	}
	if state.LastSeen.IsZero() {
		t.Errorf("expected LastSeen to be set")
	}
}

func TestEdgeServerRejectsMissingEdgeID(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)
	defer client.Close()

	req := &iop.EdgeHelloRequest{
		EdgeName: "no id edge",
	}

	res, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		req,
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}

	if res.Accepted {
		t.Errorf("expected response.Accepted to be false for missing edge_id")
	}
	if res.Reason == "" {
		t.Errorf("expected a rejection reason for missing edge_id")
	}
	if server.Registry().Len() != 0 {
		t.Errorf("expected registry to stay empty, got %d entries", server.Registry().Len())
	}
}

func TestEdgeServerDisconnectMarksEdgeDisconnected(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)

	req := &iop.EdgeHelloRequest{
		EdgeId:   "edge-dgx-group",
		EdgeName: "DGX Group",
	}
	if _, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		req,
		2*time.Second,
	); err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}

	if err := client.Close(); err != nil {
		t.Fatalf("client Close failed: %v", err)
	}

	// The server detects the closed connection asynchronously; poll until the
	// disconnect listener flips the registry state.
	deadline := time.Now().Add(2 * time.Second)
	for {
		state, ok := server.Registry().Snapshot("edge-dgx-group")
		if ok && !state.Connected {
			if state.DisconnectInfo.Reason == "" {
				t.Errorf("expected a disconnect reason to be recorded")
			}
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("edge was not marked disconnected in time (ok=%v)", ok)
		}
		time.Sleep(20 * time.Millisecond)
	}
}

func sendEdgeHello(t *testing.T, client *proto_socket.TcpClient, edgeID string) {
	t.Helper()
	if _, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		&iop.EdgeHelloRequest{EdgeId: edgeID, EdgeName: "DGX Group"},
		2*time.Second,
	); err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}
}

func TestEdgeServerStaleDisconnectDoesNotClearReconnect(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	const edgeID = "edge-dgx-group"

	// First connection enrolls the edge.
	conn1 := dialEdge(t, ctx, port)
	sendEdgeHello(t, conn1, edgeID)

	// Second connection with the same edge_id reconnects and becomes current.
	conn2 := dialEdge(t, ctx, port)
	defer conn2.Close()
	sendEdgeHello(t, conn2, edgeID)

	// The stale first connection closing must not clear the live reconnection.
	if err := conn1.Close(); err != nil {
		t.Fatalf("conn1 Close failed: %v", err)
	}

	// Give the server time to process the stale disconnect, then assert the edge
	// stays connected under the second connection's token.
	time.Sleep(300 * time.Millisecond)
	state, ok := server.Registry().Snapshot(edgeID)
	if !ok {
		t.Fatalf("edge missing from registry")
	}
	if !state.Connected {
		t.Fatalf("stale disconnect cleared the live reconnection; Connected=false")
	}
}

func TestEdgeServerStaleConnectionNodeEventDoesNotPolluteReconnect(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	const edgeID = "edge-dgx-group"

	conn1 := dialEdge(t, ctx, port)
	defer conn1.Close()
	sendEdgeHello(t, conn1, edgeID)

	conn2 := dialEdge(t, ctx, port)
	defer conn2.Close()
	sendEdgeHello(t, conn2, edgeID)

	if err := conn1.Send(&iop.EdgeNodeEvent{
		EventId: "evt-stale",
		Type:    "node.connected",
		NodeId:  "node-stale",
	}); err != nil {
		t.Fatalf("send stale EdgeNodeEvent: %v", err)
	}
	time.Sleep(300 * time.Millisecond)
	if got := server.Registry().NodeEvents(edgeID, ""); len(got) != 0 {
		t.Fatalf("stale connection event polluted registry: %+v", got)
	}

	if err := conn2.Send(&iop.EdgeNodeEvent{
		EventId: "evt-current",
		Type:    "node.connected",
		NodeId:  "node-current",
	}); err != nil {
		t.Fatalf("send current EdgeNodeEvent: %v", err)
	}

	var records []EdgeNodeEventRecord
	deadline := time.Now().Add(2 * time.Second)
	for {
		records = server.Registry().NodeEvents(edgeID, "")
		if len(records) == 1 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("current node event was not recorded in time; got %d records", len(records))
		}
		time.Sleep(20 * time.Millisecond)
	}
	if records[0].Event.GetEventId() != "evt-current" {
		t.Fatalf("recorded event_id: got %q want evt-current", records[0].Event.GetEventId())
	}
	if records[0].Event.GetNodeId() != "node-current" {
		t.Fatalf("recorded node_id: got %q want node-current", records[0].Event.GetNodeId())
	}
}

func TestEdgeServerRequestsStatusFromConnectedEdge(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)
	defer client.Close()

	const edgeID = "edge-dgx-group"

	// The test acts as the Edge: it answers a Control Plane status request with a
	// service-backed node snapshot.
	proto_socket.AddRequestListenerTyped(&client.Communicator, func(req *iop.EdgeStatusRequest) (*iop.EdgeStatusResponse, error) {
		return &iop.EdgeStatusResponse{
			RequestId:            req.GetRequestId(),
			EdgeId:               edgeID,
			EdgeName:             "DGX Group",
			ObservedTimeUnixNano: time.Now().UnixNano(),
			Nodes: []*iop.EdgeNodeSnapshot{
				{NodeId: "node-1", Alias: "alpha", Label: "node0", Connected: true},
			},
		}, nil
	})

	sendEdgeHello(t, client, edgeID)

	resp, err := server.RequestStatus(edgeID, 2*time.Second)
	if err != nil {
		t.Fatalf("RequestStatus: %v", err)
	}
	if resp.GetRequestId() == "" {
		t.Errorf("expected a non-empty correlation request_id")
	}
	if resp.GetEdgeId() != edgeID {
		t.Errorf("edge_id: got %q want %q", resp.GetEdgeId(), edgeID)
	}
	if len(resp.GetNodes()) != 1 {
		t.Fatalf("nodes: got %d want 1", len(resp.GetNodes()))
	}
	if node := resp.GetNodes()[0]; node.GetNodeId() != "node-1" || node.GetLabel() != "node0" {
		t.Errorf("node snapshot: %+v", node)
	}
}

func TestEdgeServerRecordsNodeLifecycleEventFromEdge(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)
	defer client.Close()

	if err := client.Send(&iop.EdgeNodeEvent{
		EventId: "evt-before-hello",
		Type:    "node.connected",
		NodeId:  "node-early",
	}); err != nil {
		t.Fatalf("send pre-hello EdgeNodeEvent: %v", err)
	}
	time.Sleep(100 * time.Millisecond)
	if got := server.Registry().NodeEvents("edge-dgx-group", ""); len(got) != 0 {
		t.Fatalf("expected pre-hello node event to be ignored, got %+v", got)
	}

	sendEdgeHello(t, client, "edge-dgx-group")

	event := &iop.EdgeNodeEvent{
		EventId:  "evt-node-1",
		Type:     "node.connected",
		Source:   "edge",
		NodeId:   "node-1",
		Alias:    "alpha",
		Reason:   "registered",
		Metadata: map[string]string{"rack": "r1"},
	}
	if err := client.Send(event); err != nil {
		t.Fatalf("send EdgeNodeEvent: %v", err)
	}

	var records []EdgeNodeEventRecord
	deadline := time.Now().Add(2 * time.Second)
	for {
		records = server.Registry().NodeEvents("edge-dgx-group", "")
		if len(records) == 1 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("node event was not recorded in time; got %d records", len(records))
		}
		time.Sleep(20 * time.Millisecond)
	}
	got := records[0]
	if got.EdgeID != "edge-dgx-group" {
		t.Errorf("edge_id: got %q want %q", got.EdgeID, "edge-dgx-group")
	}
	if got.Event.GetEventId() != "evt-node-1" || got.Event.GetType() != "node.connected" {
		t.Fatalf("unexpected recorded event: %+v", got.Event)
	}
	if got.Event.GetNodeId() != "node-1" || got.Event.GetAlias() != "alpha" {
		t.Fatalf("unexpected node identity: %+v", got.Event)
	}
	if got.Event.GetMetadata()["rack"] != "r1" {
		t.Fatalf("metadata rack: got %q", got.Event.GetMetadata()["rack"])
	}
	if got.Event.GetTimestamp() == 0 {
		t.Fatal("expected registry to fill missing event timestamp")
	}
	if event.GetTimestamp() != 0 {
		t.Fatal("registry mutated original event timestamp")
	}

	records[0].Event.Metadata["rack"] = "mutated"
	again := server.Registry().NodeEvents("edge-dgx-group", "node-1")
	if len(again) != 1 {
		t.Fatalf("filtered node events: got %d want 1", len(again))
	}
	if again[0].Event.GetMetadata()["rack"] != "r1" {
		t.Fatalf("registry event was mutated through returned snapshot: %+v", again[0].Event.GetMetadata())
	}
}

type commandEdgeRecorder struct {
	mu       sync.Mutex
	received []string
}

func (r *commandEdgeRecorder) operations() []string {
	r.mu.Lock()
	defer r.mu.Unlock()
	return append([]string(nil), r.received...)
}

// dialCommandEdge dials the server as an Edge that answers command requests with
// an "ok" response and records the operations it was asked to run.
func dialCommandEdge(t *testing.T, ctx context.Context, port int, edgeID string) *commandEdgeRecorder {
	t.Helper()
	client := dialEdge(t, ctx, port)
	t.Cleanup(func() { _ = client.Close() })
	rec := &commandEdgeRecorder{}
	proto_socket.AddRequestListenerTyped(&client.Communicator, func(req *iop.EdgeCommandRequest) (*iop.EdgeCommandResponse, error) {
		rec.mu.Lock()
		rec.received = append(rec.received, req.GetOperation())
		rec.mu.Unlock()
		return &iop.EdgeCommandResponse{
			RequestId: req.GetRequestId(),
			CommandId: req.GetCommandId(),
			EdgeId:    edgeID,
			Status:    "ok",
			Summary:   "applied " + req.GetOperation(),
		}, nil
	})
	sendEdgeHello(t, client, edgeID)
	return rec
}

func TestEdgeServerSendsCommandToTargetEdgeOnly(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	recA := dialCommandEdge(t, ctx, port, "edge-a")
	recB := dialCommandEdge(t, ctx, port, "edge-b")

	resp, err := server.SendCommand("edge-a", "restart-node", "node-1", map[string]string{"force": "true"}, 2*time.Second)
	if err != nil {
		t.Fatalf("SendCommand: %v", err)
	}
	if resp.GetStatus() != "ok" || resp.GetEdgeId() != "edge-a" {
		t.Fatalf("unexpected response: %+v", resp)
	}
	if resp.GetCommandId() == "" {
		t.Error("expected a non-empty command_id correlation")
	}

	if got := recA.operations(); len(got) != 1 || got[0] != "restart-node" {
		t.Fatalf("edge-a received: %+v", got)
	}
	if got := recB.operations(); len(got) != 0 {
		t.Fatalf("edge-b should not receive a targeted command, got: %+v", got)
	}

	cmds := server.Registry().Commands("edge-a")
	if len(cmds) != 1 {
		t.Fatalf("edge-a audit: got %d want 1", len(cmds))
	}
	if cmds[0].Operation != "restart-node" || cmds[0].Status != "ok" {
		t.Fatalf("unexpected audit record: %+v", cmds[0])
	}
	if got := server.Registry().Commands("edge-b"); len(got) != 0 {
		t.Fatalf("edge-b audit should be empty, got %+v", got)
	}
}

func TestEdgeServerSendCommandFailsWhenEdgeNotConnected(t *testing.T) {
	server, _ := startEdgeServer(t)
	if _, err := server.SendCommand("missing-edge", "restart-node", "", nil, time.Second); err == nil {
		t.Fatal("expected error sending command to a non-connected edge")
	}
	if got := server.Registry().Commands("missing-edge"); len(got) != 0 {
		t.Fatalf("expected no audit record for a never-connected edge, got %+v", got)
	}
}

func TestEdgeServerRecordsCommandLifecycleEventsFromEdge(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)
	defer client.Close()
	const edgeID = "edge-dgx-group"

	proto_socket.AddRequestListenerTyped(&client.Communicator, func(req *iop.EdgeCommandRequest) (*iop.EdgeCommandResponse, error) {
		_ = client.Send(&iop.EdgeCommandEvent{CommandId: req.GetCommandId(), EdgeId: edgeID, Phase: "accepted", Summary: "queued"})
		_ = client.Send(&iop.EdgeCommandEvent{CommandId: req.GetCommandId(), EdgeId: edgeID, Phase: "running", Summary: "executing"})
		return &iop.EdgeCommandResponse{
			RequestId: req.GetRequestId(),
			CommandId: req.GetCommandId(),
			EdgeId:    edgeID,
			Status:    "ok",
			Summary:   "done",
		}, nil
	})

	sendEdgeHello(t, client, edgeID)

	resp, err := server.SendCommand(edgeID, "drain", "", nil, 2*time.Second)
	if err != nil {
		t.Fatalf("SendCommand: %v", err)
	}

	var rec EdgeCommandRecord
	deadline := time.Now().Add(2 * time.Second)
	for {
		got, ok := server.Registry().Command(edgeID, resp.GetCommandId())
		if ok && len(got.Events) == 2 {
			rec = got
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("command events not recorded in time; record=%+v", got)
		}
		time.Sleep(20 * time.Millisecond)
	}
	if rec.Status != "ok" || rec.Summary != "done" {
		t.Fatalf("unexpected final command record: %+v", rec)
	}
	if rec.Events[0].Phase != "accepted" || rec.Events[1].Phase != "running" {
		t.Fatalf("unexpected event phases: %+v", rec.Events)
	}
}

func TestEdgeRegistryCommandsAreBoundedAndCopied(t *testing.T) {
	registry := NewEdgeRegistry()
	at := time.Unix(1780142500, 0).UTC()

	for i := 0; i < recentCommandCapacity+2; i++ {
		cmdID := fmt.Sprintf("cmd-%03d", i)
		registry.RecordCommandRequest("edge-a", &iop.EdgeCommandRequest{CommandId: cmdID, Operation: "op"}, at.Add(time.Duration(i)*time.Second))
		registry.RecordCommandResult("edge-a", &iop.EdgeCommandResponse{CommandId: cmdID, Status: "ok", Summary: "done"}, at.Add(time.Duration(i)*time.Second).Add(time.Millisecond))
	}

	cmds := registry.Commands("edge-a")
	if len(cmds) != recentCommandCapacity {
		t.Fatalf("commands len: got %d want %d", len(cmds), recentCommandCapacity)
	}
	if cmds[0].CommandID != "cmd-002" {
		t.Fatalf("expected oldest retained cmd-002, got %q", cmds[0].CommandID)
	}
	if cmds[0].Status != "ok" {
		t.Fatalf("expected result status recorded, got %q", cmds[0].Status)
	}

	registry.RecordCommandRequest("edge-b", &iop.EdgeCommandRequest{CommandId: "cmd-b", Operation: "drain"}, at)
	for i := 0; i < commandEventCapacity+3; i++ {
		registry.RecordCommandEvent("edge-b", &iop.EdgeCommandEvent{CommandId: "cmd-b", Phase: fmt.Sprintf("phase-%03d", i)}, at.Add(time.Duration(i)*time.Nanosecond))
	}
	rec, ok := registry.Command("edge-b", "cmd-b")
	if !ok {
		t.Fatal("edge-b cmd-b missing")
	}
	if len(rec.Events) != commandEventCapacity {
		t.Fatalf("events len: got %d want %d", len(rec.Events), commandEventCapacity)
	}
	if rec.Events[0].Phase != "phase-003" {
		t.Fatalf("expected oldest retained phase-003, got %q", rec.Events[0].Phase)
	}

	rec.Events[0].Phase = "mutated"
	again, _ := registry.Command("edge-b", "cmd-b")
	if again.Events[0].Phase == "mutated" {
		t.Fatal("Command returned mutable registry storage")
	}
}

func TestEdgeServerRequestStatusFailsWhenEdgeNotConnected(t *testing.T) {
	server, _ := startEdgeServer(t)
	if _, err := server.RequestStatus("missing-edge", time.Second); err == nil {
		t.Fatal("expected error requesting status from a non-connected edge")
	}
}

func TestEdgeRegistrySnapshotsAreSortedAndCopied(t *testing.T) {
	registry := NewEdgeRegistry()
	at := time.Unix(1780142300, 0).UTC()
	registry.MarkConnected(&iop.EdgeHelloRequest{
		EdgeId:       "edge-b",
		EdgeName:     "Edge B",
		Capabilities: []string{"run"},
	}, at)
	registry.MarkConnected(&iop.EdgeHelloRequest{
		EdgeId:       "edge-a",
		EdgeName:     "Edge A",
		Capabilities: []string{"status"},
	}, at.Add(time.Second))

	snapshots := registry.Snapshots()
	if len(snapshots) != 2 {
		t.Fatalf("expected 2 snapshots, got %d", len(snapshots))
	}
	if snapshots[0].EdgeID != "edge-a" || snapshots[1].EdgeID != "edge-b" {
		t.Fatalf("expected snapshots sorted by edge_id, got %+v", snapshots)
	}

	snapshots[0].Capabilities[0] = "mutated"
	again, ok := registry.Snapshot("edge-a")
	if !ok {
		t.Fatal("edge-a missing from registry")
	}
	if again.Capabilities[0] != "status" {
		t.Fatalf("snapshot capabilities alias registry storage: %+v", again.Capabilities)
	}
}

func TestEdgeRegistryNodeEventsAreBoundedAndCopied(t *testing.T) {
	registry := NewEdgeRegistry()
	receivedAt := time.Unix(1780142400, 0).UTC()

	for i := 0; i < recentNodeEventCapacity+2; i++ {
		registry.RecordNodeEvent("edge-a", &iop.EdgeNodeEvent{
			EventId: fmt.Sprintf("evt-%03d", i),
			Type:    "node.connected",
			NodeId:  "node-1",
			Metadata: map[string]string{
				"index": fmt.Sprintf("%d", i),
			},
		}, receivedAt.Add(time.Duration(i)*time.Nanosecond))
	}

	events := registry.NodeEvents("edge-a", "")
	if len(events) != recentNodeEventCapacity {
		t.Fatalf("node events len: got %d want %d", len(events), recentNodeEventCapacity)
	}
	if events[0].Event.GetEventId() != "evt-002" {
		t.Fatalf("expected oldest retained event evt-002, got %q", events[0].Event.GetEventId())
	}
	events[0].Event.Metadata["index"] = "mutated"
	again := registry.NodeEvents("edge-a", "")
	if again[0].Event.GetMetadata()["index"] == "mutated" {
		t.Fatal("NodeEvents returned mutable registry storage")
	}
}

func TestEdgeHeartbeatWaitExceedsInterval(t *testing.T) {
	if EdgeHeartbeatWaitSec <= EdgeHeartbeatIntervalSec {
		t.Fatalf("EdgeHeartbeatWaitSec (%d) must exceed EdgeHeartbeatIntervalSec (%d)",
			EdgeHeartbeatWaitSec, EdgeHeartbeatIntervalSec)
	}
}