package bootstrap_test import ( "context" "errors" "fmt" "net" "os" "sync" "sync/atomic" "testing" "time" toki "git.toki-labs.com/toki/proto-socket/go" "git.toki-labs.com/toki/proto-socket/go/packets" "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) }, toki.TypeNameOf(&iop.NodeReadyRequest{}): func(b []byte) (proto.Message, error) { m := &iop.NodeReadyRequest{} return m, proto.Unmarshal(b, m) }, } } // bindEdgeReadyResponder answers the node's post-handler NodeReadyRequest with a // positive ack, so connectRuntime's SignalReady completes and the session is // treated as fully established. Mock edges that accept registration must install // it or the node's ready handshake times out and reconnects. When readyCh is // non-nil it fires only after the ready response has been written to the socket. // This lets a test safely close the edge side after a fully established session, // instead of racing the in-flight ready handshake. func bindEdgeReadyResponder(client *toki.TcpClient, readyCh chan<- struct{}) { client.Communicator.AddRequestListener( toki.TypeNameOf(&iop.NodeReadyRequest{}), func(message proto.Message, requestNonce int32) { if _, ok := message.(*iop.NodeReadyRequest); !ok { return } data, err := proto.Marshal(&iop.NodeReadyResponse{Ready: true}) if err != nil { return } if err := client.QueuePacket(&packets.PacketBase{ TypeName: toki.TypeNameOf(&iop.NodeReadyResponse{}), ResponseNonce: requestNonce, Data: data, }); err != nil { return } if readyCh != nil { select { case readyCh <- struct{}{}: default: } } }, ) } 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) } }) } // startAcceptingEdge starts a mock edge bound to host:port that accepts every // RegisterRequest. The returned channels fire once per registration and // disconnect, respectively. func startAcceptingEdge(t *testing.T, ctx context.Context, host string, port int, nodeID string) (chan struct{}, chan struct{}) { t.Helper() regCh := make(chan struct{}, 16) disconnectCh := make(chan struct{}, 16) server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient { client := toki.NewTcpClient(conn, 30, 10, bootstrapEdgeParserMap()) client.AddDisconnectListener(func(_ *toki.TcpClient) { disconnectCh <- struct{}{} }) toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse]( &client.Communicator, func(_ *iop.RegisterRequest) (*iop.RegisterResponse, error) { regCh <- struct{}{} return &iop.RegisterResponse{ Accepted: true, NodeId: nodeID, Config: &iop.NodeConfigPayload{Runtime: &iop.NodeRuntimeConfig{Concurrency: 1}}, }, nil }, ) bindEdgeReadyResponder(client, nil) return client }) if err := server.Start(ctx); err != nil { t.Fatalf("start accepting edge: %v", err) } t.Cleanup(func() { server.Stop() }) return regCh, disconnectCh } // startRejectingEdge starts a mock edge bound to host:port that rejects every // RegisterRequest with the given reason. func startRejectingEdge(t *testing.T, ctx context.Context, host string, port int, reason string) { t.Helper() 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: false, Reason: reason}, nil }, ) return client }) if err := server.Start(ctx); err != nil { t.Fatalf("start rejecting edge: %v", err) } t.Cleanup(func() { server.Stop() }) } // TestSupervisorConnectsWhenEdgeStartsAfterNode verifies the actual DialEdge // connection-refused path: Node starts first against a closed port, keeps // retrying in the same process, and establishes exactly one session after Edge // begins listening on that same port. SDD S16. func TestSupervisorConnectsWhenEdgeStartsAfterNode(t *testing.T) { useTempCwd(t) host, port, addr := freeAddrForBootstrap(t) const refusedAttempts = 3 var dialCount int32 retryReady := make(chan struct{}, refusedAttempts) releaseRetry := make(chan struct{}) observedDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) { atomic.AddInt32(&dialCount, 1) return transport.DialEdge(ctx, edgeAddr, token, logger) }) fakeSleeper := func(ctx context.Context, _ time.Duration) { retryReady <- struct{}{} if atomic.LoadInt32(&dialCount) >= refusedAttempts { select { case <-releaseRetry: case <-ctx.Done(): } } } cfg := &config.NodeConfig{ Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"}, Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1}, Logging: config.LoggingConf{Level: "error"}, } app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(observedDialer), 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()) }() // Edge is not listening yet. Each completed sleep entry follows a real // DialEdge connection-refused result. Hold the third retry until the listener // is live so the next actual dial succeeds deterministically. for attempt := 1; attempt <= refusedAttempts; attempt++ { select { case <-retryReady: case <-time.After(5 * time.Second): t.Fatalf("timeout waiting for refused DialEdge attempt %d", attempt) } } if count := atomic.LoadInt32(&dialCount); count != refusedAttempts { t.Fatalf("expected %d real refused dials before Edge startup, got %d", refusedAttempts, count) } serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second) t.Cleanup(serverCancel) regCh, _ := startAcceptingEdge(t, serverCtx, host, port, "delayed-node") close(releaseRetry) select { case <-regCh: case <-time.After(5 * time.Second): t.Fatal("timeout waiting for delayed initial registration") } if count := atomic.LoadInt32(&dialCount); count != refusedAttempts+1 { t.Fatalf("expected %d dials (3 real refusals + 1 success), got %d", refusedAttempts+1, count) } // Exactly one session must be established. select { case <-regCh: t.Fatal("unexpected second registration; supervisor should hold one session") case <-time.After(200 * time.Millisecond): } } // TestSupervisorUnlimitedAttemptsPassesTen verifies that an explicit // max_attempts=0 (unlimited) keeps retrying past the former hard-coded 10-attempt // cap and eventually connects. SDD S16/S17. func TestSupervisorUnlimitedAttemptsPassesTen(t *testing.T) { useTempCwd(t) host, port, addr := freeAddrForBootstrap(t) serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second) t.Cleanup(serverCancel) regCh, _ := startAcceptingEdge(t, serverCtx, host, port, "unlimited-node") const failFirst = 15 // exceeds the old hard-coded 10-attempt cap var dialCount int32 fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) { if atomic.AddInt32(&dialCount, 1) <= failFirst { return nil, errors.New("edge unavailable") } return transport.DialEdge(ctx, edgeAddr, token, logger) }) fakeSleeper := func(context.Context, time.Duration) {} cfg := &config.NodeConfig{ Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"}, Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1}, 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()) }() select { case <-regCh: case <-time.After(5 * time.Second): t.Fatal("timeout waiting for registration after >10 unlimited attempts") } count := atomic.LoadInt32(&dialCount) if count != failFirst+1 { t.Fatalf("expected %d dials, got %d", failFirst+1, count) } if count <= 10 { t.Fatalf("unlimited policy must retry past 10 attempts, got %d", count) } } // TestSupervisorFiniteExhaustionExitsOne verifies that a finite reconnect policy // makes exactly max_attempts initial-connect attempts and then terminates the // node with exit code 1. SDD S17. func TestSupervisorFiniteExhaustionExitsOne(t *testing.T) { useTempCwd(t) _, _, addr := freeAddrForBootstrap(t) // no server listening var dialCount int32 fakeDialer := bootstrap.DialFunc(func(_ context.Context, _, _ string, _ *zap.Logger) (*transport.RegisterResult, error) { atomic.AddInt32(&dialCount, 1) return nil, errors.New("edge unavailable") }) fakeSleeper := func(context.Context, time.Duration) {} 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), 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()) }() 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 finite exhaustion shutdown") } if count := atomic.LoadInt32(&dialCount); count != 3 { t.Fatalf("expected exactly 3 initial-connect attempts, got %d", count) } } // TestSupervisorFatalFailureDoesNotRetry verifies that an authoritative // registration rejection is non-retryable: the supervisor exits with code 1 // after a single attempt without any retry. SDD S17. func TestSupervisorFatalFailureDoesNotRetry(t *testing.T) { useTempCwd(t) host, port, addr := freeAddrForBootstrap(t) serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second) t.Cleanup(serverCancel) startRejectingEdge(t, serverCtx, host, port, "bad-token") var dialCount int32 fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) { atomic.AddInt32(&dialCount, 1) return transport.DialEdge(ctx, edgeAddr, token, logger) }) 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, 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()) }() select { case sig := <-app.Wait(): if sig.ExitCode != 1 { t.Fatalf("expected exit code 1 on fatal failure, got %d", sig.ExitCode) } case <-time.After(5 * time.Second): t.Fatal("timeout waiting for fatal shutdown") } if count := atomic.LoadInt32(&dialCount); count != 1 { t.Fatalf("fatal failure must not retry: expected 1 dial, got %d", count) } } // TestSupervisorShutdownCancelsRetry verifies that a local shutdown cancels a // pending retry sleep, ends cleanly with no error, and runs no additional // attempt. SDD S17. func TestSupervisorShutdownCancelsRetry(t *testing.T) { useTempCwd(t) _, _, addr := freeAddrForBootstrap(t) // no server listening var dialCount int32 fakeDialer := bootstrap.DialFunc(func(_ context.Context, _, _ string, _ *zap.Logger) (*transport.RegisterResult, error) { atomic.AddInt32(&dialCount, 1) return nil, errors.New("edge unavailable") }) sleeping := make(chan struct{}, 1) fakeSleeper := func(ctx context.Context, _ time.Duration) { select { case sleeping <- struct{}{}: default: } <-ctx.Done() // block until the supervisor is cancelled by shutdown } cfg := &config.NodeConfig{ Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"}, Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1}, 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) } // Wait until the supervisor has dialed once and parked in the retry sleep. select { case <-sleeping: case <-time.After(5 * time.Second): t.Fatal("timeout waiting for supervisor to enter retry sleep") } countAtSleep := atomic.LoadInt32(&dialCount) if countAtSleep != 1 { t.Fatalf("expected exactly 1 dial before parking in sleep, got %d", countAtSleep) } // Local shutdown must cancel the sleep and stop cleanly with no error. stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second) defer stopCancel() if err := app.Stop(stopCtx); err != nil { t.Fatalf("app stop must be clean, got: %v", err) } if count := atomic.LoadInt32(&dialCount); count != countAtSleep { t.Fatalf("shutdown must not trigger extra dial: had %d, got %d", countAtSleep, count) } } // TestSupervisorDoesNotStartMetricsWhenPortDisabled guards the Node default // against treating port zero as an ephemeral wildcard listener. It waits for // the reconnect supervisor to be active before asserting the injected starter // was never called, then verifies normal shutdown still cancels the retry. func TestSupervisorDoesNotStartMetricsWhenPortDisabled(t *testing.T) { useTempCwd(t) _, _, addr := freeAddrForBootstrap(t) fakeDialer := bootstrap.DialFunc(func(_ context.Context, _, _ string, _ *zap.Logger) (*transport.RegisterResult, error) { return nil, errors.New("edge unavailable") }) sleeping := make(chan struct{}, 1) fakeSleeper := func(ctx context.Context, _ time.Duration) { select { case sleeping <- struct{}{}: default: } <-ctx.Done() } metricsStarted := make(chan struct{}, 1) metricsStarter := func(int) error { metricsStarted <- struct{}{} return nil } cfg := &config.NodeConfig{ Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"}, Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1}, Logging: config.LoggingConf{Level: "error"}, Metrics: config.MetricsConf{Port: 0}, } app := fx.New( bootstrap.Module( cfg, bootstrap.WithDialer(fakeDialer), bootstrap.WithSleeper(fakeSleeper), bootstrap.WithMetricsStarter(metricsStarter), ), 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) } select { case <-sleeping: case <-time.After(5 * time.Second): t.Fatal("timeout waiting for supervisor to enter retry sleep") } select { case <-metricsStarted: t.Fatal("metrics starter must not run when metrics port is disabled") case <-time.After(100 * time.Millisecond): } stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second) defer stopCancel() if err := app.Stop(stopCtx); err != nil { t.Fatalf("app stop must be clean, got: %v", err) } } // TestSupervisorShutdownWinsInFlightDialFatal fixes the ordering between a // local stop and a dialer that observes cancellation but returns a fatal result // afterward. Cancellation must produce a clean stop, no exit 1, and no retry. func TestSupervisorShutdownWinsInFlightDialFatal(t *testing.T) { useTempCwd(t) _, _, addr := freeAddrForBootstrap(t) started := make(chan struct{}, 1) var dialCount int32 dialer := bootstrap.DialFunc(func(ctx context.Context, _, _ string, logger *zap.Logger) (*transport.RegisterResult, error) { atomic.AddInt32(&dialCount, 1) started <- struct{}{} <-ctx.Done() return transport.DialEdge(context.Background(), "127.0.0.1:0", "tok", logger) }) cfg := &config.NodeConfig{ Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"}, Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1}, Logging: config.LoggingConf{Level: "error"}, } app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(dialer)), 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) } select { case <-started: case <-time.After(5 * time.Second): t.Fatal("timeout waiting for in-flight dial") } stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second) defer stopCancel() if err := app.Stop(stopCtx); err != nil { t.Fatalf("app stop must be clean, got: %v", err) } if count := atomic.LoadInt32(&dialCount); count != 1 { t.Fatalf("shutdown must not retry fatal late result: got %d dials", count) } select { case sig := <-app.Wait(): t.Fatalf("shutdown must not request failure exit, got %+v", sig) case <-time.After(100 * time.Millisecond): } } // TestSupervisorShutdownClosesInFlightDialSuccess verifies that a dialer which // observes cancellation and then returns a successful registered session cannot // install or leak that late owner. The one registered session is closed by the // post-dial cancellation fence and no failure exit or retry is emitted. func TestSupervisorShutdownClosesInFlightDialSuccess(t *testing.T) { useTempCwd(t) host, port, addr := freeAddrForBootstrap(t) serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second) t.Cleanup(serverCancel) regCh, disconnectCh := startAcceptingEdge(t, serverCtx, host, port, "late-success-node") started := make(chan struct{}, 1) var dialCount int32 dialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) { atomic.AddInt32(&dialCount, 1) started <- struct{}{} <-ctx.Done() lateCtx, lateCancel := context.WithTimeout(context.Background(), 5*time.Second) defer lateCancel() return transport.DialEdge(lateCtx, edgeAddr, token, logger) }) cfg := &config.NodeConfig{ Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"}, Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1}, Logging: config.LoggingConf{Level: "error"}, } app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(dialer)), 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) } select { case <-started: case <-time.After(5 * time.Second): t.Fatal("timeout waiting for in-flight dial") } stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second) defer stopCancel() stopDone := make(chan error, 1) go func() { stopDone <- app.Stop(stopCtx) }() select { case <-regCh: case <-time.After(5 * time.Second): t.Fatal("timeout waiting for late successful registration") } select { case err := <-stopDone: if err != nil { t.Fatalf("app stop must be clean, got: %v", err) } case <-time.After(5 * time.Second): t.Fatal("timeout waiting for app stop") } select { case <-disconnectCh: case <-time.After(5 * time.Second): t.Fatal("late successful session leaked after supervisor shutdown") } if count := atomic.LoadInt32(&dialCount); count != 1 { t.Fatalf("shutdown must not retry late success: got %d dials", count) } select { case <-regCh: t.Fatal("unexpected second registration after shutdown") case sig := <-app.Wait(): t.Fatalf("shutdown must not request failure exit, got %+v", sig) case <-time.After(100 * time.Millisecond): } } // 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; readyCh fires when a // connection completes its ready handshake so the test can close it only after // the session is fully established. regCh := make(chan struct{}, 8) readyCh := 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 }, ) bindEdgeReadyResponder(client, readyCh) 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"}, } // Wrap the real dialer so the test can wait for the initial session to be // fully established before inducing a disconnect (OnStart is non-blocking). established := make(chan struct{}, 1) var establishOnce sync.Once fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) { res, err := transport.DialEdge(ctx, edgeAddr, token, logger) if err == nil { establishOnce.Do(func() { established <- struct{}{} }) } return res, err }) 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 registration and full establishment. select { case <-regCh: case <-time.After(3 * time.Second): t.Fatal("timeout waiting for first registration") } select { case <-established: case <-time.After(3 * time.Second): t.Fatal("timeout waiting for initial session establishment") } var firstClient *toki.TcpClient select { case firstClient = <-acceptedCh: case <-time.After(time.Second): t.Fatal("timeout waiting for first accepted client") } // Wait for the ready handshake before closing so the disconnect is induced on a // fully-established session, not mid-SignalReady. select { case <-readyCh: case <-time.After(3 * time.Second): t.Fatal("timeout waiting for initial ready handshake") } // 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) readyCh := make(chan struct{}, 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 }, ) bindEdgeReadyResponder(client, readyCh) 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. // established fires once the initial session is fully established so the test // can induce the disconnect without racing the initial registration. established := make(chan struct{}, 1) 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 { res, err := transport.DialEdge(ctx, edgeAddr, token, logger) if err == nil { established <- struct{}{} } return res, err } 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 the initial session to establish, then close it from the edge side. select { case <-established: case <-time.After(3 * time.Second): t.Fatal("timeout waiting for initial session establishment") } var firstClient *toki.TcpClient select { case firstClient = <-acceptedCh: case <-time.After(3 * time.Second): t.Fatal("timeout waiting for first connection") } // Close only after the ready handshake so the first session counts as an // established connection, not a mid-SignalReady failure. select { case <-readyCh: case <-time.After(3 * time.Second): t.Fatal("timeout waiting for initial ready handshake") } _ = 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) readyCh := make(chan struct{}, 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 }, ) bindEdgeReadyResponder(client, readyCh) 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. // established fires once the initial session is fully established so the test // can induce the disconnect without racing the initial registration. established := make(chan struct{}, 1) 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 { res, err := transport.DialEdge(ctx, edgeAddr, token, logger) if err == nil { established <- struct{}{} } return res, err } 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 the initial session to establish, then close it from the edge side. select { case <-established: case <-time.After(3 * time.Second): t.Fatal("timeout waiting for initial session establishment") } var firstClient *toki.TcpClient select { case firstClient = <-acceptedCh: case <-time.After(3 * time.Second): t.Fatal("timeout waiting for first connection") } // Close only after the ready handshake so the first session counts as an // established connection, not a mid-SignalReady failure. select { case <-readyCh: case <-time.After(3 * time.Second): t.Fatal("timeout waiting for initial ready handshake") } _ = 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) } } }