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

fix: limit calls to GetWorkspaceAgentByID in agentapi (#23015)

Browse Source 
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:
Jon Ayers
2026-03-20 15:42:05 -05:00
committed by GitHub
parent 32021b3ac2
commit f135ffdb3a
23 changed files with 173 additions and 228 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/agentapi/stats.go
+8 -25
View File
@@ -4,20 +4,21 @@ import (
"context"
"time"
"github.com/google/uuid"
"golang.org/x/xerrors"
"google.golang.org/protobuf/types/known/durationpb"
"cdr.dev/slog/v3"
agentproto "github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/workspacestats"
"github.com/coder/coder/v2/codersdk"
)
type StatsAPI struct {
AgentFn func(context.Context) (database.WorkspaceAgent, error)
AgentID uuid.UUID
AgentName string
Workspace *CachedWorkspaceFields
Database database.Store
Log slog.Logger
@@ -44,32 +45,13 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
return res, nil
}
// Inject RBAC object into context for dbauthz fast path, avoid having to
// call GetWorkspaceAgentByID on every stats update.
rbacCtx := ctx
if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
var err error
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))
}
}
workspaceAgent, err := a.AgentFn(rbacCtx)
if err != nil {
return nil, err
}
// If cache is empty (prebuild or invalid), fall back to DB
var ws database.WorkspaceIdentity
var ok bool
if ws, ok = a.Workspace.AsWorkspaceIdentity(); !ok {
w, err := a.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID)
w, err := a.Database.GetWorkspaceByAgentID(ctx, a.AgentID)
if err != nil {
return nil, xerrors.Errorf("get workspace by agent ID %q: %w", workspaceAgent.ID, err)
return nil, xerrors.Errorf("get workspace by agent ID %q: %w", a.AgentID, err)
}
ws = database.WorkspaceIdentityFromWorkspace(w)
}
@@ -90,11 +72,12 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
req.Stats.SessionCountReconnectingPty = 0
}
err = a.StatsReporter.ReportAgentStats(
err := a.StatsReporter.ReportAgentStats(
ctx,
a.now(),
ws,
workspaceAgent,
a.AgentID,
a.AgentName,
req.Stats,
false,
)