Skip to content

Commit 5634f14

Browse files
ammarioammario
committed
Merge remote-tracking branch 'origin/main' into server-yaml
2 parents 0c1a02f + 5876dc1 commit 5634f14

20 files changed

+297
-106
lines changed

‎agent/agent.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -601,7 +601,9 @@ func (a *agent) runCoordinator(ctx context.Context, network *tailnet.Conn) error
601601
}
602602
defer coordinator.Close()
603603
a.logger.Info(ctx, "connected to coordination server")
604-
sendNodes, errChan := tailnet.ServeCoordinator(coordinator, network.UpdateNodes)
604+
sendNodes, errChan := tailnet.ServeCoordinator(coordinator, func(nodes []*tailnet.Node) error {
605+
return network.UpdateNodes(nodes, false)
606+
})
605607
network.SetNodeCallback(sendNodes)
606608
select {
607609
case <-ctx.Done():

‎agent/agent_test.go

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1179,12 +1179,21 @@ func setupAgent(t *testing.T, metadata agentsdk.Metadata, ptyTimeout time.Durati
11791179
coordinator.ServeClient(serverConn, uuid.New(), agentID)
11801180
}()
11811181
sendNode, _ := tailnet.ServeCoordinator(clientConn, func(node []*tailnet.Node) error {
1182-
return conn.UpdateNodes(node)
1182+
return conn.UpdateNodes(node, false)
11831183
})
11841184
conn.SetNodeCallback(sendNode)
1185-
return &codersdk.WorkspaceAgentConn{
1185+
agentConn := &codersdk.WorkspaceAgentConn{
11861186
Conn: conn,
1187-
}, c, statsCh, fs
1187+
}
1188+
t.Cleanup(func() {
1189+
_ = agentConn.Close()
1190+
})
1191+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
1192+
defer cancel()
1193+
if !agentConn.AwaitReachable(ctx) {
1194+
t.Fatal("agent not reachable")
1195+
}
1196+
return agentConn, c, statsCh, fs
11881197
}
11891198

11901199
var dialTestPayload = []byte("dean-was-here123")

‎cli/root.go

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ import (
1919

2020
"golang.org/x/xerrors"
2121

22+
"cdr.dev/slog"
23+
2224
"github.com/charmbracelet/lipgloss"
2325
"github.com/kirsle/configdir"
2426
"github.com/mattn/go-isatty"
@@ -179,6 +181,21 @@ func Root(subcommands []*cobra.Command) *cobra.Command {
179181
return cmd
180182
}
181183

184+
type contextKey int
185+
186+
const (
187+
contextKeyLogger contextKey = iota
188+
)
189+
190+
func ContextWithLogger(ctx context.Context, l slog.Logger) context.Context {
191+
return context.WithValue(ctx, contextKeyLogger, l)
192+
}
193+
194+
func LoggerFromContext(ctx context.Context) (slog.Logger, bool) {
195+
l, ok := ctx.Value(contextKeyLogger).(slog.Logger)
196+
return l, ok
197+
}
198+
182199
// fixUnknownSubcommandError modifies the provided commands so that the
183200
// ones with subcommands output the correct error message when an
184201
// unknown subcommand is invoked.

‎cli/speedtest.go

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,10 @@ func speedtest() *cobra.Command {
5151
if err != nil && !xerrors.Is(err, cliui.AgentStartError) {
5252
return xerrors.Errorf("await agent: %w", err)
5353
}
54-
logger := slog.Make(sloghuman.Sink(cmd.ErrOrStderr()))
54+
logger, ok := LoggerFromContext(ctx)
55+
if !ok {
56+
logger = slog.Make(sloghuman.Sink(cmd.ErrOrStderr()))
57+
}
5558
if cliflag.IsSetBool(cmd, varVerbose) {
5659
logger = logger.Leveled(slog.LevelDebug)
5760
}

‎cli/speedtest_test.go

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@ import (
77
"github.com/stretchr/testify/assert"
88
"github.com/stretchr/testify/require"
99

10+
"cdr.dev/slog"
1011
"cdr.dev/slog/sloggers/slogtest"
1112
"github.com/coder/coder/agent"
13+
"github.com/coder/coder/cli"
1214
"github.com/coder/coder/cli/clitest"
1315
"github.com/coder/coder/coderd/coderdtest"
1416
"github.com/coder/coder/codersdk"
@@ -28,7 +30,7 @@ func TestSpeedtest(t *testing.T) {
2830
agentClient.SetSessionToken(agentToken)
2931
agentCloser := agent.New(agent.Options{
3032
Client: agentClient,
31-
Logger: slogtest.Make(t, nil).Named("agent"),
33+
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
3234
})
3335
defer agentCloser.Close()
3436
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
@@ -50,10 +52,12 @@ func TestSpeedtest(t *testing.T) {
5052
clitest.SetupConfig(t, client, root)
5153
pty := ptytest.New(t)
5254
cmd.SetOut(pty.Output())
55+
cmd.SetErr(pty.Output())
5356

5457
ctx, cancel = context.WithTimeout(context.Background(), testutil.WaitLong)
5558
defer cancel()
5659

60+
ctx = cli.ContextWithLogger(ctx, slogtest.Make(t, nil).Named("speedtest").Leveled(slog.LevelDebug))
5761
cmdDone := tGo(t, func() {
5862
err := cmd.ExecuteContext(ctx)
5963
assert.NoError(t, err)

‎cli/ssh_test.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"golang.org/x/crypto/ssh"
2525
gosshagent "golang.org/x/crypto/ssh/agent"
2626

27+
"cdr.dev/slog"
2728
"cdr.dev/slog/sloggers/slogtest"
2829

2930
"github.com/coder/coder/agent"
@@ -47,6 +48,7 @@ func setupWorkspaceForAgent(t *testing.T, mutate func([]*proto.Agent) []*proto.A
4748
}
4849
}
4950
client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
51+
client.Logger = slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug)
5052
user := coderdtest.CreateFirstUser(t, client)
5153
agentToken := uuid.NewString()
5254
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{

‎coderd/coderd_test.go

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"net/http"
77
"net/netip"
88
"strconv"
9+
"sync"
910
"testing"
1011

1112
"github.com/stretchr/testify/assert"
@@ -78,11 +79,17 @@ func TestDERP(t *testing.T) {
7879
DERPMap: derpMap,
7980
})
8081
require.NoError(t, err)
82+
83+
w2Ready := make(chan struct{})
84+
w2ReadyOnce := sync.Once{}
8185
w1.SetNodeCallback(func(node *tailnet.Node) {
82-
w2.UpdateNodes([]*tailnet.Node{node})
86+
w2.UpdateNodes([]*tailnet.Node{node}, false)
87+
w2ReadyOnce.Do(func() {
88+
close(w2Ready)
89+
})
8390
})
8491
w2.SetNodeCallback(func(node *tailnet.Node) {
85-
w1.UpdateNodes([]*tailnet.Node{node})
92+
w1.UpdateNodes([]*tailnet.Node{node}, false)
8693
})
8794

8895
conn := make(chan struct{})
@@ -98,6 +105,7 @@ func TestDERP(t *testing.T) {
98105
}()
99106

100107
<-conn
108+
<-w2Ready
101109
nc, err := w2.DialContextTCP(context.Background(), netip.AddrPortFrom(w1IP, 35565))
102110
require.NoError(t, err)
103111
_ = nc.Close()

‎coderd/workspaceagents.go

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -404,6 +404,7 @@ func (api *API) workspaceAgentListeningPorts(rw http.ResponseWriter, r *http.Req
404404
}
405405

406406
func (api *API) dialWorkspaceAgentTailnet(r *http.Request, agentID uuid.UUID) (*codersdk.WorkspaceAgentConn, error) {
407+
ctx := r.Context()
407408
clientConn, serverConn := net.Pipe()
408409

409410
derpMap := api.DERPMap.Clone()
@@ -453,32 +454,32 @@ func (api *API) dialWorkspaceAgentTailnet(r *http.Request, agentID uuid.UUID) (*
453454
}
454455

455456
sendNodes, _ := tailnet.ServeCoordinator(clientConn, func(node []*tailnet.Node) error {
456-
err := conn.RemoveAllPeers()
457-
if err != nil {
458-
return xerrors.Errorf("remove all peers: %w", err)
459-
}
460-
461-
err = conn.UpdateNodes(node)
457+
err = conn.UpdateNodes(node, true)
462458
if err != nil {
463459
return xerrors.Errorf("update nodes: %w", err)
464460
}
465461
return nil
466462
})
467463
conn.SetNodeCallback(sendNodes)
464+
agentConn := &codersdk.WorkspaceAgentConn{
465+
Conn: conn,
466+
CloseFunc: func() {
467+
_ = clientConn.Close()
468+
_ = serverConn.Close()
469+
},
470+
}
468471
go func() {
469472
err := (*api.TailnetCoordinator.Load()).ServeClient(serverConn, uuid.New(), agentID)
470473
if err != nil {
471474
api.Logger.Warn(r.Context(), "tailnet coordinator client error", slog.Error(err))
472-
_ = conn.Close()
475+
_ = agentConn.Close()
473476
}
474477
}()
475-
return &codersdk.WorkspaceAgentConn{
476-
Conn: conn,
477-
CloseFunc: func() {
478-
_ = clientConn.Close()
479-
_ = serverConn.Close()
480-
},
481-
}, nil
478+
if !agentConn.AwaitReachable(ctx) {
479+
_ = agentConn.Close()
480+
return nil, xerrors.Errorf("agent not reachable")
481+
}
482+
return agentConn, nil
482483
}
483484

484485
// @Summary Get connection info for workspace agent

‎coderd/workspaceagents_test.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -469,6 +469,8 @@ func TestWorkspaceAgentListeningPorts(t *testing.T) {
469469
t.Parallel()
470470

471471
setup := func(t *testing.T, apps []*proto.App) (*codersdk.Client, uint16, uuid.UUID) {
472+
t.Helper()
473+
472474
client := coderdtest.New(t, &coderdtest.Options{
473475
IncludeProvisionerDaemon: true,
474476
})

‎coderd/wsconncache/wsconncache_test.go

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import (
2929
"github.com/coder/coder/codersdk/agentsdk"
3030
"github.com/coder/coder/tailnet"
3131
"github.com/coder/coder/tailnet/tailnettest"
32+
"github.com/coder/coder/testutil"
3233
)
3334

3435
func TestMain(m *testing.M) {
@@ -131,6 +132,14 @@ func TestCache(t *testing.T) {
131132
return
132133
}
133134
defer release()
135+
136+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
137+
defer cancel()
138+
if !conn.AwaitReachable(ctx) {
139+
t.Error("agent not reachable")
140+
return
141+
}
142+
134143
transport := conn.HTTPTransport()
135144
defer transport.CloseIdleConnections()
136145
proxy.Transport = transport
@@ -146,6 +155,8 @@ func TestCache(t *testing.T) {
146155
}
147156

148157
func setupAgent(t *testing.T, metadata agentsdk.Metadata, ptyTimeout time.Duration) *codersdk.WorkspaceAgentConn {
158+
t.Helper()
159+
149160
metadata.DERPMap = tailnettest.RunDERPAndSTUN(t)
150161

151162
coordinator := tailnet.NewCoordinator()
@@ -180,12 +191,21 @@ func setupAgent(t *testing.T, metadata agentsdk.Metadata, ptyTimeout time.Durati
180191
})
181192
go coordinator.ServeClient(serverConn, uuid.New(), agentID)
182193
sendNode, _ := tailnet.ServeCoordinator(clientConn, func(node []*tailnet.Node) error {
183-
return conn.UpdateNodes(node)
194+
return conn.UpdateNodes(node, false)
184195
})
185196
conn.SetNodeCallback(sendNode)
186-
return &codersdk.WorkspaceAgentConn{
197+
agentConn := &codersdk.WorkspaceAgentConn{
187198
Conn: conn,
188199
}
200+
t.Cleanup(func() {
201+
_ = agentConn.Close()
202+
})
203+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
204+
defer cancel()
205+
if !agentConn.AwaitReachable(ctx) {
206+
t.Fatal("agent not reachable")
207+
}
208+
return agentConn
189209
}
190210

191211
type client struct {

‎codersdk/workspaceagentconn.go

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,9 @@ type ReconnectingPTYRequest struct {
176176
func (c *WorkspaceAgentConn) ReconnectingPTY(ctx context.Context, id uuid.UUID, height, width uint16, command string) (net.Conn, error) {
177177
ctx, span := tracing.StartSpan(ctx)
178178
defer span.End()
179-
179+
if !c.AwaitReachable(ctx) {
180+
return nil, xerrors.Errorf("workspace agent not reachable in time: %v", ctx.Err())
181+
}
180182
conn, err := c.DialContextTCP(ctx, netip.AddrPortFrom(WorkspaceAgentIP, WorkspaceAgentReconnectingPTYPort))
181183
if err != nil {
182184
return nil, err
@@ -207,6 +209,9 @@ func (c *WorkspaceAgentConn) ReconnectingPTY(ctx context.Context, id uuid.UUID,
207209
func (c *WorkspaceAgentConn) SSH(ctx context.Context) (net.Conn, error) {
208210
ctx, span := tracing.StartSpan(ctx)
209211
defer span.End()
212+
if !c.AwaitReachable(ctx) {
213+
return nil, xerrors.Errorf("workspace agent not reachable in time: %v", ctx.Err())
214+
}
210215
return c.DialContextTCP(ctx, netip.AddrPortFrom(WorkspaceAgentIP, WorkspaceAgentSSHPort))
211216
}
212217

@@ -235,6 +240,9 @@ func (c *WorkspaceAgentConn) SSHClient(ctx context.Context) (*ssh.Client, error)
235240
func (c *WorkspaceAgentConn) Speedtest(ctx context.Context, direction speedtest.Direction, duration time.Duration) ([]speedtest.Result, error) {
236241
ctx, span := tracing.StartSpan(ctx)
237242
defer span.End()
243+
if !c.AwaitReachable(ctx) {
244+
return nil, xerrors.Errorf("workspace agent not reachable in time: %v", ctx.Err())
245+
}
238246
speedConn, err := c.DialContextTCP(ctx, netip.AddrPortFrom(WorkspaceAgentIP, WorkspaceAgentSpeedtestPort))
239247
if err != nil {
240248
return nil, xerrors.Errorf("dial speedtest: %w", err)
@@ -257,6 +265,9 @@ func (c *WorkspaceAgentConn) DialContext(ctx context.Context, network string, ad
257265
_, rawPort, _ := net.SplitHostPort(addr)
258266
port, _ := strconv.ParseUint(rawPort, 10, 16)
259267
ipp := netip.AddrPortFrom(WorkspaceAgentIP, uint16(port))
268+
if !c.AwaitReachable(ctx) {
269+
return nil, xerrors.Errorf("workspace agent not reachable in time: %v", ctx.Err())
270+
}
260271
if network == "udp" {
261272
return c.Conn.DialContextUDP(ctx, ipp)
262273
}
@@ -317,7 +328,7 @@ func (c *WorkspaceAgentConn) apiClient() *http.Client {
317328
// Disable keep alives as we're usually only making a single
318329
// request, and this triggers goleak in tests
319330
DisableKeepAlives: true,
320-
DialContext: func(_ context.Context, network, addr string) (net.Conn, error) {
331+
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
321332
if network != "tcp" {
322333
return nil, xerrors.Errorf("network must be tcp")
323334
}
@@ -331,7 +342,11 @@ func (c *WorkspaceAgentConn) apiClient() *http.Client {
331342
return nil, xerrors.Errorf("request %q does not appear to be for http api", addr)
332343
}
333344

334-
conn, err := c.DialContextTCP(context.Background(), netip.AddrPortFrom(WorkspaceAgentIP, WorkspaceAgentHTTPAPIServerPort))
345+
if !c.AwaitReachable(ctx) {
346+
return nil, xerrors.Errorf("workspace agent not reachable in time: %v", ctx.Err())
347+
}
348+
349+
conn, err := c.DialContextTCP(ctx, netip.AddrPortFrom(WorkspaceAgentIP, WorkspaceAgentHTTPAPIServerPort))
335350
if err != nil {
336351
return nil, xerrors.Errorf("dial http api: %w", err)
337352
}

0 commit comments

Comments
 (0)