mirror of
https://github.com/coder/coder.git
synced 2026-06-03 13:08:25 +00:00
Spike Curtis
bddb808b25
Fixes all our Go file imports to match the preferred spec that we've _mostly_ been using. For example: ``` import ( "context" "time" "github.com/prometheus/client_golang/prometheus" "golang.org/x/xerrors" "gopkg.in/natefinch/lumberjack.v2" "cdr.dev/slog/v3" "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/serpent" ) ``` 3 groups: standard library, 3rd partly libs, Coder libs. This PR makes the change across the codebase. The PR in the stack above modifies our formatting to maintain this state of affairs, and is a separate PR so it's possible to review that one in detail.
107 lines
3.0 KiB
Go
107 lines
3.0 KiB
Go
package agentapi
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"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)
|
|
Workspace *CachedWorkspaceFields
|
|
Database database.Store
|
|
Log slog.Logger
|
|
StatsReporter *workspacestats.Reporter
|
|
AgentStatsRefreshInterval time.Duration
|
|
Experiments codersdk.Experiments
|
|
|
|
TimeNowFn func() time.Time // defaults to dbtime.Now()
|
|
}
|
|
|
|
func (a *StatsAPI) now() time.Time {
|
|
if a.TimeNowFn != nil {
|
|
return a.TimeNowFn()
|
|
}
|
|
return dbtime.Now()
|
|
}
|
|
|
|
func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsRequest) (*agentproto.UpdateStatsResponse, error) {
|
|
res := &agentproto.UpdateStatsResponse{
|
|
ReportInterval: durationpb.New(a.AgentStatsRefreshInterval),
|
|
}
|
|
// An empty stat means it's just looking for the report interval.
|
|
if req.Stats == nil {
|
|
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)
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("get workspace by agent ID %q: %w", workspaceAgent.ID, err)
|
|
}
|
|
ws = database.WorkspaceIdentityFromWorkspace(w)
|
|
}
|
|
|
|
a.Log.Debug(ctx, "read stats report",
|
|
slog.F("interval", a.AgentStatsRefreshInterval),
|
|
slog.F("workspace_id", ws.ID),
|
|
slog.F("payload", req),
|
|
)
|
|
|
|
if a.Experiments.Enabled(codersdk.ExperimentWorkspaceUsage) {
|
|
// while the experiment is enabled we will not report
|
|
// session stats from the agent. This is because it is
|
|
// being handled by the CLI and the postWorkspaceUsage route.
|
|
req.Stats.SessionCountSsh = 0
|
|
req.Stats.SessionCountJetbrains = 0
|
|
req.Stats.SessionCountVscode = 0
|
|
req.Stats.SessionCountReconnectingPty = 0
|
|
}
|
|
|
|
err = a.StatsReporter.ReportAgentStats(
|
|
ctx,
|
|
a.now(),
|
|
ws,
|
|
workspaceAgent,
|
|
req.Stats,
|
|
false,
|
|
)
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("report agent stats: %w", err)
|
|
}
|
|
|
|
return res, nil
|
|
}
|