From c8f173e9f93759912460bb19c6140d7597f33ca0 Mon Sep 17 00:00:00 2001 From: Garrett Delfosse Date: Sat, 22 Jun 2024 15:29:29 +0000 Subject: [PATCH 1/8] chore: switch ssh session stats based on experiment --- cli/ssh.go | 34 ++++ cli/ssh_test.go | 104 ++++++++++ coderd/agentapi/api.go | 3 + coderd/agentapi/stats.go | 12 ++ coderd/agentapi/stats_test.go | 183 ++++++++++++++---- coderd/coderd.go | 2 +- coderd/coderdtest/coderdtest.go | 2 +- coderd/workspaceagentsrpc.go | 1 + .../workspacestatstest/batcher.go | 38 ++++ 9 files changed, 337 insertions(+), 42 deletions(-) create mode 100644 coderd/workspacestats/workspacestatstest/batcher.go diff --git a/cli/ssh.go b/cli/ssh.go index ac849649b9184..3f837234d5d67 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -12,6 +12,7 @@ import ( "os" "os/exec" "path/filepath" + "slices" "strings" "sync" "time" @@ -57,6 +58,7 @@ func (r *RootCmd) ssh() *serpent.Command { logDirPath string remoteForwards []string env []string + usageApp string disableAutostart bool ) client := new(codersdk.Client) @@ -251,6 +253,15 @@ func (r *RootCmd) ssh() *serpent.Command { stopPolling := tryPollWorkspaceAutostop(ctx, client, workspace) defer stopPolling() + usageAppName := getUsageAppName(usageApp) + if usageAppName != "" { + closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{ + AgentID: workspaceAgent.ID, + AppName: usageAppName, + }) + defer closeUsage() + } + if stdio { rawSSH, err := conn.SSH(ctx) if err != nil { @@ -509,6 +520,12 @@ func (r *RootCmd) ssh() *serpent.Command { FlagShorthand: "e", Value: serpent.StringArrayOf(&env), }, + { + Flag: "usage-app", + Description: "Specifies the usage app to use for workspace activity tracking.", + Env: "CODER_SSH_USAGE_APP", + Value: serpent.StringOf(&usageApp), + }, sshDisableAutostartOption(serpent.BoolOf(&disableAutostart)), } return cmd @@ -1044,3 +1061,20 @@ func (r stdioErrLogReader) Read(_ []byte) (int, error) { r.l.Error(context.Background(), "reading from stdin in stdio mode is not allowed") return 0, io.EOF } + +func getUsageAppName(usageApp string) codersdk.UsageAppName { + if usageApp == "" { + usageApp = "ssh" + } + + allowedUsageApps := []string{ + string(codersdk.UsageAppNameSSH), + string(codersdk.UsageAppNameVscode), + string(codersdk.UsageAppNameJetbrains), + } + if slices.Contains(allowedUsageApps, usageApp) { + return codersdk.UsageAppName(usageApp) + } + + return "" +} diff --git a/cli/ssh_test.go b/cli/ssh_test.go index 8c3c1a4e40fd1..501f25338fc7d 100644 --- a/cli/ssh_test.go +++ b/cli/ssh_test.go @@ -36,6 +36,7 @@ import ( "github.com/coder/coder/v2/agent" "github.com/coder/coder/v2/agent/agentssh" "github.com/coder/coder/v2/agent/agenttest" + agentproto "github.com/coder/coder/v2/agent/proto" "github.com/coder/coder/v2/cli/clitest" "github.com/coder/coder/v2/cli/cliui" "github.com/coder/coder/v2/coderd/coderdtest" @@ -43,6 +44,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbfake" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/rbac" + "github.com/coder/coder/v2/coderd/workspacestats/workspacestatstest" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/provisioner/echo" "github.com/coder/coder/v2/provisionersdk/proto" @@ -1292,6 +1294,108 @@ func TestSSH(t *testing.T) { require.NoError(t, err) require.Len(t, ents, 1, "expected one file in logdir %s", logDir) }) + t.Run("UpdateUsage", func(t *testing.T) { + t.Parallel() + + type testCase struct { + name string + experiment bool + usageAppName string + expectedCalls int + expectedCountSSH int + expectedCountJetbrains int + expectedCountVscode int + } + tcs := []testCase{ + { + name: "NoExperiment", + }, + { + name: "Empty", + experiment: true, + expectedCalls: 1, + expectedCountSSH: 1, + }, + { + name: "SSH", + experiment: true, + usageAppName: "ssh", + expectedCalls: 1, + expectedCountSSH: 1, + }, + { + name: "Jetbrains", + experiment: true, + usageAppName: "jetbrains", + expectedCalls: 1, + expectedCountJetbrains: 1, + }, + { + name: "Vscode", + experiment: true, + usageAppName: "vscode", + expectedCalls: 1, + expectedCountVscode: 1, + }, + { + name: "Invalid", + experiment: true, + usageAppName: "invalid", + }, + } + + for _, tc := range tcs { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + if tc.experiment { + dv.Experiments = []string{string(codersdk.ExperimentWorkspaceUsage)} + } + batcher := &workspacestatstest.StatsBatcher{ + LastStats: &agentproto.Stats{}, + } + admin, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{ + DeploymentValues: dv, + StatsBatcher: batcher, + }) + admin.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug)) + first := coderdtest.CreateFirstUser(t, admin) + client, user := coderdtest.CreateAnotherUser(t, admin, first.OrganizationID) + r := dbfake.WorkspaceBuild(t, store, database.Workspace{ + OrganizationID: first.OrganizationID, + OwnerID: user.ID, + }).WithAgent().Do() + workspace := r.Workspace + agentToken := r.AgentToken + inv, root := clitest.New(t, "ssh", workspace.Name, fmt.Sprintf("--usage-app=%s", tc.usageAppName)) + clitest.SetupConfig(t, client, root) + pty := ptytest.New(t).Attach(inv) + + ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) + defer cancel() + + cmdDone := tGo(t, func() { + err := inv.WithContext(ctx).Run() + assert.NoError(t, err) + }) + pty.ExpectMatch("Waiting") + + _ = agenttest.New(t, client.URL, agentToken) + coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID) + + // Shells on Mac, Windows, and Linux all exit shells with the "exit" command. + pty.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) + }) + } + }) } //nolint:paralleltest // This test uses t.Setenv, parent test MUST NOT be parallel. diff --git a/coderd/agentapi/api.go b/coderd/agentapi/api.go index ae0d594314e66..4e5e30ad9c761 100644 --- a/coderd/agentapi/api.go +++ b/coderd/agentapi/api.go @@ -25,6 +25,7 @@ import ( "github.com/coder/coder/v2/coderd/schedule" "github.com/coder/coder/v2/coderd/tracing" "github.com/coder/coder/v2/coderd/workspacestats" + "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/tailnet" tailnetproto "github.com/coder/coder/v2/tailnet/proto" @@ -72,6 +73,7 @@ type Options struct { DerpForceWebSockets bool DerpMapUpdateFrequency time.Duration ExternalAuthConfigs []*externalauth.Config + Experiments codersdk.Experiments // Optional: // WorkspaceID avoids a future lookup to find the workspace ID by setting @@ -118,6 +120,7 @@ func New(opts Options) *API { Log: opts.Log, StatsReporter: opts.StatsReporter, AgentStatsRefreshInterval: opts.AgentStatsRefreshInterval, + Experiments: opts.Experiments, } api.LifecycleAPI = &LifecycleAPI{ diff --git a/coderd/agentapi/stats.go b/coderd/agentapi/stats.go index a167fb5d6f275..4f6a6da1c8c66 100644 --- a/coderd/agentapi/stats.go +++ b/coderd/agentapi/stats.go @@ -12,6 +12,7 @@ import ( "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/coderd/workspacestats" + "github.com/coder/coder/v2/codersdk" ) type StatsAPI struct { @@ -20,6 +21,7 @@ type StatsAPI struct { Log slog.Logger StatsReporter *workspacestats.Reporter AgentStatsRefreshInterval time.Duration + Experiments codersdk.Experiments TimeNowFn func() time.Time // defaults to dbtime.Now() } @@ -55,6 +57,16 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR 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(), diff --git a/coderd/agentapi/stats_test.go b/coderd/agentapi/stats_test.go index 8b4d72fc1d579..5c016f8beee99 100644 --- a/coderd/agentapi/stats_test.go +++ b/coderd/agentapi/stats_test.go @@ -3,7 +3,6 @@ package agentapi_test import ( "context" "database/sql" - "sync" "sync/atomic" "testing" "time" @@ -23,37 +22,11 @@ import ( "github.com/coder/coder/v2/coderd/prometheusmetrics" "github.com/coder/coder/v2/coderd/schedule" "github.com/coder/coder/v2/coderd/workspacestats" + "github.com/coder/coder/v2/coderd/workspacestats/workspacestatstest" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" ) -type statsBatcher struct { - mu sync.Mutex - - called int64 - lastTime time.Time - lastAgentID uuid.UUID - lastTemplateID uuid.UUID - lastUserID uuid.UUID - lastWorkspaceID uuid.UUID - lastStats *agentproto.Stats -} - -var _ workspacestats.Batcher = &statsBatcher{} - -func (b *statsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error { - b.mu.Lock() - defer b.mu.Unlock() - b.called++ - b.lastTime = now - b.lastAgentID = agentID - b.lastTemplateID = templateID - b.lastUserID = userID - b.lastWorkspaceID = workspaceID - b.lastStats = st - return nil -} - func TestUpdateStates(t *testing.T) { t.Parallel() @@ -94,7 +67,7 @@ func TestUpdateStates(t *testing.T) { panic("not implemented") }, } - batcher = &statsBatcher{} + batcher = &workspacestatstest.StatsBatcher{} updateAgentMetricsFnCalled = false req = &agentproto.UpdateStatsRequest{ @@ -188,15 +161,15 @@ func TestUpdateStates(t *testing.T) { ReportInterval: durationpb.New(10 * time.Second), }, resp) - batcher.mu.Lock() - defer batcher.mu.Unlock() - require.Equal(t, int64(1), batcher.called) - require.Equal(t, now, batcher.lastTime) - require.Equal(t, agent.ID, batcher.lastAgentID) - require.Equal(t, template.ID, batcher.lastTemplateID) - require.Equal(t, user.ID, batcher.lastUserID) - require.Equal(t, workspace.ID, batcher.lastWorkspaceID) - require.Equal(t, req.Stats, batcher.lastStats) + batcher.Mu.Lock() + defer batcher.Mu.Unlock() + require.Equal(t, int64(1), batcher.Called) + require.Equal(t, now, batcher.LastTime) + require.Equal(t, agent.ID, batcher.LastAgentID) + require.Equal(t, template.ID, batcher.LastTemplateID) + require.Equal(t, user.ID, batcher.LastUserID) + require.Equal(t, workspace.ID, batcher.LastWorkspaceID) + require.Equal(t, req.Stats, batcher.LastStats) ctx := testutil.Context(t, testutil.WaitShort) select { case <-ctx.Done(): @@ -222,7 +195,7 @@ func TestUpdateStates(t *testing.T) { panic("not implemented") }, } - batcher = &statsBatcher{} + batcher = &workspacestatstest.StatsBatcher{} req = &agentproto.UpdateStatsRequest{ Stats: &agentproto.Stats{ @@ -336,7 +309,7 @@ func TestUpdateStates(t *testing.T) { panic("not implemented") }, } - batcher = &statsBatcher{} + batcher = &workspacestatstest.StatsBatcher{} updateAgentMetricsFnCalled = false req = &agentproto.UpdateStatsRequest{ @@ -406,6 +379,136 @@ func TestUpdateStates(t *testing.T) { require.True(t, updateAgentMetricsFnCalled) }) + + t.Run("WorkspaceUsageExperiment", func(t *testing.T) { + t.Parallel() + + var ( + now = dbtime.Now() + dbM = dbmock.NewMockStore(gomock.NewController(t)) + ps = pubsub.NewInMemory() + + templateScheduleStore = schedule.MockTemplateScheduleStore{ + GetFn: func(context.Context, database.Store, uuid.UUID) (schedule.TemplateScheduleOptions, error) { + panic("should not be called") + }, + SetFn: func(context.Context, database.Store, database.Template, schedule.TemplateScheduleOptions) (database.Template, error) { + panic("not implemented") + }, + } + batcher = &workspacestatstest.StatsBatcher{} + updateAgentMetricsFnCalled = false + + req = &agentproto.UpdateStatsRequest{ + Stats: &agentproto.Stats{ + ConnectionsByProto: map[string]int64{ + "tcp": 1, + "dean": 2, + }, + ConnectionCount: 3, + ConnectionMedianLatencyMs: 23, + RxPackets: 120, + RxBytes: 1000, + TxPackets: 130, + TxBytes: 2000, + SessionCountVscode: 1, + SessionCountJetbrains: 2, + SessionCountReconnectingPty: 3, + SessionCountSsh: 4, + Metrics: []*agentproto.Stats_Metric{ + { + Name: "awesome metric", + Value: 42, + }, + { + Name: "uncool metric", + Value: 0, + }, + }, + }, + } + ) + api := agentapi.StatsAPI{ + AgentFn: func(context.Context) (database.WorkspaceAgent, error) { + return agent, nil + }, + Database: dbM, + StatsReporter: workspacestats.NewReporter(workspacestats.ReporterOptions{ + Database: dbM, + Pubsub: ps, + StatsBatcher: batcher, + TemplateScheduleStore: templateScheduleStorePtr(templateScheduleStore), + UpdateAgentMetricsFn: func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) { + updateAgentMetricsFnCalled = true + assert.Equal(t, prometheusmetrics.AgentMetricLabels{ + Username: user.Username, + WorkspaceName: workspace.Name, + AgentName: agent.Name, + TemplateName: template.Name, + }, labels) + assert.Equal(t, req.Stats.Metrics, metrics) + }, + }), + AgentStatsRefreshInterval: 10 * time.Second, + TimeNowFn: func() time.Time { + return now + }, + Experiments: codersdk.Experiments{ + codersdk.ExperimentWorkspaceUsage, + }, + } + + // Workspace gets fetched. + dbM.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(database.GetWorkspaceByAgentIDRow{ + Workspace: workspace, + TemplateName: template.Name, + }, nil) + + // We expect an activity bump because ConnectionCount > 0. + dbM.EXPECT().ActivityBumpWorkspace(gomock.Any(), database.ActivityBumpWorkspaceParams{ + WorkspaceID: workspace.ID, + NextAutostart: time.Time{}.UTC(), + }).Return(nil) + + // Workspace last used at gets bumped. + dbM.EXPECT().UpdateWorkspaceLastUsedAt(gomock.Any(), database.UpdateWorkspaceLastUsedAtParams{ + ID: workspace.ID, + LastUsedAt: now, + }).Return(nil) + + // User gets fetched to hit the UpdateAgentMetricsFn. + dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil) + + // Ensure that pubsub notifications are sent. + notifyDescription := make(chan []byte) + ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, description []byte) { + go func() { + notifyDescription <- description + }() + }) + + resp, err := api.UpdateStats(context.Background(), req) + require.NoError(t, err) + require.Equal(t, &agentproto.UpdateStatsResponse{ + ReportInterval: durationpb.New(10 * time.Second), + }, resp) + + batcher.Mu.Lock() + defer batcher.Mu.Unlock() + require.EqualValues(t, 1, batcher.Called) + require.EqualValues(t, 0, batcher.LastStats.SessionCountSsh) + require.EqualValues(t, 0, batcher.LastStats.SessionCountJetbrains) + require.EqualValues(t, 0, batcher.LastStats.SessionCountVscode) + require.EqualValues(t, 0, batcher.LastStats.SessionCountReconnectingPty) + ctx := testutil.Context(t, testutil.WaitShort) + select { + case <-ctx.Done(): + t.Error("timed out while waiting for pubsub notification") + case description := <-notifyDescription: + require.Equal(t, description, []byte{}) + } + require.True(t, updateAgentMetricsFnCalled) + }) } func templateScheduleStorePtr(store schedule.TemplateScheduleStore) *atomic.Pointer[schedule.TemplateScheduleStore] { diff --git a/coderd/coderd.go b/coderd/coderd.go index cca4faf36a203..f92e3008604c3 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -187,7 +187,7 @@ type Options struct { HTTPClient *http.Client UpdateAgentMetrics func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric) - StatsBatcher *workspacestats.DBBatcher + StatsBatcher workspacestats.Batcher WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go index 49388aa3537a5..fbeaed43bd8be 100644 --- a/coderd/coderdtest/coderdtest.go +++ b/coderd/coderdtest/coderdtest.go @@ -145,7 +145,7 @@ type Options struct { // Logger should only be overridden if you expect errors // as part of your test. Logger *slog.Logger - StatsBatcher *workspacestats.DBBatcher + StatsBatcher workspacestats.Batcher WorkspaceAppsStatsCollectorOptions workspaceapps.StatsCollectorOptions AllowWorkspaceRenames bool diff --git a/coderd/workspaceagentsrpc.go b/coderd/workspaceagentsrpc.go index ec8dcd8a0e3fc..b413db264feac 100644 --- a/coderd/workspaceagentsrpc.go +++ b/coderd/workspaceagentsrpc.go @@ -143,6 +143,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) { DerpForceWebSockets: api.DeploymentValues.DERP.Config.ForceWebSockets.Value(), DerpMapUpdateFrequency: api.Options.DERPMapUpdateFrequency, ExternalAuthConfigs: api.ExternalAuthConfigs, + Experiments: api.Experiments, // Optional: WorkspaceID: build.WorkspaceID, // saves the extra lookup later diff --git a/coderd/workspacestats/workspacestatstest/batcher.go b/coderd/workspacestats/workspacestatstest/batcher.go new file mode 100644 index 0000000000000..ad5ba60ad16d0 --- /dev/null +++ b/coderd/workspacestats/workspacestatstest/batcher.go @@ -0,0 +1,38 @@ +package workspacestatstest + +import ( + "sync" + "time" + + "github.com/google/uuid" + + agentproto "github.com/coder/coder/v2/agent/proto" + "github.com/coder/coder/v2/coderd/workspacestats" +) + +type StatsBatcher struct { + Mu sync.Mutex + + Called int64 + LastTime time.Time + LastAgentID uuid.UUID + LastTemplateID uuid.UUID + LastUserID uuid.UUID + LastWorkspaceID uuid.UUID + LastStats *agentproto.Stats +} + +var _ workspacestats.Batcher = &StatsBatcher{} + +func (b *StatsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error { + b.Mu.Lock() + defer b.Mu.Unlock() + b.Called++ + b.LastTime = now + b.LastAgentID = agentID + b.LastTemplateID = templateID + b.LastUserID = userID + b.LastWorkspaceID = workspaceID + b.LastStats = st + return nil +} From aaad0753d9dd3a7e60015466b275ca67c26ac894 Mon Sep 17 00:00:00 2001 From: Garrett Delfosse Date: Sat, 22 Jun 2024 16:41:42 +0000 Subject: [PATCH 2/8] make gen --- cli/testdata/coder_ssh_--help.golden | 3 +++ docs/cli/ssh.md | 9 +++++++++ 2 files changed, 12 insertions(+) diff --git a/cli/testdata/coder_ssh_--help.golden b/cli/testdata/coder_ssh_--help.golden index 80aaa3c204fda..08be664586e25 100644 --- a/cli/testdata/coder_ssh_--help.golden +++ b/cli/testdata/coder_ssh_--help.golden @@ -42,6 +42,9 @@ OPTIONS: --stdio bool, $CODER_SSH_STDIO Specifies whether to emit SSH output over stdin/stdout. + --usage-app string, $CODER_SSH_USAGE_APP + Specifies the usage app to use for workspace activity tracking. + --wait yes|no|auto, $CODER_SSH_WAIT (default: auto) Specifies whether or not to wait for the startup script to finish executing. Auto means that the agent startup script behavior diff --git a/docs/cli/ssh.md b/docs/cli/ssh.md index d2110628feec4..decea998189df 100644 --- a/docs/cli/ssh.md +++ b/docs/cli/ssh.md @@ -104,6 +104,15 @@ Enable remote port forwarding (remote_port:local_address:local_port). Set environment variable(s) for session (key1=value1,key2=value2,...). +### --usage-app + +| | | +| ----------- | --------------------------------- | +| Type | string | +| Environment | $CODER_SSH_USAGE_APP | + +Specifies the usage app to use for workspace activity tracking. + ### --disable-autostart | | | From 438b7a60666fb3f87866861553bc414729f9d5f6 Mon Sep 17 00:00:00 2001 From: Garrett Delfosse Date: Sat, 22 Jun 2024 23:37:45 +0000 Subject: [PATCH 3/8] add support to vscodessh command --- cli/vscodessh.go | 9 ++++++++- cli/vscodessh_test.go | 30 +++++++++++++++++++++++++++++- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/cli/vscodessh.go b/cli/vscodessh.go index 147436374b1f6..558b50c00fe95 100644 --- a/cli/vscodessh.go +++ b/cli/vscodessh.go @@ -110,7 +110,7 @@ func (r *RootCmd) vscodeSSH() *serpent.Command { // will call this command after the workspace is started. autostart := false - _, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, autostart, fmt.Sprintf("%s/%s", owner, name)) + workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, autostart, fmt.Sprintf("%s/%s", owner, name)) if err != nil { return xerrors.Errorf("find workspace and agent: %w", err) } @@ -176,6 +176,13 @@ func (r *RootCmd) vscodeSSH() *serpent.Command { defer agentConn.Close() agentConn.AwaitReachable(ctx) + + closeUsage := client.UpdateWorkspaceUsageWithBodyContext(ctx, workspace.ID, codersdk.PostWorkspaceUsageRequest{ + AgentID: workspaceAgent.ID, + AppName: codersdk.UsageAppNameVscode, + }) + defer closeUsage() + rawSSH, err := agentConn.SSH(ctx) if err != nil { return err diff --git a/cli/vscodessh_test.go b/cli/vscodessh_test.go index a4f6ca19132c6..f80b6b0b6029e 100644 --- a/cli/vscodessh_test.go +++ b/cli/vscodessh_test.go @@ -9,9 +9,16 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "cdr.dev/slog" + "cdr.dev/slog/sloggers/slogtest" + "github.com/coder/coder/v2/agent/agenttest" + agentproto "github.com/coder/coder/v2/agent/proto" "github.com/coder/coder/v2/cli/clitest" "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbfake" + "github.com/coder/coder/v2/coderd/workspacestats/workspacestatstest" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/pty/ptytest" "github.com/coder/coder/v2/testutil" @@ -22,7 +29,25 @@ import ( func TestVSCodeSSH(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client, workspace, agentToken := setupWorkspaceForAgent(t) + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentWorkspaceUsage)} + batcher := &workspacestatstest.StatsBatcher{ + LastStats: &agentproto.Stats{}, + } + admin, store := coderdtest.NewWithDatabase(t, &coderdtest.Options{ + DeploymentValues: dv, + StatsBatcher: batcher, + }) + admin.SetLogger(slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug)) + first := coderdtest.CreateFirstUser(t, admin) + client, user := coderdtest.CreateAnotherUser(t, admin, first.OrganizationID) + r := dbfake.WorkspaceBuild(t, store, database.Workspace{ + OrganizationID: first.OrganizationID, + OwnerID: user.ID, + }).WithAgent().Do() + workspace := r.Workspace + agentToken := r.AgentToken + user, err := client.User(ctx, codersdk.Me) require.NoError(t, err) @@ -65,4 +90,7 @@ func TestVSCodeSSH(t *testing.T) { if err := waiter.Wait(); err != nil { waiter.RequireIs(context.Canceled) } + + require.EqualValues(t, 1, batcher.Called) + require.EqualValues(t, 1, batcher.LastStats.SessionCountVscode) } From 61648cbca9078d66e86d9c99f3385f7b6ec1ad55 Mon Sep 17 00:00:00 2001 From: Garrett Delfosse Date: Sun, 23 Jun 2024 01:09:46 +0000 Subject: [PATCH 4/8] add support for web terminal --- site/src/api/api.ts | 10 ++++++ site/src/api/queries/workspaces.ts | 34 ++++++++++++++++++++ site/src/pages/TerminalPage/TerminalPage.tsx | 12 ++++++- 3 files changed, 55 insertions(+), 1 deletion(-) diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 50dbc32a1867d..0627852627e9c 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -24,6 +24,7 @@ import type dayjs from "dayjs"; import userAgentParser from "ua-parser-js"; import { delay } from "../utils/delay"; import * as TypesGen from "./typesGenerated"; +import { PostWorkspaceUsageRequest } from "./typesGenerated"; const getMissingParameters = ( oldBuildParameters: TypesGen.WorkspaceBuildParameter[], @@ -1879,6 +1880,15 @@ class ApiMethods { throw error; } }; + + postWorkspaceUsage = async (workspaceID: string, options: PostWorkspaceUsageRequest) => { + const response = await this.axios.post( + `/api/v2/workspaces/${workspaceID}/usage`, + options, + ); + + return response.data; + }; } // This is a hard coded CSRF token/cookie pair for local development. In prod, diff --git a/site/src/api/queries/workspaces.ts b/site/src/api/queries/workspaces.ts index 809c8a4c2862f..8e0ae33bd70da 100644 --- a/site/src/api/queries/workspaces.ts +++ b/site/src/api/queries/workspaces.ts @@ -8,6 +8,7 @@ import { type DeleteWorkspaceOptions, API } from "api/api"; import type { CreateWorkspaceRequest, ProvisionerLogLevel, + UsageAppName, Workspace, WorkspaceBuild, WorkspaceBuildParameter, @@ -16,6 +17,7 @@ import type { } from "api/typesGenerated"; import { disabledRefetchOptions } from "./util"; import { workspaceBuildsKey } from "./workspaceBuilds"; +import { Terminal } from "xterm"; export const workspaceByOwnerAndNameKey = (owner: string, name: string) => [ "workspace", @@ -313,3 +315,35 @@ export const agentLogs = (workspaceId: string, agentId: string) => { ...disabledRefetchOptions, }; }; + +// workspace usage options +export interface WorkspaceUsageOptions { + usageApp: UsageAppName; + terminal: Terminal | undefined; + workspaceId: string | undefined; + agentId: string | undefined; +} + +export const workspaceUsage = (options: WorkspaceUsageOptions) => { + return { + queryKey: [ + "workspaces", options.workspaceId, + "agents", options.agentId, + "usage", options.usageApp, + ], + enabled: options.terminal !== undefined && options.workspaceId !== undefined && options.agentId !== undefined, + queryFn: () => { + if (options.terminal === undefined || options.workspaceId === undefined || options.agentId === undefined) { + return Promise.reject(); + } + + return API.postWorkspaceUsage(options.workspaceId, { + agent_id: options.agentId, + app_name: options.usageApp, + }) + }, + // ...disabledRefetchOptions, + refetchInterval: 60000, + refetchIntervalInBackground: true, + }; +} diff --git a/site/src/pages/TerminalPage/TerminalPage.tsx b/site/src/pages/TerminalPage/TerminalPage.tsx index 73ddc1ceebe06..85936ca9ecd44 100644 --- a/site/src/pages/TerminalPage/TerminalPage.tsx +++ b/site/src/pages/TerminalPage/TerminalPage.tsx @@ -12,7 +12,7 @@ import { Unicode11Addon } from "xterm-addon-unicode11"; import { WebLinksAddon } from "xterm-addon-web-links"; import { WebglAddon } from "xterm-addon-webgl"; import { deploymentConfig } from "api/queries/deployment"; -import { workspaceByOwnerAndName } from "api/queries/workspaces"; +import { workspaceByOwnerAndName, workspaceUsage } from "api/queries/workspaces"; import { useProxy } from "contexts/ProxyContext"; import { ThemeOverride } from "contexts/ThemeProvider"; import themes from "theme"; @@ -67,6 +67,16 @@ const TerminalPage: FC = () => { const config = useQuery(deploymentConfig()); const renderer = config.data?.config.web_terminal_renderer; + // Periodically report workspace usage. + useQuery( + workspaceUsage({ + usageApp: "reconnecting-pty", + terminal, + workspaceId: workspace.data?.id, + agentId: workspaceAgent?.id, + }) + ); + // handleWebLink handles opening of URLs in the terminal! const handleWebLink = useCallback( (uri: string) => { From 338d2b1225206091485afb5e7ec9cbe11cad40c7 Mon Sep 17 00:00:00 2001 From: Garrett Delfosse Date: Mon, 24 Jun 2024 15:40:37 +0000 Subject: [PATCH 5/8] hide flag --- cli/ssh.go | 11 ++++++++--- cli/ssh_test.go | 11 +++++++++-- cli/testdata/coder_ssh_--help.golden | 3 --- docs/cli/ssh.md | 9 --------- 4 files changed, 17 insertions(+), 17 deletions(-) diff --git a/cli/ssh.go b/cli/ssh.go index 3f837234d5d67..e4e9fadf5e8e8 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -41,6 +41,10 @@ import ( "github.com/coder/serpent" ) +const ( + disableUsageApp = "disable" +) + var ( workspacePollInterval = time.Minute autostopNotifyCountdown = []time.Duration{30 * time.Minute} @@ -525,6 +529,7 @@ func (r *RootCmd) ssh() *serpent.Command { Description: "Specifies the usage app to use for workspace activity tracking.", Env: "CODER_SSH_USAGE_APP", Value: serpent.StringOf(&usageApp), + Hidden: true, }, sshDisableAutostartOption(serpent.BoolOf(&disableAutostart)), } @@ -1063,8 +1068,8 @@ func (r stdioErrLogReader) Read(_ []byte) (int, error) { } func getUsageAppName(usageApp string) codersdk.UsageAppName { - if usageApp == "" { - usageApp = "ssh" + if usageApp == disableUsageApp { + return "" } allowedUsageApps := []string{ @@ -1076,5 +1081,5 @@ func getUsageAppName(usageApp string) codersdk.UsageAppName { return codersdk.UsageAppName(usageApp) } - return "" + return codersdk.UsageAppNameSSH } diff --git a/cli/ssh_test.go b/cli/ssh_test.go index 501f25338fc7d..ae93c4b0cea05 100644 --- a/cli/ssh_test.go +++ b/cli/ssh_test.go @@ -1338,9 +1338,16 @@ func TestSSH(t *testing.T) { expectedCountVscode: 1, }, { - name: "Invalid", + name: "InvalidDefaultsToSSH", + experiment: true, + usageAppName: "invalid", + expectedCalls: 1, + expectedCountSSH: 1, + }, + { + name: "Disable", experiment: true, - usageAppName: "invalid", + usageAppName: "disable", }, } diff --git a/cli/testdata/coder_ssh_--help.golden b/cli/testdata/coder_ssh_--help.golden index 08be664586e25..80aaa3c204fda 100644 --- a/cli/testdata/coder_ssh_--help.golden +++ b/cli/testdata/coder_ssh_--help.golden @@ -42,9 +42,6 @@ OPTIONS: --stdio bool, $CODER_SSH_STDIO Specifies whether to emit SSH output over stdin/stdout. - --usage-app string, $CODER_SSH_USAGE_APP - Specifies the usage app to use for workspace activity tracking. - --wait yes|no|auto, $CODER_SSH_WAIT (default: auto) Specifies whether or not to wait for the startup script to finish executing. Auto means that the agent startup script behavior diff --git a/docs/cli/ssh.md b/docs/cli/ssh.md index decea998189df..d2110628feec4 100644 --- a/docs/cli/ssh.md +++ b/docs/cli/ssh.md @@ -104,15 +104,6 @@ Enable remote port forwarding (remote_port:local_address:local_port). Set environment variable(s) for session (key1=value1,key2=value2,...). -### --usage-app - -| | | -| ----------- | --------------------------------- | -| Type | string | -| Environment | $CODER_SSH_USAGE_APP | - -Specifies the usage app to use for workspace activity tracking. - ### --disable-autostart | | | From b3a240b161d8068cee5490f87ea52ae561732c98 Mon Sep 17 00:00:00 2001 From: Garrett Delfosse Date: Mon, 24 Jun 2024 15:58:07 +0000 Subject: [PATCH 6/8] fe lint --- site/src/api/api.ts | 7 ++++-- site/src/api/queries/workspaces.ts | 26 ++++++++++++++------ site/src/pages/TerminalPage/TerminalPage.tsx | 7 ++++-- 3 files changed, 28 insertions(+), 12 deletions(-) diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 0627852627e9c..01b4ccec76b9d 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -24,7 +24,7 @@ import type dayjs from "dayjs"; import userAgentParser from "ua-parser-js"; import { delay } from "../utils/delay"; import * as TypesGen from "./typesGenerated"; -import { PostWorkspaceUsageRequest } from "./typesGenerated"; +import type { PostWorkspaceUsageRequest } from "./typesGenerated"; const getMissingParameters = ( oldBuildParameters: TypesGen.WorkspaceBuildParameter[], @@ -1881,7 +1881,10 @@ class ApiMethods { } }; - postWorkspaceUsage = async (workspaceID: string, options: PostWorkspaceUsageRequest) => { + postWorkspaceUsage = async ( + workspaceID: string, + options: PostWorkspaceUsageRequest, + ) => { const response = await this.axios.post( `/api/v2/workspaces/${workspaceID}/usage`, options, diff --git a/site/src/api/queries/workspaces.ts b/site/src/api/queries/workspaces.ts index 8e0ae33bd70da..7b58f47c26e7a 100644 --- a/site/src/api/queries/workspaces.ts +++ b/site/src/api/queries/workspaces.ts @@ -4,6 +4,7 @@ import type { QueryOptions, UseMutationOptions, } from "react-query"; +import type { Terminal } from "xterm"; import { type DeleteWorkspaceOptions, API } from "api/api"; import type { CreateWorkspaceRequest, @@ -17,7 +18,6 @@ import type { } from "api/typesGenerated"; import { disabledRefetchOptions } from "./util"; import { workspaceBuildsKey } from "./workspaceBuilds"; -import { Terminal } from "xterm"; export const workspaceByOwnerAndNameKey = (owner: string, name: string) => [ "workspace", @@ -327,23 +327,33 @@ export interface WorkspaceUsageOptions { export const workspaceUsage = (options: WorkspaceUsageOptions) => { return { queryKey: [ - "workspaces", options.workspaceId, - "agents", options.agentId, - "usage", options.usageApp, + "workspaces", + options.workspaceId, + "agents", + options.agentId, + "usage", + options.usageApp, ], - enabled: options.terminal !== undefined && options.workspaceId !== undefined && options.agentId !== undefined, + enabled: + options.terminal !== undefined && + options.workspaceId !== undefined && + options.agentId !== undefined, queryFn: () => { - if (options.terminal === undefined || options.workspaceId === undefined || options.agentId === undefined) { + if ( + options.terminal === undefined || + options.workspaceId === undefined || + options.agentId === undefined + ) { return Promise.reject(); } return API.postWorkspaceUsage(options.workspaceId, { agent_id: options.agentId, app_name: options.usageApp, - }) + }); }, // ...disabledRefetchOptions, refetchInterval: 60000, refetchIntervalInBackground: true, }; -} +}; diff --git a/site/src/pages/TerminalPage/TerminalPage.tsx b/site/src/pages/TerminalPage/TerminalPage.tsx index 85936ca9ecd44..d498fd32eff50 100644 --- a/site/src/pages/TerminalPage/TerminalPage.tsx +++ b/site/src/pages/TerminalPage/TerminalPage.tsx @@ -12,7 +12,10 @@ import { Unicode11Addon } from "xterm-addon-unicode11"; import { WebLinksAddon } from "xterm-addon-web-links"; import { WebglAddon } from "xterm-addon-webgl"; import { deploymentConfig } from "api/queries/deployment"; -import { workspaceByOwnerAndName, workspaceUsage } from "api/queries/workspaces"; +import { + workspaceByOwnerAndName, + workspaceUsage, +} from "api/queries/workspaces"; import { useProxy } from "contexts/ProxyContext"; import { ThemeOverride } from "contexts/ThemeProvider"; import themes from "theme"; @@ -74,7 +77,7 @@ const TerminalPage: FC = () => { terminal, workspaceId: workspace.data?.id, agentId: workspaceAgent?.id, - }) + }), ); // handleWebLink handles opening of URLs in the terminal! From 645bc5a74d747540c357be13fccf0793af7a818d Mon Sep 17 00:00:00 2001 From: Garrett Delfosse Date: Mon, 24 Jun 2024 16:25:04 +0000 Subject: [PATCH 7/8] use connection status --- site/src/api/queries/workspaces.ts | 14 +++++--------- site/src/pages/TerminalPage/TerminalPage.tsx | 2 +- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/site/src/api/queries/workspaces.ts b/site/src/api/queries/workspaces.ts index 7b58f47c26e7a..aa5f8f29a9783 100644 --- a/site/src/api/queries/workspaces.ts +++ b/site/src/api/queries/workspaces.ts @@ -4,7 +4,6 @@ import type { QueryOptions, UseMutationOptions, } from "react-query"; -import type { Terminal } from "xterm"; import { type DeleteWorkspaceOptions, API } from "api/api"; import type { CreateWorkspaceRequest, @@ -16,6 +15,7 @@ import type { WorkspacesRequest, WorkspacesResponse, } from "api/typesGenerated"; +import type { ConnectionStatus } from "pages/TerminalPage/types"; import { disabledRefetchOptions } from "./util"; import { workspaceBuildsKey } from "./workspaceBuilds"; @@ -319,7 +319,7 @@ export const agentLogs = (workspaceId: string, agentId: string) => { // workspace usage options export interface WorkspaceUsageOptions { usageApp: UsageAppName; - terminal: Terminal | undefined; + connectionStatus: ConnectionStatus; workspaceId: string | undefined; agentId: string | undefined; } @@ -335,15 +335,11 @@ export const workspaceUsage = (options: WorkspaceUsageOptions) => { options.usageApp, ], enabled: - options.terminal !== undefined && options.workspaceId !== undefined && - options.agentId !== undefined, + options.agentId !== undefined && + options.connectionStatus === "connected", queryFn: () => { - if ( - options.terminal === undefined || - options.workspaceId === undefined || - options.agentId === undefined - ) { + if (options.workspaceId === undefined || options.agentId === undefined) { return Promise.reject(); } diff --git a/site/src/pages/TerminalPage/TerminalPage.tsx b/site/src/pages/TerminalPage/TerminalPage.tsx index d498fd32eff50..4d99b3e4862f7 100644 --- a/site/src/pages/TerminalPage/TerminalPage.tsx +++ b/site/src/pages/TerminalPage/TerminalPage.tsx @@ -74,7 +74,7 @@ const TerminalPage: FC = () => { useQuery( workspaceUsage({ usageApp: "reconnecting-pty", - terminal, + connectionStatus, workspaceId: workspace.data?.id, agentId: workspaceAgent?.id, }), From 1370ab318af6c8113454968c66c772e09b4ac15b Mon Sep 17 00:00:00 2001 From: Garrett Delfosse Date: Tue, 25 Jun 2024 14:37:31 +0000 Subject: [PATCH 8/8] use t.fatal --- coderd/agentapi/stats_test.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/coderd/agentapi/stats_test.go b/coderd/agentapi/stats_test.go index 5c016f8beee99..57534208be110 100644 --- a/coderd/agentapi/stats_test.go +++ b/coderd/agentapi/stats_test.go @@ -390,10 +390,12 @@ func TestUpdateStates(t *testing.T) { templateScheduleStore = schedule.MockTemplateScheduleStore{ GetFn: func(context.Context, database.Store, uuid.UUID) (schedule.TemplateScheduleOptions, error) { - panic("should not be called") + t.Fatal("getfn should not be called") + return schedule.TemplateScheduleOptions{}, nil }, SetFn: func(context.Context, database.Store, database.Template, schedule.TemplateScheduleOptions) (database.Template, error) { - panic("not implemented") + t.Fatal("setfn not implemented") + return database.Template{}, nil }, } batcher = &workspacestatstest.StatsBatcher{}