Skip to content

Commit fed668b

Browse files
f0sself0ssel
authored
chore: switch ssh session stats based on experiment (#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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy