Skip to content

Commit 1deaf15

Browse files
BrunoQuaresmaBrunoQuaresma
committed
Merge branch 'main' into bq/refactor-agent
2 parents 2d882ad + 9abfe97 commit 1deaf15

Some content is hidden

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

59 files changed

+2480
-353
lines changed

‎Makefile

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -515,7 +515,7 @@ docs/admin/audit-logs.md: scripts/auditdocgen/main.go enterprise/audit/table.go
515515
cd site
516516
yarn run format:write:only ../docs/admin/audit-logs.md
517517

518-
coderd/apidoc/swagger.json: $(shell find ./scripts/apidocgen $(FIND_EXCLUSIONS) -type f) $(wildcard coderd/*.go) $(wildcard enterprise/coderd/*.go) $(wildcard codersdk/*.go) .swaggo docs/manifest.json coderd/rbac/object_gen.go
518+
coderd/apidoc/swagger.json: $(shell find ./scripts/apidocgen $(FIND_EXCLUSIONS) -type f) $(wildcard coderd/*.go) $(wildcard enterprise/coderd/*.go) $(wildcard codersdk/*.go) $(wildcard enterprise/wsproxy/wsproxysdk/*.go) coderd/database/querier.go .swaggo docs/manifest.json coderd/rbac/object_gen.go
519519
./scripts/apidocgen/generate.sh
520520
yarn run --cwd=site format:write:only ../docs/api ../docs/manifest.json ../coderd/apidoc/swagger.json
521521

‎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/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)