Dec 7, 2023 · Nov 30, 2023 · Nov 30, 2023 · Nov 30, 2023 · Nov 30, 2023 · Dec 1, 2023
diff --git a/agent/agent_test.go b/agent/agent_test.go
 remotePort := tcpAddr.Port
 go echoOnce(t, rl)

 sshClient := setupAgentSSHClient(ctx, t)
 //nolint:dogsled
 conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 sshClient, err := conn.SSHClient(ctx)
 require.NoError(t, err)

 conn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
 tunneledConn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
 require.NoError(t, err)
 defer conn.Close()
 requireEcho(t, conn)
 t.Cleanup(func() {
 // always close on failure of test
 _ = conn.Close()
 _ = tunneledConn.Close()
 })

 var s *agentsdk.Stats
 require.Eventuallyf(t, func() bool {
 var ok bool
 s, ok = <-stats
 return ok && s.ConnectionCount > 0 &&
 s.SessionCountJetBrains == 1
 }, testutil.WaitLong, testutil.IntervalFast,
 "never saw stats with conn open: %+v", s,
 )

 // Manually closing the connection
 requireEcho(t, tunneledConn)
 _ = rl.Close()

 require.Eventuallyf(t, func() bool {
 var ok bool
 s, ok = <-stats
 return ok && s.ConnectionCount == 0 &&
 s.SessionCountJetBrains == 0
 }, testutil.WaitLong, testutil.IntervalFast,
 "never saw stats after conn closes: %+v", s,
 )
 })
 }

diff --git a/agent/agentssh/jetbrainstrack.go b/agent/agentssh/jetbrainstrack.go
 package agentssh

 import (
 "sync"

 "cdr.dev/slog"
 "go.uber.org/atomic"
 gossh "golang.org/x/crypto/ssh"

 type ChannelOnClose struct {
 gossh.Channel
 // once ensures close only decrements the counter once.
 // Because close can be called multiple times.
 once sync.Once
 done func()
 }

 func (c *ChannelOnClose) Close() error {
 c.done()
 c.once.Do(c.done)
 return c.Channel.Close()
 }
Original file line number	Diff line number	Diff line change
Expand Up		@@ -206,12 +206,41 @@ func TestAgent_Stats_Magic(t *testing.T) {
		remotePort := tcpAddr.Port
		go echoOnce(t, rl)

		sshClient := setupAgentSSHClient(ctx, t)
		//nolint:dogsled
		conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
		sshClient, err := conn.SSHClient(ctx)
		require.NoError(t, err)

		conn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
		tunneledConn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
		require.NoError(t, err)
		defer conn.Close()
		requireEcho(t, conn)
		t.Cleanup(func() {
		// always close on failure of test
		_ = conn.Close()
		_ = tunneledConn.Close()
		})

		var s *agentsdk.Stats
		require.Eventuallyf(t, func() bool {
		var ok bool
		s, ok = <-stats
		return ok && s.ConnectionCount > 0 &&
		s.SessionCountJetBrains == 1
		}, testutil.WaitLong, testutil.IntervalFast,
		"never saw stats with conn open: %+v", s,
		)

		// Manually closing the connection
		requireEcho(t, tunneledConn)
		_ = rl.Close()

		require.Eventuallyf(t, func() bool {
		var ok bool
		s, ok = <-stats
		return ok && s.ConnectionCount == 0 &&
		s.SessionCountJetBrains == 0
		}, testutil.WaitLong, testutil.IntervalFast,
		"never saw stats after conn closes: %+v", s,
		)
		})
		}

Expand Down
Original file line number	Diff line number	Diff line change
		@@ -1,6 +1,8 @@
		package agentssh

		import (
		"sync"

		"cdr.dev/slog"
		"go.uber.org/atomic"
		gossh "golang.org/x/crypto/ssh"
Expand DownExpand Up		@@ -54,10 +56,13 @@ func (w ChannelAcceptWatcher) Accept() (gossh.Channel, <-chan gossh.Request, e

		type ChannelOnClose struct {
		gossh.Channel
		// once ensures close only decrements the counter once.
		// Because close can be called multiple times.
		once sync.Once
		done func()
		}

		func (c *ChannelOnClose) Close() error {
		c.done()
		c.once.Do(c.done)
		return c.Channel.Close()
		}