fix test

f0ssel · f0ssel · commit 218f4b0c28bd · 2024-10-28T16:07:56.000Z
diff --git a/agent/agent.go b/agent/agent.go
@@ -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
diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
@@ -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
 }
@@ -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,
@@ -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,
 	}
 }
 
@@ -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))
diff --git a/agent/agentssh/jetbrainstrack.go b/agent/agentssh/jetbrainstrack.go
@@ -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.
@@ -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)),
 	}
 }
@@ -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")
 
diff --git a/coderd/activitybump_test.go b/coderd/activitybump_test.go
@@ -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)
