iop/apps/node/internal/bootstrap/module_test.go

package bootstrap_test

import (
	"context"
	"errors"
	"fmt"
	"net"
	"os"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

	"iop/apps/node/internal/bootstrap"
	"iop/apps/node/internal/transport"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

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

func freeAddrForBootstrap(t *testing.T) (host string, port int, addr 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()
	h, portStr, _ := net.SplitHostPort(addr)
	fmt.Sscanf(portStr, "%d", &port)
	return h, port, addr
}

func useTempCwd(t *testing.T) {
	t.Helper()
	previous, err := os.Getwd()
	if err != nil {
		t.Fatalf("get cwd: %v", err)
	}
	if err := os.Chdir(t.TempDir()); err != nil {
		t.Fatalf("chdir temp: %v", err)
	}
	t.Cleanup(func() {
		if err := os.Chdir(previous); err != nil {
			t.Fatalf("restore cwd: %v", err)
		}
	})
}

// TestReconnectSupervisorReestablishesSession verifies that after the edge-side
// client closes the connection, the supervisor reconnects and re-registers.
func TestReconnectSupervisorReestablishesSession(t *testing.T) {
	useTempCwd(t)
	host, port, addr := freeAddrForBootstrap(t)

	serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
	t.Cleanup(serverCancel)

	// regCh fires each time a RegisterRequest is handled.
	regCh := make(chan struct{}, 8)
	acceptedCh := make(chan *toki.TcpClient, 8)

	server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
		client := toki.NewTcpClient(conn, 30, 10, bootstrapEdgeParserMap())
		toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
			&client.Communicator,
			func(_ *iop.RegisterRequest) (*iop.RegisterResponse, error) {
				regCh <- struct{}{}
				return &iop.RegisterResponse{
					Accepted: true,
					NodeId:   "reconnect-test-node",
					Config:   &iop.NodeConfigPayload{Runtime: &iop.NodeRuntimeConfig{Concurrency: 1}},
				}, nil
			},
		)
		acceptedCh <- client
		return client
	})
	if err := server.Start(serverCtx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	t.Cleanup(func() { server.Stop() })

	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 5, IntervalSec: 0},
		Logging:   config.LoggingConf{Level: "error"},
	}

	app := fx.New(bootstrap.Module(cfg), fx.NopLogger)
	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()

	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}
	defer func() { _ = app.Stop(context.Background()) }()

	// Wait for first registration.
	select {
	case <-regCh:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for first registration")
	}
	var firstClient *toki.TcpClient
	select {
	case firstClient = <-acceptedCh:
	case <-time.After(time.Second):
		t.Fatal("timeout waiting for first accepted client")
	}

	// Force edge-side close → remote disconnect on node.
	_ = firstClient.Close()

	// Wait for second registration (supervisor reconnected and re-registered).
	select {
	case <-regCh:
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for supervisor reconnect registration")
	}
}

// TestReconnectSupervisorExhaustsRetries verifies that when all reconnect
// attempts fail, the node cleans up resources and calls Shutdown(ExitCode(1)).
func TestReconnectSupervisorExhaustsRetries(t *testing.T) {
	useTempCwd(t)
	host, port, addr := freeAddrForBootstrap(t)

	serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
	t.Cleanup(serverCancel)

	acceptedCh := make(chan *toki.TcpClient, 4)
	server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
		client := toki.NewTcpClient(conn, 30, 10, bootstrapEdgeParserMap())
		toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
			&client.Communicator,
			func(_ *iop.RegisterRequest) (*iop.RegisterResponse, error) {
				return &iop.RegisterResponse{
					Accepted: true,
					NodeId:   "exhaust-test-node",
					Config:   &iop.NodeConfigPayload{Runtime: &iop.NodeRuntimeConfig{Concurrency: 1}},
				}, nil
			},
		)
		acceptedCh <- client
		return client
	})
	if err := server.Start(serverCtx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	t.Cleanup(func() { server.Stop() })

	// Fake dialer: first call uses real DialEdge; subsequent calls fail immediately.
	var dialCount int32
	fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
		if atomic.AddInt32(&dialCount, 1) == 1 {
			return transport.DialEdge(ctx, edgeAddr, token, logger)
		}
		return nil, errors.New("fake reconnect failure")
	})

	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 3, IntervalSec: 0},
		Logging:   config.LoggingConf{Level: "error"},
	}

	app := fx.New(
		bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer)),
		fx.NopLogger,
	)

	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()

	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}
	defer func() { _ = app.Stop(context.Background()) }()

	// Wait for first connection then close it from the edge side.
	var firstClient *toki.TcpClient
	select {
	case firstClient = <-acceptedCh:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for first connection")
	}
	_ = firstClient.Close()

	// Supervisor exhausts retries → Shutdown(ExitCode(1)).
	select {
	case sig := <-app.Wait():
		if sig.ExitCode != 1 {
			t.Fatalf("expected exit code 1, got %d", sig.ExitCode)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for supervisor exhaustion shutdown")
	}

	// 1 initial dial + 3 retry attempts = 4 total.
	if count := atomic.LoadInt32(&dialCount); count != 4 {
		t.Fatalf("expected 4 dial attempts (1 initial + 3 retries), got %d", count)
	}
}

// TestReconnectSupervisorPolicyTimingAndLimit verifies that the supervisor
// respects MaxAttempts=10 and that the fake sleeper receives IntervalSec=10s
// before each reconnect attempt (SDD reconnect_waiting state).
func TestReconnectSupervisorPolicyTimingAndLimit(t *testing.T) {
	useTempCwd(t)
	host, port, addr := freeAddrForBootstrap(t)

	serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
	t.Cleanup(serverCancel)

	acceptedCh := make(chan *toki.TcpClient, 4)
	server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
		client := toki.NewTcpClient(conn, 30, 10, bootstrapEdgeParserMap())
		toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
			&client.Communicator,
			func(_ *iop.RegisterRequest) (*iop.RegisterResponse, error) {
				return &iop.RegisterResponse{
					Accepted: true,
					NodeId:   "policy-test-node",
					Config:   &iop.NodeConfigPayload{Runtime: &iop.NodeRuntimeConfig{Concurrency: 1}},
				}, nil
			},
		)
		acceptedCh <- client
		return client
	})
	if err := server.Start(serverCtx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	t.Cleanup(func() { server.Stop() })

	// Fake dialer: first call uses real DialEdge; subsequent calls fail immediately.
	var dialCount int32
	fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
		if atomic.AddInt32(&dialCount, 1) == 1 {
			return transport.DialEdge(ctx, edgeAddr, token, logger)
		}
		return nil, errors.New("fake reconnect failure")
	})

	// Fake sleeper records received durations without blocking.
	var slMu sync.Mutex
	var sleptDurations []time.Duration
	fakeSleeper := func(_ context.Context, d time.Duration) {
		slMu.Lock()
		sleptDurations = append(sleptDurations, d)
		slMu.Unlock()
	}

	cfg := &config.NodeConfig{
		Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
		Reconnect: config.ReconnectConf{MaxAttempts: 10, IntervalSec: 10},
		Logging:   config.LoggingConf{Level: "error"},
	}

	app := fx.New(
		bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer), bootstrap.WithSleeper(fakeSleeper)),
		fx.NopLogger,
	)

	startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer startCancel()

	if err := app.Start(startCtx); err != nil {
		t.Fatalf("app start: %v", err)
	}
	defer func() { _ = app.Stop(context.Background()) }()

	// Wait for first connection then close it from the edge side.
	var firstClient *toki.TcpClient
	select {
	case firstClient = <-acceptedCh:
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for first connection")
	}
	_ = firstClient.Close()

	// Supervisor exhausts all 10 retries → Shutdown(ExitCode(1)).
	select {
	case sig := <-app.Wait():
		if sig.ExitCode != 1 {
			t.Fatalf("expected exit code 1, got %d", sig.ExitCode)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for supervisor exhaustion shutdown")
	}

	// 1 initial dial + 10 retry attempts = 11 total.
	if count := atomic.LoadInt32(&dialCount); count != 11 {
		t.Fatalf("expected 11 dial attempts (1 initial + 10 retries), got %d", count)
	}

	// Each retry is preceded by a sleep of IntervalSec=10s (reconnect_waiting state).
	slMu.Lock()
	durations := append([]time.Duration(nil), sleptDurations...)
	slMu.Unlock()
	if len(durations) != 10 {
		t.Fatalf("expected 10 sleep calls (one before each retry), got %d: %v", len(durations), durations)
	}
	for i, d := range durations {
		if d != 10*time.Second {
			t.Errorf("sleep[%d] = %v, want 10s", i, d)
		}
	}
}