fix(cli/cliui): handle never ending startup log stream in Agent

mafredri · mafredri · commit 506dce5ee5d6 · 2023-07-17T09:03:20.000Z
diff --git a/cli/cliui/agent.go b/cli/cliui/agent.go
@@ -137,26 +137,44 @@ func Agent(ctx context.Context, writer io.Writer, agentID uuid.UUID, opts AgentO
 				}
 				defer logsCloser.Close()
 
+				var lastLog codersdk.WorkspaceAgentStartupLog
+				fetchedAgentWhileFollowing := fetchedAgent
+				if !follow {
+					fetchedAgentWhileFollowing = nil
+				}
 				for {
 					// This select is essentially and inline `fetch()`.
 					select {
 					case <-ctx.Done():
 						return ctx.Err()
-					case f := <-fetchedAgent:
+					case f := <-fetchedAgentWhileFollowing:
 						if f.err != nil {
 							return xerrors.Errorf("fetch: %w", f.err)
 						}
-						// We could handle changes in the agent status here, like
-						// if the agent becomes disconnected, we may want to stop.
-						// But for now, we'll just keep going, hopefully the agent
-						// will reconnect and update its status.
 						agent = f.agent
+
+						// If the agent is no longer starting, stop following
+						// logs because FetchLogs will keep streaming forever.
+						// We do one last non-follow request to ensure we have
+						// fetched all logs.
+						if !agent.LifecycleState.Starting() {
+							_ = logsCloser.Close()
+							fetchedAgentWhileFollowing = nil
+
+							logStream, logsCloser, err = opts.FetchLogs(ctx, agent.ID, lastLog.ID, false)
+							if err != nil {
+								return xerrors.Errorf("fetch workspace agent startup logs: %w", err)
+							}
+							// Logs are already primed, so we can call close.
+							_ = logsCloser.Close()
+						}
 					case logs, ok := <-logStream:
 						if !ok {
 							return nil
 						}
 						for _, log := range logs {
 							sw.Log(log.CreatedAt, log.Level, log.Output)
+							lastLog = log
 						}
 					}
 				}
diff --git a/cli/cliui/agent_test.go b/cli/cliui/agent_test.go
@@ -46,7 +46,6 @@ func TestAgent(t *testing.T) {
 				func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentStartupLog) error {
 					agent.Status = codersdk.WorkspaceAgentConnected
 					agent.FirstConnectedAt = ptr.Ref(time.Now())
-					close(logs)
 					return nil
 				},
 			},
@@ -79,7 +78,6 @@ func TestAgent(t *testing.T) {
 					agent.FirstConnectedAt = ptr.Ref(time.Now())
 					agent.LifecycleState = codersdk.WorkspaceAgentLifecycleReady
 					agent.ReadyAt = ptr.Ref(time.Now())
-					close(logs)
 					return nil
 				},
 			},
@@ -113,10 +111,6 @@ func TestAgent(t *testing.T) {
 					agent.LastConnectedAt = ptr.Ref(time.Now())
 					return nil
 				},
-				func(_ context.Context, _ *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentStartupLog) error {
-					close(logs)
-					return nil
-				},
 			},
 			want: []string{
 				"⧗ The workspace agent lost connection",
@@ -154,7 +148,6 @@ func TestAgent(t *testing.T) {
 							Output:    "Bye now",
 						},
 					}
-					close(logs)
 					return nil
 				},
 			},
@@ -184,7 +177,6 @@ func TestAgent(t *testing.T) {
 							Output:    "Hello world",
 						},
 					}
-					close(logs)
 					return nil
 				},
 			},
@@ -205,7 +197,6 @@ func TestAgent(t *testing.T) {
 				func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentStartupLog) error {
 					agent.Status = codersdk.WorkspaceAgentDisconnected
 					agent.LifecycleState = codersdk.WorkspaceAgentLifecycleOff
-					close(logs)
 					return nil
 				},
 			},
@@ -234,7 +225,6 @@ func TestAgent(t *testing.T) {
 				func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentStartupLog) error {
 					agent.ReadyAt = ptr.Ref(time.Now())
 					agent.LifecycleState = codersdk.WorkspaceAgentLifecycleShuttingDown
-					close(logs)
 					return nil
 				},
 			},
@@ -316,8 +306,21 @@ func TestAgent(t *testing.T) {
 						}
 						return agent, err
 					}
-					tc.opts.FetchLogs = func(_ context.Context, _ uuid.UUID, _ int64, _ bool) (<-chan []codersdk.WorkspaceAgentStartupLog, io.Closer, error) {
-						return logs, closeFunc(func() error { return nil }), nil
+					tc.opts.FetchLogs = func(ctx context.Context, _ uuid.UUID, _ int64, follow bool) (<-chan []codersdk.WorkspaceAgentStartupLog, io.Closer, error) {
+						if follow {
+							return logs, closeFunc(func() error { return nil }), nil
+						}
+
+						fetchLogs := make(chan []codersdk.WorkspaceAgentStartupLog, 1)
+						select {
+						case <-ctx.Done():
+							return nil, nil, ctx.Err()
+						case l := <-logs:
+							fetchLogs <- l
+						default:
+						}
+						close(fetchLogs)
+						return fetchLogs, closeFunc(func() error { return nil }), nil
 					}
 					err := cliui.Agent(inv.Context(), &buf, uuid.Nil, tc.opts)
 					return err
