opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: refactor agent routines that use the v2 API (#12223)

Browse Source 
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.
This commit is contained in:
Spike Curtis
2024-02-23 11:04:23 +04:00
committed by GitHub
parent 66585f042f
commit af3fdc68c3
9 changed files with 501 additions and 259 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/coderd/coderd_test.go
+4 -3
View File
@@ -231,13 +231,14 @@ func TestAuditLogging(t *testing.T) {
},
DontAddLicense: true,
})
workspace, agent := setupWorkspaceAgent(t, client, user, 0)
conn, err := client.DialWorkspaceAgent(ctx, agent.ID, nil) //nolint:gocritic // RBAC is not the purpose of this test
r := setupWorkspaceAgent(t, client, user, 0)
conn, err := client.DialWorkspaceAgent(ctx, r.sdkAgent.ID, nil) //nolint:gocritic // RBAC is not the purpose of this test
require.NoError(t, err)
defer conn.Close()
connected := conn.AwaitReachable(ctx)
require.True(t, connected)
build := coderdtest.CreateWorkspaceBuild(t, client, workspace, database.WorkspaceTransitionStop)
_ = r.agent.Close() // close first so we don't drop error logs from outdated build
build := coderdtest.CreateWorkspaceBuild(t, client, r.workspace, database.WorkspaceTransitionStop)
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, build.ID)
})
}
enterprise/coderd/replicas_test.go
+4 -4
View File
@@ -81,8 +81,8 @@ func TestReplicas(t *testing.T) {
require.NoError(t, err)
require.Len(t, replicas, 2)
_, agent := setupWorkspaceAgent(t, firstClient, firstUser, 0)
conn, err := secondClient.DialWorkspaceAgent(context.Background(), agent.ID, &codersdk.DialWorkspaceAgentOptions{
r := setupWorkspaceAgent(t, firstClient, firstUser, 0)
conn, err := secondClient.DialWorkspaceAgent(context.Background(), r.sdkAgent.ID, &codersdk.DialWorkspaceAgentOptions{
BlockEndpoints: true,
Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
})
@@ -127,8 +127,8 @@ func TestReplicas(t *testing.T) {
require.NoError(t, err)
require.Len(t, replicas, 2)
_, agent := setupWorkspaceAgent(t, firstClient, firstUser, 0)
conn, err := secondClient.DialWorkspaceAgent(context.Background(), agent.ID, &codersdk.DialWorkspaceAgentOptions{
r := setupWorkspaceAgent(t, firstClient, firstUser, 0)
conn, err := secondClient.DialWorkspaceAgent(context.Background(), r.sdkAgent.ID, &codersdk.DialWorkspaceAgentOptions{
BlockEndpoints: true,
Logger: slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug),
})
enterprise/coderd/workspaceagents_test.go
+15 -9
View File
@@ -44,9 +44,9 @@ func TestBlockNonBrowser(t *testing.T) {
},
},
})
_, agent := setupWorkspaceAgent(t, client, user, 0)
r := setupWorkspaceAgent(t, client, user, 0)
//nolint:gocritic // Testing that even the owner gets blocked.
_, err := client.DialWorkspaceAgent(context.Background(), agent.ID, nil)
_, err := client.DialWorkspaceAgent(context.Background(), r.sdkAgent.ID, nil)
var apiErr *codersdk.Error
require.ErrorAs(t, err, &apiErr)
require.Equal(t, http.StatusConflict, apiErr.StatusCode())
@@ -63,15 +63,21 @@ func TestBlockNonBrowser(t *testing.T) {
},
},
})
_, agent := setupWorkspaceAgent(t, client, user, 0)
r := setupWorkspaceAgent(t, client, user, 0)
//nolint:gocritic // Testing RBAC is not the point of this test.
conn, err := client.DialWorkspaceAgent(context.Background(), agent.ID, nil)
conn, err := client.DialWorkspaceAgent(context.Background(), r.sdkAgent.ID, nil)
require.NoError(t, err)
_ = conn.Close()
})
}
func setupWorkspaceAgent(t *testing.T, client *codersdk.Client, user codersdk.CreateFirstUserResponse, appPort uint16) (codersdk.Workspace, codersdk.WorkspaceAgent) {
type setupResp struct {
workspace codersdk.Workspace
sdkAgent codersdk.WorkspaceAgent
agent agent.Agent
}
func setupWorkspaceAgent(t *testing.T, client *codersdk.Client, user codersdk.CreateFirstUserResponse, appPort uint16) setupResp {
authToken := uuid.NewString()
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
Parse: echo.ParseComplete,
@@ -127,20 +133,20 @@ func setupWorkspaceAgent(t *testing.T, client *codersdk.Client, user codersdk.Cr
},
}
agentClient.SetSessionToken(authToken)
agentCloser := agent.New(agent.Options{
agnt := agent.New(agent.Options{
Client: agentClient,
Logger: slogtest.Make(t, nil).Named("agent"),
})
t.Cleanup(func() {
_ = agentCloser.Close()
_ = agnt.Close()
})
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
agnt, err := client.WorkspaceAgent(ctx, resources[0].Agents[0].ID)
sdkAgent, err := client.WorkspaceAgent(ctx, resources[0].Agents[0].ID)
require.NoError(t, err)
return workspace, agnt
return setupResp{workspace, sdkAgent, agnt}
}
enterprise/coderd/workspaceportshare_test.go
+6 -6
View File
@@ -31,7 +31,7 @@ func TestWorkspacePortShare(t *testing.T) {
},
})
client, user := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleTemplateAdmin())
workspace, agent := setupWorkspaceAgent(t, client, codersdk.CreateFirstUserResponse{
r := setupWorkspaceAgent(t, client, codersdk.CreateFirstUserResponse{
UserID: user.ID,
OrganizationID: owner.OrganizationID,
}, 0)
@@ -39,8 +39,8 @@ func TestWorkspacePortShare(t *testing.T) {
defer cancel()
// try to update port share with template max port share level owner
_, err := client.UpsertWorkspaceAgentPortShare(ctx, workspace.ID, codersdk.UpsertWorkspaceAgentPortShareRequest{
AgentName: agent.Name,
_, err := client.UpsertWorkspaceAgentPortShare(ctx, r.workspace.ID, codersdk.UpsertWorkspaceAgentPortShareRequest{
AgentName: r.sdkAgent.Name,
Port: 8080,
ShareLevel: codersdk.WorkspaceAgentPortShareLevelPublic,
})
@@ -48,13 +48,13 @@ func TestWorkspacePortShare(t *testing.T) {
// update the template max port share level to public
var level codersdk.WorkspaceAgentPortShareLevel = codersdk.WorkspaceAgentPortShareLevelPublic
client.UpdateTemplateMeta(ctx, workspace.TemplateID, codersdk.UpdateTemplateMeta{
client.UpdateTemplateMeta(ctx, r.workspace.TemplateID, codersdk.UpdateTemplateMeta{
MaxPortShareLevel: &level,
})
// OK
ps, err := client.UpsertWorkspaceAgentPortShare(ctx, workspace.ID, codersdk.UpsertWorkspaceAgentPortShareRequest{
AgentName: agent.Name,
ps, err := client.UpsertWorkspaceAgentPortShare(ctx, r.workspace.ID, codersdk.UpsertWorkspaceAgentPortShareRequest{
AgentName: r.sdkAgent.Name,
Port: 8080,
ShareLevel: codersdk.WorkspaceAgentPortShareLevelPublic,
})