Skip to content

Commit a753703

Browse files
mafredrimafredri
authored
feat(cli): Add support for delay_login_until_ready (#5851)
1 parent cf93fbd commit a753703

File tree

6 files changed

+427
-60
lines changed

6 files changed

+427
-60
lines changed

‎cli/cliui/agent.go

+136-41
Original file line numberDiff line numberDiff line change
@@ -10,16 +10,21 @@ import (
1010
"time"
1111

1212
"github.com/briandowns/spinner"
13+
"github.com/muesli/reflow/indent"
14+
"github.com/muesli/reflow/wordwrap"
1315
"golang.org/x/xerrors"
1416

1517
"github.com/coder/coder/codersdk"
1618
)
1719

20+
var AgentStartError = xerrors.New("agent startup exited with non-zero exit status")
21+
1822
type AgentOptions struct {
1923
WorkspaceName string
2024
Fetch func(context.Context) (codersdk.WorkspaceAgent, error)
2125
FetchInterval time.Duration
2226
WarnInterval time.Duration
27+
NoWait bool // If true, don't wait for the agent to be ready.
2328
}
2429

2530
// Agent displays a spinning indicator that waits for a workspace agent to connect.
@@ -36,48 +41,33 @@ func Agent(ctx context.Context, writer io.Writer, opts AgentOptions) error {
3641
return xerrors.Errorf("fetch: %w", err)
3742
}
3843

39-
if agent.Status == codersdk.WorkspaceAgentConnected {
44+
// Fast path if the agent is ready (avoid showing connecting prompt).
45+
// We don't take the fast path for opts.NoWait yet because we want to
46+
// show the message.
47+
if agent.Status == codersdk.WorkspaceAgentConnected &&
48+
(!agent.DelayLoginUntilReady || agent.LifecycleState == codersdk.WorkspaceAgentLifecycleReady) {
4049
return nil
4150
}
4251

52+
ctx, cancel := signal.NotifyContext(ctx, os.Interrupt)
53+
defer cancel()
54+
4355
spin := spinner.New(spinner.CharSets[78], 100*time.Millisecond, spinner.WithColor("fgHiGreen"))
4456
spin.Writer = writer
4557
spin.ForceOutput = true
46-
spin.Suffix = " Waiting for connection from " + Styles.Field.Render(agent.Name) + "..."
47-
spin.Start()
48-
defer spin.Stop()
58+
spin.Suffix = waitingMessage(agent, opts).Spin
4959

50-
ctx, cancelFunc := context.WithCancel(ctx)
51-
defer cancelFunc()
52-
stopSpin := make(chan os.Signal, 1)
53-
signal.Notify(stopSpin, os.Interrupt)
54-
defer signal.Stop(stopSpin)
55-
go func() {
56-
select {
57-
case <-ctx.Done():
58-
return
59-
case <-stopSpin:
60-
}
61-
cancelFunc()
62-
signal.Stop(stopSpin)
63-
spin.Stop()
64-
// nolint:revive
65-
os.Exit(1)
66-
}()
67-
68-
var waitMessage string
69-
messageAfter := time.NewTimer(opts.WarnInterval)
70-
defer messageAfter.Stop()
60+
waitMessage := &message{}
7161
showMessage := func() {
7262
resourceMutex.Lock()
7363
defer resourceMutex.Unlock()
7464

75-
m := waitingMessage(agent)
76-
if m == waitMessage {
65+
m := waitingMessage(agent, opts)
66+
if m.Prompt == waitMessage.Prompt {
7767
return
7868
}
7969
moveUp := ""
80-
if waitMessage != "" {
70+
if waitMessage.Prompt != "" {
8171
// If this is an update, move a line up
8272
// to keep it tidy and aligned.
8373
moveUp = "\033[1A"
@@ -86,20 +76,43 @@ func Agent(ctx context.Context, writer io.Writer, opts AgentOptions) error {
8676

8777
// Stop the spinner while we write our message.
8878
spin.Stop()
79+
spin.Suffix = waitMessage.Spin
8980
// Clear the line and (if necessary) move up a line to write our message.
90-
_, _ = fmt.Fprintf(writer, "\033[2K%s%s\n\n", moveUp, Styles.Paragraph.Render(Styles.Prompt.String()+waitMessage))
81+
_, _ = fmt.Fprintf(writer, "\033[2K%s\n%s\n", moveUp, waitMessage.Prompt)
9182
select {
9283
case <-ctx.Done():
9384
default:
9485
// Safe to resume operation.
95-
spin.Start()
86+
if spin.Suffix != "" {
87+
spin.Start()
88+
}
9689
}
9790
}
91+
92+
// Fast path for showing the error message even when using no wait,
93+
// we do this just before starting the spinner to avoid needless
94+
// spinning.
95+
if agent.Status == codersdk.WorkspaceAgentConnected &&
96+
agent.DelayLoginUntilReady && opts.NoWait {
97+
showMessage()
98+
return nil
99+
}
100+
101+
// Start spinning after fast paths are handled.
102+
if spin.Suffix != "" {
103+
spin.Start()
104+
}
105+
defer spin.Stop()
106+
107+
warnAfter := time.NewTimer(opts.WarnInterval)
108+
defer warnAfter.Stop()
109+
warningShown := make(chan struct{})
98110
go func() {
99111
select {
100112
case <-ctx.Done():
101-
case <-messageAfter.C:
102-
messageAfter.Stop()
113+
close(warningShown)
114+
case <-warnAfter.C:
115+
close(warningShown)
103116
showMessage()
104117
}
105118
}()
@@ -121,26 +134,108 @@ func Agent(ctx context.Context, writer io.Writer, opts AgentOptions) error {
121134
resourceMutex.Unlock()
122135
switch agent.Status {
123136
case codersdk.WorkspaceAgentConnected:
137+
// NOTE(mafredri): Once we have access to the workspace agent's
138+
// startup script logs, we can show them here.
139+
// https://github.com/coder/coder/issues/2957
140+
if agent.DelayLoginUntilReady && !opts.NoWait {
141+
switch agent.LifecycleState {
142+
case codersdk.WorkspaceAgentLifecycleReady:
143+
return nil
144+
case codersdk.WorkspaceAgentLifecycleStartTimeout:
145+
showMessage()
146+
case codersdk.WorkspaceAgentLifecycleStartError:
147+
showMessage()
148+
return AgentStartError
149+
default:
150+
select {
151+
case <-warningShown:
152+
showMessage()
153+
default:
154+
// This state is normal, we don't want
155+
// to show a message prematurely.
156+
}
157+
}
158+
continue
159+
}
124160
return nil
125161
case codersdk.WorkspaceAgentTimeout, codersdk.WorkspaceAgentDisconnected:
126162
showMessage()
127163
}
128164
}
129165
}
130166

131-
func waitingMessage(agent codersdk.WorkspaceAgent) string {
132-
var m string
167+
type message struct {
168+
Spin string
169+
Prompt string
170+
Troubleshoot bool
171+
}
172+
173+
func waitingMessage(agent codersdk.WorkspaceAgent, opts AgentOptions) (m *message) {
174+
m = &message{
175+
Spin: fmt.Sprintf("Waiting for connection from %s...", Styles.Field.Render(agent.Name)),
176+
Prompt: "Don't panic, your workspace is booting up!",
177+
}
178+
defer func() {
179+
if opts.NoWait {
180+
m.Spin = ""
181+
}
182+
if m.Spin != "" {
183+
m.Spin = " " + m.Spin
184+
}
185+
186+
// We don't want to wrap the troubleshooting URL, so we'll handle word
187+
// wrapping ourselves (vs using lipgloss).
188+
w := wordwrap.NewWriter(Styles.Paragraph.GetWidth() - Styles.Paragraph.GetMarginLeft()*2)
189+
w.Breakpoints = []rune{' ', '\n'}
190+
191+
_, _ = fmt.Fprint(w, m.Prompt)
192+
if m.Troubleshoot {
193+
if agent.TroubleshootingURL != "" {
194+
_, _ = fmt.Fprintf(w, " See troubleshooting instructions at:\n%s", agent.TroubleshootingURL)
195+
} else {
196+
_, _ = fmt.Fprint(w, " Wait for it to (re)connect or restart your workspace.")
197+
}
198+
}
199+
_, _ = fmt.Fprint(w, "\n")
200+
201+
// We want to prefix the prompt with a caret, but we want text on the
202+
// following lines to align with the text on the first line (i.e. added
203+
// spacing).
204+
ind := " " + Styles.Prompt.String()
205+
iw := indent.NewWriter(1, func(w io.Writer) {
206+
_, _ = w.Write([]byte(ind))
207+
ind = " " // Set indentation to space after initial prompt.
208+
})
209+
_, _ = fmt.Fprint(iw, w.String())
210+
m.Prompt = iw.String()
211+
}()
212+
133213
switch agent.Status {
134214
case codersdk.WorkspaceAgentTimeout:
135-
m = "The workspace agent is having trouble connecting."
215+
m.Prompt = "The workspace agent is having trouble connecting."
136216
case codersdk.WorkspaceAgentDisconnected:
137-
m = "The workspace agent lost connection!"
217+
m.Prompt = "The workspace agent lost connection!"
218+
case codersdk.WorkspaceAgentConnected:
219+
m.Spin = fmt.Sprintf("Waiting for %s to become ready...", Styles.Field.Render(agent.Name))
220+
m.Prompt = "Don't panic, your workspace agent has connected and the workspace is getting ready!"
221+
if opts.NoWait {
222+
m.Prompt = "Your workspace is still getting ready, it may be in an incomplete state."
223+
}
224+
225+
switch agent.LifecycleState {
226+
case codersdk.WorkspaceAgentLifecycleStartTimeout:
227+
m.Prompt = "The workspace is taking longer than expected to get ready, the agent startup script is still executing."
228+
case codersdk.WorkspaceAgentLifecycleStartError:
229+
m.Spin = ""
230+
m.Prompt = "The workspace ran into a problem while getting ready, the agent startup script exited with non-zero status."
231+
default:
232+
// Not a failure state, no troubleshooting necessary.
233+
return m
234+
}
138235
default:
139236
// Not a failure state, no troubleshooting necessary.
140-
return "Don't panic, your workspace is booting up!"
141-
}
142-
if agent.TroubleshootingURL != "" {
143-
return fmt.Sprintf("%s See troubleshooting instructions at: %s", m, agent.TroubleshootingURL)
237+
return m
144238
}
145-
return fmt.Sprintf("%s Wait for it to (re)connect or restart your workspace.", m)
239+
m.Troubleshoot = true
240+
return m
146241
}

0 commit comments

Comments
 (0)