mirror of
https://github.com/coder/coder.git
synced 2025-07-18 14:17:22 +00:00
Daily Active User Metrics (#3735)
* agent: add StatsReporter * Stabilize protoc
This commit is contained in:
@ -19,6 +19,7 @@ import (
|
||||
"time"
|
||||
|
||||
"golang.org/x/xerrors"
|
||||
"tailscale.com/tailcfg"
|
||||
|
||||
scp "github.com/bramvdbogaerde/go-scp"
|
||||
"github.com/google/uuid"
|
||||
@ -51,6 +52,67 @@ func TestMain(m *testing.M) {
|
||||
|
||||
func TestAgent(t *testing.T) {
|
||||
t.Parallel()
|
||||
t.Run("Stats", func(t *testing.T) {
|
||||
for _, tailscale := range []bool{true, false} {
|
||||
t.Run(fmt.Sprintf("tailscale=%v", tailscale), func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
setupAgent := func(t *testing.T) (agent.Conn, <-chan *agent.Stats) {
|
||||
var derpMap *tailcfg.DERPMap
|
||||
if tailscale {
|
||||
derpMap = tailnettest.RunDERPAndSTUN(t)
|
||||
}
|
||||
conn, stats := setupAgent(t, agent.Metadata{
|
||||
DERPMap: derpMap,
|
||||
}, 0)
|
||||
assert.Empty(t, <-stats)
|
||||
return conn, stats
|
||||
}
|
||||
|
||||
t.Run("SSH", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
conn, stats := setupAgent(t)
|
||||
|
||||
sshClient, err := conn.SSHClient()
|
||||
require.NoError(t, err)
|
||||
session, err := sshClient.NewSession()
|
||||
require.NoError(t, err)
|
||||
defer session.Close()
|
||||
|
||||
assert.EqualValues(t, 1, (<-stats).NumConns)
|
||||
assert.Greater(t, (<-stats).RxBytes, int64(0))
|
||||
assert.Greater(t, (<-stats).TxBytes, int64(0))
|
||||
})
|
||||
|
||||
t.Run("ReconnectingPTY", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
conn, stats := setupAgent(t)
|
||||
|
||||
ptyConn, err := conn.ReconnectingPTY(uuid.NewString(), 128, 128, "/bin/bash")
|
||||
require.NoError(t, err)
|
||||
defer ptyConn.Close()
|
||||
|
||||
data, err := json.Marshal(agent.ReconnectingPTYRequest{
|
||||
Data: "echo test\r\n",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
_, err = ptyConn.Write(data)
|
||||
require.NoError(t, err)
|
||||
|
||||
var s *agent.Stats
|
||||
require.Eventuallyf(t, func() bool {
|
||||
var ok bool
|
||||
s, ok = (<-stats)
|
||||
return ok && s.NumConns > 0 && s.RxBytes > 0 && s.TxBytes > 0
|
||||
}, testutil.WaitShort, testutil.IntervalFast,
|
||||
"never saw stats: %+v", s,
|
||||
)
|
||||
})
|
||||
})
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("SessionExec", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
session := setupSSHSession(t, agent.Metadata{})
|
||||
@ -169,7 +231,8 @@ func TestAgent(t *testing.T) {
|
||||
|
||||
t.Run("SFTP", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
sshClient, err := setupAgent(t, agent.Metadata{}, 0).SSHClient()
|
||||
conn, _ := setupAgent(t, agent.Metadata{}, 0)
|
||||
sshClient, err := conn.SSHClient()
|
||||
require.NoError(t, err)
|
||||
client, err := sftp.NewClient(sshClient)
|
||||
require.NoError(t, err)
|
||||
@ -184,7 +247,9 @@ func TestAgent(t *testing.T) {
|
||||
|
||||
t.Run("SCP", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
sshClient, err := setupAgent(t, agent.Metadata{}, 0).SSHClient()
|
||||
|
||||
conn, _ := setupAgent(t, agent.Metadata{}, 0)
|
||||
sshClient, err := conn.SSHClient()
|
||||
require.NoError(t, err)
|
||||
scpClient, err := scp.NewClientBySSH(sshClient)
|
||||
require.NoError(t, err)
|
||||
@ -318,7 +383,7 @@ func TestAgent(t *testing.T) {
|
||||
t.Skip("ConPTY appears to be inconsistent on Windows.")
|
||||
}
|
||||
|
||||
conn := setupAgent(t, agent.Metadata{
|
||||
conn, _ := setupAgent(t, agent.Metadata{
|
||||
DERPMap: tailnettest.RunDERPAndSTUN(t),
|
||||
}, 0)
|
||||
id := uuid.NewString()
|
||||
@ -431,7 +496,7 @@ func TestAgent(t *testing.T) {
|
||||
}()
|
||||
|
||||
// Dial the listener over WebRTC twice and test out of order
|
||||
conn := setupAgent(t, agent.Metadata{}, 0)
|
||||
conn, _ := setupAgent(t, agent.Metadata{}, 0)
|
||||
conn1, err := conn.DialContext(context.Background(), l.Addr().Network(), l.Addr().String())
|
||||
require.NoError(t, err)
|
||||
defer conn1.Close()
|
||||
@ -462,7 +527,7 @@ func TestAgent(t *testing.T) {
|
||||
})
|
||||
|
||||
// Try to dial the non-existent Unix socket over WebRTC
|
||||
conn := setupAgent(t, agent.Metadata{}, 0)
|
||||
conn, _ := setupAgent(t, agent.Metadata{}, 0)
|
||||
netConn, err := conn.DialContext(context.Background(), "unix", filepath.Join(tmpDir, "test.sock"))
|
||||
require.Error(t, err)
|
||||
require.ErrorContains(t, err, "remote dial error")
|
||||
@ -473,7 +538,7 @@ func TestAgent(t *testing.T) {
|
||||
t.Run("Tailnet", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
derpMap := tailnettest.RunDERPAndSTUN(t)
|
||||
conn := setupAgent(t, agent.Metadata{
|
||||
conn, _ := setupAgent(t, agent.Metadata{
|
||||
DERPMap: derpMap,
|
||||
}, 0)
|
||||
defer conn.Close()
|
||||
@ -485,7 +550,7 @@ func TestAgent(t *testing.T) {
|
||||
}
|
||||
|
||||
func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
|
||||
agentConn := setupAgent(t, agent.Metadata{}, 0)
|
||||
agentConn, _ := setupAgent(t, agent.Metadata{}, 0)
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
waitGroup := sync.WaitGroup{}
|
||||
@ -523,7 +588,8 @@ func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exe
|
||||
}
|
||||
|
||||
func setupSSHSession(t *testing.T, options agent.Metadata) *ssh.Session {
|
||||
sshClient, err := setupAgent(t, options, 0).SSHClient()
|
||||
conn, _ := setupAgent(t, options, 0)
|
||||
sshClient, err := conn.SSHClient()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() {
|
||||
_ = sshClient.Close()
|
||||
@ -533,11 +599,21 @@ func setupSSHSession(t *testing.T, options agent.Metadata) *ssh.Session {
|
||||
return session
|
||||
}
|
||||
|
||||
func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration) agent.Conn {
|
||||
type closeFunc func() error
|
||||
|
||||
func (c closeFunc) Close() error {
|
||||
return c()
|
||||
}
|
||||
|
||||
func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration) (
|
||||
agent.Conn,
|
||||
<-chan *agent.Stats,
|
||||
) {
|
||||
client, server := provisionersdk.TransportPipe()
|
||||
tailscale := metadata.DERPMap != nil
|
||||
coordinator := tailnet.NewCoordinator()
|
||||
agentID := uuid.New()
|
||||
statsCh := make(chan *agent.Stats)
|
||||
closer := agent.New(agent.Options{
|
||||
FetchMetadata: func(ctx context.Context) (agent.Metadata, error) {
|
||||
return metadata, nil
|
||||
@ -557,6 +633,38 @@ func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration)
|
||||
},
|
||||
Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
|
||||
ReconnectingPTYTimeout: ptyTimeout,
|
||||
StatsReporter: func(ctx context.Context, log slog.Logger, statsFn func() *agent.Stats) (io.Closer, error) {
|
||||
doneCh := make(chan struct{})
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
|
||||
go func() {
|
||||
defer close(doneCh)
|
||||
|
||||
t := time.NewTicker(time.Millisecond * 100)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
}
|
||||
select {
|
||||
case statsCh <- statsFn():
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
// We don't want to send old stats.
|
||||
continue
|
||||
}
|
||||
}
|
||||
}()
|
||||
return closeFunc(func() error {
|
||||
cancel()
|
||||
<-doneCh
|
||||
close(statsCh)
|
||||
return nil
|
||||
}), nil
|
||||
},
|
||||
})
|
||||
t.Cleanup(func() {
|
||||
_ = client.Close()
|
||||
@ -586,7 +694,7 @@ func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration)
|
||||
conn.SetNodeCallback(sendNode)
|
||||
return &agent.TailnetConn{
|
||||
Conn: conn,
|
||||
}
|
||||
}, statsCh
|
||||
}
|
||||
conn, err := peerbroker.Dial(stream, []webrtc.ICEServer{}, &peer.ConnOptions{
|
||||
Logger: slogtest.Make(t, nil),
|
||||
@ -599,7 +707,7 @@ func setupAgent(t *testing.T, metadata agent.Metadata, ptyTimeout time.Duration)
|
||||
return &agent.WebRTCConn{
|
||||
Negotiator: api,
|
||||
Conn: conn,
|
||||
}
|
||||
}, statsCh
|
||||
}
|
||||
|
||||
var dialTestPayload = []byte("dean-was-here123")
|
||||
|
Reference in New Issue
Block a user