Oct 29, 2024 · Oct 25, 2024 · Oct 25, 2024 · Oct 28, 2024 · Oct 28, 2024 · Oct 28, 2024
diff --git a/agent/agent.go b/agent/agent.go

 stats.SessionCountReconnectingPty = a.connCountReconnectingPTY.Load()

 // if we've seen sessions but currently have no connections we
 // just count the sum of the sessions as connections
 if stats.ConnectionCount == 0 {
 stats.ConnectionCount = stats.SessionCountSsh + stats.SessionCountVscode + stats.SessionCountJetbrains + stats.SessionCountReconnectingPty
 }

 // Compute the median connection latency!
 a.logger.Debug(ctx, "starting peer latency measurement for stats")
 var wg sync.WaitGroup
diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
 connCountVSCode     atomic.Int64
 connCountJetBrains  atomic.Int64
 connCountSSHSession atomic.Int64
 seenVSCode          atomic.Bool
 seenJetBrains       atomic.Bool
 seenSSHSession      atomic.Bool

 metrics *sshServerMetrics
 }
 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, newChan, &s.connCountJetBrains)
 wrapped := NewJetbrainsChannelWatcher(ctx, s.logger, newChan, &s.connCountJetBrains, &s.seenJetBrains)
 ssh.DirectTCPIPHandler(srv, conn, wrapped, ctx)
 },
 "direct-streamlocal@openssh.com": directStreamLocalHandler,
 }

 func (s *Server) ConnStats() ConnStats {
 // if we have 0 active connections, but we have seen a connection
 // since the last time we collected, count it as 1 so that workspace
 // activity is properly counted.
 sshCount := s.connCountSSHSession.Load()
 if sshCount == 0 && s.seenSSHSession.Load() {
 sshCount = 1
 }
 vscode := s.connCountVSCode.Load()
 if vscode == 0 && s.seenVSCode.Load() {
 vscode = 1
 }
 jetbrains := s.connCountJetBrains.Load()
 if jetbrains == 0 && s.seenJetBrains.Load() {
 jetbrains = 1
 }

 // Reset the seen trackers for the next collection.
 s.seenSSHSession.Store(false)
 s.seenVSCode.Store(false)
 s.seenJetBrains.Store(false)

 return ConnStats{
 Sessions:s.connCountSSHSession.Load(),
 VSCode:s.connCountVSCode.Load(),
 JetBrains:s.connCountJetBrains.Load(),
 Sessions:sshCount,
 VSCode:vscode,
 JetBrains:jetbrains,
 }
 }

 switch magicType {
 case MagicSessionTypeVSCode:
 s.connCountVSCode.Add(1)
 s.seenVSCode.Store(true)
 defer s.connCountVSCode.Add(-1)
 case MagicSessionTypeJetBrains:
 // Do nothing here because JetBrains launches hundreds of ssh sessions.
 // We instead track JetBrains in the single persistent tcp forwarding channel.
 case "":
 s.connCountSSHSession.Add(1)
 s.seenSSHSession.Store(true)
 defer s.connCountSSHSession.Add(-1)
 default:
 logger.Warn(ctx, "invalid magic ssh session type specified", slog.F("type", magicType))
diff --git a/agent/agentssh/jetbrainstrack.go b/agent/agentssh/jetbrainstrack.go
 type JetbrainsChannelWatcher struct {
 gossh.NewChannel
 jetbrainsCounter *atomic.Int64
 jetbrainsSeen    *atomic.Bool
 logger           slog.Logger
 }

 func NewJetbrainsChannelWatcher(ctx ssh.Context, logger slog.Logger, newChannel gossh.NewChannel, counter *atomic.Int64) gossh.NewChannel {
 func NewJetbrainsChannelWatcher(ctx ssh.Context, logger slog.Logger, newChannel gossh.NewChannel, counter *atomic.Int64, seen *atomic.Bool) gossh.NewChannel {
 d := localForwardChannelData{}
 if err := gossh.Unmarshal(newChannel.ExtraData(), &d); err != nil {
 // If the data fails to unmarshal, do nothing.
 return &JetbrainsChannelWatcher{
 NewChannel:       newChannel,
 jetbrainsCounter: counter,
 jetbrainsSeen:    seen,
 logger:           logger.With(slog.F("destination_port", d.DestPort)),
 }
 }
 return c, r, err
 }
 w.jetbrainsCounter.Add(1)
 w.jetbrainsSeen.Store(true)
 // nolint: gocritic // JetBrains is a proper noun and should be capitalized
 w.logger.Debug(context.Background(), "JetBrains watcher accepted channel")

diff --git a/coderd/activitybump_test.go b/coderd/activitybump_test.go
 time.Sleep(time.Second * 3)
 sshConn, err := conn.SSHClient(ctx)
 require.NoError(t, err)
 sess, err := sshConn.NewSession()
 require.NoError(t, err)
 err = sess.Shell()
 require.NoError(t, err)
 err = sess.Close()
 require.NoError(t, err)
 _ = sshConn.Close()

 assertBumped(true)
Original file line number	Diff line number	Diff line change
Expand Up		@@ -1503,6 +1503,12 @@ func (a *agent) Collect(ctx context.Context, networkStats map[netlogtype.Connect

		stats.SessionCountReconnectingPty = a.connCountReconnectingPTY.Load()

		// if we've seen sessions but currently have no connections we
		// just count the sum of the sessions as connections
		if stats.ConnectionCount == 0 {
		stats.ConnectionCount = stats.SessionCountSsh + stats.SessionCountVscode + stats.SessionCountJetbrains + stats.SessionCountReconnectingPty
		}

		// Compute the median connection latency!
		a.logger.Debug(ctx, "starting peer latency measurement for stats")
		var wg sync.WaitGroup
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -105,6 +105,9 @@ type Server struct {
		connCountVSCode atomic.Int64
		connCountJetBrains atomic.Int64
		connCountSSHSession atomic.Int64
		seenVSCode atomic.Bool
		seenJetBrains atomic.Bool
		seenSSHSession atomic.Bool

		metrics *sshServerMetrics
		}
Expand DownExpand Up		@@ -167,7 +170,7 @@ func NewServer(ctx context.Context, logger slog.Logger, prometheusRegistry *prom
		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, newChan, &s.connCountJetBrains)
		wrapped := NewJetbrainsChannelWatcher(ctx, s.logger, newChan, &s.connCountJetBrains, &s.seenJetBrains)
		ssh.DirectTCPIPHandler(srv, conn, wrapped, ctx)
		},
		"direct-streamlocal@openssh.com": directStreamLocalHandler,
Expand DownExpand Up		@@ -245,10 +248,31 @@ type ConnStats struct {
		}

		func (s *Server) ConnStats() ConnStats {
		// if we have 0 active connections, but we have seen a connection
		// since the last time we collected, count it as 1 so that workspace
		// activity is properly counted.
		sshCount := s.connCountSSHSession.Load()
		if sshCount == 0 && s.seenSSHSession.Load() {
		sshCount = 1
		}
		vscode := s.connCountVSCode.Load()
		if vscode == 0 && s.seenVSCode.Load() {
		vscode = 1
		}
		jetbrains := s.connCountJetBrains.Load()
		if jetbrains == 0 && s.seenJetBrains.Load() {
		jetbrains = 1
		}

		// Reset the seen trackers for the next collection.
		s.seenSSHSession.Store(false)
		s.seenVSCode.Store(false)
		s.seenJetBrains.Store(false)

		return ConnStats{
		Sessions:s.connCountSSHSession.Load(),
		VSCode:s.connCountVSCode.Load(),
		JetBrains:s.connCountJetBrains.Load(),
		Sessions:sshCount,
		VSCode:vscode,
		JetBrains:jetbrains,
		}
		}

Expand DownExpand Up		@@ -392,12 +416,14 @@ func (s *Server) sessionStart(logger slog.Logger, session ssh.Session, extraEnv
		switch magicType {
		case MagicSessionTypeVSCode:
		s.connCountVSCode.Add(1)
		s.seenVSCode.Store(true)
		defer s.connCountVSCode.Add(-1)
		case MagicSessionTypeJetBrains:
		// Do nothing here because JetBrains launches hundreds of ssh sessions.
		// We instead track JetBrains in the single persistent tcp forwarding channel.
		case "":
		s.connCountSSHSession.Add(1)
		s.seenSSHSession.Store(true)
		defer s.connCountSSHSession.Add(-1)
		default:
		logger.Warn(ctx, "invalid magic ssh session type specified", slog.F("type", magicType))
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -27,10 +27,11 @@ type localForwardChannelData struct {
		type JetbrainsChannelWatcher struct {
		gossh.NewChannel
		jetbrainsCounter *atomic.Int64
		jetbrainsSeen *atomic.Bool
		logger slog.Logger
		}

		func NewJetbrainsChannelWatcher(ctx ssh.Context, logger slog.Logger, newChannel gossh.NewChannel, counter *atomic.Int64) gossh.NewChannel {
		func NewJetbrainsChannelWatcher(ctx ssh.Context, logger slog.Logger, newChannel gossh.NewChannel, counter atomic.Int64, seen atomic.Bool) gossh.NewChannel {
		d := localForwardChannelData{}
		if err := gossh.Unmarshal(newChannel.ExtraData(), &d); err != nil {
		// If the data fails to unmarshal, do nothing.
Expand DownExpand Up		@@ -60,6 +61,7 @@ func NewJetbrainsChannelWatcher(ctx ssh.Context, logger slog.Logger, newChannel
		return &JetbrainsChannelWatcher{
		NewChannel: newChannel,
		jetbrainsCounter: counter,
		jetbrainsSeen: seen,
		logger: logger.With(slog.F("destination_port", d.DestPort)),
		}
		}
Expand All		@@ -70,6 +72,7 @@ func (w JetbrainsChannelWatcher) Accept() (gossh.Channel, <-chan gossh.Request
		return c, r, err
		}
		w.jetbrainsCounter.Add(1)
		w.jetbrainsSeen.Store(true)
		// nolint: gocritic // JetBrains is a proper noun and should be capitalized
		w.logger.Debug(context.Background(), "JetBrains watcher accepted channel")

Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -212,6 +212,12 @@ func TestWorkspaceActivityBump(t *testing.T) {
		time.Sleep(time.Second * 3)
		sshConn, err := conn.SSHClient(ctx)
		require.NoError(t, err)
		sess, err := sshConn.NewSession()
		require.NoError(t, err)
		err = sess.Shell()
		require.NoError(t, err)
		err = sess.Close()
		require.NoError(t, err)
		_ = sshConn.Close()

		assertBumped(true)
Expand Down