mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix: limit calls to GetWorkspaceAgentByID in agentapi (#23015)
We currently call GetWorkspaceAgentByID millions of times at scale unnecessarily. This PR embeds immutable fields into the relevant services instead of fetching for them every time. resolves https://github.com/coder/scaletest/issues/84 Confirmed with a 10k scaletest that this changeset takes the query from 10M+ queries down to 39k
This commit is contained in:
@@ -14,11 +14,11 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/connectionlog"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/db2sdk"
|
||||
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
||||
)
|
||||
|
||||
type ConnLogAPI struct {
|
||||
AgentFn func(context.Context) (database.WorkspaceAgent, error)
|
||||
AgentID uuid.UUID
|
||||
AgentName string
|
||||
ConnectionLogger *atomic.Pointer[connectionlog.ConnectionLogger]
|
||||
Workspace *CachedWorkspaceFields
|
||||
Database database.Store
|
||||
@@ -53,27 +53,12 @@ func (a *ConnLogAPI) ReportConnection(ctx context.Context, req *agentproto.Repor
|
||||
}
|
||||
}
|
||||
|
||||
// Inject RBAC object into context for dbauthz fast path, avoid having to
|
||||
// call GetWorkspaceByAgentID on every metadata update.
|
||||
rbacCtx := ctx
|
||||
var ws database.WorkspaceIdentity
|
||||
if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
|
||||
ws = dbws
|
||||
rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject())
|
||||
if err != nil {
|
||||
// Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID.
|
||||
//nolint:gocritic
|
||||
a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err))
|
||||
}
|
||||
}
|
||||
|
||||
// Fetch contextual data for this connection log event.
|
||||
workspaceAgent, err := a.AgentFn(rbacCtx)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("get agent: %w", err)
|
||||
}
|
||||
if ws.Equal(database.WorkspaceIdentity{}) {
|
||||
workspace, err := a.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID)
|
||||
workspace, err := a.Database.GetWorkspaceByAgentID(ctx, a.AgentID)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("get workspace by agent id: %w", err)
|
||||
}
|
||||
@@ -97,7 +82,7 @@ func (a *ConnLogAPI) ReportConnection(ctx context.Context, req *agentproto.Repor
|
||||
WorkspaceOwnerID: ws.OwnerID,
|
||||
WorkspaceID: ws.ID,
|
||||
WorkspaceName: ws.Name,
|
||||
AgentName: workspaceAgent.Name,
|
||||
AgentName: a.AgentName,
|
||||
Type: connectionType,
|
||||
Code: code,
|
||||
Ip: logIP,
|
||||
|
||||
Reference in New Issue
Block a user