mirror of
https://github.com/coder/coder.git
synced 2025-07-09 11:45:56 +00:00
bf3224e373
* fix: Refactor agent to consume API client This simplifies a lot of code by creating an interface for the codersdk client into the agent. It also moves agent authentication code so instance identity will work between restarts. Fixes #3485 and #4082. * Fix client reconnections
59 lines
1.2 KiB
Go
59 lines
1.2 KiB
Go
package agent
|
|
|
|
import (
|
|
"net"
|
|
"sync/atomic"
|
|
|
|
"github.com/coder/coder/codersdk"
|
|
)
|
|
|
|
// 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() *codersdk.AgentStats {
|
|
return &codersdk.AgentStats{
|
|
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
|
|
}
|