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

package transport

import (
	"context"
	"net"
	"strconv"
	"sync/atomic"
	"testing"
	"time"

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

	iop "iop/proto/gen/iop"
)

// TestHeartbeatSurvivesIdleAfterNodeCommand validates that, after a RegisterRequest
// and NodeCommandRequest round-trip, both edge and node TcpClients stay alive across
// multiple heartbeat intervals of idle traffic. This guards against the heartbeat
// timeout regression reproduced in agent-task/12 (real Gemini idle).
func TestHeartbeatSurvivesIdleAfterNodeCommand(t *testing.T) {
	// 1s interval, 2s wait — wait > interval mirrors the production ratio
	// (interval 30s, wait 45s) at a faster test timescale. The library-level
	// wait-timer state check (proto-socket base_client.go) is the primary
	// defence against false heartbeat_timeout; this iop test exercises the
	// idle path end-to-end through DialTcp + NodeCommandRequest round-trip.
	const (
		intervalSec = 1
		waitSec     = 2
		idleWindow  = 5 * time.Second
	)

	ctx, cancel := context.WithTimeout(context.Background(), 7*time.Second)
	defer cancel()

	listenAddr := freeAddr(t)
	host, portStr, _ := net.SplitHostPort(listenAddr)
	port, _ := strconv.Atoi(portStr)

	// Mock edge server with 1s/1s heartbeat.
	acceptedCh := make(chan *toki.TcpClient, 1)
	server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
		client := toki.NewTcpClient(conn, intervalSec, waitSec, heartbeatTestEdgeParserMap())
		toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
			&client.Communicator,
			func(req *iop.RegisterRequest) (*iop.RegisterResponse, error) {
				return &iop.RegisterResponse{
					Accepted: true,
					NodeId:   "hb-node",
					Alias:    "hb-alias",
				}, nil
			},
		)
		acceptedCh <- client
		return client
	})
	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	// Node client with 1s/1s heartbeat. We bypass DialEdge to allow custom intervals
	// while keeping the rest of the node parser map / wire format identical.
	nodeClient, err := toki.DialTcp(ctx, host, port, intervalSec, waitSec, nodeParserMap())
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	defer nodeClient.Close()

	resp, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&nodeClient.Communicator,
		&iop.RegisterRequest{Token: "hb-token"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("register: %v", err)
	}
	if !resp.GetAccepted() {
		t.Fatalf("register rejected: %s", resp.GetReason())
	}

	edgeClient := waitForHeartbeatTestClient(t, acceptedCh)

	// NodeCommandRequest handler on node side, request driven from edge.
	var cmdCount atomic.Int32
	toki.AddRequestListenerTyped[*iop.NodeCommandRequest, *iop.NodeCommandResponse](
		&nodeClient.Communicator,
		func(req *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
			cmdCount.Add(1)
			return &iop.NodeCommandResponse{RequestId: req.GetRequestId()}, nil
		},
	)

	cmdResp, err := toki.SendRequestTyped[*iop.NodeCommandRequest, *iop.NodeCommandResponse](
		&edgeClient.Communicator,
		&iop.NodeCommandRequest{RequestId: "cmd-1", Type: iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("node command: %v", err)
	}
	if cmdResp.GetRequestId() != "cmd-1" {
		t.Fatalf("cmd response id: %q", cmdResp.GetRequestId())
	}

	// Idle window: span several heartbeat intervals with no app traffic.
	deadline := time.Now().Add(idleWindow)
	for time.Now().Before(deadline) {
		time.Sleep(200 * time.Millisecond)
		if !nodeClient.IsAlive() {
			t.Fatalf("node client died during idle window after %v", time.Since(deadline.Add(-idleWindow)))
		}
		if !edgeClient.IsAlive() {
			t.Fatalf("edge client died during idle window after %v", time.Since(deadline.Add(-idleWindow)))
		}
	}

	if !nodeClient.IsAlive() {
		t.Fatal("node client not alive after idle window")
	}
	if !edgeClient.IsAlive() {
		t.Fatal("edge client not alive after idle window")
	}
}

func freeAddr(t *testing.T) string {
	t.Helper()
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	addr := l.Addr().String()
	l.Close()
	return addr
}

func waitForHeartbeatTestClient(t *testing.T, ch <-chan *toki.TcpClient) *toki.TcpClient {
	t.Helper()
	select {
	case c := <-ch:
		return c
	case <-time.After(2 * time.Second):
		t.Fatal("server did not accept connection")
		return nil
	}
}

// TestHeartbeatWaitExceedsInterval guards the invariant that protects against the
// library's sendHeartBeat race: the wait timer can be installed even after the
// heartbeat response was already processed, so the peer's next heartbeat (one
// interval later) must arrive before the stray wait timer fires.
func TestHeartbeatWaitExceedsInterval(t *testing.T) {
	if heartbeatWaitSec <= heartbeatIntervalSec {
		t.Fatalf("heartbeatWaitSec (%d) must exceed heartbeatIntervalSec (%d)", heartbeatWaitSec, heartbeatIntervalSec)
	}
}

func TestHeartbeatUsesLowLatencyNodeLivenessProfile(t *testing.T) {
	if heartbeatIntervalSec != 2 {
		t.Fatalf("heartbeatIntervalSec: got %d want 2", heartbeatIntervalSec)
	}
	if heartbeatWaitSec != 5 {
		t.Fatalf("heartbeatWaitSec: got %d want 5", heartbeatWaitSec)
	}
}

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