forked from coder/coder
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathstats.go
More file actions
106 lines (91 loc) · 3.04 KB
/
stats.go
File metadata and controls
106 lines (91 loc) · 3.04 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
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
}