Skip to content

Commit e0f806f

Browse files
ThomasK33ThomasK33
committed
fix: remove unnecessary user lookup in agent API calls
Change-Id: I3b6e7581134e2374b364ee059e3b18ece3d98b41 Signed-off-by: Thomas Kosiewski <tk@coder.com>
1 parent e5758a1 commit e0f806f

File tree

4 files changed

+144
-102
lines changed

4 files changed

+144
-102
lines changed

‎coderd/agentapi/manifest.go

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,6 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
4747
scripts []database.WorkspaceAgentScript
4848
metadata []database.WorkspaceAgentMetadatum
4949
workspace database.Workspace
50-
owner database.User
5150
devcontainers []database.WorkspaceAgentDevcontainer
5251
)
5352

@@ -76,10 +75,6 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
7675
if err != nil {
7776
return xerrors.Errorf("getting workspace by id: %w", err)
7877
}
79-
owner, err = a.Database.GetUserByID(ctx, workspace.OwnerID)
80-
if err != nil {
81-
return xerrors.Errorf("getting workspace owner by id: %w", err)
82-
}
8378
return err
8479
})
8580
eg.Go(func() (err error) {
@@ -98,7 +93,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
9893
AppSlugOrPort: "{{port}}",
9994
AgentName: workspaceAgent.Name,
10095
WorkspaceName: workspace.Name,
101-
Username: owner.Username,
96+
Username: workspace.OwnerUsername,
10297
}
10398

10499
vscodeProxyURI := vscodeProxyURI(appSlug, a.AccessURL, a.AppHostname)
@@ -115,7 +110,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
115110
}
116111
}
117112

118-
apps, err := dbAppsToProto(dbApps, workspaceAgent, owner.Username, workspace)
113+
apps, err := dbAppsToProto(dbApps, workspaceAgent, workspace.OwnerUsername, workspace)
119114
if err != nil {
120115
return nil, xerrors.Errorf("converting workspace apps: %w", err)
121116
}
@@ -128,7 +123,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
128123
return &agentproto.Manifest{
129124
AgentId: workspaceAgent.ID[:],
130125
AgentName: workspaceAgent.Name,
131-
OwnerUsername: owner.Username,
126+
OwnerUsername: workspace.OwnerUsername,
132127
WorkspaceId: workspace.ID[:],
133128
WorkspaceName: workspace.Name,
134129
GitAuthConfigs: gitAuthConfigs,

‎coderd/agentapi/manifest_test.go

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,9 +46,10 @@ func TestGetManifest(t *testing.T) {
4646
Username: "cool-user",
4747
}
4848
workspace = database.Workspace{
49-
ID: uuid.New(),
50-
OwnerID: owner.ID,
51-
Name: "cool-workspace",
49+
ID: uuid.New(),
50+
OwnerID: owner.ID,
51+
OwnerUsername: owner.Username,
52+
Name: "cool-workspace",
5253
}
5354
agent = database.WorkspaceAgent{
5455
ID: uuid.New(),
@@ -336,7 +337,6 @@ func TestGetManifest(t *testing.T) {
336337
}).Return(metadata, nil)
337338
mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), agent.ID).Return(devcontainers, nil)
338339
mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil)
339-
mDB.EXPECT().GetUserByID(gomock.Any(), workspace.OwnerID).Return(owner, nil)
340340

341341
got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{})
342342
require.NoError(t, err)
@@ -404,7 +404,6 @@ func TestGetManifest(t *testing.T) {
404404
}).Return([]database.WorkspaceAgentMetadatum{}, nil)
405405
mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), childAgent.ID).Return([]database.WorkspaceAgentDevcontainer{}, nil)
406406
mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil)
407-
mDB.EXPECT().GetUserByID(gomock.Any(), workspace.OwnerID).Return(owner, nil)
408407

409408
got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{})
410409
require.NoError(t, err)
@@ -468,7 +467,6 @@ func TestGetManifest(t *testing.T) {
468467
}).Return(metadata, nil)
469468
mDB.EXPECT().GetWorkspaceAgentDevcontainersByAgentID(gomock.Any(), agent.ID).Return(devcontainers, nil)
470469
mDB.EXPECT().GetWorkspaceByID(gomock.Any(), workspace.ID).Return(workspace, nil)
471-
mDB.EXPECT().GetUserByID(gomock.Any(), workspace.OwnerID).Return(owner, nil)
472470

473471
got, err := api.GetManifest(context.Background(), &agentproto.GetManifestRequest{})
474472
require.NoError(t, err)

‎coderd/workspaceagentsrpc.go

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -76,17 +76,8 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
7676
return
7777
}
7878

79-
owner, err := api.Database.GetUserByID(ctx, workspace.OwnerID)
80-
if err != nil {
81-
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
82-
Message: "Internal error fetching user.",
83-
Detail: err.Error(),
84-
})
85-
return
86-
}
87-
8879
logger = logger.With(
89-
slog.F("owner", owner.Username),
80+
slog.F("owner", workspace.OwnerUsername),
9081
slog.F("workspace_name", workspace.Name),
9182
slog.F("agent_name", workspaceAgent.Name),
9283
)
@@ -170,7 +161,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
170161
})
171162

172163
streamID := tailnet.StreamID{
173-
Name: fmt.Sprintf("%s-%s-%s", owner.Username, workspace.Name, workspaceAgent.Name),
164+
Name: fmt.Sprintf("%s-%s-%s", workspace.OwnerUsername, workspace.Name, workspaceAgent.Name),
174165
ID: workspaceAgent.ID,
175166
Auth: tailnet.AgentCoordinateeAuth{ID: workspaceAgent.ID},
176167
}

‎coderd/workspaceagentsrpc_test.go

Lines changed: 135 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -13,98 +13,156 @@ import (
1313
"github.com/coder/coder/v2/coderd/database"
1414
"github.com/coder/coder/v2/coderd/database/dbfake"
1515
"github.com/coder/coder/v2/coderd/database/dbtime"
16+
"github.com/coder/coder/v2/coderd/rbac"
1617
"github.com/coder/coder/v2/codersdk/agentsdk"
1718
"github.com/coder/coder/v2/provisionersdk/proto"
1819
"github.com/coder/coder/v2/testutil"
1920
)
2021

2122
// Ported to RPC API from coderd/workspaceagents_test.go
2223
func TestWorkspaceAgentReportStats(t *testing.T) {
23-
t.Parallel()
24+
for _, tc := range []struct {
25+
name string
26+
apiKeyScope rbac.ScopeName
27+
}{
28+
{
29+
name: "empty (backwards compat)",
30+
apiKeyScope: "",
31+
},
32+
{
33+
name: "all",
34+
apiKeyScope: rbac.ScopeAll,
35+
},
36+
{
37+
name: "no_user_data",
38+
apiKeyScope: rbac.ScopeNoUserData,
39+
},
40+
{
41+
name: "application_connect",
42+
apiKeyScope: rbac.ScopeApplicationConnect,
43+
},
44+
} {
45+
t.Run(tc.name, func(t *testing.T) {
46+
t.Parallel()
2447

25-
tickCh := make(chan time.Time)
26-
flushCh := make(chan int, 1)
27-
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
28-
WorkspaceUsageTrackerFlush: flushCh,
29-
WorkspaceUsageTrackerTick: tickCh,
30-
})
31-
user := coderdtest.CreateFirstUser(t, client)
32-
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
33-
OrganizationID: user.OrganizationID,
34-
OwnerID: user.UserID,
35-
LastUsedAt: dbtime.Now().Add(-time.Minute),
36-
}).WithAgent().Do()
48+
tickCh := make(chan time.Time)
49+
flushCh := make(chan int, 1)
50+
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
51+
WorkspaceUsageTrackerFlush: flushCh,
52+
WorkspaceUsageTrackerTick: tickCh,
53+
})
54+
user := coderdtest.CreateFirstUser(t, client)
55+
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
56+
OrganizationID: user.OrganizationID,
57+
OwnerID: user.UserID,
58+
LastUsedAt: dbtime.Now().Add(-time.Minute),
59+
}).WithAgent(
60+
func(agent []*proto.Agent) []*proto.Agent {
61+
for _, a := range agent {
62+
a.ApiKeyScope = string(tc.apiKeyScope)
63+
}
3764

38-
ac := agentsdk.New(client.URL)
39-
ac.SetSessionToken(r.AgentToken)
40-
conn, err := ac.ConnectRPC(context.Background())
41-
require.NoError(t, err)
42-
defer func() {
43-
_ = conn.Close()
44-
}()
45-
agentAPI := agentproto.NewDRPCAgentClient(conn)
65+
return agent
66+
},
67+
).Do()
4668

47-
_, err = agentAPI.UpdateStats(context.Background(), &agentproto.UpdateStatsRequest{
48-
Stats: &agentproto.Stats{
49-
ConnectionsByProto: map[string]int64{"TCP": 1},
50-
ConnectionCount: 1,
51-
RxPackets: 1,
52-
RxBytes: 1,
53-
TxPackets: 1,
54-
TxBytes: 1,
55-
SessionCountVscode: 1,
56-
SessionCountJetbrains: 0,
57-
SessionCountReconnectingPty: 0,
58-
SessionCountSsh: 0,
59-
ConnectionMedianLatencyMs: 10,
60-
},
61-
})
62-
require.NoError(t, err)
69+
ac := agentsdk.New(client.URL)
70+
ac.SetSessionToken(r.AgentToken)
71+
conn, err := ac.ConnectRPC(context.Background())
72+
require.NoError(t, err)
73+
defer func() {
74+
_ = conn.Close()
75+
}()
76+
agentAPI := agentproto.NewDRPCAgentClient(conn)
77+
78+
_, err = agentAPI.UpdateStats(context.Background(), &agentproto.UpdateStatsRequest{
79+
Stats: &agentproto.Stats{
80+
ConnectionsByProto: map[string]int64{"TCP": 1},
81+
ConnectionCount: 1,
82+
RxPackets: 1,
83+
RxBytes: 1,
84+
TxPackets: 1,
85+
TxBytes: 1,
86+
SessionCountVscode: 1,
87+
SessionCountJetbrains: 0,
88+
SessionCountReconnectingPty: 0,
89+
SessionCountSsh: 0,
90+
ConnectionMedianLatencyMs: 10,
91+
},
92+
})
93+
require.NoError(t, err)
6394

64-
tickCh <- dbtime.Now()
65-
count := <-flushCh
66-
require.Equal(t, 1, count, "expected one flush with one id")
95+
tickCh <- dbtime.Now()
96+
count := <-flushCh
97+
require.Equal(t, 1, count, "expected one flush with one id")
6798

68-
newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID)
69-
require.NoError(t, err)
99+
newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID)
100+
require.NoError(t, err)
70101

71-
assert.True(t,
72-
newWorkspace.LastUsedAt.After(r.Workspace.LastUsedAt),
73-
"%s is not after %s", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt,
74-
)
102+
assert.True(t,
103+
newWorkspace.LastUsedAt.After(r.Workspace.LastUsedAt),
104+
"%s is not after %s", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt,
105+
)
106+
})
107+
}
75108
}
76109

77110
func TestAgentAPI_LargeManifest(t *testing.T) {
78-
t.Parallel()
79-
ctx := testutil.Context(t, testutil.WaitLong)
80-
client, store := coderdtest.NewWithDatabase(t, nil)
81-
adminUser := coderdtest.CreateFirstUser(t, client)
82-
n := 512000
83-
longScript := make([]byte, n)
84-
for i := range longScript {
85-
longScript[i] = 'q'
111+
for _, tc := range []struct {
112+
name string
113+
apiKeyScope rbac.ScopeName
114+
}{
115+
{
116+
name: "empty (backwards compat)",
117+
apiKeyScope: "",
118+
},
119+
{
120+
name: "all",
121+
apiKeyScope: rbac.ScopeAll,
122+
},
123+
{
124+
name: "no_user_data",
125+
apiKeyScope: rbac.ScopeNoUserData,
126+
},
127+
{
128+
name: "application_connect",
129+
apiKeyScope: rbac.ScopeApplicationConnect,
130+
},
131+
} {
132+
t.Run(tc.name, func(t *testing.T) {
133+
t.Parallel()
134+
ctx := testutil.Context(t, testutil.WaitLong)
135+
client, store := coderdtest.NewWithDatabase(t, nil)
136+
adminUser := coderdtest.CreateFirstUser(t, client)
137+
n := 512000
138+
longScript := make([]byte, n)
139+
for i := range longScript {
140+
longScript[i] = 'q'
141+
}
142+
r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{
143+
OrganizationID: adminUser.OrganizationID,
144+
OwnerID: adminUser.UserID,
145+
}).WithAgent(func(agents []*proto.Agent) []*proto.Agent {
146+
agents[0].Scripts = []*proto.Script{
147+
{
148+
Script: string(longScript),
149+
},
150+
}
151+
agents[0].ApiKeyScope = string(tc.apiKeyScope)
152+
return agents
153+
}).Do()
154+
ac := agentsdk.New(client.URL)
155+
ac.SetSessionToken(r.AgentToken)
156+
conn, err := ac.ConnectRPC(ctx)
157+
defer func() {
158+
_ = conn.Close()
159+
}()
160+
require.NoError(t, err)
161+
agentAPI := agentproto.NewDRPCAgentClient(conn)
162+
manifest, err := agentAPI.GetManifest(ctx, &agentproto.GetManifestRequest{})
163+
require.NoError(t, err)
164+
require.Len(t, manifest.Scripts, 1)
165+
require.Len(t, manifest.Scripts[0].Script, n)
166+
})
86167
}
87-
r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{
88-
OrganizationID: adminUser.OrganizationID,
89-
OwnerID: adminUser.UserID,
90-
}).WithAgent(func(agents []*proto.Agent) []*proto.Agent {
91-
agents[0].Scripts = []*proto.Script{
92-
{
93-
Script: string(longScript),
94-
},
95-
}
96-
return agents
97-
}).Do()
98-
ac := agentsdk.New(client.URL)
99-
ac.SetSessionToken(r.AgentToken)
100-
conn, err := ac.ConnectRPC(ctx)
101-
defer func() {
102-
_ = conn.Close()
103-
}()
104-
require.NoError(t, err)
105-
agentAPI := agentproto.NewDRPCAgentClient(conn)
106-
manifest, err := agentAPI.GetManifest(ctx, &agentproto.GetManifestRequest{})
107-
require.NoError(t, err)
108-
require.Len(t, manifest.Scripts, 1)
109-
require.Len(t, manifest.Scripts[0].Script, n)
110168
}

0 commit comments

Comments
 (0)