opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

Daily Active User Metrics (#3735)

Browse Source 
* agent: add StatsReporter

* Stabilize protoc
This commit is contained in:
Ammar Bandukwala
2022-09-01 14:58:23 -05:00
committed by GitHub
parent e0cb52ceea
commit 30f8fd9b95
47 changed files with 2006 additions and 279 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent/agent.go
+39 -8
View File
@@ -65,6 +65,7 @@ type Options struct {
WebRTCDialer WebRTCDialer
FetchMetadata FetchMetadata
StatsReporter StatsReporter
ReconnectingPTYTimeout time.Duration
EnvironmentVariables map[string]string
Logger slog.Logger
@@ -100,6 +101,8 @@ func New(options Options) io.Closer {
envVars: options.EnvironmentVariables,
coordinatorDialer: options.CoordinatorDialer,
fetchMetadata: options.FetchMetadata,
stats: &Stats{},
statsReporter: options.StatsReporter,
}
server.init(ctx)
return server
@@ -125,6 +128,8 @@ type agent struct {
network *tailnet.Conn
coordinatorDialer CoordinatorDialer
stats *Stats
statsReporter StatsReporter
}
func (a *agent) run(ctx context.Context) {
@@ -194,6 +199,13 @@ func (a *agent) runTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) {
a.logger.Critical(ctx, "create tailnet", slog.Error(err))
return
}
a.network.SetForwardTCPCallback(func(conn net.Conn, listenerExists bool) net.Conn {
if listenerExists {
// If a listener already exists, we would double-wrap the conn.
return conn
}
return a.stats.wrapConn(conn)
})
go a.runCoordinator(ctx)
sshListener, err := a.network.Listen("tcp", ":"+strconv.Itoa(tailnetSSHPort))
@@ -207,7 +219,7 @@ func (a *agent) runTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) {
if err != nil {
return
}
go a.sshServer.HandleConn(conn)
a.sshServer.HandleConn(a.stats.wrapConn(conn))
}
}()
reconnectingPTYListener, err := a.network.Listen("tcp", ":"+strconv.Itoa(tailnetReconnectingPTYPort))
@@ -219,8 +231,10 @@ func (a *agent) runTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) {
for {
conn, err := reconnectingPTYListener.Accept()
if err != nil {
a.logger.Debug(ctx, "accept pty failed", slog.Error(err))
return
}
conn = a.stats.wrapConn(conn)
// This cannot use a JSON decoder, since that can
// buffer additional data that is required for the PTY.
rawLen := make([]byte, 2)
@@ -364,17 +378,17 @@ func (a *agent) runStartupScript(ctx context.Context, script string) error {
return nil
}
func (a *agent) handlePeerConn(ctx context.Context, conn *peer.Conn) {
func (a *agent) handlePeerConn(ctx context.Context, peerConn *peer.Conn) {
go func() {
select {
case <-a.closed:
case <-conn.Closed():
case <-peerConn.Closed():
}
_ = conn.Close()
_ = peerConn.Close()
a.connCloseWait.Done()
}()
for {
channel, err := conn.Accept(ctx)
channel, err := peerConn.Accept(ctx)
if err != nil {
if errors.Is(err, peer.ErrClosed) || a.isClosed() {
return
@@ -383,9 +397,11 @@ func (a *agent) handlePeerConn(ctx context.Context, conn *peer.Conn) {
return
}
conn := channel.NetConn()
switch channel.Protocol() {
case ProtocolSSH:
go a.sshServer.HandleConn(channel.NetConn())
go a.sshServer.HandleConn(a.stats.wrapConn(conn))
case ProtocolReconnectingPTY:
rawID := channel.Label()
// The ID format is referenced in conn.go.
@@ -418,9 +434,9 @@ func (a *agent) handlePeerConn(ctx context.Context, conn *peer.Conn) {
Height: uint16(height),
Width: uint16(width),
Command: idParts[3],
}, channel.NetConn())
}, a.stats.wrapConn(conn))
case ProtocolDial:
go a.handleDial(ctx, channel.Label(), channel.NetConn())
go a.handleDial(ctx, channel.Label(), a.stats.wrapConn(conn))
default:
a.logger.Warn(ctx, "unhandled protocol from channel",
slog.F("protocol", channel.Protocol()),
@@ -514,6 +530,21 @@ func (a *agent) init(ctx context.Context) {
}
go a.run(ctx)
if a.statsReporter != nil {
cl, err := a.statsReporter(ctx, a.logger, func() *Stats {
return a.stats.Copy()
})
if err != nil {
a.logger.Error(ctx, "report stats", slog.Error(err))
return
}
a.connCloseWait.Add(1)
go func() {
defer a.connCloseWait.Done()
<-a.closed
cl.Close()
}()
}
}
// createCommand processes raw command input with OpenSSH-like behavior.
agent/agent_test.go
+119 -11
View File
@@ -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")
agent/stats.go
+67
View File
@@ -0,0 +1,67 @@
package agent
import (
"context"
"io"
"net"
"sync/atomic"
"cdr.dev/slog"
)
// statsConn wraps a net.Conn with statistics.
type statsConn struct {
*Stats
net.Conn `json:"-"`
}
var _ net.Conn = new(statsConn)
func (c *statsConn) Read(b []byte) (n int, err error) {
n, err = c.Conn.Read(b)
atomic.AddInt64(&c.RxBytes, int64(n))
return n, err
}
func (c *statsConn) Write(b []byte) (n int, err error) {
n, err = c.Conn.Write(b)
atomic.AddInt64(&c.TxBytes, int64(n))
return n, err
}
var _ net.Conn = new(statsConn)
// Stats records the Agent's network connection statistics for use in
// user-facing metrics and debugging.
// Each member value must be written and read with atomic.
type Stats struct {
NumConns int64 `json:"num_comms"`
RxBytes int64 `json:"rx_bytes"`
TxBytes int64 `json:"tx_bytes"`
}
func (s *Stats) Copy() *Stats {
return &Stats{
NumConns: atomic.LoadInt64(&s.NumConns),
RxBytes: atomic.LoadInt64(&s.RxBytes),
TxBytes: atomic.LoadInt64(&s.TxBytes),
}
}
// wrapConn returns a new connection that records statistics.
func (s *Stats) wrapConn(conn net.Conn) net.Conn {
atomic.AddInt64(&s.NumConns, 1)
cs := &statsConn{
Stats: s,
Conn: conn,
}
return cs
}
// StatsReporter periodically accept and records agent stats.
type StatsReporter func(
ctx context.Context,
log slog.Logger,
stats func() *Stats,
) (io.Closer, error)