Skip to content

Commit fed668b

Browse files
f0sself0ssel
authored
chore: switch ssh session stats based on experiment (coder#13637)
1 parent d7eadee commit fed668b

File tree

14 files changed

+455
-45
lines changed

14 files changed

+455
-45
lines changed

‎cli/ssh.go

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"os"
1313
"os/exec"
1414
"path/filepath"
15+
"slices"
1516
"strings"
1617
"sync"
1718
"time"
@@ -40,6 +41,10 @@ import (
4041
"github.com/coder/serpent"
4142
)
4243

44+
const (
45+
disableUsageApp = "disable"
46+
)
47+
4348
var (
4449
workspacePollInterval = time.Minute
4550
autostopNotifyCountdown = []time.Duration{30 * time.Minute}
@@ -57,6 +62,7 @@ func (r *RootCmd) ssh() *serpent.Command {
5762
logDirPath string
5863
remoteForwards []string
5964
env []string
65+
usageApp string
6066
disableAutostart bool
6167
)
6268
client := new(codersdk.Client)
@@ -251,6 +257,15 @@ func (r *RootCmd) ssh() *serpent.Command {
251257
stopPolling := tryPollWorkspaceAutostop(ctx, client, workspace)
252258
defer stopPolling()
253259

260+
usageAppName := getUsageAppName(usageApp)
261+
if usageAppName != "" {
262+
closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{
263+
AgentID: workspaceAgent.ID,
264+
AppName: usageAppName,
265+
})
266+
defer closeUsage()
267+
}
268+
254269
if stdio {
255270
rawSSH, err := conn.SSH(ctx)
256271
if err != nil {
@@ -509,6 +524,13 @@ func (r *RootCmd) ssh() *serpent.Command {
509524
FlagShorthand: "e",
510525
Value: serpent.StringArrayOf(&env),
511526
},
527+
{
528+
Flag: "usage-app",
529+
Description: "Specifies the usage app to use for workspace activity tracking.",
530+
Env: "CODER_SSH_USAGE_APP",
531+
Value: serpent.StringOf(&usageApp),
532+
Hidden: true,
533+
},
512534
sshDisableAutostartOption(serpent.BoolOf(&disableAutostart)),
513535
}
514536
return cmd
@@ -1044,3 +1066,20 @@ func (r stdioErrLogReader) Read(_ []byte) (int, error) {
10441066
r.l.Error(context.Background(), "reading from stdin in stdio mode is not allowed")
10451067
return 0, io.EOF
10461068
}
1069+
1070+
func getUsageAppName(usageApp string) codersdk.UsageAppName {
1071+
if usageApp == disableUsageApp {
1072+
return ""
1073+
}
1074+
1075+
allowedUsageApps := []string{
1076+
string(codersdk.UsageAppNameSSH),
1077+
string(codersdk.UsageAppNameVscode),
1078+
string(codersdk.UsageAppNameJetbrains),
1079+
}
1080+
if slices.Contains(allowedUsageApps, usageApp) {
1081+
return codersdk.UsageAppName(usageApp)
1082+
}
1083+
1084+
return codersdk.UsageAppNameSSH
1085+
}

‎cli/ssh_test.go

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,13 +36,15 @@ import (
3636
"github.com/coder/coder/v2/agent"
3737
"github.com/coder/coder/v2/agent/agentssh"
3838
"github.com/coder/coder/v2/agent/agenttest"
39+
agentproto "github.com/coder/coder/v2/agent/proto"
3940
"github.com/coder/coder/v2/cli/clitest"
4041
"github.com/coder/coder/v2/cli/cliui"
4142
"github.com/coder/coder/v2/coderd/coderdtest"
4243
"github.com/coder/coder/v2/coderd/database"
4344
"github.com/coder/coder/v2/coderd/database/dbfake"
4445
"github.com/coder/coder/v2/coderd/database/dbtestutil"
4546
"github.com/coder/coder/v2/coderd/rbac"
47+
"github.com/coder/coder/v2/coderd/workspacestats/workspacestatstest"
4648
"github.com/coder/coder/v2/codersdk"
4749
"github.com/coder/coder/v2/provisioner/echo"
4850
"github.com/coder/coder/v2/provisionersdk/proto"
@@ -1292,6 +1294,115 @@ func TestSSH(t *testing.T) {
12921294
require.NoError(t, err)
12931295
require.Len(t, ents, 1, "expected one file in logdir %s", logDir)
12941296
})
1297+
t.Run("UpdateUsage", func(t *testing.T) {
1298+
t.Parallel()
1299+
1300+
type testCase struct {
1301+
name string
1302+
experiment bool
1303+
usageAppName string
1304+
expectedCalls int
1305+
expectedCountSSH int
1306+
expectedCountJetbrains int
1307+
expectedCountVscode int
1308+
}
1309+
tcs := []testCase{
1310+
{
1311+
name: "NoExperiment",
1312+
},
1313+
{
1314+
name: "Empty",
1315+
experiment: true,
1316+
expectedCalls: 1,
1317+
expectedCountSSH: 1,
1318+
},
1319+
{
1320+
name: "SSH",
1321+
experiment: true,
1322+
usageAppName: "ssh",
1323+
expectedCalls: 1,
1324+
expectedCountSSH: 1,
1325+
},
1326+
{
1327+
name: "Jetbrains",
1328+
experiment: true,
1329+
usageAppName: "jetbrains",
1330+
expectedCalls: 1,
1331+
expectedCountJetbrains: 1,
1332+
},
1333+
{
1334+
name: "Vscode",
1335+
experiment: true,
1336+
usageAppName: "vscode",
1337+
expectedCalls: 1,
1338+
expectedCountVscode: 1,
1339+
},
1340+
{
1341+
name: "InvalidDefaultsToSSH",
1342+
experiment: true,
1343+
usageAppName: "invalid",
1344+
expectedCalls: 1,
1345+
expectedCountSSH: 1,
1346+
},
1347+
{
1348+
name: "Disable",
1349+
experiment: true,
1350+
usageAppName: "disable",
1351+
},
1352+
}
1353+
1354+
for _, tc := range tcs {
1355+
tc := tc
1356+
t.Run(tc.name, func(t *testing.T) {
1357+
t.Parallel()
1358+
1359+
dv := coderdtest.DeploymentValues(t)
1360+
if tc.experiment {
1361+
dv.Experiments = []string{string(codersdk.ExperimentWorkspaceUsage)}
1362+
}
1363+
batcher := &workspacestatstest.StatsBatcher{
1364+
LastStats: &agentproto.Stats{},
1365+
}
1366+
admin, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{
1367+
DeploymentValues: dv,
1368+
StatsBatcher: batcher,
1369+
})
1370+
admin.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug))
1371+
first := coderdtest.CreateFirstUser(t, admin)
1372+
client, user := coderdtest.CreateAnotherUser(t, admin, first.OrganizationID)
1373+
r := dbfake.WorkspaceBuild(t, store, database.Workspace{
1374+
OrganizationID: first.OrganizationID,
1375+
OwnerID: user.ID,
1376+
}).WithAgent().Do()
1377+
workspace := r.Workspace
1378+
agentToken := r.AgentToken
1379+
inv, root := clitest.New(t, "ssh", workspace.Name, fmt.Sprintf("--usage-app=%s", tc.usageAppName))
1380+
clitest.SetupConfig(t, client, root)
1381+
pty := ptytest.New(t).Attach(inv)
1382+
1383+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
1384+
defer cancel()
1385+
1386+
cmdDone := tGo(t, func() {
1387+
err := inv.WithContext(ctx).Run()
1388+
assert.NoError(t, err)
1389+
})
1390+
pty.ExpectMatch("Waiting")
1391+
1392+
_ = agenttest.New(t, client.URL, agentToken)
1393+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
1394+
1395+
// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
1396+
pty.WriteLine("exit")
1397+
<-cmdDone
1398+
1399+
require.EqualValues(t, tc.expectedCalls, batcher.Called)
1400+
require.EqualValues(t, tc.expectedCountSSH, batcher.LastStats.SessionCountSsh)
1401+
require.EqualValues(t, tc.expectedCountJetbrains, batcher.LastStats.SessionCountJetbrains)
1402+
require.EqualValues(t, tc.expectedCountVscode, batcher.LastStats.SessionCountVscode)
1403+
})
1404+
}
1405+
})
12951406
}
12961407

12971408
//nolint:paralleltest // This test uses t.Setenv, parent test MUST NOT be parallel.

‎cli/vscodessh.go

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ func (r *RootCmd) vscodeSSH() *serpent.Command {
110110
// will call this command after the workspace is started.
111111
autostart := false
112112

113-
_, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, autostart, fmt.Sprintf("%s/%s", owner, name))
113+
workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, autostart, fmt.Sprintf("%s/%s", owner, name))
114114
if err != nil {
115115
return xerrors.Errorf("find workspace and agent: %w", err)
116116
}
@@ -176,6 +176,13 @@ func (r *RootCmd) vscodeSSH() *serpent.Command {
176176
defer agentConn.Close()
177177

178178
agentConn.AwaitReachable(ctx)
179+
180+
closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{
181+
AgentID: workspaceAgent.ID,
182+
AppName: codersdk.UsageAppNameVscode,
183+
})
184+
defer closeUsage()
185+
179186
rawSSH, err := agentConn.SSH(ctx)
180187
if err != nil {
181188
return err

‎cli/vscodessh_test.go

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,16 @@ import (
99
"github.com/stretchr/testify/assert"
1010
"github.com/stretchr/testify/require"
1111

12+
"cdr.dev/slog"
13+
"cdr.dev/slog/sloggers/slogtest"
14+
1215
"github.com/coder/coder/v2/agent/agenttest"
16+
agentproto "github.com/coder/coder/v2/agent/proto"
1317
"github.com/coder/coder/v2/cli/clitest"
1418
"github.com/coder/coder/v2/coderd/coderdtest"
19+
"github.com/coder/coder/v2/coderd/database"
20+
"github.com/coder/coder/v2/coderd/database/dbfake"
21+
"github.com/coder/coder/v2/coderd/workspacestats/workspacestatstest"
1522
"github.com/coder/coder/v2/codersdk"
1623
"github.com/coder/coder/v2/pty/ptytest"
1724
"github.com/coder/coder/v2/testutil"
@@ -22,7 +29,25 @@ import (
2229
func TestVSCodeSSH(t *testing.T) {
2330
t.Parallel()
2431
ctx := testutil.Context(t, testutil.WaitLong)
25-
client, workspace, agentToken := setupWorkspaceForAgent(t)
32+
dv := coderdtest.DeploymentValues(t)
33+
dv.Experiments = []string{string(codersdk.ExperimentWorkspaceUsage)}
34+
batcher := &workspacestatstest.StatsBatcher{
35+
LastStats: &agentproto.Stats{},
36+
}
37+
admin, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{
38+
DeploymentValues: dv,
39+
StatsBatcher: batcher,
40+
})
41+
admin.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug))
42+
first := coderdtest.CreateFirstUser(t, admin)
43+
client, user := coderdtest.CreateAnotherUser(t, admin, first.OrganizationID)
44+
r := dbfake.WorkspaceBuild(t, store, database.Workspace{
45+
OrganizationID: first.OrganizationID,
46+
OwnerID: user.ID,
47+
}).WithAgent().Do()
48+
workspace := r.Workspace
49+
agentToken := r.AgentToken
50+
2651
user, err := client.User(ctx, codersdk.Me)
2752
require.NoError(t, err)
2853

@@ -65,4 +90,7 @@ func TestVSCodeSSH(t *testing.T) {
6590
if err := waiter.Wait(); err != nil {
6691
waiter.RequireIs(context.Canceled)
6792
}
93+
94+
require.EqualValues(t, 1, batcher.Called)
95+
require.EqualValues(t, 1, batcher.LastStats.SessionCountVscode)
6896
}

‎coderd/agentapi/api.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import (
2525
"github.com/coder/coder/v2/coderd/schedule"
2626
"github.com/coder/coder/v2/coderd/tracing"
2727
"github.com/coder/coder/v2/coderd/workspacestats"
28+
"github.com/coder/coder/v2/codersdk"
2829
"github.com/coder/coder/v2/codersdk/agentsdk"
2930
"github.com/coder/coder/v2/tailnet"
3031
tailnetproto "github.com/coder/coder/v2/tailnet/proto"
@@ -72,6 +73,7 @@ type Options struct {
7273
DerpForceWebSockets bool
7374
DerpMapUpdateFrequency time.Duration
7475
ExternalAuthConfigs []*externalauth.Config
76+
Experiments codersdk.Experiments
7577

7678
// Optional:
7779
// WorkspaceID avoids a future lookup to find the workspace ID by setting
@@ -118,6 +120,7 @@ func New(opts Options) *API {
118120
Log: opts.Log,
119121
StatsReporter: opts.StatsReporter,
120122
AgentStatsRefreshInterval: opts.AgentStatsRefreshInterval,
123+
Experiments: opts.Experiments,
121124
}
122125

123126
api.LifecycleAPI = &LifecycleAPI{

‎coderd/agentapi/stats.go

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"github.com/coder/coder/v2/coderd/database"
1313
"github.com/coder/coder/v2/coderd/database/dbtime"
1414
"github.com/coder/coder/v2/coderd/workspacestats"
15+
"github.com/coder/coder/v2/codersdk"
1516
)
1617

1718
type StatsAPI struct {
@@ -20,6 +21,7 @@ type StatsAPI struct {
2021
Log slog.Logger
2122
StatsReporter *workspacestats.Reporter
2223
AgentStatsRefreshInterval time.Duration
24+
Experiments codersdk.Experiments
2325

2426
TimeNowFn func() time.Time // defaults to dbtime.Now()
2527
}
@@ -55,6 +57,16 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
5557
slog.F("payload", req),
5658
)
5759

60+
if a.Experiments.Enabled(codersdk.ExperimentWorkspaceUsage) {
61+
// while the experiment is enabled we will not report
62+
// session stats from the agent. This is because it is
63+
// being handled by the CLI and the postWorkspaceUsage route.
64+
req.Stats.SessionCountSsh = 0
65+
req.Stats.SessionCountJetbrains = 0
66+
req.Stats.SessionCountVscode = 0
67+
req.Stats.SessionCountReconnectingPty = 0
68+
}
69+
5870
err = a.StatsReporter.ReportAgentStats(
5971
ctx,
6072
a.now(),

0 commit comments

Comments
 (0)