Skip to content

Commit 0850463

Browse files
deansheatherdeansheather
committed
Merge branch 'main' into dreamteam/moon-terminal
2 parents 0f77db3 + 745868f commit 0850463

File tree

85 files changed

+3348
-1346
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

85 files changed

+3348
-1346
lines changed

‎agent/agent_test.go

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -733,12 +733,15 @@ func TestAgent_UnixRemoteForwarding(t *testing.T) {
733733

734734
// It's possible that the socket is created but the server is not ready to
735735
// accept connections yet. We need to retry until we can connect.
736+
//
737+
// Note that we wait long here because if the tailnet connection has trouble
738+
// connecting, it could take 5 seconds or more to reconnect.
736739
var conn net.Conn
737740
require.Eventually(t, func() bool {
738741
var err error
739742
conn, err = net.Dial("unix", remoteSocketPath)
740743
return err == nil
741-
}, testutil.WaitShort, testutil.IntervalFast)
744+
}, testutil.WaitLong, testutil.IntervalFast)
742745
defer conn.Close()
743746
_, err = conn.Write([]byte("test"))
744747
require.NoError(t, err)

‎cli/agent.go

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import (
1818
"cloud.google.com/go/compute/metadata"
1919
"golang.org/x/xerrors"
2020
"gopkg.in/natefinch/lumberjack.v2"
21+
"tailscale.com/util/clientmetric"
2122

2223
"cdr.dev/slog"
2324
"cdr.dev/slog/sloggers/sloghuman"
@@ -36,6 +37,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
3637
noReap bool
3738
sshMaxTimeout time.Duration
3839
tailnetListenPort int64
40+
prometheusAddress string
3941
)
4042
cmd := &clibase.Cmd{
4143
Use: "agent",
@@ -126,6 +128,13 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
126128
agentPorts[port] = "pprof"
127129
}
128130

131+
prometheusSrvClose := ServeHandler(ctx, logger, prometheusMetricsHandler(), prometheusAddress, "prometheus")
132+
defer prometheusSrvClose()
133+
// Do a best effort here. If this fails, it's not a big deal.
134+
if port, err := urlPort(prometheusAddress); err == nil {
135+
agentPorts[port] = "prometheus"
136+
}
137+
129138
// exchangeToken returns a session token.
130139
// This is abstracted to allow for the same looping condition
131140
// regardless of instance identity auth type.
@@ -257,6 +266,13 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
257266
Description: "Specify a static port for Tailscale to use for listening.",
258267
Value: clibase.Int64Of(&tailnetListenPort),
259268
},
269+
{
270+
Flag: "prometheus-address",
271+
Default: "127.0.0.1:2112",
272+
Env: "CODER_AGENT_PROMETHEUS_ADDRESS",
273+
Value: clibase.StringOf(&prometheusAddress),
274+
Description: "The bind address to serve Prometheus metrics.",
275+
},
260276
}
261277

262278
return cmd
@@ -343,3 +359,12 @@ func urlPort(u string) (int, error) {
343359
}
344360
return -1, xerrors.Errorf("invalid port: %s", u)
345361
}
362+
363+
func prometheusMetricsHandler() http.Handler {
364+
// We don't have any other internal metrics so far, so it's safe to expose metrics this way.
365+
// Based on: https://github.com/tailscale/tailscale/blob/280255acae604796a1113861f5a84e6fa2dc6121/ipn/localapi/localapi.go#L489
366+
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
367+
w.Header().Set("Content-Type", "text/plain")
368+
clientmetric.WritePrometheusExpositionFormat(w)
369+
})
370+
}

‎cli/clibase/option.go

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,17 @@ func (s *OptionSet) Add(opts ...Option) {
8080
*s = append(*s, opts...)
8181
}
8282

83+
// Filter will only return options that match the given filter. (return true)
84+
func (s OptionSet) Filter(filter func(opt Option) bool) OptionSet {
85+
cpy := make(OptionSet, 0)
86+
for _, opt := range s {
87+
if filter(opt) {
88+
cpy = append(cpy, opt)
89+
}
90+
}
91+
return cpy
92+
}
93+
8394
// FlagSet returns a pflag.FlagSet for the OptionSet.
8495
func (s *OptionSet) FlagSet() *pflag.FlagSet {
8596
if s == nil {

‎cli/clitest/clitest.go

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -127,8 +127,8 @@ func extractTar(t *testing.T, data []byte, directory string) {
127127
}
128128
}
129129

130-
// Start runs the command in a goroutine and cleans it up when
131-
// the test completed.
130+
// Start runs the command in a goroutine and cleans it up when the test
131+
// completed.
132132
func Start(t *testing.T, inv *clibase.Invocation) {
133133
t.Helper()
134134

@@ -170,7 +170,7 @@ func (w *ErrorWaiter) Wait() error {
170170
var ok bool
171171
w.cachedError, ok = <-w.c
172172
if !ok {
173-
panic("unexpoected channel close")
173+
panic("unexpected channel close")
174174
}
175175
})
176176
return w.cachedError
@@ -196,18 +196,18 @@ func (w *ErrorWaiter) RequireAs(want interface{}) {
196196
require.ErrorAs(w.t, w.Wait(), want)
197197
}
198198

199-
// StartWithWaiter runs the command in a goroutine but returns the error
200-
// instead of asserting it. This is useful for testing error cases.
199+
// StartWithWaiter runs the command in a goroutine but returns the error instead
200+
// of asserting it. This is useful for testing error cases.
201201
func StartWithWaiter(t *testing.T, inv *clibase.Invocation) *ErrorWaiter {
202202
t.Helper()
203203

204-
errCh := make(chan error, 1)
205-
206-
var cleaningUp atomic.Bool
207-
208204
var (
209205
ctx = inv.Context()
210206
cancel func()
207+
208+
cleaningUp atomic.Bool
209+
errCh = make(chan error, 1)
210+
doneCh = make(chan struct{})
211211
)
212212
if _, ok := ctx.Deadline(); !ok {
213213
ctx, cancel = context.WithDeadline(ctx, time.Now().Add(testutil.WaitMedium))
@@ -218,12 +218,13 @@ func StartWithWaiter(t *testing.T, inv *clibase.Invocation) *ErrorWaiter {
218218
inv = inv.WithContext(ctx)
219219

220220
go func() {
221+
defer close(doneCh)
221222
defer close(errCh)
222223
err := inv.Run()
223224
if cleaningUp.Load() && errors.Is(err, context.DeadlineExceeded) {
224-
// If we're cleaning up, this error is likely related to the
225-
// CLI teardown process. E.g., the server could be slow to shut
226-
// down Postgres.
225+
// If we're cleaning up, this error is likely related to the CLI
226+
// teardown process. E.g., the server could be slow to shut down
227+
// Postgres.
227228
t.Logf("command %q timed out during test cleanup", inv.Command.FullName())
228229
}
229230
// Whether or not this fails the test is left to the caller.
@@ -235,7 +236,7 @@ func StartWithWaiter(t *testing.T, inv *clibase.Invocation) *ErrorWaiter {
235236
t.Cleanup(func() {
236237
cancel()
237238
cleaningUp.Store(true)
238-
<-errCh
239+
<-doneCh
239240
})
240241
return &ErrorWaiter{c: errCh, t: t}
241242
}

‎cli/cliui/output.go

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,3 +192,35 @@ func (textFormat) AttachOptions(_ *clibase.OptionSet) {}
192192
func (textFormat) Format(_ context.Context, data any) (string, error) {
193193
return fmt.Sprintf("%s", data), nil
194194
}
195+
196+
// DataChangeFormat allows manipulating the data passed to an output format.
197+
// This is because sometimes the data needs to be manipulated before it can be
198+
// passed to the output format.
199+
// For example, you may want to pass something different to the text formatter
200+
// than what you pass to the json formatter.
201+
type DataChangeFormat struct {
202+
format OutputFormat
203+
change func(data any) (any, error)
204+
}
205+
206+
// ChangeFormatterData allows manipulating the data passed to an output
207+
// format.
208+
func ChangeFormatterData(format OutputFormat, change func(data any) (any, error)) *DataChangeFormat {
209+
return &DataChangeFormat{format: format, change: change}
210+
}
211+
212+
func (d *DataChangeFormat) ID() string {
213+
return d.format.ID()
214+
}
215+
216+
func (d *DataChangeFormat) AttachOptions(opts *clibase.OptionSet) {
217+
d.format.AttachOptions(opts)
218+
}
219+
220+
func (d *DataChangeFormat) Format(ctx context.Context, data any) (string, error) {
221+
newData, err := d.change(data)
222+
if err != nil {
223+
return "", err
224+
}
225+
return d.format.Format(ctx, newData)
226+
}

‎cli/root.go

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -364,6 +364,13 @@ func (r *RootCmd) Command(subcommands []*clibase.Cmd) (*clibase.Cmd, error) {
364364
Value: clibase.BoolOf(&r.verbose),
365365
Group: globalGroup,
366366
},
367+
{
368+
Flag: "debug-http",
369+
Description: "Debug codersdk HTTP requests.",
370+
Value: clibase.BoolOf(&r.debugHTTP),
371+
Group: globalGroup,
372+
Hidden: true,
373+
},
367374
{
368375
Flag: config.FlagName,
369376
Env: "CODER_CONFIG_DIR",
@@ -412,6 +419,7 @@ type RootCmd struct {
412419
forceTTY bool
413420
noOpen bool
414421
verbose bool
422+
debugHTTP bool
415423

416424
noVersionCheck bool
417425
noFeatureWarning bool
@@ -464,6 +472,11 @@ func (r *RootCmd) InitClient(client *codersdk.Client) clibase.MiddlewareFunc {
464472

465473
client.SetSessionToken(r.token)
466474

475+
if r.debugHTTP {
476+
client.PlainLogger = os.Stderr
477+
client.LogBodies = true
478+
}
479+
467480
// We send these requests in parallel to minimize latency.
468481
var (
469482
versionErr = make(chan error)

‎cli/server.go

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1740,24 +1740,24 @@ func connectToPostgres(ctx context.Context, logger slog.Logger, driver string, d
17401740
}
17411741

17421742
// Ensure the PostgreSQL version is >=13.0.0!
1743-
version, err := sqlDB.QueryContext(ctx, "SHOW server_version;")
1743+
version, err := sqlDB.QueryContext(ctx, "SHOW server_version_num;")
17441744
if err != nil {
17451745
return nil, xerrors.Errorf("get postgres version: %w", err)
17461746
}
17471747
if !version.Next() {
17481748
return nil, xerrors.Errorf("no rows returned for version select")
17491749
}
1750-
var versionStr string
1751-
err = version.Scan(&versionStr)
1750+
var versionNum int
1751+
err = version.Scan(&versionNum)
17521752
if err != nil {
17531753
return nil, xerrors.Errorf("scan version: %w", err)
17541754
}
17551755
_ = version.Close()
1756-
versionStr = strings.Split(versionStr, " ")[0]
1757-
if semver.Compare("v"+versionStr, "v13") < 0 {
1758-
return nil, xerrors.New("PostgreSQL version must be v13.0.0 or higher!")
1756+
1757+
if versionNum < 130000 {
1758+
return nil, xerrors.Errorf("PostgreSQL version must be v13.0.0 or higher! Got: %d", versionNum)
17591759
}
1760-
logger.Debug(ctx, "connected to postgresql", slog.F("version", versionStr))
1760+
logger.Debug(ctx, "connected to postgresql", slog.F("version", versionNum))
17611761

17621762
err = migrations.Up(sqlDB)
17631763
if err != nil {

‎cli/ssh.go

Lines changed: 72 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import (
3030
"github.com/coder/coder/coderd/util/ptr"
3131
"github.com/coder/coder/codersdk"
3232
"github.com/coder/coder/cryptorand"
33+
"github.com/coder/retry"
3334
)
3435

3536
var (
@@ -100,17 +101,82 @@ func (r *RootCmd) ssh() *clibase.Cmd {
100101
stopPolling := tryPollWorkspaceAutostop(ctx, client, workspace)
101102
defer stopPolling()
102103

104+
// Enure connection is closed if the context is canceled or
105+
// the workspace reaches the stopped state.
106+
//
107+
// Watching the stopped state is a work-around for cases
108+
// where the agent is not gracefully shut down and the
109+
// connection is left open. If, for instance, the networking
110+
// is stopped before the agent is shut down, the disconnect
111+
// will usually not propagate.
112+
//
113+
// See: https://github.com/coder/coder/issues/6180
114+
watchAndClose := func(closer func() error) {
115+
// Ensure session is ended on both context cancellation
116+
// and workspace stop.
117+
defer func() {
118+
_ = closer()
119+
}()
120+
121+
startWatchLoop:
122+
for {
123+
// (Re)connect to the coder server and watch workspace events.
124+
var wsWatch <-chan codersdk.Workspace
125+
var err error
126+
for r := retry.New(time.Second, 15*time.Second); r.Wait(ctx); {
127+
wsWatch, err = client.WatchWorkspace(ctx, workspace.ID)
128+
if err == nil {
129+
break
130+
}
131+
if ctx.Err() != nil {
132+
return
133+
}
134+
}
135+
136+
for {
137+
select {
138+
case <-ctx.Done():
139+
return
140+
case w, ok := <-wsWatch:
141+
if !ok {
142+
continue startWatchLoop
143+
}
144+
145+
// Transitioning to stop or delete could mean that
146+
// the agent will still gracefully stop. If a new
147+
// build is starting, there's no reason to wait for
148+
// the agent, it should be long gone.
149+
if workspace.LatestBuild.ID != w.LatestBuild.ID && w.LatestBuild.Transition == codersdk.WorkspaceTransitionStart {
150+
return
151+
}
152+
// Note, we only react to the stopped state here because we
153+
// want to give the agent a chance to gracefully shut down
154+
// during "stopping".
155+
if w.LatestBuild.Status == codersdk.WorkspaceStatusStopped {
156+
return
157+
}
158+
}
159+
}
160+
}
161+
}
162+
103163
if stdio {
104164
rawSSH, err := conn.SSH(ctx)
105165
if err != nil {
106166
return err
107167
}
108168
defer rawSSH.Close()
169+
go watchAndClose(rawSSH.Close)
109170

110171
go func() {
111-
_, _ = io.Copy(inv.Stdout, rawSSH)
172+
// Ensure stdout copy closes incase stdin is closed
173+
// unexpectedly. Typically we wouldn't worry about
174+
// this since OpenSSH should kill the proxy command.
175+
defer rawSSH.Close()
176+
177+
_, _ = io.Copy(rawSSH, inv.Stdin)
112178
}()
113-
_, _ = io.Copy(rawSSH, inv.Stdin)
179+
_, _ = io.Copy(inv.Stdout, rawSSH)
114180
return nil
115181
}
116182

@@ -125,13 +191,11 @@ func (r *RootCmd) ssh() *clibase.Cmd {
125191
return err
126192
}
127193
defer sshSession.Close()
128-
129-
// Ensure context cancellation is propagated to the
130-
// SSH session, e.g. to cancel `Wait()` at the end.
131-
go func() {
132-
<-ctx.Done()
194+
go watchAndClose(func() error {
133195
_ = sshSession.Close()
134-
}()
196+
_ = sshClient.Close()
197+
return nil
198+
})
135199

136200
if identityAgent == "" {
137201
identityAgent = os.Getenv("SSH_AUTH_SOCK")

0 commit comments

Comments
 (0)