mirror of
https://github.com/coder/coder.git
synced 2026-10-16 14:48:56 +08:00
feat: report arbitrary workspace agent session types (#28338)
Report session counts as a dynamic app-name map from the agent, the SSH CLI, and the workspace usage API, so a new IDE flows end to end without a schema, proto, or server change. - Report through `Stats.session_counts` on Agent API v2.11. Older agents keep reporting through the deprecated count fields. - Accept arbitrary names from `coder ssh --usage-app` and the workspace usage API. A name the registry does not recognize keeps its identity and is attributed to the `unknown` family. - Add reconnecting PTY counts to the map without overwriting the counts SSH reported. - Normalize and merge names before applying the deterministic 64-name cap, ordered by known apps, then count, then name. Overflow sums into `unknown`, and non-positive counts cannot cancel valid activity. - Strip control characters, treat whitespace-only HTTP app names as absent, and require the exact environment variable prefix when you extract an SSH session type. - Count overflow entries in `coderd_agentstats_session_counts_overflow_total` and keep raw payloads out of logs. - Change the usage API's `app_name` to a string, and update the SDK types and API reference. closes #27410
This commit is contained in:
Ehab Younes
+12
-25
@@ -127,23 +127,10 @@ type Options struct {
|
||||
}
|
||||
|
||||
type Client interface {
|
||||
ConnectRPC29(ctx context.Context) (
|
||||
proto.DRPCAgentClient29, tailnetproto.DRPCTailnetClient28, error,
|
||||
)
|
||||
// ConnectRPC29WithRole is like ConnectRPC29 but sends an explicit
|
||||
// role query parameter to the server. The workspace agent should
|
||||
// use role "agent" to enable connection monitoring.
|
||||
ConnectRPC29WithRole(ctx context.Context, role string) (
|
||||
proto.DRPCAgentClient29, tailnetproto.DRPCTailnetClient28, error,
|
||||
)
|
||||
ConnectRPC210(ctx context.Context) (
|
||||
proto.DRPCAgentClient210, tailnetproto.DRPCTailnetClient28, error,
|
||||
)
|
||||
// ConnectRPC210WithRole is like ConnectRPC210 but sends an explicit
|
||||
// role query parameter to the server. The workspace agent should
|
||||
// use role "agent" to enable connection monitoring.
|
||||
ConnectRPC210WithRole(ctx context.Context, role string) (
|
||||
proto.DRPCAgentClient210, tailnetproto.DRPCTailnetClient28, error,
|
||||
// ConnectRPC211WithRole connects to the Agent API v2.11. The workspace
|
||||
// agent should use role "agent" to enable connection monitoring.
|
||||
ConnectRPC211WithRole(ctx context.Context, role string) (
|
||||
proto.DRPCAgentClient211, tailnetproto.DRPCTailnetClient28, error,
|
||||
)
|
||||
tailnet.DERPMapRewriter
|
||||
agentsdk.RefreshableSessionTokenProvider
|
||||
@@ -1166,7 +1153,7 @@ func (a *agent) run() (retErr error) {
|
||||
// ConnectRPC returns the dRPC connection we use for the Agent and Tailnet v2+ APIs.
|
||||
// We pass role "agent" to enable connection monitoring on the server, which tracks
|
||||
// the agent's connectivity state (first_connected_at, last_connected_at, disconnected_at).
|
||||
aAPI, tAPI, err := a.client.ConnectRPC210WithRole(a.hardCtx, "agent")
|
||||
aAPI, tAPI, err := a.client.ConnectRPC211WithRole(a.hardCtx, "agent")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -2152,13 +2139,13 @@ func (a *agent) Collect(ctx context.Context, networkStats map[netlogtype.Connect
|
||||
stats.TxPackets += int64(counts.TxPackets)
|
||||
}
|
||||
|
||||
// The count of active sessions; types without a protocol field are dropped.
|
||||
sessionCounts := a.sshServer.SessionCounts()
|
||||
stats.SessionCountSsh = sessionCounts[string(codersdk.AppFamilySSH)]
|
||||
stats.SessionCountVscode = sessionCounts[string(codersdk.AppFamilyVSCode)]
|
||||
stats.SessionCountJetbrains = sessionCounts[string(codersdk.AppFamilyJetBrains)]
|
||||
|
||||
stats.SessionCountReconnectingPty = a.reconnectingPTYServer.ConnCount()
|
||||
// A client picks its own app name and can pick reconnecting_pty, so add
|
||||
// rather than overwrite. A non-empty map suppresses the deprecated
|
||||
// session_count_* fields on ingest, so leave them zero.
|
||||
stats.SessionCounts = a.sshServer.SessionCounts()
|
||||
if count := a.reconnectingPTYServer.ConnCount(); count > 0 {
|
||||
stats.SessionCounts[string(codersdk.AppFamilyReconnectingPTY)] += count
|
||||
}
|
||||
|
||||
// Compute the median connection latency!
|
||||
a.logger.Debug(ctx, "starting peer latency measurement for stats")
|
||||
|
||||
+10
-10
@@ -229,7 +229,7 @@ func assertSSHStats(t *testing.T, stats <-chan *proto.Stats) {
|
||||
return false
|
||||
}
|
||||
t.Logf("got stats: ConnectionCount=%d, RxBytes=%d, TxBytes=%d, SessionCountSsh=%d",
|
||||
s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCountSsh)
|
||||
s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCounts["ssh"])
|
||||
if s.ConnectionCount > 0 {
|
||||
connectionCountSeen = true
|
||||
}
|
||||
@@ -239,7 +239,7 @@ func assertSSHStats(t *testing.T, stats <-chan *proto.Stats) {
|
||||
if s.TxBytes > 0 {
|
||||
txBytesSeen = true
|
||||
}
|
||||
if s.SessionCountSsh == 1 {
|
||||
if s.SessionCounts["ssh"] == 1 {
|
||||
sessionCountSSHSeen = true
|
||||
}
|
||||
return connectionCountSeen && rxBytesSeen && txBytesSeen && sessionCountSSHSeen
|
||||
@@ -287,7 +287,7 @@ func TestAgent_Stats_ReconnectingPTY(t *testing.T) {
|
||||
if s.TxBytes > 0 {
|
||||
txBytesSeen = true
|
||||
}
|
||||
if s.SessionCountReconnectingPty == 1 {
|
||||
if s.SessionCounts["reconnecting_pty"] == 1 {
|
||||
sessionCountReconnectingPTYSeen = true
|
||||
}
|
||||
return connectionCountSeen && rxBytesSeen && txBytesSeen && sessionCountReconnectingPTYSeen
|
||||
@@ -346,12 +346,12 @@ func TestAgent_Stats_Magic(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
require.Eventuallyf(t, func() bool {
|
||||
s, ok := <-stats
|
||||
t.Logf("got stats: ok=%t, ConnectionCount=%d, RxBytes=%d, TxBytes=%d, SessionCountVSCode=%d, ConnectionMedianLatencyMS=%f",
|
||||
ok, s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCountVscode, s.ConnectionMedianLatencyMs)
|
||||
t.Logf("got stats: ok=%t, ConnectionCount=%d, RxBytes=%d, TxBytes=%d, SessionCounts[vscode]=%d, ConnectionMedianLatencyMS=%f",
|
||||
ok, s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCounts["vscode"], s.ConnectionMedianLatencyMs)
|
||||
return ok &&
|
||||
// Ensure that the connection didn't count as a "normal" SSH session.
|
||||
// This was a special one, so it should be labeled specially in the stats!
|
||||
s.SessionCountVscode == 1 &&
|
||||
s.SessionCounts["vscode"] == 1 &&
|
||||
// Ensure that connection latency is being counted!
|
||||
// If it isn't, it's set to -1.
|
||||
s.ConnectionMedianLatencyMs >= 0
|
||||
@@ -417,8 +417,8 @@ func TestAgent_Stats_Magic(t *testing.T) {
|
||||
require.Eventuallyf(t, func() bool {
|
||||
s, ok := <-stats
|
||||
t.Logf("got stats with conn open: ok=%t, ConnectionCount=%d, SessionCountJetBrains=%d",
|
||||
ok, s.ConnectionCount, s.SessionCountJetbrains)
|
||||
return ok && s.SessionCountJetbrains == 1
|
||||
ok, s.ConnectionCount, s.SessionCounts["jetbrains"])
|
||||
return ok && s.SessionCounts["jetbrains"] == 1
|
||||
}, testutil.WaitLong, testutil.IntervalFast,
|
||||
"never saw stats with conn open",
|
||||
)
|
||||
@@ -431,9 +431,9 @@ func TestAgent_Stats_Magic(t *testing.T) {
|
||||
require.Eventuallyf(t, func() bool {
|
||||
s, ok := <-stats
|
||||
t.Logf("got stats after disconnect %t, %d",
|
||||
ok, s.SessionCountJetbrains)
|
||||
ok, s.SessionCounts["jetbrains"])
|
||||
return ok &&
|
||||
s.SessionCountJetbrains == 0
|
||||
s.SessionCounts["jetbrains"] == 0
|
||||
}, testutil.WaitLong, testutil.IntervalFast,
|
||||
"never saw stats after conn closes",
|
||||
)
|
||||
|
||||
@@ -340,13 +340,14 @@ func (s *Server) SessionCounts() map[string]int64 {
|
||||
}
|
||||
|
||||
func extractAppName(env []string) (appName, rawAppName string, filteredEnv []string) {
|
||||
// Match the full assignment so a longer variable such as
|
||||
// CODER_SSH_SESSION_TYPE_FOO=bar is not mistaken for this one.
|
||||
prefix := AppNameEnvironmentVariable + "="
|
||||
for _, kv := range env {
|
||||
if !strings.HasPrefix(kv, AppNameEnvironmentVariable) {
|
||||
continue
|
||||
if value, ok := strings.CutPrefix(kv, prefix); ok {
|
||||
// The last instance of the variable wins.
|
||||
rawAppName = value
|
||||
}
|
||||
|
||||
rawAppName = strings.TrimPrefix(kv, AppNameEnvironmentVariable+"=")
|
||||
// Keep going, we'll use the last instance of the env.
|
||||
}
|
||||
|
||||
if rawAppName == "" {
|
||||
@@ -357,7 +358,7 @@ func extractAppName(env []string) (appName, rawAppName string, filteredEnv []str
|
||||
}
|
||||
|
||||
return appName, rawAppName, slices.DeleteFunc(env, func(kv string) bool {
|
||||
return strings.HasPrefix(kv, AppNameEnvironmentVariable+"=")
|
||||
return strings.HasPrefix(kv, prefix)
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -21,14 +21,30 @@ func TestExtractAppName(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
env []string
|
||||
want string
|
||||
// keepEnv, when set, must survive filtering.
|
||||
keepEnv string
|
||||
want string
|
||||
}{
|
||||
{"NoEnvDefaultsToSSH", []string{"FOO=bar"}, "ssh"},
|
||||
{"EmptyValueDefaultsToSSH", envWith(""), "ssh"},
|
||||
{"VSCode", envWith("vscode"), "vscode"},
|
||||
{"JetBrainsLegacyCasing", envWith("JetBrains"), "jetbrains"},
|
||||
{"UnknownTypeNormalized", envWith("Cursor-Nightly"), "cursor_nightly"},
|
||||
{"LastInstanceWins", append(envWith("vscode"), AppNameEnvironmentVariable+"=cursor"), "cursor"},
|
||||
{"NoEnvDefaultsToSSH", []string{"FOO=bar"}, "FOO=bar", "ssh"},
|
||||
{"EmptyValueDefaultsToSSH", envWith(""), "", "ssh"},
|
||||
{"VSCode", envWith("vscode"), "", "vscode"},
|
||||
{"JetBrainsLegacyCasing", envWith("JetBrains"), "", "jetbrains"},
|
||||
{"UnknownTypeNormalized", envWith("Cursor-Nightly"), "", "cursor_nightly"},
|
||||
{"LastInstanceWins", append(envWith("vscode"), AppNameEnvironmentVariable+"=cursor"), "", "cursor"},
|
||||
// A variable that merely starts with the same characters is a
|
||||
// different variable, not an app name.
|
||||
{
|
||||
"LongerVariableIgnored",
|
||||
[]string{AppNameEnvironmentVariable + "_FOO=bar"},
|
||||
AppNameEnvironmentVariable + "_FOO=bar",
|
||||
"ssh",
|
||||
},
|
||||
{
|
||||
"LongerVariableDoesNotShadowAppName",
|
||||
append(envWith("cursor"), AppNameEnvironmentVariable+"_FOO=bar"),
|
||||
AppNameEnvironmentVariable + "_FOO=bar",
|
||||
"cursor",
|
||||
},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
@@ -37,6 +53,9 @@ func TestExtractAppName(t *testing.T) {
|
||||
for _, kv := range filteredEnv {
|
||||
require.NotContains(t, kv, AppNameEnvironmentVariable+"=")
|
||||
}
|
||||
if tc.keepEnv != "" {
|
||||
require.Contains(t, filteredEnv, tc.keepEnv)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -153,6 +153,7 @@ func (c *Client) Close() {
|
||||
c.derpMapOnce.Do(func() { close(c.derpMapUpdates) })
|
||||
}
|
||||
|
||||
// Required by the agentfake scaletest dialer.
|
||||
func (c *Client) ConnectRPC29WithRole(ctx context.Context, _ string) (
|
||||
agentproto.DRPCAgentClient29, proto.DRPCTailnetClient28, error,
|
||||
) {
|
||||
@@ -177,8 +178,9 @@ func (c *Client) ConnectRPC210(ctx context.Context) (
|
||||
return client, tAPI, nil
|
||||
}
|
||||
|
||||
func (c *Client) ConnectRPC210WithRole(ctx context.Context, _ string) (
|
||||
agentproto.DRPCAgentClient210, proto.DRPCTailnetClient28, error,
|
||||
// v2.11 adds no RPCs, so the v2.10 client satisfies it as-is.
|
||||
func (c *Client) ConnectRPC211WithRole(ctx context.Context, _ string) (
|
||||
agentproto.DRPCAgentClient211, proto.DRPCTailnetClient28, error,
|
||||
) {
|
||||
return c.ConnectRPC210(ctx)
|
||||
}
|
||||
|
||||
+1
-1
@@ -168,7 +168,7 @@ func handleRPTY(inv *serpent.Invocation, client *codersdk.Client, args handleRPT
|
||||
|
||||
closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, ws.ID, codersdk.PostWorkspaceUsageRequest{
|
||||
AgentID: agt.ID,
|
||||
AppName: codersdk.UsageAppNameReconnectingPty,
|
||||
AppName: string(codersdk.UsageAppNameReconnectingPty),
|
||||
})
|
||||
defer closeUsage()
|
||||
|
||||
|
||||
@@ -1115,6 +1115,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
|
||||
batcher, closeBatcher, err := workspacestats.NewBatcher(ctx,
|
||||
workspacestats.BatcherWithLogger(options.Logger.Named("batchstats")),
|
||||
workspacestats.BatcherWithStore(options.Database),
|
||||
workspacestats.BatcherWithRegisterer(options.PrometheusRegistry),
|
||||
)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("failed to create agent stats batcher: %w", err)
|
||||
|
||||
+10
-13
@@ -1532,21 +1532,18 @@ func (r stdioErrLogReader) Read(_ []byte) (int, error) {
|
||||
return 0, io.EOF
|
||||
}
|
||||
|
||||
func getUsageAppName(usageApp string) codersdk.UsageAppName {
|
||||
if usageApp == disableUsageApp {
|
||||
// getUsageAppName returns the app name to report usage under, or the empty
|
||||
// string to report none. Any name is valid because the server normalizes it
|
||||
// at ingestion.
|
||||
func getUsageAppName(usageApp string) string {
|
||||
switch usageApp {
|
||||
case disableUsageApp:
|
||||
return ""
|
||||
case "":
|
||||
return string(codersdk.UsageAppNameSSH)
|
||||
default:
|
||||
return usageApp
|
||||
}
|
||||
|
||||
allowedUsageApps := []string{
|
||||
string(codersdk.UsageAppNameSSH),
|
||||
string(codersdk.UsageAppNameVscode),
|
||||
string(codersdk.UsageAppNameJetbrains),
|
||||
}
|
||||
if slices.Contains(allowedUsageApps, usageApp) {
|
||||
return codersdk.UsageAppName(usageApp)
|
||||
}
|
||||
|
||||
return codersdk.UsageAppNameSSH
|
||||
}
|
||||
|
||||
func setStatsCallback(
|
||||
|
||||
+43
-35
@@ -1630,51 +1630,51 @@ func TestSSH(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
type testCase struct {
|
||||
name string
|
||||
experiment bool
|
||||
usageAppName string
|
||||
expectedCalls int
|
||||
expectedCountSSH int
|
||||
expectedCountJetbrains int
|
||||
expectedCountVscode int
|
||||
name string
|
||||
experiment bool
|
||||
usageAppName string
|
||||
expectedCalls int
|
||||
expectedCounts map[string]int64
|
||||
}
|
||||
tcs := []testCase{
|
||||
{
|
||||
name: "NoExperiment",
|
||||
},
|
||||
{
|
||||
name: "Empty",
|
||||
experiment: true,
|
||||
expectedCalls: 1,
|
||||
expectedCountSSH: 1,
|
||||
name: "Empty",
|
||||
experiment: true,
|
||||
expectedCalls: 1,
|
||||
expectedCounts: map[string]int64{"ssh": 1},
|
||||
},
|
||||
{
|
||||
name: "SSH",
|
||||
experiment: true,
|
||||
usageAppName: "ssh",
|
||||
expectedCalls: 1,
|
||||
expectedCountSSH: 1,
|
||||
name: "SSH",
|
||||
experiment: true,
|
||||
usageAppName: "ssh",
|
||||
expectedCalls: 1,
|
||||
expectedCounts: map[string]int64{"ssh": 1},
|
||||
},
|
||||
{
|
||||
name: "Jetbrains",
|
||||
experiment: true,
|
||||
usageAppName: "jetbrains",
|
||||
expectedCalls: 1,
|
||||
expectedCountJetbrains: 1,
|
||||
name: "Jetbrains",
|
||||
experiment: true,
|
||||
usageAppName: "jetbrains",
|
||||
expectedCalls: 1,
|
||||
expectedCounts: map[string]int64{"jetbrains": 1},
|
||||
},
|
||||
{
|
||||
name: "Vscode",
|
||||
experiment: true,
|
||||
usageAppName: "vscode",
|
||||
expectedCalls: 1,
|
||||
expectedCountVscode: 1,
|
||||
name: "Vscode",
|
||||
experiment: true,
|
||||
usageAppName: "vscode",
|
||||
expectedCalls: 1,
|
||||
expectedCounts: map[string]int64{"vscode": 1},
|
||||
},
|
||||
{
|
||||
name: "InvalidDefaultsToSSH",
|
||||
experiment: true,
|
||||
usageAppName: "invalid",
|
||||
expectedCalls: 1,
|
||||
expectedCountSSH: 1,
|
||||
// Arbitrary app names are accepted, normalized by the server
|
||||
// at ingestion, so new IDEs need no CLI changes.
|
||||
name: "ArbitraryNamePassthrough",
|
||||
experiment: true,
|
||||
usageAppName: "Some-Future-IDE",
|
||||
expectedCalls: 1,
|
||||
expectedCounts: map[string]int64{"some_future_ide": 1},
|
||||
},
|
||||
{
|
||||
name: "Disable",
|
||||
@@ -1729,10 +1729,18 @@ func TestSSH(t *testing.T) {
|
||||
stdin.WriteLine("exit")
|
||||
<-cmdDone
|
||||
|
||||
require.EqualValues(t, tc.expectedCalls, batcher.Called)
|
||||
require.EqualValues(t, tc.expectedCountSSH, batcher.LastStats.SessionCountSsh)
|
||||
require.EqualValues(t, tc.expectedCountJetbrains, batcher.LastStats.SessionCountJetbrains)
|
||||
require.EqualValues(t, tc.expectedCountVscode, batcher.LastStats.SessionCountVscode)
|
||||
// The agent may still be reporting stats concurrently, so
|
||||
// the batcher fields are only safe to read under its mutex.
|
||||
batcher.Mu.Lock()
|
||||
called := batcher.Called
|
||||
sessionCounts := batcher.LastStats.GetSessionCounts()
|
||||
batcher.Mu.Unlock()
|
||||
require.EqualValues(t, tc.expectedCalls, called)
|
||||
if len(tc.expectedCounts) == 0 {
|
||||
require.Empty(t, sessionCounts)
|
||||
} else {
|
||||
require.EqualValues(t, tc.expectedCounts, sessionCounts)
|
||||
}
|
||||
})
|
||||
}
|
||||
})
|
||||
|
||||
+1
-1
@@ -178,7 +178,7 @@ func (r *RootCmd) vscodeSSH() *serpent.Command {
|
||||
|
||||
closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{
|
||||
AgentID: workspaceAgent.ID,
|
||||
AppName: codersdk.UsageAppNameVscode,
|
||||
AppName: string(codersdk.UsageAppNameVscode),
|
||||
})
|
||||
defer closeUsage()
|
||||
|
||||
|
||||
@@ -87,5 +87,5 @@ func TestVSCodeSSH(t *testing.T) {
|
||||
}
|
||||
|
||||
require.EqualValues(t, 1, batcher.Called)
|
||||
require.EqualValues(t, 1, batcher.LastStats.SessionCountVscode)
|
||||
require.EqualValues(t, 1, batcher.LastStats.GetSessionCounts()["vscode"])
|
||||
}
|
||||
|
||||
@@ -56,10 +56,14 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
|
||||
ws = database.WorkspaceIdentityFromWorkspace(w)
|
||||
}
|
||||
|
||||
// The payload is unnormalized and unbounded until the batcher caps it, so
|
||||
// log scalars rather than its contents.
|
||||
a.Log.Debug(ctx, "read stats report",
|
||||
slog.F("interval", a.AgentStatsRefreshInterval),
|
||||
slog.F("workspace_id", ws.ID),
|
||||
slog.F("payload", req),
|
||||
slog.F("connection_count", req.Stats.GetConnectionCount()),
|
||||
slog.F("session_count_keys", len(req.Stats.GetSessionCounts())),
|
||||
slog.F("connections_by_proto_keys", len(req.Stats.GetConnectionsByProto())),
|
||||
)
|
||||
|
||||
if a.Experiments.Enabled(codersdk.ExperimentWorkspaceUsage) {
|
||||
|
||||
@@ -3,6 +3,9 @@ package agentapi_test
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -666,6 +669,133 @@ func TestUpdateStats(t *testing.T) {
|
||||
}
|
||||
require.True(t, updateAgentMetricsFnCalled)
|
||||
})
|
||||
|
||||
t.Run("SessionCounts", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Every case starts from 300 whitespace aliases of one app, which
|
||||
// normalize to a single name.
|
||||
vscodeAliases := func() map[string]int64 {
|
||||
counts := make(map[string]int64, 300)
|
||||
for i := range 300 {
|
||||
counts["VSCode"+strings.Repeat(" ", i)] = 1
|
||||
}
|
||||
return counts
|
||||
}
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
setup func() (counts, want map[string]int64)
|
||||
}{
|
||||
{
|
||||
name: "AliasesFitWithoutFolding",
|
||||
setup: func() (map[string]int64, map[string]int64) {
|
||||
return vscodeAliases(), map[string]int64{"vscode": 300}
|
||||
},
|
||||
},
|
||||
{
|
||||
// Known apps win even with smaller counts. Unknown aliases
|
||||
// merge before ranking, with name breaking count ties, and
|
||||
// the kept names fill the 64 cap.
|
||||
name: "RankAfterMerging",
|
||||
setup: func() (map[string]int64, map[string]int64) {
|
||||
counts := vscodeAliases()
|
||||
counts["SSH"] = 1
|
||||
counts["unknown"] = -1_000
|
||||
counts[" VSCODE"] = -1_000
|
||||
counts["idle"] = 0
|
||||
for i := range 300 {
|
||||
counts[fmt.Sprintf("app-%03d", i)] = 2
|
||||
counts["zzz-busy"+strings.Repeat(" ", i)] = 1
|
||||
}
|
||||
want := map[string]int64{"vscode": 300, "ssh": 1, "zzz_busy": 300, "unknown": 478}
|
||||
for i := range 61 {
|
||||
want[fmt.Sprintf("app_%03d", i)] = 2
|
||||
}
|
||||
return counts, want
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "NonPositiveSuppressesFallback",
|
||||
setup: func() (map[string]int64, map[string]int64) {
|
||||
counts := vscodeAliases()
|
||||
for key := range counts {
|
||||
counts[key] = -1
|
||||
}
|
||||
counts["unknown"] = 0
|
||||
return counts, map[string]int64{}
|
||||
},
|
||||
},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
counts, wantCounts := tc.setup()
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
dbM := dbmock.NewMockStore(gomock.NewController(t))
|
||||
inserted := make(chan database.InsertWorkspaceAgentStatsParams, 1)
|
||||
dbM.EXPECT().InsertWorkspaceAgentStats(gomock.Any(), gomock.Any()).Times(1).DoAndReturn(
|
||||
func(_ context.Context, params database.InsertWorkspaceAgentStatsParams) error {
|
||||
inserted <- params
|
||||
return nil
|
||||
})
|
||||
batcher, closeBatcher, err := workspacestats.NewBatcher(ctx,
|
||||
workspacestats.BatcherWithStore(dbM),
|
||||
workspacestats.BatcherWithLogger(testutil.Logger(t)),
|
||||
workspacestats.BatcherWithInterval(time.Hour),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(closeBatcher)
|
||||
api := agentapi.StatsAPI{
|
||||
AgentID: agent.ID,
|
||||
AgentName: agent.Name,
|
||||
Workspace: &workspaceAsCacheFields,
|
||||
Database: dbM,
|
||||
StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{
|
||||
StatsBatcher: batcher,
|
||||
}),
|
||||
AgentStatsRefreshInterval: 10 * time.Second,
|
||||
}
|
||||
_, err = api.UpdateStats(ctx, &agentproto.UpdateStatsRequest{Stats: &agentproto.Stats{
|
||||
SessionCounts: counts,
|
||||
// A present map must suppress deprecated fields even when
|
||||
// every reported count is non-positive.
|
||||
SessionCountVscode: 9,
|
||||
ConnectionsByProto: map[string]int64{"tcp": 1},
|
||||
ConnectionMedianLatencyMs: 23,
|
||||
RxPackets: 120,
|
||||
RxBytes: 1000,
|
||||
TxPackets: 130,
|
||||
TxBytes: 2000,
|
||||
}})
|
||||
require.NoError(t, err)
|
||||
closeBatcher() // Flush synchronously without waiting for a timer.
|
||||
|
||||
params := testutil.RequireReceive(ctx, t, inserted)
|
||||
wantSessionCounts, err := json.Marshal([]map[string]int64{wantCounts})
|
||||
require.NoError(t, err)
|
||||
// ID and CreatedAt are generated inside the batcher; the rest
|
||||
// must survive capping untouched.
|
||||
require.Equal(t, database.InsertWorkspaceAgentStatsParams{
|
||||
ID: params.ID,
|
||||
CreatedAt: params.CreatedAt,
|
||||
UserID: []uuid.UUID{workspace.OwnerID},
|
||||
WorkspaceID: []uuid.UUID{workspace.ID},
|
||||
TemplateID: []uuid.UUID{workspace.TemplateID},
|
||||
AgentID: []uuid.UUID{agent.ID},
|
||||
ConnectionsByProto: json.RawMessage(`[{"tcp":1}]`),
|
||||
ConnectionCount: []int64{0},
|
||||
RxPackets: []int64{120},
|
||||
RxBytes: []int64{1000},
|
||||
TxPackets: []int64{130},
|
||||
TxBytes: []int64{2000},
|
||||
SessionCounts: wantSessionCounts,
|
||||
ConnectionMedianLatencyMS: []float64{23},
|
||||
Usage: []bool{false},
|
||||
}, params)
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func templateScheduleStorePtr(store schedule.TemplateScheduleStore) *atomic.Pointer[schedule.TemplateScheduleStore] {
|
||||
|
||||
Generated
+2
-16
@@ -26596,7 +26596,8 @@ const docTemplate = `{
|
||||
"format": "uuid"
|
||||
},
|
||||
"app_name": {
|
||||
"$ref": "#/definitions/codersdk.UsageAppName"
|
||||
"description": "AppName is any name for the app reporting usage. The server normalizes\nit at ingestion, so a new app needs no server change. The UsageAppName\nconstants are the well-known names.",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -30411,21 +30412,6 @@ const docTemplate = `{
|
||||
}
|
||||
}
|
||||
},
|
||||
"codersdk.UsageAppName": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"vscode",
|
||||
"jetbrains",
|
||||
"reconnecting-pty",
|
||||
"ssh"
|
||||
],
|
||||
"x-enum-varnames": [
|
||||
"UsageAppNameVscode",
|
||||
"UsageAppNameJetbrains",
|
||||
"UsageAppNameReconnectingPty",
|
||||
"UsageAppNameSSH"
|
||||
]
|
||||
},
|
||||
"codersdk.UsagePeriod": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
|
||||
Generated
+2
-11
@@ -24374,7 +24374,8 @@
|
||||
"format": "uuid"
|
||||
},
|
||||
"app_name": {
|
||||
"$ref": "#/definitions/codersdk.UsageAppName"
|
||||
"description": "AppName is any name for the app reporting usage. The server normalizes\nit at ingestion, so a new app needs no server change. The UsageAppName\nconstants are the well-known names.",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -28001,16 +28002,6 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"codersdk.UsageAppName": {
|
||||
"type": "string",
|
||||
"enum": ["vscode", "jetbrains", "reconnecting-pty", "ssh"],
|
||||
"x-enum-varnames": [
|
||||
"UsageAppNameVscode",
|
||||
"UsageAppNameJetbrains",
|
||||
"UsageAppNameReconnectingPty",
|
||||
"UsageAppNameSSH"
|
||||
]
|
||||
},
|
||||
"codersdk.UsagePeriod": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
|
||||
+22
-27
@@ -11,6 +11,7 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode"
|
||||
|
||||
"github.com/dustin/go-humanize"
|
||||
"github.com/go-chi/chi/v5"
|
||||
@@ -1766,7 +1767,11 @@ func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
if req.AgentID == uuid.Nil && req.AppName == "" {
|
||||
// Normalize at the edge so storage and lookup agree on the key, and so a
|
||||
// name that carries no information reads the same as an absent one.
|
||||
appName := normalizeUsageAppName(req.AppName)
|
||||
|
||||
if req.AgentID == uuid.Nil && appName == "" {
|
||||
// Continue previous behavior if body is empty.
|
||||
rw.WriteHeader(http.StatusNoContent)
|
||||
return
|
||||
@@ -1781,7 +1786,7 @@ func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
|
||||
})
|
||||
return
|
||||
}
|
||||
if req.AppName == "" {
|
||||
if appName == "" {
|
||||
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
|
||||
Message: "Invalid request",
|
||||
Validations: []codersdk.ValidationError{{
|
||||
@@ -1791,34 +1796,10 @@ func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
|
||||
})
|
||||
return
|
||||
}
|
||||
if !slices.Contains(codersdk.AllowedAppNames, req.AppName) {
|
||||
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
|
||||
Message: "Invalid request",
|
||||
Validations: []codersdk.ValidationError{{
|
||||
Field: "app_name",
|
||||
Detail: fmt.Sprintf("must be one of %v", codersdk.AllowedAppNames),
|
||||
}},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
stat := &proto.Stats{
|
||||
ConnectionCount: 1,
|
||||
}
|
||||
switch req.AppName {
|
||||
case codersdk.UsageAppNameVscode:
|
||||
stat.SessionCountVscode = 1
|
||||
case codersdk.UsageAppNameJetbrains:
|
||||
stat.SessionCountJetbrains = 1
|
||||
case codersdk.UsageAppNameReconnectingPty:
|
||||
stat.SessionCountReconnectingPty = 1
|
||||
case codersdk.UsageAppNameSSH:
|
||||
stat.SessionCountSsh = 1
|
||||
default:
|
||||
// This means the app_name is in the codersdk.AllowedAppNames but not being
|
||||
// handled by this switch statement.
|
||||
httpapi.InternalServerError(rw, xerrors.Errorf("unknown app_name %q", req.AppName))
|
||||
return
|
||||
SessionCounts: map[string]int64{appName: 1},
|
||||
}
|
||||
|
||||
agent, err := api.Database.GetWorkspaceAgentByID(ctx, req.AgentID)
|
||||
@@ -1846,6 +1827,20 @@ func (api *API) postWorkspaceUsage(rw http.ResponseWriter, r *http.Request) {
|
||||
rw.WriteHeader(http.StatusNoContent)
|
||||
}
|
||||
|
||||
// normalizeUsageAppName prepares a client-supplied app name for storage. A
|
||||
// name of only whitespace and control characters carries no app, so it
|
||||
// returns the empty string and the caller rejects the request rather than
|
||||
// counting a session under the unknown family.
|
||||
func normalizeUsageAppName(appName string) string {
|
||||
named := strings.ContainsFunc(appName, func(r rune) bool {
|
||||
return !unicode.IsControl(r) && !unicode.IsSpace(r)
|
||||
})
|
||||
if !named {
|
||||
return ""
|
||||
}
|
||||
return codersdk.NormalizeAppName(appName)
|
||||
}
|
||||
|
||||
// @Summary Favorite workspace by ID.
|
||||
// @ID favorite-workspace-by-id
|
||||
// @Security CoderSessionToken
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
package coderd
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestNormalizeUsageAppName(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
appName string
|
||||
want string
|
||||
}{
|
||||
{"Empty", "", ""},
|
||||
{"Whitespace", " \t\n ", ""},
|
||||
{"NullBytes", "\x00\x00", ""},
|
||||
// Control characters carry no name, so they must read as an absent
|
||||
// app name rather than a session under the unknown family.
|
||||
{"ControlOnly", "\x1b\x07\x01", ""},
|
||||
{"ControlAndWhitespace", " \x1b \x07\t", ""},
|
||||
{"EscapeSequenceKeepsPrintableText", "\x1b[31m", "[31m"},
|
||||
{"Normalized", "Some-Future-IDE", "some_future_ide"},
|
||||
{"ControlStripped", "cur\x1bsor", "cursor"},
|
||||
{"LongWhitespace", strings.Repeat(" ", 300), ""},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
require.Equal(t, tc.want, normalizeUsageAppName(tc.appName))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -45,6 +45,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/schedule/cron"
|
||||
"github.com/coder/coder/v2/coderd/util/ptr"
|
||||
"github.com/coder/coder/v2/coderd/util/slice"
|
||||
"github.com/coder/coder/v2/coderd/workspacestats"
|
||||
"github.com/coder/coder/v2/coderd/wsbuilder"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/cryptorand"
|
||||
@@ -5000,12 +5001,29 @@ func TestWorkspaceUsageTracking(t *testing.T) {
|
||||
AppName: "ssh",
|
||||
})
|
||||
require.ErrorContains(t, err, "app_name")
|
||||
// unknown app name fails
|
||||
// unknown app names are accepted
|
||||
err = client.PostWorkspaceUsageWithBody(ctx, r.Workspace.ID, codersdk.PostWorkspaceUsageRequest{
|
||||
AgentID: workspace.LatestBuild.Resources[0].Agents[0].ID,
|
||||
AppName: "unknown",
|
||||
AppName: "SomeFutureIDE",
|
||||
})
|
||||
require.ErrorContains(t, err, "app_name")
|
||||
require.NoError(t, err)
|
||||
|
||||
// an app name of only whitespace or control characters, at any
|
||||
// length, reads as an absent app name: the usage bump still happens
|
||||
// and no session is counted
|
||||
for _, appName := range []string{" ", strings.Repeat(" ", 300), "\x00", "\x1b\x07", " \x1b \t"} {
|
||||
err = client.PostWorkspaceUsageWithBody(ctx, r.Workspace.ID, codersdk.PostWorkspaceUsageRequest{
|
||||
AppName: appName,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// with an agent set, those names fail like an empty one
|
||||
err = client.PostWorkspaceUsageWithBody(ctx, r.Workspace.ID, codersdk.PostWorkspaceUsageRequest{
|
||||
AgentID: workspace.LatestBuild.Resources[0].Agents[0].ID,
|
||||
AppName: appName,
|
||||
})
|
||||
require.ErrorContains(t, err, "app_name")
|
||||
}
|
||||
|
||||
// vscode works
|
||||
err = client.PostWorkspaceUsageWithBody(ctx, r.Workspace.ID, codersdk.PostWorkspaceUsageRequest{
|
||||
@@ -5041,6 +5059,70 @@ func TestWorkspaceUsageTracking(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
// TestWorkspaceUsageArbitraryAppNameNormalized posts an arbitrary app name
|
||||
// through a real batcher and asserts the database stores the normalized key,
|
||||
// so readers of session_counts only ever see canonical names.
|
||||
func TestWorkspaceUsageArbitraryAppNameNormalized(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
store, ps, sqlDB := dbtestutil.NewDBWithSQLDB(t)
|
||||
batcher, closeBatcher, err := workspacestats.NewBatcher(ctx,
|
||||
workspacestats.BatcherWithStore(store),
|
||||
workspacestats.BatcherWithLogger(testutil.Logger(t).Named("batcher")),
|
||||
workspacestats.BatcherWithInterval(testutil.IntervalFast),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(closeBatcher)
|
||||
|
||||
dv := coderdtest.DeploymentValues(t)
|
||||
dv.Experiments = []string{string(codersdk.ExperimentWorkspaceUsage)}
|
||||
client := coderdtest.New(t, &coderdtest.Options{
|
||||
Database: store,
|
||||
Pubsub: ps,
|
||||
DeploymentValues: dv,
|
||||
StatsBatcher: batcher,
|
||||
})
|
||||
user := coderdtest.CreateFirstUser(t, client)
|
||||
r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{
|
||||
OrganizationID: user.OrganizationID,
|
||||
OwnerID: user.UserID,
|
||||
Name: fmt.Sprintf("usage-%d", time.Now().UnixNano()),
|
||||
}).WithAgent().Do()
|
||||
require.Len(t, r.Agents, 1)
|
||||
agentID := r.Agents[0].ID
|
||||
|
||||
err = client.PostWorkspaceUsageWithBody(ctx, r.Workspace.ID, codersdk.PostWorkspaceUsageRequest{
|
||||
AgentID: agentID,
|
||||
AppName: "Some-Future-IDE",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// The batcher flushes on its own schedule, so poll until the row lands.
|
||||
var count int64
|
||||
require.True(t, testutil.Eventually(ctx, t, func(ctx context.Context) bool {
|
||||
err := sqlDB.QueryRowContext(ctx, `
|
||||
SELECT coalesce(sum((session_counts->>'some_future_ide')::bigint), 0)
|
||||
FROM workspace_agent_stats
|
||||
WHERE agent_id = $1`, agentID).Scan(&count)
|
||||
if err != nil {
|
||||
t.Logf("query session counts: %s", err)
|
||||
return false
|
||||
}
|
||||
return count > 0
|
||||
}, testutil.IntervalFast), "expected the normalized app name to reach the database")
|
||||
require.EqualValues(t, 1, count)
|
||||
|
||||
// The raw name must never reach storage as a key.
|
||||
var rawKeyRows int64
|
||||
err = sqlDB.QueryRowContext(ctx, `
|
||||
SELECT count(*)
|
||||
FROM workspace_agent_stats
|
||||
WHERE agent_id = $1 AND session_counts ? 'Some-Future-IDE'`, agentID).Scan(&rawKeyRows)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, rawKeyRows)
|
||||
}
|
||||
|
||||
func TestWorkspaceNotifications(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
@@ -51,6 +52,10 @@ type DBBatcher struct {
|
||||
flushForced atomic.Bool
|
||||
// flushed is used during testing to signal that a flush has completed.
|
||||
flushed chan<- int
|
||||
|
||||
metrics batcherMetrics
|
||||
// registerer is nil unless BatcherWithRegisterer is passed.
|
||||
registerer prometheus.Registerer
|
||||
}
|
||||
|
||||
// Option is a functional option for configuring a Batcher.
|
||||
@@ -84,15 +89,25 @@ func BatcherWithLogger(log slog.Logger) BatcherOption {
|
||||
}
|
||||
}
|
||||
|
||||
// BatcherWithRegisterer sets the Prometheus registerer for batcher metrics.
|
||||
func BatcherWithRegisterer(reg prometheus.Registerer) BatcherOption {
|
||||
return func(b *DBBatcher) {
|
||||
b.registerer = reg
|
||||
}
|
||||
}
|
||||
|
||||
// NewBatcher creates a new Batcher and starts it.
|
||||
func NewBatcher(ctx context.Context, opts ...BatcherOption) (*DBBatcher, func(), error) {
|
||||
b := &DBBatcher{}
|
||||
b.log = slog.Make(sloghuman.Sink(os.Stderr))
|
||||
b.metrics = newBatcherMetrics()
|
||||
b.flushLever = make(chan struct{}, 1) // Buffered so that it doesn't block.
|
||||
for _, opt := range opts {
|
||||
opt(b)
|
||||
}
|
||||
|
||||
b.metrics.register(b.registerer)
|
||||
|
||||
if b.store == nil {
|
||||
return nil, nil, xerrors.Errorf("no store configured for batcher")
|
||||
}
|
||||
@@ -141,15 +156,17 @@ func (b *DBBatcher) Add(
|
||||
usage bool,
|
||||
) {
|
||||
// Normalize and cap outside the lock.
|
||||
sessionCounts := normalizedSessionCounts(st)
|
||||
if len(sessionCounts) > maxSessionCountEntries {
|
||||
b.log.Warn(context.Background(), "too many distinct session types, overflow counted under unknown",
|
||||
sessionCounts, overflow := capSessionCounts(normalizedSessionCounts(st))
|
||||
if overflow > 0 {
|
||||
// A misbehaving agent hits this on every report, so alert on the
|
||||
// counter and keep the log at debug.
|
||||
b.log.Debug(context.Background(), "too many distinct session types, overflow counted under unknown",
|
||||
slog.F("agent_id", agentID),
|
||||
slog.F("reported", len(sessionCounts)),
|
||||
slog.F("overflow", overflow),
|
||||
slog.F("max", maxSessionCountEntries),
|
||||
)
|
||||
b.metrics.SessionCountsOverflowTotal.Add(float64(overflow))
|
||||
}
|
||||
sessionCounts = capSessionCounts(sessionCounts)
|
||||
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
|
||||
@@ -3,10 +3,13 @@ package workspacestats
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
prom_testutil "github.com/prometheus/client_golang/prometheus/testutil"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
@@ -19,6 +22,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/database/pubsub"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/cryptorand"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
func TestBatchStats(t *testing.T) {
|
||||
@@ -146,6 +150,46 @@ func TestBatchStats(t *testing.T) {
|
||||
require.Equal(t, defaultBufferSize, cap(b.buf.ID), "buffer grew beyond expected capacity")
|
||||
}
|
||||
|
||||
// TestBatchStatsSessionCountFold asserts that a report over the session count
|
||||
// cap folds the excess into the overflow counter and logs it at debug, never
|
||||
// at warn, so one misbehaving agent cannot flood the logs.
|
||||
func TestBatchStatsSessionCountFold(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
sink := testutil.NewFakeSink(t)
|
||||
store, ps := dbtestutil.NewDB(t)
|
||||
deps := setupDeps(t, store, ps)
|
||||
|
||||
b, closer, err := NewBatcher(ctx,
|
||||
BatcherWithStore(store),
|
||||
BatcherWithLogger(sink.Logger(slog.LevelDebug)),
|
||||
func(b *DBBatcher) {
|
||||
// Take control of flushes so they do not interleave with the
|
||||
// assertions below.
|
||||
b.tickCh = make(chan time.Time)
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(closer)
|
||||
|
||||
const extra = 6
|
||||
st := randStats(t, func(s *agentproto.Stats) {
|
||||
s.SessionCounts = make(map[string]int64, maxSessionCountEntries+extra)
|
||||
for i := range maxSessionCountEntries + extra {
|
||||
s.SessionCounts[fmt.Sprintf("app_%d", i)] = 1
|
||||
}
|
||||
})
|
||||
b.Add(dbtime.Now(), deps.Agent.ID, deps.Template.ID, deps.User.ID, deps.Workspace.ID, st, false)
|
||||
|
||||
overcap := sink.Entries(func(e slog.SinkEntry) bool {
|
||||
return strings.Contains(e.Message, "too many distinct session types")
|
||||
})
|
||||
require.Len(t, overcap, 1)
|
||||
require.Equal(t, slog.LevelDebug, overcap[0].Level)
|
||||
require.Equal(t, float64(extra), prom_testutil.ToFloat64(b.metrics.SessionCountsOverflowTotal))
|
||||
}
|
||||
|
||||
// randStats returns a random agentproto.Stats
|
||||
func randStats(t *testing.T, opts ...func(*agentproto.Stats)) *agentproto.Stats {
|
||||
t.Helper()
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
package workspacestats
|
||||
|
||||
import (
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
)
|
||||
|
||||
// batcherMetrics collects Prometheus metrics for the stats batcher.
|
||||
type batcherMetrics struct {
|
||||
SessionCountsOverflowTotal prometheus.Counter
|
||||
}
|
||||
|
||||
func newBatcherMetrics() batcherMetrics {
|
||||
return batcherMetrics{
|
||||
SessionCountsOverflowTotal: prometheus.NewCounter(prometheus.CounterOpts{
|
||||
Namespace: "coderd",
|
||||
Subsystem: "agentstats",
|
||||
Name: "session_counts_overflow_total",
|
||||
Help: "Total number of reported session count entries summed into the unknown app after exceeding the per-report cap.",
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
func (m batcherMetrics) register(reg prometheus.Registerer) {
|
||||
if reg != nil {
|
||||
reg.MustRegister(m.SessionCountsOverflowTotal)
|
||||
}
|
||||
}
|
||||
@@ -41,19 +41,22 @@ func normalizedSessionCounts(st *agentproto.Stats) map[string]int64 {
|
||||
// aggregates under AppFamilyUnknown.
|
||||
const maxSessionCountEntries = 64
|
||||
|
||||
// capSessionCounts keeps the busiest maxSessionCountEntries names, preferring
|
||||
// known apps, and sums the rest into AppFamilyUnknown, so the result can hold
|
||||
// one name past the cap.
|
||||
func capSessionCounts(counts map[string]int64) map[string]int64 {
|
||||
// capSessionCounts keeps the busiest maxSessionCountEntries normalized names,
|
||||
// preferring known apps, and sums the rest into AppFamilyUnknown, so the
|
||||
// result can hold one name past the cap. The second return counts the names
|
||||
// summed away.
|
||||
func capSessionCounts(counts map[string]int64) (map[string]int64, int) {
|
||||
if len(counts) <= maxSessionCountEntries {
|
||||
return counts
|
||||
return counts, 0
|
||||
}
|
||||
// Known apps rank 0, unknown 1, so known apps win the cap.
|
||||
// Names are already normalized, so look up families directly rather
|
||||
// than normalizing again on every sort comparison.
|
||||
families := codersdk.SessionCountAppFamilies()
|
||||
rank := func(name string) int {
|
||||
if codersdk.AppNameFamily(name) == codersdk.AppFamilyUnknown {
|
||||
return 1
|
||||
if family, ok := families[name]; ok && family != codersdk.AppFamilyUnknown {
|
||||
return 0
|
||||
}
|
||||
return 0
|
||||
return 1
|
||||
}
|
||||
ranked := slices.SortedFunc(maps.Keys(counts), func(a, b string) int {
|
||||
return cmp.Or(
|
||||
@@ -66,10 +69,11 @@ func capSessionCounts(counts map[string]int64) map[string]int64 {
|
||||
for _, name := range ranked[:maxSessionCountEntries] {
|
||||
capped[name] = counts[name]
|
||||
}
|
||||
for _, name := range ranked[maxSessionCountEntries:] {
|
||||
overflow := ranked[maxSessionCountEntries:]
|
||||
for _, name := range overflow {
|
||||
capped[string(codersdk.AppFamilyUnknown)] += counts[name]
|
||||
}
|
||||
return capped
|
||||
return capped, len(overflow)
|
||||
}
|
||||
|
||||
// HasSessionCounts reports whether the stats contain any active session.
|
||||
|
||||
@@ -11,7 +11,8 @@ import (
|
||||
|
||||
// sessionCountsFromProto is what the batcher applies on ingest.
|
||||
func sessionCountsFromProto(st *agentproto.Stats) map[string]int64 {
|
||||
return capSessionCounts(normalizedSessionCounts(st))
|
||||
counts, _ := capSessionCounts(normalizedSessionCounts(st))
|
||||
return counts
|
||||
}
|
||||
|
||||
func TestSessionCountsFromProto(t *testing.T) {
|
||||
|
||||
@@ -349,19 +349,6 @@ func (c *Client) ConnectRPC210(ctx context.Context) (
|
||||
return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil
|
||||
}
|
||||
|
||||
// ConnectRPC210WithRole is like ConnectRPC210 but sends an explicit role
|
||||
// query parameter to the server. Use "agent" for workspace agents to
|
||||
// enable connection monitoring.
|
||||
func (c *Client) ConnectRPC210WithRole(ctx context.Context, role string) (
|
||||
proto.DRPCAgentClient210, tailnetproto.DRPCTailnetClient28, error,
|
||||
) {
|
||||
conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 10), role)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil
|
||||
}
|
||||
|
||||
// ConnectRPC211WithRole returns a dRPC client to the Agent API v2.11, which
|
||||
// reports per-app session counts on Stats. Pass role "agent" for workspace
|
||||
// agents to enable connection monitoring.
|
||||
|
||||
+11
-4
@@ -5,6 +5,7 @@ import (
|
||||
"maps"
|
||||
"slices"
|
||||
"strings"
|
||||
"unicode"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/xerrors"
|
||||
@@ -120,11 +121,17 @@ func AppNameFamily(appName string) AppFamilyName {
|
||||
}
|
||||
|
||||
// NormalizeAppName prepares a client-supplied app name for storage and
|
||||
// lookup: it strips the null bytes Postgres TEXT rejects, trims, truncates,
|
||||
// lowercases, and folds hyphens to underscores. Empty becomes
|
||||
// AppFamilyUnknown.
|
||||
// lookup: it strips control characters, then trims, truncates, lowercases,
|
||||
// and folds hyphens to underscores. Empty becomes AppFamilyUnknown.
|
||||
// Stripping control characters covers both the null bytes Postgres TEXT
|
||||
// rejects and escape sequences that would otherwise reach logs and terminals.
|
||||
func NormalizeAppName(appName string) string {
|
||||
appName = strings.ReplaceAll(appName, "\x00", "")
|
||||
appName = strings.Map(func(r rune) rune {
|
||||
if unicode.IsControl(r) {
|
||||
return -1
|
||||
}
|
||||
return r
|
||||
}, appName)
|
||||
// Trim before truncating so padding does not spend the budget, and after
|
||||
// in case the cut lands in whitespace.
|
||||
appName = strings.TrimSpace(appName)
|
||||
|
||||
@@ -31,6 +31,11 @@ func TestNormalizeAppName(t *testing.T) {
|
||||
{"UnknownLowercased", "Cursor Nightly", "cursor nightly"},
|
||||
{"UnknownPreservesUnicode", "エディタ", "エディタ"},
|
||||
{"StripsNullBytes", "cur\x00sor", "cursor"},
|
||||
{"StripsControlCharacters", "vs\ncode\r\t", "vscode"},
|
||||
// Only the escape byte itself is a control character, so the rest of
|
||||
// an ANSI sequence survives as ordinary text.
|
||||
{"StripsANSIEscape", "\x1b[31mvscode\x1b[0m", "[31mvscode[0m"},
|
||||
{"OnlyControlCharacters", "\n\r\t\x1b", "unknown"},
|
||||
{"TrimsWhitespace", " vscode\t", "vscode"},
|
||||
{"Empty", "", "unknown"},
|
||||
{"OnlyNullBytes", "\x00\x00", "unknown"},
|
||||
|
||||
@@ -376,12 +376,18 @@ func (c *Client) PutExtendWorkspace(ctx context.Context, id uuid.UUID, req PutEx
|
||||
}
|
||||
|
||||
type PostWorkspaceUsageRequest struct {
|
||||
AgentID uuid.UUID `json:"agent_id" format:"uuid"`
|
||||
AppName UsageAppName `json:"app_name"`
|
||||
AgentID uuid.UUID `json:"agent_id" format:"uuid"`
|
||||
// AppName is any name for the app reporting usage. The server normalizes
|
||||
// it at ingestion, so a new app needs no server change. The UsageAppName
|
||||
// constants are the well-known names.
|
||||
AppName string `json:"app_name"`
|
||||
}
|
||||
|
||||
type UsageAppName string
|
||||
|
||||
// Well-known usage app names. The values are wire format and cannot change.
|
||||
// UsageAppNameReconnectingPty keeps its hyphen, which the server folds to the
|
||||
// canonical reconnecting_pty.
|
||||
const (
|
||||
UsageAppNameVscode UsageAppName = "vscode"
|
||||
UsageAppNameJetbrains UsageAppName = "jetbrains"
|
||||
@@ -389,13 +395,6 @@ const (
|
||||
UsageAppNameSSH UsageAppName = "ssh"
|
||||
)
|
||||
|
||||
var AllowedAppNames = []UsageAppName{
|
||||
UsageAppNameVscode,
|
||||
UsageAppNameJetbrains,
|
||||
UsageAppNameReconnectingPty,
|
||||
UsageAppNameSSH,
|
||||
}
|
||||
|
||||
// PostWorkspaceUsage marks the workspace as having been used recently and records an app stat.
|
||||
func (c *Client) PostWorkspaceUsageWithBody(ctx context.Context, id uuid.UUID, req PostWorkspaceUsageRequest) error {
|
||||
path := fmt.Sprintf("/api/v2/workspaces/%s/usage", id.String())
|
||||
|
||||
@@ -218,6 +218,7 @@ The `coder_ai_gateway_cost_control_*` metrics are exported only by `coderd`.
|
||||
| `coderd_agentstats_session_count_reconnecting_pty` | gauge | The number of session established by reconnecting PTY | `agent_name` `username` `workspace_name` |
|
||||
| `coderd_agentstats_session_count_ssh` | gauge | The number of session established by SSH | `agent_name` `username` `workspace_name` |
|
||||
| `coderd_agentstats_session_count_vscode` | gauge | The number of session established by VSCode | `agent_name` `username` `workspace_name` |
|
||||
| `coderd_agentstats_session_counts_overflow_total` | counter | Total number of reported session count entries summed into the unknown app after exceeding the per-report cap. | |
|
||||
| `coderd_agentstats_startup_script_seconds` | gauge | The number of seconds the startup script took to execute. | `agent_name` `success` `template_name` `username` `workspace_name` |
|
||||
| `coderd_agentstats_tx_bytes` | gauge | Agent Tx bytes | `agent_name` `username` `workspace_name` |
|
||||
| `coderd_api_active_users_duration_hour` | gauge | The number of users that have been active within the last hour. | |
|
||||
|
||||
Generated
+5
-19
@@ -12210,16 +12210,16 @@ Git clone makes use of this by parsing the URL from: 'Username for "https://gith
|
||||
```json
|
||||
{
|
||||
"agent_id": "2b1e3b65-2c04-4fa2-a2d7-467901e98978",
|
||||
"app_name": "vscode"
|
||||
"app_name": "string"
|
||||
}
|
||||
```
|
||||
|
||||
### Properties
|
||||
|
||||
| Name | Type | Required | Restrictions | Description |
|
||||
|------------|------------------------------------------------|----------|--------------|-------------|
|
||||
| `agent_id` | string | false | | |
|
||||
| `app_name` | [codersdk.UsageAppName](#codersdkusageappname) | false | | |
|
||||
| Name | Type | Required | Restrictions | Description |
|
||||
|------------|--------|----------|--------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| `agent_id` | string | false | | |
|
||||
| `app_name` | string | false | | App name is any name for the app reporting usage. The server normalizes it at ingestion, so a new app needs no server change. The UsageAppName constants are the well-known names. |
|
||||
|
||||
## codersdk.PprofConfig
|
||||
|
||||
@@ -16477,20 +16477,6 @@ If the schedule is empty, the user will be updated to use the default schedule.|
|
||||
| `protocol` | `http`, `https` |
|
||||
| `share_level` | `authenticated`, `organization`, `owner`, `public` |
|
||||
|
||||
## codersdk.UsageAppName
|
||||
|
||||
```json
|
||||
"vscode"
|
||||
```
|
||||
|
||||
### Properties
|
||||
|
||||
#### Enumerated Values
|
||||
|
||||
| Value(s) |
|
||||
|--------------------------------------------------|
|
||||
| `jetbrains`, `reconnecting-pty`, `ssh`, `vscode` |
|
||||
|
||||
## codersdk.UsagePeriod
|
||||
|
||||
```json
|
||||
|
||||
Generated
+1
-1
@@ -2568,7 +2568,7 @@ curl -X POST http://coder-server:8080/api/v2/workspaces/{workspace}/usage \
|
||||
```json
|
||||
{
|
||||
"agent_id": "2b1e3b65-2c04-4fa2-a2d7-467901e98978",
|
||||
"app_name": "vscode"
|
||||
"app_name": "string"
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
@@ -226,6 +226,9 @@ coderd_agentstats_session_count_ssh 0
|
||||
# HELP coderd_agentstats_session_count_vscode The number of session established by VSCode
|
||||
# TYPE coderd_agentstats_session_count_vscode gauge
|
||||
coderd_agentstats_session_count_vscode 0
|
||||
# HELP coderd_agentstats_session_counts_overflow_total Total number of reported session count entries summed into the unknown app after exceeding the per-report cap.
|
||||
# TYPE coderd_agentstats_session_counts_overflow_total counter
|
||||
coderd_agentstats_session_counts_overflow_total 0
|
||||
# HELP coderd_agentstats_startup_script_seconds Amount of time taken to run the startup script in seconds.
|
||||
# TYPE coderd_agentstats_startup_script_seconds gauge
|
||||
coderd_agentstats_startup_script_seconds{success=""} 0
|
||||
|
||||
Generated
+6
-1
@@ -7507,7 +7507,12 @@ export interface PostOAuth2ProviderAppRequest {
|
||||
// From codersdk/workspaces.go
|
||||
export interface PostWorkspaceUsageRequest {
|
||||
readonly agent_id: string;
|
||||
readonly app_name: UsageAppName;
|
||||
/**
|
||||
* AppName is any name for the app reporting usage. The server normalizes
|
||||
* it at ingestion, so a new app needs no server change. The UsageAppName
|
||||
* constants are the well-known names.
|
||||
*/
|
||||
readonly app_name: string;
|
||||
}
|
||||
|
||||
// From codersdk/deployment.go
|
||||
|
||||
Reference in New Issue
Block a user