Skip to content

Commit b1f8370

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 b1f8370

File tree

4 files changed

+146
-100
lines changed

4 files changed

+146
-100
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: 137 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ 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"
@@ -22,89 +23,150 @@ import (
2223
func TestWorkspaceAgentReportStats(t *testing.T) {
2324
t.Parallel()
2425

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()
26+
for _, tc := range []struct {
27+
name string
28+
apiKeyScope rbac.ScopeName
29+
}{
30+
{
31+
name: "empty (backwards compat)",
32+
apiKeyScope: "",
33+
},
34+
{
35+
name: "all",
36+
apiKeyScope: rbac.ScopeAll,
37+
},
38+
{
39+
name: "no_user_data",
40+
apiKeyScope: rbac.ScopeNoUserData,
41+
},
42+
{
43+
name: "application_connect",
44+
apiKeyScope: rbac.ScopeApplicationConnect,
45+
},
46+
} {
47+
t.Run(tc.name, func(t *testing.T) {
48+
t.Parallel()
3749

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)
50+
tickCh := make(chan time.Time)
51+
flushCh := make(chan int, 1)
52+
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
53+
WorkspaceUsageTrackerFlush: flushCh,
54+
WorkspaceUsageTrackerTick: tickCh,
55+
})
56+
user := coderdtest.CreateFirstUser(t, client)
57+
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
58+
OrganizationID: user.OrganizationID,
59+
OwnerID: user.UserID,
60+
LastUsedAt: dbtime.Now().Add(-time.Minute),
61+
}).WithAgent(
62+
func(agent []*proto.Agent) []*proto.Agent {
63+
for _, a := range agent {
64+
a.ApiKeyScope = string(tc.apiKeyScope)
65+
}
4666

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)
67+
return agent
68+
},
69+
).Do()
70+
71+
ac := agentsdk.New(client.URL)
72+
ac.SetSessionToken(r.AgentToken)
73+
conn, err := ac.ConnectRPC(context.Background())
74+
require.NoError(t, err)
75+
defer func() {
76+
_ = conn.Close()
77+
}()
78+
agentAPI := agentproto.NewDRPCAgentClient(conn)
79+
80+
_, err = agentAPI.UpdateStats(context.Background(), &agentproto.UpdateStatsRequest{
81+
Stats: &agentproto.Stats{
82+
ConnectionsByProto: map[string]int64{"TCP": 1},
83+
ConnectionCount: 1,
84+
RxPackets: 1,
85+
RxBytes: 1,
86+
TxPackets: 1,
87+
TxBytes: 1,
88+
SessionCountVscode: 1,
89+
SessionCountJetbrains: 0,
90+
SessionCountReconnectingPty: 0,
91+
SessionCountSsh: 0,
92+
ConnectionMedianLatencyMs: 10,
93+
},
94+
})
95+
require.NoError(t, err)
6396

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

68-
newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID)
69-
require.NoError(t, err)
101+
newWorkspace, err := client.Workspace(context.Background(), r.Workspace.ID)
102+
require.NoError(t, err)
70103

71-
assert.True(t,
72-
newWorkspace.LastUsedAt.After(r.Workspace.LastUsedAt),
73-
"%s is not after %s", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt,
74-
)
104+
assert.True(t,
105+
newWorkspace.LastUsedAt.After(r.Workspace.LastUsedAt),
106+
"%s is not after %s", newWorkspace.LastUsedAt, r.Workspace.LastUsedAt,
107+
)
108+
})
109+
}
75110
}
76111

77112
func TestAgentAPI_LargeManifest(t *testing.T) {
78113
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'
114+
115+
for _, tc := range []struct {
116+
name string
117+
apiKeyScope rbac.ScopeName
118+
}{
119+
{
120+
name: "empty (backwards compat)",
121+
apiKeyScope: "",
122+
},
123+
{
124+
name: "all",
125+
apiKeyScope: rbac.ScopeAll,
126+
},
127+
{
128+
name: "no_user_data",
129+
apiKeyScope: rbac.ScopeNoUserData,
130+
},
131+
{
132+
name: "application_connect",
133+
apiKeyScope: rbac.ScopeApplicationConnect,
134+
},
135+
} {
136+
t.Run(tc.name, func(t *testing.T) {
137+
t.Parallel()
138+
ctx := testutil.Context(t, testutil.WaitLong)
139+
client, store := coderdtest.NewWithDatabase(t, nil)
140+
adminUser := coderdtest.CreateFirstUser(t, client)
141+
n := 512000
142+
longScript := make([]byte, n)
143+
for i := range longScript {
144+
longScript[i] = 'q'
145+
}
146+
r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{
147+
OrganizationID: adminUser.OrganizationID,
148+
OwnerID: adminUser.UserID,
149+
}).WithAgent(func(agents []*proto.Agent) []*proto.Agent {
150+
agents[0].Scripts = []*proto.Script{
151+
{
152+
Script: string(longScript),
153+
},
154+
}
155+
agents[0].ApiKeyScope = string(tc.apiKeyScope)
156+
return agents
157+
}).Do()
158+
ac := agentsdk.New(client.URL)
159+
ac.SetSessionToken(r.AgentToken)
160+
conn, err := ac.ConnectRPC(ctx)
161+
defer func() {
162+
_ = conn.Close()
163+
}()
164+
require.NoError(t, err)
165+
agentAPI := agentproto.NewDRPCAgentClient(conn)
166+
manifest, err := agentAPI.GetManifest(ctx, &agentproto.GetManifestRequest{})
167+
require.NoError(t, err)
168+
require.Len(t, manifest.Scripts, 1)
169+
require.Len(t, manifest.Scripts[0].Script, n)
170+
})
86171
}
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)
110172
}

0 commit comments

Comments
 (0)