feat: Add workspace agent lifecycle state reporting (#5785)

agent/agent.go

@@ -71,6 +71,7 @@ type Client interface {
 	WorkspaceAgentMetadata(ctx context.Context) (codersdk.WorkspaceAgentMetadata, error)
 	ListenWorkspaceAgent(ctx context.Context) (net.Conn, error)
 	AgentReportStats(ctx context.Context, log slog.Logger, stats func() *codersdk.AgentStats) (io.Closer, error)
+	PostWorkspaceAgentLifecycle(ctx context.Context, state codersdk.PostWorkspaceAgentLifecycleRequest) error
 	PostWorkspaceAgentAppHealth(ctx context.Context, req codersdk.PostWorkspaceAppHealthsRequest) error
 	PostWorkspaceAgentVersion(ctx context.Context, version string) error
 }
@@ -101,6 +102,7 @@ func New(options Options) io.Closer {
 		exchangeToken:          options.ExchangeToken,
 		filesystem:             options.Filesystem,
 		tempDir:                options.TempDir,
+		lifecycleUpdate:        make(chan struct{}, 1),
 	}
 	a.init(ctx)
 	return a
@@ -127,6 +129,10 @@ type agent struct {
 	sessionToken atomic.Pointer[string]
 	sshServer    *ssh.Server
 
+	lifecycleUpdate chan struct{}
+	lifecycleMu     sync.Mutex // Protects following.
+	lifecycleState  codersdk.WorkspaceAgentLifecycle
+
 	network *tailnet.Conn
 }
 
@@ -135,6 +141,8 @@ type agent struct {
 // may be happening, but regardless after the intermittent
 // failure, you'll want the agent to reconnect.
 func (a *agent) runLoop(ctx context.Context) {
+	go a.reportLifecycleLoop(ctx)
+
 	for retrier := retry.New(100*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
 		a.logger.Info(ctx, "running loop")
 		err := a.run(ctx)
@@ -156,6 +164,58 @@ func (a *agent) runLoop(ctx context.Context) {
 	}
 }
 
+// reportLifecycleLoop reports the current lifecycle state once.
+// Only the latest state is reported, intermediate states may be
+// lost if the agent can't communicate with the API.
+func (a *agent) reportLifecycleLoop(ctx context.Context) {
+	var lastReported codersdk.WorkspaceAgentLifecycle
+	for {
+		select {
+		case <-a.lifecycleUpdate:
+		case <-ctx.Done():
+			return
+		}
+
+		for r := retry.New(time.Second, 15*time.Second); r.Wait(ctx); {
+			a.lifecycleMu.Lock()
+			state := a.lifecycleState
+			a.lifecycleMu.Unlock()
+
+			if state == lastReported {
+				break
+			}
+
+			a.logger.Debug(ctx, "post lifecycle state", slog.F("state", state))
+
+			err := a.client.PostWorkspaceAgentLifecycle(ctx, codersdk.PostWorkspaceAgentLifecycleRequest{
+				State: state,
+			})
+			if err == nil {
+				lastReported = state
+				break
+			}
+			if xerrors.Is(err, context.Canceled) || xerrors.Is(err, context.DeadlineExceeded) {
+				return
+			}
+			// If we fail to report the state we probably shouldn't exit, log only.
+			a.logger.Error(ctx, "post state", slog.Error(err))
+		}
+	}
+}
+
+func (a *agent) setLifecycle(ctx context.Context, state codersdk.WorkspaceAgentLifecycle) {
+	a.lifecycleMu.Lock()
+	defer a.lifecycleMu.Unlock()
+
+	a.logger.Debug(ctx, "set lifecycle state", slog.F("state", state), slog.F("previous", a.lifecycleState))
+
+	a.lifecycleState = state
+	select {
+	case a.lifecycleUpdate <- struct{}{}:
+	default:
+	}
+}
+
 func (a *agent) run(ctx context.Context) error {
 	// This allows the agent to refresh it's token if necessary.
 	// For instance identity this is required, since the instance
@@ -180,22 +240,60 @@ func (a *agent) run(ctx context.Context) error {
 
 	// The startup script should only execute on the first run!
 	if oldMetadata == nil {
+		a.setLifecycle(ctx, codersdk.WorkspaceAgentLifecycleStarting)
+
+		// Perform overrides early so that Git auth can work even if users
+		// connect to a workspace that is not yet ready. We don't run this
+		// concurrently with the startup script to avoid conflicts between
+		// them.
+		if metadata.GitAuthConfigs > 0 {
+			// If this fails, we should consider surfacing the error in the
+			// startup log and setting the lifecycle state to be "start_error"
+			// (after startup script completion), but for now we'll just log it.
+			err := gitauth.OverrideVSCodeConfigs(a.filesystem)
+			if err != nil {
+				a.logger.Warn(ctx, "failed to override vscode git auth configs", slog.Error(err))
+			}
+		}
+
+		scriptDone := make(chan error, 1)
+		scriptStart := time.Now()
 		go func() {
-			err := a.runStartupScript(ctx, metadata.StartupScript)
+			defer close(scriptDone)
+			scriptDone <- a.runStartupScript(ctx, metadata.StartupScript)
+		}()
+		go func() {
+			var timeout <-chan time.Time
+			// If timeout is zero, an older version of the coder
+			// provider was used. Otherwise a timeout is always > 0.
+			if metadata.StartupScriptTimeout > 0 {
+				t := time.NewTimer(metadata.StartupScriptTimeout)
+				defer t.Stop()
+				timeout = t.C
+			}
+
+			var err error
+			select {
+			case err = <-scriptDone:
+			case <-timeout:
+				a.logger.Warn(ctx, "startup script timed out")
+				a.setLifecycle(ctx, codersdk.WorkspaceAgentLifecycleStartTimeout)
+				err = <-scriptDone // The script can still complete after a timeout.
+			}
 			if errors.Is(err, context.Canceled) {
 				return
 			}
+			execTime := time.Since(scriptStart)
+			lifecycleStatus := codersdk.WorkspaceAgentLifecycleReady
 			if err != nil {
-				a.logger.Warn(ctx, "agent script failed", slog.Error(err))
+				a.logger.Warn(ctx, "startup script failed", slog.F("execution_time", execTime), slog.Error(err))
+				lifecycleStatus = codersdk.WorkspaceAgentLifecycleStartError
+			} else {
+				a.logger.Info(ctx, "startup script completed", slog.F("execution_time", execTime))
 			}
-		}()
-	}
 
-	if metadata.GitAuthConfigs > 0 {
-		err = gitauth.OverrideVSCodeConfigs(a.filesystem)
-		if err != nil {
-			return xerrors.Errorf("override vscode configuration for git auth: %w", err)
-		}
+			a.setLifecycle(ctx, lifecycleStatus)
+		}()
 	}
 
 	// This automatically closes when the context ends!

agent/agent_test.go

@@ -57,7 +57,7 @@ func TestAgent_Stats_SSH(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 
-	conn, stats, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
+	conn, _, stats, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
 
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
@@ -83,7 +83,7 @@ func TestAgent_Stats_ReconnectingPTY(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 
-	conn, stats, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
+	conn, _, stats, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
 
 	ptyConn, err := conn.ReconnectingPTY(ctx, uuid.New(), 128, 128, "/bin/bash")
 	require.NoError(t, err)
@@ -529,7 +529,8 @@ func TestAgent_SFTP(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		home = "/" + strings.ReplaceAll(home, "\\", "/")
 	}
-	conn, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
+	//nolint:dogsled
+	conn, _, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	defer sshClient.Close()
@@ -560,7 +561,8 @@ func TestAgent_SCP(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 
-	conn, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
+	//nolint:dogsled
+	conn, _, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	defer sshClient.Close()
@@ -664,7 +666,8 @@ func TestAgent_StartupScript(t *testing.T) {
 		t.Skip("This test doesn't work on Windows for some reason...")
 	}
 	content := "output"
-	_, _, fs := setupAgent(t, codersdk.WorkspaceAgentMetadata{
+	//nolint:dogsled
+	_, _, _, fs := setupAgent(t, codersdk.WorkspaceAgentMetadata{
 		StartupScript: "echo " + content,
 	}, 0)
 	var gotContent string
@@ -692,6 +695,97 @@ func TestAgent_StartupScript(t *testing.T) {
 	require.Equal(t, content, strings.TrimSpace(gotContent))
 }
 
+func TestAgent_Lifecycle(t *testing.T) {
+	t.Parallel()
+
+	t.Run("Timeout", func(t *testing.T) {
+		t.Parallel()
+
+		_, client, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{
+			StartupScript:        "sleep 10",
+			StartupScriptTimeout: time.Nanosecond,
+		}, 0)
+
+		want := []codersdk.WorkspaceAgentLifecycle{
+			codersdk.WorkspaceAgentLifecycleStarting,
+			codersdk.WorkspaceAgentLifecycleStartTimeout,
+		}
+
+		var got []codersdk.WorkspaceAgentLifecycle
+		assert.Eventually(t, func() bool {
+			got = client.getLifecycleStates()
+			return len(got) > 0 && got[len(got)-1] == want[len(want)-1]
+		}, testutil.WaitShort, testutil.IntervalMedium)
+		switch len(got) {
+		case 1:
+			// This can happen if lifecycle state updates are
+			// too fast, only the latest one is reported.
+			require.Equal(t, want[1:], got)
+		default:
+			// This is the expected case.
+			require.Equal(t, want, got)
+		}
+	})
+
+	t.Run("Error", func(t *testing.T) {
+		t.Parallel()
+
+		_, client, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{
+			StartupScript:        "false",
+			StartupScriptTimeout: 30 * time.Second,
+		}, 0)
+
+		want := []codersdk.WorkspaceAgentLifecycle{
+			codersdk.WorkspaceAgentLifecycleStarting,
+			codersdk.WorkspaceAgentLifecycleStartError,
+		}
+
+		var got []codersdk.WorkspaceAgentLifecycle
+		assert.Eventually(t, func() bool {
+			got = client.getLifecycleStates()
+			return len(got) > 0 && got[len(got)-1] == want[len(want)-1]
+		}, testutil.WaitShort, testutil.IntervalMedium)
+		switch len(got) {
+		case 1:
+			// This can happen if lifecycle state updates are
+			// too fast, only the latest one is reported.
+			require.Equal(t, want[1:], got)
+		default:
+			// This is the expected case.
+			require.Equal(t, want, got)
+		}
+	})
+
+	t.Run("Ready", func(t *testing.T) {
+		t.Parallel()
+
+		_, client, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{
+			StartupScript:        "true",
+			StartupScriptTimeout: 30 * time.Second,
+		}, 0)
+
+		want := []codersdk.WorkspaceAgentLifecycle{
+			codersdk.WorkspaceAgentLifecycleStarting,
+			codersdk.WorkspaceAgentLifecycleReady,
+		}
+
+		var got []codersdk.WorkspaceAgentLifecycle
+		assert.Eventually(t, func() bool {
+			got = client.getLifecycleStates()
+			return len(got) > 0 && got[len(got)-1] == want[len(want)-1]
+		}, testutil.WaitShort, testutil.IntervalMedium)
+		switch len(got) {
+		case 1:
+			// This can happen if lifecycle state updates are
+			// too fast, only the latest one is reported.
+			require.Equal(t, want[1:], got)
+		default:
+			// This is the expected case.
+			require.Equal(t, want, got)
+		}
+	})
+}
+
 func TestAgent_ReconnectingPTY(t *testing.T) {
 	t.Parallel()
 	if runtime.GOOS == "windows" {
@@ -704,7 +798,8 @@ func TestAgent_ReconnectingPTY(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 
-	conn, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
+	//nolint:dogsled
+	conn, _, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
 	id := uuid.New()
 	netConn, err := conn.ReconnectingPTY(ctx, id, 100, 100, "/bin/bash")
 	require.NoError(t, err)
@@ -805,7 +900,8 @@ func TestAgent_Dial(t *testing.T) {
 				}
 			}()
 
-			conn, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
+			//nolint:dogsled
+			conn, _, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
 			require.True(t, conn.AwaitReachable(context.Background()))
 			conn1, err := conn.DialContext(context.Background(), l.Addr().Network(), l.Addr().String())
 			require.NoError(t, err)
@@ -826,7 +922,8 @@ func TestAgent_Speedtest(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 	derpMap := tailnettest.RunDERPAndSTUN(t)
-	conn, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{
+	//nolint:dogsled
+	conn, _, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{
 		DERPMap: derpMap,
 	}, 0)
 	defer conn.Close()
@@ -911,7 +1008,8 @@ func TestAgent_WriteVSCodeConfigs(t *testing.T) {
 }
 
 func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
-	agentConn, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
+	//nolint:dogsled
+	agentConn, _, _, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
 	waitGroup := sync.WaitGroup{}
@@ -957,7 +1055,8 @@ func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exe
 func setupSSHSession(t *testing.T, options codersdk.WorkspaceAgentMetadata) *ssh.Session {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
-	conn, _, _ := setupAgent(t, options, 0)
+	//nolint:dogsled
+	conn, _, _, _ := setupAgent(t, options, 0)
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	t.Cleanup(func() {
@@ -979,6 +1078,7 @@ func (c closeFunc) Close() error {
 
 func setupAgent(t *testing.T, metadata codersdk.WorkspaceAgentMetadata, ptyTimeout time.Duration) (
 	*codersdk.AgentConn,
+	*client,
 	<-chan *codersdk.AgentStats,
 	afero.Fs,
 ) {
@@ -992,14 +1092,15 @@ func setupAgent(t *testing.T, metadata codersdk.WorkspaceAgentMetadata, ptyTimeo
 	agentID := uuid.New()
 	statsCh := make(chan *codersdk.AgentStats, 50)
 	fs := afero.NewMemMapFs()
+	c := &client{
+		t:           t,
+		agentID:     agentID,
+		metadata:    metadata,
+		statsChan:   statsCh,
+		coordinator: coordinator,
+	}
 	closer := agent.New(agent.Options{
-		Client: &client{
-			t:           t,
-			agentID:     agentID,
-			metadata:    metadata,
-			statsChan:   statsCh,
-			coordinator: coordinator,
-		},
+		Client:                 c,
 		Filesystem:             fs,
 		Logger:                 slogtest.Make(t, nil).Leveled(slog.LevelDebug),
 		ReconnectingPTYTimeout: ptyTimeout,
@@ -1032,7 +1133,7 @@ func setupAgent(t *testing.T, metadata codersdk.WorkspaceAgentMetadata, ptyTimeo
 	conn.SetNodeCallback(sendNode)
 	return &codersdk.AgentConn{
 		Conn: conn,
-	}, statsCh, fs
+	}, c, statsCh, fs
 }
 
 var dialTestPayload = []byte("dean-was-here123")
@@ -1073,6 +1174,9 @@ type client struct {
 	statsChan          chan *codersdk.AgentStats
 	coordinator        tailnet.Coordinator
 	lastWorkspaceAgent func()
+
+	mu              sync.Mutex // Protects following.
+	lifecycleStates []codersdk.WorkspaceAgentLifecycle
 }
 
 func (c *client) WorkspaceAgentMetadata(_ context.Context) (codersdk.WorkspaceAgentMetadata, error) {
@@ -1128,6 +1232,19 @@ func (c *client) AgentReportStats(ctx context.Context, _ slog.Logger, stats func
 	}), nil
 }
 
+func (c *client) getLifecycleStates() []codersdk.WorkspaceAgentLifecycle {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.lifecycleStates
+}
+
+func (c *client) PostWorkspaceAgentLifecycle(_ context.Context, req codersdk.PostWorkspaceAgentLifecycleRequest) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.lifecycleStates = append(c.lifecycleStates, req.State)
+	return nil
+}
+
 func (*client) PostWorkspaceAgentAppHealth(_ context.Context, _ codersdk.PostWorkspaceAppHealthsRequest) error {
 	return nil
 }