Skip to content

Commit af3fdc6

Browse files
spikecurtisspikecurtis
authored
chore: refactor agent routines that use the v2 API (#12223)
In anticipation of needing the `LogSender` to run on a context that doesn't get immediately canceled when you `Close()` the agent, I've undertaken a little refactor to manage the goroutines that get run against the Tailnet and Agent API connection. This handles controlling two contexts, one that gets canceled right away at the start of graceful shutdown, and another that stays up to allow graceful shutdown to complete.
1 parent 66585f0 commit af3fdc6

File tree

9 files changed

+491
-249
lines changed

9 files changed

+491
-249
lines changed

‎agent/agent.go

+450-222
Large diffs are not rendered by default.

‎cli/ssh_test.go

+8-2
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,13 @@ func TestSSH(t *testing.T) {
162162
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspaceBuild.ID)
163163

164164
// Update template version
165-
version = coderdtest.UpdateTemplateVersion(t, ownerClient, owner.OrganizationID, echoResponses, template.ID)
165+
authToken2 := uuid.NewString()
166+
echoResponses2 := &echo.Responses{
167+
Parse: echo.ParseComplete,
168+
ProvisionPlan: echo.PlanComplete,
169+
ProvisionApply: echo.ProvisionApplyWithAgent(authToken2),
170+
}
171+
version = coderdtest.UpdateTemplateVersion(t, ownerClient, owner.OrganizationID, echoResponses2, template.ID)
166172
coderdtest.AwaitTemplateVersionJobCompleted(t, ownerClient, version.ID)
167173
err := ownerClient.UpdateActiveTemplateVersion(context.Background(), template.ID, codersdk.UpdateActiveTemplateVersion{
168174
ID: version.ID,
@@ -184,7 +190,7 @@ func TestSSH(t *testing.T) {
184190

185191
// When the agent connects, the workspace was started, and we should
186192
// have access to the shell.
187-
_ = agenttest.New(t, client.URL, authToken)
193+
_ = agenttest.New(t, client.URL, authToken2)
188194
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
189195

190196
// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.

‎coderd/coderdtest/coderdtest.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
193193
options = &Options{}
194194
}
195195
if options.Logger == nil {
196-
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
196+
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug).Named("coderd")
197197
options.Logger = &logger
198198
}
199199
if options.GoogleTokenValidator == nil {

‎codersdk/workspaceagents.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -534,7 +534,7 @@ func (tac *tailnetAPIConnector) coordinate(client proto.DRPCTailnetClient) {
534534
tac.logger.Debug(tac.ctx, "main context canceled; do graceful disconnect")
535535
crdErr := coordination.Close()
536536
if crdErr != nil {
537-
tac.logger.Error(tac.ctx, "failed to close remote coordination", slog.Error(err))
537+
tac.logger.Warn(tac.ctx, "failed to close remote coordination", slog.Error(err))
538538
}
539539
case err = <-coordination.Error():
540540
if err != nil &&

‎enterprise/coderd/coderd_test.go

+4-3
Original file line numberDiff line numberDiff line change
@@ -231,13 +231,14 @@ func TestAuditLogging(t *testing.T) {
231231
},
232232
DontAddLicense: true,
233233
})
234-
workspace, agent := setupWorkspaceAgent(t, client, user, 0)
235-
conn, err := client.DialWorkspaceAgent(ctx, agent.ID, nil) //nolint:gocritic // RBAC is not the purpose of this test
234+
r := setupWorkspaceAgent(t, client, user, 0)
235+
conn, err := client.DialWorkspaceAgent(ctx, r.sdkAgent.ID, nil) //nolint:gocritic // RBAC is not the purpose of this test
236236
require.NoError(t, err)
237237
defer conn.Close()
238238
connected := conn.AwaitReachable(ctx)
239239
require.True(t, connected)
240-
build := coderdtest.CreateWorkspaceBuild(t, client, workspace, database.WorkspaceTransitionStop)
240+
_ = r.agent.Close() // close first so we don't drop error logs from outdated build
241+
build := coderdtest.CreateWorkspaceBuild(t, client, r.workspace, database.WorkspaceTransitionStop)
241242
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, build.ID)
242243
})
243244
}

‎enterprise/coderd/replicas_test.go

+4-4
Original file line numberDiff line numberDiff line change
@@ -81,8 +81,8 @@ func TestReplicas(t *testing.T) {
8181
require.NoError(t, err)
8282
require.Len(t, replicas, 2)
8383

84-
_, agent := setupWorkspaceAgent(t, firstClient, firstUser, 0)
85-
conn, err := secondClient.DialWorkspaceAgent(context.Background(), agent.ID, &codersdk.DialWorkspaceAgentOptions{
84+
r := setupWorkspaceAgent(t, firstClient, firstUser, 0)
85+
conn, err := secondClient.DialWorkspaceAgent(context.Background(), r.sdkAgent.ID, &codersdk.DialWorkspaceAgentOptions{
8686
BlockEndpoints: true,
8787
Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
8888
})
@@ -127,8 +127,8 @@ func TestReplicas(t *testing.T) {
127127
require.NoError(t, err)
128128
require.Len(t, replicas, 2)
129129

130-
_, agent := setupWorkspaceAgent(t, firstClient, firstUser, 0)
131-
conn, err := secondClient.DialWorkspaceAgent(context.Background(), agent.ID, &codersdk.DialWorkspaceAgentOptions{
130+
r := setupWorkspaceAgent(t, firstClient, firstUser, 0)
131+
conn, err := secondClient.DialWorkspaceAgent(context.Background(), r.sdkAgent.ID, &codersdk.DialWorkspaceAgentOptions{
132132
BlockEndpoints: true,
133133
Logger: slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug),
134134
})

‎enterprise/coderd/workspaceagents_test.go

+15-9
Original file line numberDiff line numberDiff line change
@@ -44,9 +44,9 @@ func TestBlockNonBrowser(t *testing.T) {
4444
},
4545
},
4646
})
47-
_, agent := setupWorkspaceAgent(t, client, user, 0)
47+
r := setupWorkspaceAgent(t, client, user, 0)
4848
//nolint:gocritic // Testing that even the owner gets blocked.
49-
_, err := client.DialWorkspaceAgent(context.Background(), agent.ID, nil)
49+
_, err := client.DialWorkspaceAgent(context.Background(), r.sdkAgent.ID, nil)
5050
var apiErr *codersdk.Error
5151
require.ErrorAs(t, err, &apiErr)
5252
require.Equal(t, http.StatusConflict, apiErr.StatusCode())
@@ -63,15 +63,21 @@ func TestBlockNonBrowser(t *testing.T) {
6363
},
6464
},
6565
})
66-
_, agent := setupWorkspaceAgent(t, client, user, 0)
66+
r := setupWorkspaceAgent(t, client, user, 0)
6767
//nolint:gocritic // Testing RBAC is not the point of this test.
68-
conn, err := client.DialWorkspaceAgent(context.Background(), agent.ID, nil)
68+
conn, err := client.DialWorkspaceAgent(context.Background(), r.sdkAgent.ID, nil)
6969
require.NoError(t, err)
7070
_ = conn.Close()
7171
})
7272
}
7373

74-
func setupWorkspaceAgent(t *testing.T, client *codersdk.Client, user codersdk.CreateFirstUserResponse, appPort uint16) (codersdk.Workspace, codersdk.WorkspaceAgent) {
74+
type setupResp struct {
75+
workspace codersdk.Workspace
76+
sdkAgent codersdk.WorkspaceAgent
77+
agent agent.Agent
78+
}
79+
80+
func setupWorkspaceAgent(t *testing.T, client *codersdk.Client, user codersdk.CreateFirstUserResponse, appPort uint16) setupResp {
7581
authToken := uuid.NewString()
7682
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
7783
Parse: echo.ParseComplete,
@@ -127,20 +133,20 @@ func setupWorkspaceAgent(t *testing.T, client *codersdk.Client, user codersdk.Cr
127133
},
128134
}
129135
agentClient.SetSessionToken(authToken)
130-
agentCloser := agent.New(agent.Options{
136+
agnt := agent.New(agent.Options{
131137
Client: agentClient,
132138
Logger: slogtest.Make(t, nil).Named("agent"),
133139
})
134140
t.Cleanup(func() {
135-
_ = agentCloser.Close()
141+
_ = agnt.Close()
136142
})
137143

138144
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
139145
defer cancel()
140146

141147
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
142-
agnt, err := client.WorkspaceAgent(ctx, resources[0].Agents[0].ID)
148+
sdkAgent, err := client.WorkspaceAgent(ctx, resources[0].Agents[0].ID)
143149
require.NoError(t, err)
144150

145-
return workspace, agnt
151+
return setupResp{workspace, sdkAgent, agnt}
146152
}

‎enterprise/coderd/workspaceportshare_test.go

+6-6
Original file line numberDiff line numberDiff line change
@@ -31,30 +31,30 @@ func TestWorkspacePortShare(t *testing.T) {
3131
},
3232
})
3333
client, user := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleTemplateAdmin())
34-
workspace, agent := setupWorkspaceAgent(t, client, codersdk.CreateFirstUserResponse{
34+
r := setupWorkspaceAgent(t, client, codersdk.CreateFirstUserResponse{
3535
UserID: user.ID,
3636
OrganizationID: owner.OrganizationID,
3737
}, 0)
3838
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
3939
defer cancel()
4040

4141
// try to update port share with template max port share level owner
42-
_, err := client.UpsertWorkspaceAgentPortShare(ctx, workspace.ID, codersdk.UpsertWorkspaceAgentPortShareRequest{
43-
AgentName: agent.Name,
42+
_, err := client.UpsertWorkspaceAgentPortShare(ctx, r.workspace.ID, codersdk.UpsertWorkspaceAgentPortShareRequest{
43+
AgentName: r.sdkAgent.Name,
4444
Port: 8080,
4545
ShareLevel: codersdk.WorkspaceAgentPortShareLevelPublic,
4646
})
4747
require.Error(t, err, "Port sharing level not allowed")
4848

4949
// update the template max port share level to public
5050
var level codersdk.WorkspaceAgentPortShareLevel = codersdk.WorkspaceAgentPortShareLevelPublic
51-
client.UpdateTemplateMeta(ctx, workspace.TemplateID, codersdk.UpdateTemplateMeta{
51+
client.UpdateTemplateMeta(ctx, r.workspace.TemplateID, codersdk.UpdateTemplateMeta{
5252
MaxPortShareLevel: &level,
5353
})
5454

5555
// OK
56-
ps, err := client.UpsertWorkspaceAgentPortShare(ctx, workspace.ID, codersdk.UpsertWorkspaceAgentPortShareRequest{
57-
AgentName: agent.Name,
56+
ps, err := client.UpsertWorkspaceAgentPortShare(ctx, r.workspace.ID, codersdk.UpsertWorkspaceAgentPortShareRequest{
57+
AgentName: r.sdkAgent.Name,
5858
Port: 8080,
5959
ShareLevel: codersdk.WorkspaceAgentPortShareLevelPublic,
6060
})

‎tailnet/coordinator.go

+2-1
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,8 @@ func (c *remoteCoordination) Close() (retErr error) {
131131
}
132132
}()
133133
err := c.protocol.Send(&proto.CoordinateRequest{Disconnect: &proto.CoordinateRequest_Disconnect{}})
134-
if err != nil {
134+
if err != nil && !xerrors.Is(err, io.EOF) {
135+
// Coordinator RPC hangs up when it gets disconnect, so EOF is expected.
135136
return xerrors.Errorf("send disconnect: %w", err)
136137
}
137138
c.logger.Debug(context.Background(), "sent disconnect")

0 commit comments

Comments
 (0)