chore: add support for tailnet updates to Tunneler FSM (#23875)

<!--

If you have used AI to produce some or all of this PR, please ensure you have read our [AI Contribution guidelines](https://coder.com/docs/about/contributing/AI_CONTRIBUTING) before submitting.

-->

relates to GRU-18  
  
Adds support for tailnet updates to Tunneler FSM.

codersdk/workspacesdk/tunneler/tunneler.go

@@ -20,7 +20,7 @@ type NetworkedApplication interface {
 	// Closer is used to gracefully tear down the application prior to stopping the tunnel.
 	io.Closer
 	// Start the NetworkedApplication, using the provided AgentConn to connect.
-	Start(conn workspacesdk.AgentConn)
+	Start(conn workspacesdk.AgentConn) error
 }
 
 // WorkspaceStarter is used to create a start build of the workspace. It is an interface here because the CLI has lots
@@ -63,6 +63,33 @@ const (
 	maxState // used for testing
 )
 
+func (s state) String() string {
+	switch s {
+	case stateInit:
+		return "init"
+	case exit:
+		return "exit"
+	case waitToStart:
+		return "waitToStart"
+	case waitForWorkspaceStarted:
+		return "waitForWorkspaceStarted"
+	case waitForAgent:
+		return "waitForAgent"
+	case establishTailnet:
+		return "establishTailnet"
+	case tailnetUp:
+		return "tailnetUp"
+	case applicationUp:
+		return "applicationUp"
+	case shutdownApplication:
+		return "shutdownApplication"
+	case shutdownTailnet:
+		return "shutdownTailnet"
+	default:
+		return fmt.Sprintf("unknown(%d)", s)
+	}
+}
+
 type Tunneler struct {
 	config    Config
 	ctx       context.Context
@@ -179,10 +206,12 @@ func (t *Tunneler) eventLoop() {
 		case e.tailnetUpdate != nil:
 			t.handleTailnetUpdate(e.tailnetUpdate)
 		}
+		t.config.DebugLogger.Debug(t.ctx, "handled event", slog.F("state", t.state))
 	}
 }
 
 func (t *Tunneler) handleSignal() {
+	t.config.DebugLogger.Debug(t.ctx, "got shutdown signal")
 	switch t.state {
 	case exit, shutdownTailnet, shutdownApplication:
 		return
@@ -313,6 +342,10 @@ func (*Tunneler) handleProvisionerJobLog(*codersdk.ProvisionerJobLog) {
 }
 
 func (t *Tunneler) handleAgentUpdate(update *agentUpdate) {
+	t.config.DebugLogger.Debug(t.ctx, "handling agent update",
+		slog.F("state", t.state),
+		slog.F("lifecycle", update.lifecycle),
+		slog.F("agent_id", update.id))
 	if t.state != waitForAgent {
 		return
 	}
@@ -399,10 +432,91 @@ func (t *Tunneler) handleAppUpdate(update *networkedApplicationUpdate) {
 		slog.F("state", t.state), slog.F("app_up", update.up))
 }
 
-func (*Tunneler) handleTailnetUpdate(*tailnetUpdate) {
+func (t *Tunneler) handleTailnetUpdate(update *tailnetUpdate) {
+	switch t.state {
+	case exit:
+		return
+	case stateInit, waitToStart, waitForAgent, waitForWorkspaceStarted:
+		t.config.DebugLogger.Error(t.ctx, "unexpected: tailnet update before we started it",
+			slog.F("state", t.state), slog.F("app_up", update.up), slog.Error(update.err))
+		return
+	}
+	if update.up {
+		t.config.DebugLogger.Debug(t.ctx, "got tailnet 'up' update", slog.F("state", t.state))
+		switch t.state {
+		case establishTailnet:
+			t.agentConn = update.conn
+			t.state = tailnetUp
+			t.wg.Add(1)
+			go t.startApp()
+			return
+		case shutdownTailnet:
+			// this means we were notified to shut down while we were starting the tailnet. We need to tear it down.
+			t.config.DebugLogger.Debug(t.ctx, "gracefully shutting down tailnet after it started")
+			t.agentConn = update.conn
+			t.wg.Add(1)
+			go t.shutdownTailnet()
+			return
+		case tailnetUp:
+			t.config.DebugLogger.Error(t.ctx, "unexpected: got tailnet 'up' update when it is already up")
+			if update.conn != nil && update.conn != t.agentConn {
+				// somehow we have two updates with different connections. Something very bad has happened so we are
+				// going to just bail, rather than try to gracefully tear them both down.
+				t.config.DebugLogger.Fatal(t.ctx, "unexpected: got two different connections")
+			}
+			return
+		case shutdownApplication:
+			t.config.DebugLogger.Error(t.ctx, "unexpected: got tailnet 'up' update when we expected application update")
+			return
+		}
+	}
+	t.config.DebugLogger.Debug(t.ctx, "got tailnet 'down' update", slog.F("state", t.state))
+	switch t.state {
+	case establishTailnet, shutdownTailnet:
+		// Either we failed to establish, or we successfully shut down. In the former case, the error has already been
+		// logged. Nothing else to do now that tailnet is down, since it implies the application is also down.
+		t.cancel()
+		t.state = exit
+		return
+	case tailnetUp:
+		t.config.DebugLogger.Error(t.ctx,
+			"unexpected: got tailnet 'down' update when we were starting the application")
+		return
+	case shutdownApplication:
+		t.config.DebugLogger.Error(t.ctx,
+			"unexpected: got tailnet 'down' update when we were stopping the application")
+		return
+	}
+	t.config.DebugLogger.Critical(t.ctx, "unhandled tailnet update",
+		slog.F("state", t.state), slog.F("app_up", update.up))
+}
+
+func (t *Tunneler) startApp() {
+	t.config.DebugLogger.Debug(t.ctx, "starting networked application")
+	defer t.wg.Done()
+	err := t.config.App.Start(t.agentConn)
+	if err != nil {
+		t.config.DebugLogger.Error(t.ctx, "failed to start application", slog.Error(err))
+		if t.config.LogWriter != nil {
+			_, _ = fmt.Fprintf(t.config.LogWriter, "failed to start: %s", err.Error())
+		}
+		select {
+		case <-t.ctx.Done():
+			t.config.DebugLogger.Info(t.ctx,
+				"context expired before sending event after failed network application start")
+		case t.events <- tunnelerEvent{appUpdate: &networkedApplicationUpdate{up: false, err: err}}:
+		}
+		return
+	}
+	select {
+	case <-t.ctx.Done():
+		t.config.DebugLogger.Info(t.ctx, "context expired before sending network application start update")
+	case t.events <- tunnelerEvent{appUpdate: &networkedApplicationUpdate{up: true}}:
+	}
 }
 
 func (t *Tunneler) closeApp() {
+	t.config.DebugLogger.Info(t.ctx, "closing networked application")
 	defer t.wg.Done()
 	err := t.config.App.Close()
 	if err != nil {
@@ -416,6 +530,7 @@ func (t *Tunneler) closeApp() {
 }
 
 func (t *Tunneler) startWorkspace() {
+	t.config.DebugLogger.Info(t.ctx, "starting workspace")
 	defer t.wg.Done()
 	err := t.config.WorkspaceStarter.StartWorkspace()
 	if err != nil {
@@ -428,10 +543,12 @@ func (t *Tunneler) startWorkspace() {
 			t.config.DebugLogger.Info(t.ctx, "context expired before sending signal after failed workspace start")
 		case t.events <- tunnelerEvent{appUpdate: &networkedApplicationUpdate{up: false}}:
 		}
+		return
 	}
 }
 
 func (t *Tunneler) connectTailnet(id uuid.UUID) {
+	t.config.DebugLogger.Info(t.ctx, "connecting tailnet")
 	defer t.wg.Done()
 	conn, err := t.client.DialAgent(t.ctx, id, &workspacesdk.DialAgentOptions{
 		Logger: t.config.DebugLogger.Named("dialer"),
@@ -446,6 +563,7 @@ func (t *Tunneler) connectTailnet(id uuid.UUID) {
 			t.config.DebugLogger.Info(t.ctx, "context expired before sending event after failed agent dial")
 		case t.events <- tunnelerEvent{tailnetUpdate: &tailnetUpdate{up: false, err: err}}:
 		}
+		return
 	}
 	select {
 	case <-t.ctx.Done():
@@ -455,6 +573,7 @@ func (t *Tunneler) connectTailnet(id uuid.UUID) {
 }
 
 func (t *Tunneler) shutdownTailnet() {
+	t.config.DebugLogger.Info(t.ctx, "shutting down tailnet")
 	defer t.wg.Done()
 	err := t.agentConn.Close()
 	if err != nil {

codersdk/workspacesdk/tunneler/tunneler_internal_test.go

@@ -390,11 +390,10 @@ func TestAppUpdate(t *testing.T) {
 			ctx, cancel := context.WithCancel(testCtx)
 			uut := &Tunneler{
 				config: Config{
-					WorkspaceID:      workspaceID,
-					AgentName:        "test",
-					DebugLogger:      logger.Named("tunneler"),
-					NoWaitForScripts: true,
-					App:              fApp,
+					WorkspaceID: workspaceID,
+					AgentName:   "test",
+					DebugLogger: logger.Named("tunneler"),
+					App:         fApp,
 				},
 				events:    make(chan tunnelerEvent),
 				ctx:       ctx,
@@ -415,6 +414,171 @@ func TestAppUpdate(t *testing.T) {
 	}
 }
 
+func TestTailnetUpdate(t *testing.T) {
+	t.Parallel()
+	testCases := []struct {
+		name                                  string
+		up                                    bool
+		initState, expected                   state
+		expectStartApp, expectShutdownTailnet bool
+	}{
+		{
+			name:           "mainline_up",
+			up:             true,
+			initState:      establishTailnet,
+			expected:       tailnetUp,
+			expectStartApp: true,
+		},
+		{
+			name:      "mainline_down",
+			up:        false,
+			initState: shutdownTailnet,
+			expected:  exit,
+		},
+		{
+			name:      "failed_tailnet_start",
+			up:        false,
+			initState: establishTailnet,
+			expected:  exit,
+		},
+		{
+			name:                  "graceful_shutdown_while_starting",
+			up:                    true,
+			initState:             shutdownTailnet,
+			expected:              shutdownTailnet,
+			expectShutdownTailnet: true,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			workspaceID := uuid.UUID{1}
+			logger := testutil.Logger(t)
+
+			ctrl := gomock.NewController(t)
+			mAgentConn := agentconnmock.NewMockAgentConn(ctrl)
+			fApp := &fakeApp{}
+
+			testCtx := testutil.Context(t, testutil.WaitShort)
+			ctx, cancel := context.WithCancel(testCtx)
+			uut := &Tunneler{
+				config: Config{
+					WorkspaceID: workspaceID,
+					AgentName:   "test",
+					DebugLogger: logger.Named("tunneler"),
+					App:         fApp,
+				},
+				events: make(chan tunnelerEvent),
+				ctx:    ctx,
+				cancel: cancel,
+				state:  tc.initState,
+			}
+			if tc.expectShutdownTailnet {
+				mAgentConn.EXPECT().Close().Return(nil).Times(1)
+			}
+
+			update := &tailnetUpdate{up: tc.up}
+			if tc.up {
+				update.conn = mAgentConn
+			}
+			uut.handleTailnetUpdate(update)
+			require.Equal(t, tc.expected, uut.state)
+			cancel() // so that any goroutines can complete without an event loop
+			waitForGoroutines(testCtx, t, uut)
+			require.Equal(t, tc.expectStartApp, fApp.started)
+		})
+	}
+}
+
+func TestTunneler_EventLoop_Signal(t *testing.T) {
+	t.Parallel()
+
+	workspaceID := uuid.UUID{1}
+	agentID := uuid.UUID{2}
+	logger := testutil.Logger(t)
+
+	ctrl := gomock.NewController(t)
+	mAgentConn := agentconnmock.NewMockAgentConn(ctrl)
+	fApp := &fakeApp{
+		starts: make(chan appStartRequest),
+		closes: make(chan errorResult),
+	}
+	fClient := &fakeClient{
+		dials: make(chan dialRequest),
+	}
+
+	testCtx := testutil.Context(t, testutil.WaitShort)
+	ctx, cancel := context.WithCancel(testCtx)
+	uut := &Tunneler{
+		client: fClient,
+		config: Config{
+			WorkspaceID: workspaceID,
+			AgentName:   "test",
+			DebugLogger: logger.Named("tunneler"),
+			App:         fApp,
+		},
+		events: make(chan tunnelerEvent),
+		ctx:    ctx,
+		cancel: cancel,
+		state:  stateInit,
+	}
+	uut.wg.Add(1)
+	go uut.eventLoop()
+
+	testutil.RequireSend(testCtx, t, uut.events, tunnelerEvent{
+		buildUpdate: &buildUpdate{
+			transition: codersdk.WorkspaceTransitionStart,
+			jobStatus:  codersdk.ProvisionerJobPending,
+		},
+	})
+	testutil.RequireSend(testCtx, t, uut.events, tunnelerEvent{
+		buildUpdate: &buildUpdate{
+			transition: codersdk.WorkspaceTransitionStart,
+			jobStatus:  codersdk.ProvisionerJobRunning,
+		},
+	})
+	testutil.RequireSend(testCtx, t, uut.events, tunnelerEvent{
+		buildUpdate: &buildUpdate{
+			transition: codersdk.WorkspaceTransitionStart,
+			jobStatus:  codersdk.ProvisionerJobSucceeded,
+		},
+	})
+	testutil.RequireSend(testCtx, t, uut.events, tunnelerEvent{
+		agentUpdate: &agentUpdate{
+			lifecycle: codersdk.WorkspaceAgentLifecycleReady,
+			id:        agentID,
+		},
+	})
+
+	// Workspace started, agent ready. Should connect the tailnet.
+	tailnetDial := testutil.RequireReceive(testCtx, t, fClient.dials)
+	testutil.RequireSend(testCtx, t, tailnetDial.result, dialResult{conn: mAgentConn})
+
+	// Tailnet up, should start App
+	appStart := testutil.RequireReceive(testCtx, t, fApp.starts)
+	require.Equal(t, mAgentConn, appStart.conn)
+	testutil.RequireSend(testCtx, t, appStart.result, nil)
+
+	connClosed := make(chan struct{})
+	mAgentConn.EXPECT().Close().Times(1).Do(func() {
+		close(connClosed)
+	}).Return(nil)
+
+	testutil.RequireSend(testCtx, t, uut.events, tunnelerEvent{
+		shutdownSignal: &shutdownSignal{},
+	})
+
+	closeReq := testutil.RequireReceive(testCtx, t, fApp.closes)
+	testutil.RequireSend(testCtx, t, closeReq.result, nil)
+
+	// next tailnet closes
+	_ = testutil.TryReceive(testCtx, t, connClosed)
+
+	// should cancel the loop and be at exit
+	waitForGoroutines(testCtx, t, uut)
+	require.Equal(t, exit, uut.state)
+}
+
 func waitForGoroutines(ctx context.Context, t *testing.T, tunneler *Tunneler) {
 	done := make(chan struct{})
 	go func() {
@@ -424,32 +588,87 @@ func waitForGoroutines(ctx context.Context, t *testing.T, tunneler *Tunneler) {
 	_ = testutil.TryReceive(ctx, t, done)
 }
 
+type errorResult struct {
+	result chan error
+}
+
 type fakeWorkspaceStarter struct {
+	starts  chan errorResult
 	started bool
 }
 
 func (f *fakeWorkspaceStarter) StartWorkspace() error {
-	f.started = true
-	return nil
+	if f.starts == nil {
+		f.started = true
+		return nil
+	}
+	result := make(chan error)
+	f.starts <- errorResult{result: result}
+	return <-result
+}
+
+type appStartRequest struct {
+	conn   workspacesdk.AgentConn
+	result chan error
 }
 
 type fakeApp struct {
-	closed bool
+	starts  chan appStartRequest
+	closes  chan errorResult
+	closed  bool
+	started bool
 }
 
 func (f *fakeApp) Close() error {
-	f.closed = true
-	return nil
+	if f.closes == nil {
+		f.closed = true
+		return nil
+	}
+	result := make(chan error)
+	f.closes <- errorResult{result: result}
+	return <-result
 }
 
-func (*fakeApp) Start(workspacesdk.AgentConn) {}
+func (f *fakeApp) Start(conn workspacesdk.AgentConn) error {
+	if f.starts == nil {
+		f.started = true
+		return nil
+	}
+	result := make(chan error)
+	f.starts <- appStartRequest{result: result, conn: conn}
+	return <-result
+}
+
+type dialRequest struct {
+	id     uuid.UUID
+	result chan dialResult
+}
+
+type dialResult struct {
+	conn workspacesdk.AgentConn
+	err  error
+}
 
 type fakeClient struct {
+	// async:
+	dials chan dialRequest
+
+	// sync:
 	conn   workspacesdk.AgentConn
 	dialed bool
 }
 
-func (f *fakeClient) DialAgent(context.Context, uuid.UUID, *workspacesdk.DialAgentOptions) (workspacesdk.AgentConn, error) {
-	f.dialed = true
-	return f.conn, nil
+func (f *fakeClient) DialAgent(
+	_ context.Context, id uuid.UUID, _ *workspacesdk.DialAgentOptions,
+) (
+	workspacesdk.AgentConn, error,
+) {
+	if f.dials == nil {
+		f.dialed = true
+		return f.conn, nil
+	}
+	results := make(chan dialResult)
+	f.dials <- dialRequest{id: id, result: results}
+	result := <-results
+	return result.conn, result.err
 }