Skip to content

Commit 5304b4e

Browse files
kylecarbskylecarbs
authored
feat: add connection statistics for workspace agents (#6469)
* fix: don't make session counts cumulative This made for some weird tracking... we want the point-in-time number of counts! * Add databasefake query for getting agent stats * Add deployment stats endpoint * The query... works?!? * Fix aggregation query * Select from multiple tables instead * Fix continuous stats * Increase period of stat refreshes * Add workspace counts to deployment stats * fmt * Add a slight bit of responsiveness * Fix template version editor overflow * Add refresh button * Fix font family on button * Fix latest stat being reported * Revert agent conn stats * Fix linting error * Fix tests * Fix gen * Fix migrations * Block on sending stat updates * Add test fixtures * Fix response structure * make gen
1 parent 9d40d2f commit 5304b4e

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

43 files changed

+1786
-170
lines changed

‎agent/agent.go

Lines changed: 21 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import (
1717
"os/exec"
1818
"os/user"
1919
"path/filepath"
20+
"reflect"
2021
"runtime"
2122
"sort"
2223
"strconv"
@@ -60,7 +61,7 @@ const (
6061

6162
// MagicSSHSessionTypeEnvironmentVariable is used to track the purpose behind an SSH connection.
6263
// This is stripped from any commands being executed, and is counted towards connection stats.
63-
MagicSSHSessionTypeEnvironmentVariable = "__CODER_SSH_SESSION_TYPE"
64+
MagicSSHSessionTypeEnvironmentVariable = "CODER_SSH_SESSION_TYPE"
6465
// MagicSSHSessionTypeVSCode is set in the SSH config by the VS Code extension to identify itself.
6566
MagicSSHSessionTypeVSCode = "vscode"
6667
// MagicSSHSessionTypeJetBrains is set in the SSH config by the JetBrains extension to identify itself.
@@ -122,9 +123,7 @@ func New(options Options) io.Closer {
122123
tempDir: options.TempDir,
123124
lifecycleUpdate: make(chan struct{}, 1),
124125
lifecycleReported: make(chan codersdk.WorkspaceAgentLifecycle, 1),
125-
// TODO: This is a temporary hack to make tests not flake.
126-
// @kylecarbs has a better solution in here: https://github.com/coder/coder/pull/6469
127-
connStatsChan: make(chan *agentsdk.Stats, 8),
126+
connStatsChan: make(chan *agentsdk.Stats, 1),
128127
}
129128
a.init(ctx)
130129
return a
@@ -159,11 +158,8 @@ type agent struct {
159158

160159
network *tailnet.Conn
161160
connStatsChan chan *agentsdk.Stats
161+
latestStat atomic.Pointer[agentsdk.Stats]
162162

163-
statRxPackets atomic.Int64
164-
statRxBytes atomic.Int64
165-
statTxPackets atomic.Int64
166-
statTxBytes atomic.Int64
167163
connCountVSCode atomic.Int64
168164
connCountJetBrains atomic.Int64
169165
connCountReconnectingPTY atomic.Int64
@@ -905,10 +901,13 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
905901
switch magicType {
906902
case MagicSSHSessionTypeVSCode:
907903
a.connCountVSCode.Add(1)
904+
defer a.connCountVSCode.Add(-1)
908905
case MagicSSHSessionTypeJetBrains:
909906
a.connCountJetBrains.Add(1)
907+
defer a.connCountJetBrains.Add(-1)
910908
case "":
911909
a.connCountSSHSession.Add(1)
910+
defer a.connCountSSHSession.Add(-1)
912911
default:
913912
a.logger.Warn(ctx, "invalid magic ssh session type specified", slog.F("type", magicType))
914913
}
@@ -1012,6 +1011,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
10121011
defer conn.Close()
10131012

10141013
a.connCountReconnectingPTY.Add(1)
1014+
defer a.connCountReconnectingPTY.Add(-1)
10151015

10161016
connectionID := uuid.NewString()
10171017
logger = logger.With(slog.F("id", msg.ID), slog.F("connection_id", connectionID))
@@ -1210,18 +1210,15 @@ func (a *agent) startReportingConnectionStats(ctx context.Context) {
12101210
ConnectionCount: int64(len(networkStats)),
12111211
ConnectionsByProto: map[string]int64{},
12121212
}
1213-
// Tailscale resets counts on every report!
1214-
// We'd rather have these compound, like Linux does!
12151213
for conn, counts := range networkStats {
12161214
stats.ConnectionsByProto[conn.Proto.String()]++
1217-
stats.RxBytes = a.statRxBytes.Add(int64(counts.RxBytes))
1218-
stats.RxPackets = a.statRxPackets.Add(int64(counts.RxPackets))
1219-
stats.TxBytes = a.statTxBytes.Add(int64(counts.TxBytes))
1220-
stats.TxPackets = a.statTxPackets.Add(int64(counts.TxPackets))
1215+
stats.RxBytes += int64(counts.RxBytes)
1216+
stats.RxPackets += int64(counts.RxPackets)
1217+
stats.TxBytes += int64(counts.TxBytes)
1218+
stats.TxPackets += int64(counts.TxPackets)
12211219
}
12221220

1223-
// Tailscale's connection stats are not cumulative, but it makes no sense to make
1224-
// ours temporary.
1221+
// The count of active sessions.
12251222
stats.SessionCountSSH = a.connCountSSHSession.Load()
12261223
stats.SessionCountVSCode = a.connCountVSCode.Load()
12271224
stats.SessionCountJetBrains = a.connCountJetBrains.Load()
@@ -1270,10 +1267,16 @@ func (a *agent) startReportingConnectionStats(ctx context.Context) {
12701267
// Convert from microseconds to milliseconds.
12711268
stats.ConnectionMedianLatencyMS /= 1000
12721269

1270+
lastStat := a.latestStat.Load()
1271+
if lastStat != nil && reflect.DeepEqual(lastStat, stats) {
1272+
a.logger.Info(ctx, "skipping stat because nothing changed")
1273+
return
1274+
}
1275+
a.latestStat.Store(stats)
1276+
12731277
select {
12741278
case a.connStatsChan <- stats:
1275-
default:
1276-
a.logger.Warn(ctx, "network stat dropped")
1279+
case <-a.closed:
12771280
}
12781281
}
12791282

‎agent/agent_test.go

Lines changed: 69 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,10 @@ func TestAgent_Stats_SSH(t *testing.T) {
6868
session, err := sshClient.NewSession()
6969
require.NoError(t, err)
7070
defer session.Close()
71-
require.NoError(t, session.Run("echo test"))
71+
stdin, err := session.StdinPipe()
72+
require.NoError(t, err)
73+
err = session.Shell()
74+
require.NoError(t, err)
7275

7376
var s *agentsdk.Stats
7477
require.Eventuallyf(t, func() bool {
@@ -78,6 +81,9 @@ func TestAgent_Stats_SSH(t *testing.T) {
7881
}, testutil.WaitLong, testutil.IntervalFast,
7982
"never saw stats: %+v", s,
8083
)
84+
_ = stdin.Close()
85+
err = session.Wait()
86+
require.NoError(t, err)
8187
}
8288

8389
func TestAgent_Stats_ReconnectingPTY(t *testing.T) {
@@ -112,43 +118,69 @@ func TestAgent_Stats_ReconnectingPTY(t *testing.T) {
112118

113119
func TestAgent_Stats_Magic(t *testing.T) {
114120
t.Parallel()
121+
t.Run("StripsEnvironmentVariable", func(t *testing.T) {
122+
t.Parallel()
123+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
124+
defer cancel()
125+
//nolint:dogsled
126+
conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
127+
sshClient, err := conn.SSHClient(ctx)
128+
require.NoError(t, err)
129+
defer sshClient.Close()
130+
session, err := sshClient.NewSession()
131+
require.NoError(t, err)
132+
session.Setenv(agent.MagicSSHSessionTypeEnvironmentVariable, agent.MagicSSHSessionTypeVSCode)
133+
defer session.Close()
115134

116-
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
117-
defer cancel()
118-
119-
//nolint:dogsled
120-
conn, _, stats, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
121-
sshClient, err := conn.SSHClient(ctx)
122-
require.NoError(t, err)
123-
defer sshClient.Close()
124-
session, err := sshClient.NewSession()
125-
require.NoError(t, err)
126-
session.Setenv(agent.MagicSSHSessionTypeEnvironmentVariable, agent.MagicSSHSessionTypeVSCode)
127-
defer session.Close()
128-
129-
command := "sh -c 'echo $" + agent.MagicSSHSessionTypeEnvironmentVariable + "'"
130-
expected := ""
131-
if runtime.GOOS == "windows" {
132-
expected = "%" + agent.MagicSSHSessionTypeEnvironmentVariable + "%"
133-
command = "cmd.exe /c echo " + expected
134-
}
135-
output, err := session.Output(command)
136-
require.NoError(t, err)
137-
require.Equal(t, expected, strings.TrimSpace(string(output)))
138-
var s *agentsdk.Stats
139-
require.Eventuallyf(t, func() bool {
140-
var ok bool
141-
s, ok = <-stats
142-
return ok && s.ConnectionCount > 0 && s.RxBytes > 0 && s.TxBytes > 0 &&
143-
// Ensure that the connection didn't count as a "normal" SSH session.
144-
// This was a special one, so it should be labeled specially in the stats!
145-
s.SessionCountVSCode == 1 &&
146-
// Ensure that connection latency is being counted!
147-
// If it isn't, it's set to -1.
148-
s.ConnectionMedianLatencyMS >= 0
149-
}, testutil.WaitLong, testutil.IntervalFast,
150-
"never saw stats: %+v", s,
151-
)
135+
command := "sh -c 'echo $" + agent.MagicSSHSessionTypeEnvironmentVariable + "'"
136+
expected := ""
137+
if runtime.GOOS == "windows" {
138+
expected = "%" + agent.MagicSSHSessionTypeEnvironmentVariable + "%"
139+
command = "cmd.exe /c echo " + expected
140+
}
141+
output, err := session.Output(command)
142+
require.NoError(t, err)
143+
require.Equal(t, expected, strings.TrimSpace(string(output)))
144+
})
145+
t.Run("Tracks", func(t *testing.T) {
146+
t.Parallel()
147+
if runtime.GOOS == "window" {
148+
t.Skip("Sleeping for infinity doesn't work on Windows")
149+
}
150+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
151+
defer cancel()
152+
//nolint:dogsled
153+
conn, _, stats, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
154+
sshClient, err := conn.SSHClient(ctx)
155+
require.NoError(t, err)
156+
defer sshClient.Close()
157+
session, err := sshClient.NewSession()
158+
require.NoError(t, err)
159+
session.Setenv(agent.MagicSSHSessionTypeEnvironmentVariable, agent.MagicSSHSessionTypeVSCode)
160+
defer session.Close()
161+
stdin, err := session.StdinPipe()
162+
require.NoError(t, err)
163+
err = session.Shell()
164+
require.NoError(t, err)
165+
var s *agentsdk.Stats
166+
require.Eventuallyf(t, func() bool {
167+
var ok bool
168+
s, ok = <-stats
169+
return ok && s.ConnectionCount > 0 && s.RxBytes > 0 && s.TxBytes > 0 &&
170+
// Ensure that the connection didn't count as a "normal" SSH session.
171+
// This was a special one, so it should be labeled specially in the stats!
172+
s.SessionCountVSCode == 1 &&
173+
// Ensure that connection latency is being counted!
174+
// If it isn't, it's set to -1.
175+
s.ConnectionMedianLatencyMS >= 0
176+
}, testutil.WaitLong, testutil.IntervalFast,
177+
"never saw stats: %+v", s,
178+
)
179+
// The shell will automatically exit if there is no stdin!
180+
_ = stdin.Close()
181+
err = session.Wait()
182+
require.NoError(t, err)
183+
})
152184
}
153185

154186
func TestAgent_SessionExec(t *testing.T) {

0 commit comments

Comments
 (0)