Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 12 additions & 23 deletions agent/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -128,23 +128,14 @@ type Options struct {
}

type Client interface {
ConnectRPC29(ctx context.Context) (
proto.DRPCAgentClient29, tailnetproto.DRPCTailnetClient28, error,
ConnectRPC211(ctx context.Context) (
proto.DRPCAgentClient211, tailnetproto.DRPCTailnetClient28, error,
)
// ConnectRPC29WithRole is like ConnectRPC29 but sends an explicit
// ConnectRPC211WithRole is like ConnectRPC211 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(ctx context.Context, role string) (
proto.DRPCAgentClient211, tailnetproto.DRPCTailnetClient28, error,
)
tailnet.DERPMapRewriter
agentsdk.RefreshableSessionTokenProvider
Expand Down Expand Up @@ -1175,7 +1166,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
}
Expand Down Expand Up @@ -2161,14 +2152,12 @@ func (a *agent) Collect(ctx context.Context, networkStats map[netlogtype.Connect
stats.TxPackets += int64(counts.TxPackets)
}

// The count of active sessions. Only the four canonical names the protocol
// has fields for are reported; other app names are dropped, as before.
sessionCounts := a.sshServer.SessionCounts()
stats.SessionCountSsh = sessionCounts[idemetadata.AppNameSSH]
stats.SessionCountVscode = sessionCounts[idemetadata.AppNameVSCode]
stats.SessionCountJetbrains = sessionCounts[idemetadata.AppNameJetBrains]

stats.SessionCountReconnectingPty = a.reconnectingPTYServer.ConnCount()
// Active sessions per app; the deprecated fields stay zero. A client may
// label an ssh session "reconnecting_pty", so add, don't overwrite.
stats.SessionCounts = a.sshServer.SessionCounts()
if count := a.reconnectingPTYServer.ConnCount(); count > 0 {
stats.SessionCounts[idemetadata.AppNameReconnectingPTY] += count
}

// Compute the median connection latency!
a.logger.Debug(ctx, "starting peer latency measurement for stats")
Expand Down
20 changes: 10 additions & 10 deletions agent/agent_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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",
)
Expand All @@ -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",
)
Expand Down
20 changes: 20 additions & 0 deletions agent/agenttest/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,26 @@ func (c *Client) ConnectRPC210WithRole(ctx context.Context, _ string) (
return c.ConnectRPC210(ctx)
}

func (c *Client) ConnectRPC211(ctx context.Context) (
agentproto.DRPCAgentClient211, proto.DRPCTailnetClient28, error,
) {
aAPI, tAPI, err := c.ConnectRPC210(ctx)
if err != nil {
return nil, nil, err
}
client, ok := aAPI.(agentproto.DRPCAgentClient211)
if !ok {
return nil, nil, xerrors.Errorf("agenttest: connection does not implement DRPCAgentClient211; got %T", aAPI)
}
return client, tAPI, nil
}

func (c *Client) ConnectRPC211WithRole(ctx context.Context, _ string) (
agentproto.DRPCAgentClient211, proto.DRPCTailnetClient28, error,
) {
return c.ConnectRPC211(ctx)
}

func (c *Client) ConnectRPC29(ctx context.Context) (
agentproto.DRPCAgentClient29, proto.DRPCTailnetClient28, error,
) {
Expand Down
13 changes: 4 additions & 9 deletions cli/ssh.go
Original file line number Diff line number Diff line change
Expand Up @@ -1532,17 +1532,12 @@ func getUsageAppName(usageApp string) codersdk.UsageAppName {
if usageApp == disableUsageApp {
return ""
}

allowedUsageApps := []string{
string(codersdk.UsageAppNameSSH),
string(codersdk.UsageAppNameVscode),
string(codersdk.UsageAppNameJetbrains),
}
if slices.Contains(allowedUsageApps, usageApp) {
return codersdk.UsageAppName(usageApp)
if usageApp == "" {
return codersdk.UsageAppNameSSH
}

return codersdk.UsageAppNameSSH
// The server accepts arbitrary app names.
return codersdk.UsageAppName(usageApp)
}

func setStatsCallback(
Expand Down
70 changes: 36 additions & 34 deletions cli/ssh_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 pass through raw (normalized at
// ingestion), so new IDEs need no CLI changes.
name: "ArbitraryNamePassthrough",
experiment: true,
usageAppName: "SomeFutureIDE",
expectedCalls: 1,
expectedCounts: map[string]int64{"SomeFutureIDE": 1},
},
{
name: "Disable",
Expand Down Expand Up @@ -1730,9 +1730,11 @@ func TestSSH(t *testing.T) {
<-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)
if len(tc.expectedCounts) == 0 {
require.Empty(t, batcher.LastStats.GetSessionCounts())
} else {
require.EqualValues(t, tc.expectedCounts, batcher.LastStats.GetSessionCounts())
}
})
}
})
Expand Down
2 changes: 1 addition & 1 deletion cli/vscodessh_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"])
}
27 changes: 2 additions & 25 deletions coderd/workspaces.go
Original file line number Diff line number Diff line change
Expand Up @@ -1783,34 +1783,11 @@ 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
}

// Any app name is accepted, normalized at ingestion.
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{string(req.AppName): 1},
}

agent, err := api.Database.GetWorkspaceAgentByID(ctx, req.AgentID)
Expand Down
6 changes: 3 additions & 3 deletions coderd/workspaces_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5242,12 +5242,12 @@ 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)

// vscode works
err = client.PostWorkspaceUsageWithBody(ctx, r.Workspace.ID, codersdk.PostWorkspaceUsageRequest{
Expand Down
5 changes: 5 additions & 0 deletions codersdk/workspaces.go
Original file line number Diff line number Diff line change
Expand Up @@ -383,13 +383,18 @@ type PostWorkspaceUsageRequest struct {

type UsageAppName string

// Well-known usage app names. The API accepts any name, normalized at
// ingestion. These are wire format and cannot change: the hyphen in
// reconnecting-pty folds to the canonical reconnecting_pty.
const (
UsageAppNameVscode UsageAppName = "vscode"
UsageAppNameJetbrains UsageAppName = "jetbrains"
UsageAppNameReconnectingPty UsageAppName = "reconnecting-pty"
UsageAppNameSSH UsageAppName = "ssh"
)

// Deprecated: the workspace usage API no longer restricts app_name to this
// list. Use it only to recognize the well-known names.
var AllowedAppNames = []UsageAppName{
UsageAppNameVscode,
UsageAppNameJetbrains,
Expand Down
Loading