Skip to content

chore: switch ssh session stats based on experiment #13637

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Jun 25, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 39 additions & 0 deletions cli/ssh.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import (
"os"
"os/exec"
"path/filepath"
"slices"
"strings"
"sync"
"time"
Expand Down Expand Up @@ -40,6 +41,10 @@ import (
"github.com/coder/serpent"
)

const (
disableUsageApp = "disable"
)

var (
workspacePollInterval = time.Minute
autostopNotifyCountdown = []time.Duration{30 * time.Minute}
Expand All @@ -57,6 +62,7 @@ func (r *RootCmd) ssh() *serpent.Command {
logDirPath string
remoteForwards []string
env []string
usageApp string
disableAutostart bool
)
client := new(codersdk.Client)
Expand Down Expand Up @@ -251,6 +257,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 {
Expand Down Expand Up @@ -509,6 +524,13 @@ 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),
Hidden: true,
},
sshDisableAutostartOption(serpent.BoolOf(&disableAutostart)),
}
return cmd
Expand Down Expand Up @@ -1044,3 +1066,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 == disableUsageApp {
return ""
}

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

return codersdk.UsageAppNameSSH
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Letting the zero value mean 'disabled' feels more idiomatic to me, but 'disabled' is way more explicit and better from a user perspective.

Copy link
Contributor Author

@f0ssel f0ssel Jun 24, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I considered having the default value on the flag be ssh and disable reporting if it's manually set to empty string but I wasn't sure if that was explicit enough behavior. I'm thinking in the extensions where we do CODER_SSH_USAGE_APP_NAME="" coder ssh ... I'm worried there may be cases where the empty string get's passed to the CLI as "unset" instead.

}
111 changes: 111 additions & 0 deletions cli/ssh_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -36,13 +36,15 @@ 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"
"github.com/coder/coder/v2/coderd/database"
"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"
Expand Down Expand Up @@ -1292,6 +1294,115 @@ 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: "InvalidDefaultsToSSH",
experiment: true,
usageAppName: "invalid",
expectedCalls: 1,
expectedCountSSH: 1,
},
{
name: "Disable",
experiment: true,
usageAppName: "disable",
},
}

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.
Expand Down
9 changes: 8 additions & 1 deletion cli/vscodessh.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down Expand Up @@ -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
Expand Down
30 changes: 29 additions & 1 deletion cli/vscodessh_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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)

Expand Down Expand Up @@ -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)
}
3 changes: 3 additions & 0 deletions coderd/agentapi/api.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -118,6 +120,7 @@ func New(opts Options) *API {
Log: opts.Log,
StatsReporter: opts.StatsReporter,
AgentStatsRefreshInterval: opts.AgentStatsRefreshInterval,
Experiments: opts.Experiments,
}

api.LifecycleAPI = &LifecycleAPI{
Expand Down
12 changes: 12 additions & 0 deletions coderd/agentapi/stats.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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()
}
Expand Down Expand Up @@ -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(),
Expand Down
Loading
Loading