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{} 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