Thanks to visit codestin.com
Credit goes to github.com

Skip to content

feat: normalize workspace agent session counts into a child table #27179

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.

Sign up for GitHub

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

Closed
EhabY wants to merge 1 commit into from
Closed

feat: normalize workspace agent session counts into a child table #27179

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
3 changes: 3 additions & 0 deletions .github/workflows/typos.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,9 @@ AIBridge = "AI Bridge"

[default.extend-words]
AKS = "AKS"
# BRIN is a postgres index type
brin = "brin"
BRIN = "BRIN"
# do as sudo replacement
doas = "doas"
darcula = "darcula"
Expand Down
48 changes: 19 additions & 29 deletions agent/agent.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ import (
"github.com/coder/coder/v2/buildinfo"
"github.com/coder/coder/v2/cli/gitauth"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/idemetadata"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/agentsdk"
"github.com/coder/coder/v2/codersdk/workspacesdk"
Expand Down Expand Up @@ -127,23 +128,14 @@ type Options struct {
}

type Client interface {
ConnectRPC29(ctx context.Context) (
proto.DRPCAgentClient29, tailnetproto.DRPCTailnetClient28, error,
ConnectRPC211(ctx context.Context) (
proto.DRPCAgentClient211, tailnetproto.DRPCTailnetClient28, error,
)
// ConnectRPC29WithRole is like ConnectRPC29 but sends an explicit
// ConnectRPC211WithRole is like ConnectRPC211 but sends an explicit
// role query parameter to the server. The workspace agent should
// use role "agent" to enable connection monitoring.
ConnectRPC29WithRole(ctx context.Context, role string) (
proto.DRPCAgentClient29, tailnetproto.DRPCTailnetClient28, error,
)
ConnectRPC210(ctx context.Context) (
proto.DRPCAgentClient210, tailnetproto.DRPCTailnetClient28, error,
)
// ConnectRPC210WithRole is like ConnectRPC210 but sends an explicit
// role query parameter to the server. The workspace agent should
// use role "agent" to enable connection monitoring.
ConnectRPC210WithRole(ctx context.Context, role string) (
proto.DRPCAgentClient210, tailnetproto.DRPCTailnetClient28, error,
ConnectRPC211WithRole(ctx context.Context, role string) (
proto.DRPCAgentClient211, tailnetproto.DRPCTailnetClient28, error,
)
tailnet.DERPMapRewriter
agentsdk.RefreshableSessionTokenProvider
Expand Down Expand Up @@ -426,17 +418,16 @@ func (a *agent) init() {
BlockLocalPortForwarding: a.blockLocalPortForwarding,
ReportConnection: func(id uuid.UUID, magicType agentssh.MagicSessionType, ip string) func(code int, reason string) {
var connectionType proto.Connection_Type
switch magicType {
case agentssh.MagicSessionTypeSSH:
// The proto enum cannot represent arbitrary session types, so
// map by family.
switch idemetadata.Family(string(magicType)) {
case idemetadata.AppNameSSH:
connectionType = proto.Connection_SSH
case agentssh.MagicSessionTypeVSCode:
case idemetadata.AppNameVSCode:
connectionType = proto.Connection_VSCODE
case agentssh.MagicSessionTypeJetBrains:
case idemetadata.AppNameJetBrains:
connectionType = proto.Connection_JETBRAINS
case agentssh.MagicSessionTypeUnknown:
connectionType = proto.Connection_TYPE_UNSPECIFIED
default:
a.logger.Error(a.hardCtx, "unhandled magic session type when reporting connection", slog.F("magic_type", magicType))
connectionType = proto.Connection_TYPE_UNSPECIFIED
}

Expand Down Expand Up @@ -1176,7 +1167,7 @@ func (a *agent) run() (retErr error) {
// ConnectRPC returns the dRPC connection we use for the Agent and Tailnet v2+ APIs.
// We pass role "agent" to enable connection monitoring on the server, which tracks
// the agent's connectivity state (first_connected_at, last_connected_at, disconnected_at).
aAPI, tAPI, err := a.client.ConnectRPC210WithRole(a.hardCtx, "agent")
aAPI, tAPI, err := a.client.ConnectRPC211WithRole(a.hardCtx, "agent")
if err != nil {
return err
}
Expand Down Expand Up @@ -2162,13 +2153,12 @@ func (a *agent) Collect(ctx context.Context, networkStats map[netlogtype.Connect
stats.TxPackets += int64(counts.TxPackets)
}

// The count of active sessions.
sshStats := a.sshServer.ConnStats()
stats.SessionCountSsh = sshStats.Sessions
stats.SessionCountVscode = sshStats.VSCode
stats.SessionCountJetbrains = sshStats.JetBrains

stats.SessionCountReconnectingPty = a.reconnectingPTYServer.ConnCount()
// Active sessions per app; the deprecated fixed fields are left zero.
stats.SessionCounts = a.sshServer.SessionCounts()
if count := a.reconnectingPTYServer.ConnCount(); count > 0 {
// += because an SSH session may already carry this type.
stats.SessionCounts[idemetadata.AppNameReconnectingPTY] += count
}

// Compute the median connection latency!
a.logger.Debug(ctx, "starting peer latency measurement for stats")
Expand Down
18 changes: 9 additions & 9 deletions agent/agent_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,7 +229,7 @@ func assertSSHStats(t *testing.T, stats <-chan *proto.Stats) {
return false
}
t.Logf("got stats: ConnectionCount=%d, RxBytes=%d, TxBytes=%d, SessionCountSsh=%d",
s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCountSsh)
s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCounts["ssh"])
if s.ConnectionCount > 0 {
connectionCountSeen = true
}
Expand All @@ -239,7 +239,7 @@ func assertSSHStats(t *testing.T, stats <-chan *proto.Stats) {
if s.TxBytes > 0 {
txBytesSeen = true
}
if s.SessionCountSsh == 1 {
if s.SessionCounts["ssh"] == 1 {
sessionCountSSHSeen = true
}
return connectionCountSeen && rxBytesSeen && txBytesSeen && sessionCountSSHSeen
Expand Down Expand Up @@ -287,7 +287,7 @@ func TestAgent_Stats_ReconnectingPTY(t *testing.T) {
if s.TxBytes > 0 {
txBytesSeen = true
}
if s.SessionCountReconnectingPty == 1 {
if s.SessionCounts["reconnecting_pty"] == 1 {
sessionCountReconnectingPTYSeen = true
}
return connectionCountSeen && rxBytesSeen && txBytesSeen && sessionCountReconnectingPTYSeen
Expand Down Expand Up @@ -347,11 +347,11 @@ func TestAgent_Stats_Magic(t *testing.T) {
require.Eventuallyf(t, func() bool {
s, ok := <-stats
t.Logf("got stats: ok=%t, ConnectionCount=%d, RxBytes=%d, TxBytes=%d, SessionCountVSCode=%d, ConnectionMedianLatencyMS=%f",
ok, s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCountVscode, s.ConnectionMedianLatencyMs)
ok, s.ConnectionCount, s.RxBytes, s.TxBytes, s.SessionCounts["vscode"], s.ConnectionMedianLatencyMs)
return ok &&
// Ensure that the connection didn't count as a "normal" SSH session.
// This was a special one, so it should be labeled specially in the stats!
s.SessionCountVscode == 1 &&
s.SessionCounts["vscode"] == 1 &&
// Ensure that connection latency is being counted!
// If it isn't, it's set to -1.
s.ConnectionMedianLatencyMs >= 0
Expand Down Expand Up @@ -417,8 +417,8 @@ func TestAgent_Stats_Magic(t *testing.T) {
require.Eventuallyf(t, func() bool {
s, ok := <-stats
t.Logf("got stats with conn open: ok=%t, ConnectionCount=%d, SessionCountJetBrains=%d",
ok, s.ConnectionCount, s.SessionCountJetbrains)
return ok && s.SessionCountJetbrains == 1
ok, s.ConnectionCount, s.SessionCounts["jetbrains"])
return ok && s.SessionCounts["jetbrains"] == 1
}, testutil.WaitLong, testutil.IntervalFast,
"never saw stats with conn open",
)
Expand All @@ -431,9 +431,9 @@ func TestAgent_Stats_Magic(t *testing.T) {
require.Eventuallyf(t, func() bool {
s, ok := <-stats
t.Logf("got stats after disconnect %t, %d",
ok, s.SessionCountJetbrains)
ok, s.SessionCounts["jetbrains"])
return ok &&
s.SessionCountJetbrains == 0
s.SessionCounts["jetbrains"] == 0
}, testutil.WaitLong, testutil.IntervalFast,
"never saw stats after conn closes",
)
Expand Down
120 changes: 71 additions & 49 deletions agent/agentssh/agentssh.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"errors"
"fmt"
"io"
"maps"
"net"
"os"
"os/exec"
Expand All @@ -31,6 +32,7 @@ import (
"github.com/coder/coder/v2/agent/agentexec"
"github.com/coder/coder/v2/agent/agentrsa"
"github.com/coder/coder/v2/agent/usershell"
"github.com/coder/coder/v2/coderd/idemetadata"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/pty"
)
Expand Down Expand Up @@ -71,17 +73,16 @@ const (
ContainerUserEnvironmentVariable = "CODER_CONTAINER_USER"
)

// MagicSessionType enums.
// Well-known magic session types, defined as canonical app names so the
// agent and server vocabularies cannot drift.
const (
// MagicSessionTypeUnknown means the session type could not be determined.
MagicSessionTypeUnknown MagicSessionType = "unknown"
// MagicSessionTypeSSH is the default session type.
MagicSessionTypeSSH MagicSessionType = "ssh"
MagicSessionTypeSSH MagicSessionType = idemetadata.AppNameSSH
// MagicSessionTypeVSCode is set in the SSH config by the VS Code extension to identify itself.
MagicSessionTypeVSCode MagicSessionType = "vscode"
MagicSessionTypeVSCode MagicSessionType = idemetadata.AppNameVSCode
// MagicSessionTypeJetBrains is set in the SSH config by the JetBrains
// extension to identify itself.
MagicSessionTypeJetBrains MagicSessionType = "jetbrains"
MagicSessionTypeJetBrains MagicSessionType = idemetadata.AppNameJetBrains
)

// BlockedFileTransferCommands contains a list of restricted file transfer commands.
Expand Down Expand Up @@ -143,7 +144,10 @@ type Server struct {
conns map[net.Conn]struct{}
sessions map[ssh.Session]struct{}
processes map[*os.Process]struct{}
closing chan struct{}
// sessionCounts tracks active sessions per session type; zero-count
// entries are removed.
sessionCounts map[string]int64
closing chan struct{}
// Wait for goroutines to exit, waited without
// a lock on mu but protected by closing.
wg sync.WaitGroup
Expand All @@ -155,10 +159,6 @@ type Server struct {

config *Config

connCountVSCode atomic.Int64
connCountJetBrains atomic.Int64
connCountSSHSession atomic.Int64

metrics *sshServerMetrics
}

Expand Down Expand Up @@ -200,13 +200,14 @@ func NewServer(ctx context.Context, logger slog.Logger, prometheusRegistry *prom

metrics := newSSHServerMetrics(prometheusRegistry)
s := &Server{
Execer: execer,
listeners: make(map[net.Listener]struct{}),
fs: fs,
conns: make(map[net.Conn]struct{}),
sessions: make(map[ssh.Session]struct{}),
processes: make(map[*os.Process]struct{}),
logger: logger,
Execer: execer,
listeners: make(map[net.Listener]struct{}),
fs: fs,
conns: make(map[net.Conn]struct{}),
sessions: make(map[ssh.Session]struct{}),
processes: make(map[*os.Process]struct{}),
sessionCounts: make(map[string]int64),
logger: logger,

config: config,

Expand All @@ -228,11 +229,14 @@ func NewServer(ctx context.Context, logger slog.Logger, prometheusRegistry *prom
},
}

startJetBrainsSession := func() (endSession func()) {
return s.startSession(MagicSessionTypeJetBrains)
}
srv := &ssh.Server{
ChannelHandlers: map[string]ssh.ChannelHandler{
"direct-tcpip": func(srv *ssh.Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx ssh.Context) {
// Wrapper is designed to find and track JetBrains Gateway connections.
wrapped := NewJetbrainsChannelWatcher(ctx, s.logger, s.config.ReportConnection, newChan, &s.connCountJetBrains)
wrapped := NewJetbrainsChannelWatcher(ctx, s.logger, s.config.ReportConnection, newChan, startJetBrainsSession)
ssh.DirectTCPIPHandler(srv, conn, wrapped, ctx)
},
"[email protected]": func(srv *ssh.Server, conn *gossh.ServerConn, newChan gossh.NewChannel, ctx ssh.Context) {
Expand Down Expand Up @@ -324,18 +328,38 @@ func NewServer(ctx context.Context, logger slog.Logger, prometheusRegistry *prom
return s, nil
}

type ConnStats struct {
Sessions int64
VSCode int64
JetBrains int64
// startSession increments the active session count for the given session
// type and returns a function that decrements it. Known-family types
// always get their own counter; unrecognized types past the cap count
// under unknown, mirroring the server-side cap.
func (s *Server) startSession(magicType MagicSessionType) (endSession func()) {
key := string(magicType)
s.mu.Lock()
defer s.mu.Unlock()
if _, ok := s.sessionCounts[key]; !ok &&
len(s.sessionCounts) >= idemetadata.MaxSessionCountEntries &&
idemetadata.Family(key) == idemetadata.AppNameUnknown {
s.logger.Debug(context.Background(), "session type counter cap reached, counting under unknown",
slog.F("session_type", key),
)
key = idemetadata.AppNameUnknown
}
s.sessionCounts[key]++
return func() {
s.mu.Lock()
defer s.mu.Unlock()
s.sessionCounts[key]--
if s.sessionCounts[key] <= 0 {
delete(s.sessionCounts, key)
}
}
}

func (s *Server) ConnStats() ConnStats {
return ConnStats{
Sessions: s.connCountSSHSession.Load(),
VSCode: s.connCountVSCode.Load(),
JetBrains: s.connCountJetBrains.Load(),
}
// SessionCounts returns a snapshot of active sessions per session type.
func (s *Server) SessionCounts() map[string]int64 {
s.mu.RLock()
defer s.mu.RUnlock()
return maps.Clone(s.sessionCounts)
}

func extractMagicSessionType(env []string) (magicType MagicSessionType, rawType string, filteredEnv []string) {
Expand All @@ -348,16 +372,12 @@ func extractMagicSessionType(env []string) (magicType MagicSessionType, rawType
// Keep going, we'll use the last instance of the env.
}

// Always force lowercase checking to be case-insensitive.
switch MagicSessionType(strings.ToLower(rawType)) {
case MagicSessionTypeVSCode:
magicType = MagicSessionTypeVSCode
case MagicSessionTypeJetBrains:
magicType = MagicSessionTypeJetBrains
case "", MagicSessionTypeSSH:
if rawType == "" {
magicType = MagicSessionTypeSSH
default:
magicType = MagicSessionTypeUnknown
} else {
// Canonicalize, do not classify: unknown names flow through so
// new IDEs need no code changes.
magicType = MagicSessionType(idemetadata.Normalize(rawType))
}

return magicType, rawType, slices.DeleteFunc(env, func(kv string) bool {
Expand Down Expand Up @@ -424,6 +444,13 @@ func (s *Server) sessionHandler(session ssh.Session) {

env := session.Environ()
magicType, magicTypeRaw, env := extractMagicSessionType(env)
magicTypeFamily := idemetadata.Family(string(magicType))
if magicTypeFamily == idemetadata.AppNameUnknown {
logger.Debug(ctx, "unrecognized ssh session type",
slog.F("magic_type", magicType),
slog.F("raw_type", magicTypeRaw),
)
}

// It's not safe to assume RemoteAddr() returns a non-nil value. slog.F usage is fine because it correctly
// handles nil.
Expand All @@ -449,19 +476,14 @@ func (s *Server) sessionHandler(session ssh.Session) {

reportSession := true

switch magicType {
case MagicSessionTypeVSCode:
s.connCountVSCode.Add(1)
defer s.connCountVSCode.Add(-1)
case MagicSessionTypeJetBrains:
if magicTypeFamily == idemetadata.AppNameJetBrains {
// Do nothing here because JetBrains launches hundreds of ssh sessions.
// We instead track JetBrains in the single persistent tcp forwarding channel.
// We instead track JetBrains in the single persistent tcp forwarding
// channel.
reportSession = false
case MagicSessionTypeSSH:
s.connCountSSHSession.Add(1)
defer s.connCountSSHSession.Add(-1)
case MagicSessionTypeUnknown:
logger.Warn(ctx, "invalid magic ssh session type specified", slog.F("raw_type", magicTypeRaw))
} else {
endSession := s.startSession(magicType)
defer endSession()
}

closeCause := func(_ string) {}
Expand Down
Loading
Loading