agent/stats_internal_test.go
Rama main · réplica de lectura del repositorio original.
package agent
import (
"context"
"net/netip"
"sync"
"testing"
"time"
"github.com/stretchr/testify/require"
"google.golang.org/protobuf/types/known/durationpb"
"tailscale.com/types/ipproto"
"tailscale.com/types/netlogtype"
"github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/testutil"
)
func TestStatsReporter(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
logger := testutil.Logger(t)
fSource := newFakeNetworkStatsSource(ctx, t)
fCollector := newFakeCollector(t)
fDest := newFakeStatsDest()
uut := newStatsReporter(logger, fSource, fCollector, DefaultStatsReportInterval)
_ = testutil.TryReceive(ctx, t, fSource.period) // drain construction-time install
loopErr := make(chan error, 1)
loopCtx, loopCancel := context.WithCancel(ctx)
go func() {
err := uut.reportLoop(loopCtx, fDest)
loopErr <- err
}()
// initial request to get duration
req := testutil.TryReceive(ctx, t, fDest.reqs)
require.NotNil(t, req)
require.Nil(t, req.Stats)
interval := time.Second * 34
testutil.RequireSend(ctx, t, fDest.resps, &proto.UpdateStatsResponse{ReportInterval: durationpb.New(interval)})
// call to source to set the callback and interval
gotInterval := testutil.TryReceive(ctx, t, fSource.period)
require.Equal(t, interval, gotInterval)
// callback returning netstats
netStats := map[netlogtype.Connection]netlogtype.Counts{
{
Proto: ipproto.TCP,
Src: netip.MustParseAddrPort("192.168.1.33:4887"),
Dst: netip.MustParseAddrPort("192.168.2.99:9999"),
}: {
TxPackets: 22,
TxBytes: 23,
RxPackets: 24,
RxBytes: 25,
},
}
fSource.callback(time.Now(), time.Now(), netStats, nil)
// collector called to complete the stats
gotNetStats := testutil.TryReceive(ctx, t, fCollector.calls)
require.Equal(t, netStats, gotNetStats)
// while we are collecting the stats, send in two new netStats to simulate
// what happens if we don't keep up. The stats should be accumulated.
netStats0 := map[netlogtype.Connection]netlogtype.Counts{
{
Proto: ipproto.TCP,
Src: netip.MustParseAddrPort("192.168.1.33:4887"),
Dst: netip.MustParseAddrPort("192.168.2.99:9999"),
}: {
TxPackets: 10,
TxBytes: 10,
RxPackets: 10,
RxBytes: 10,
},
}
fSource.callback(time.Now(), time.Now(), netStats0, nil)
netStats1 := map[netlogtype.Connection]netlogtype.Counts{
{
Proto: ipproto.TCP,
Src: netip.MustParseAddrPort("192.168.1.33:4887"),
Dst: netip.MustParseAddrPort("192.168.2.99:9999"),
}: {
TxPackets: 11,
TxBytes: 11,
RxPackets: 11,
RxBytes: 11,
},
}
fSource.callback(time.Now(), time.Now(), netStats1, nil)
// complete first collection
stats := &proto.Stats{SessionCountJetbrains: 55}
testutil.RequireSend(ctx, t, fCollector.stats, stats)
// destination called to report the first stats
update := testutil.TryReceive(ctx, t, fDest.reqs)
require.NotNil(t, update)
require.Equal(t, stats, update.Stats)
testutil.RequireSend(ctx, t, fDest.resps, &proto.UpdateStatsResponse{ReportInterval: durationpb.New(interval)})
// second update -- netStat0 and netStats1 are accumulated and reported
wantNetStats := map[netlogtype.Connection]netlogtype.Counts{
{
Proto: ipproto.TCP,
Src: netip.MustParseAddrPort("192.168.1.33:4887"),
Dst: netip.MustParseAddrPort("192.168.2.99:9999"),
}: {
TxPackets: 21,
TxBytes: 21,
RxPackets: 21,
RxBytes: 21,
},
}
gotNetStats = testutil.TryReceive(ctx, t, fCollector.calls)
require.Equal(t, wantNetStats, gotNetStats)
stats = &proto.Stats{SessionCountJetbrains: 66}
testutil.RequireSend(ctx, t, fCollector.stats, stats)
update = testutil.TryReceive(ctx, t, fDest.reqs)
require.NotNil(t, update)
require.Equal(t, stats, update.Stats)
interval2 := 27 * time.Second
testutil.RequireSend(ctx, t, fDest.resps, &proto.UpdateStatsResponse{ReportInterval: durationpb.New(interval2)})
// set the new interval
gotInterval = testutil.TryReceive(ctx, t, fSource.period)
require.Equal(t, interval2, gotInterval)
loopCancel()
err := testutil.TryReceive(ctx, t, loopErr)
require.NoError(t, err)
}
type fakeNetworkStatsSource struct {
sync.Mutex
ctx context.Context
t testing.TB
callback func(start, end time.Time, virtual, physical map[netlogtype.Connection]netlogtype.Counts)
period chan time.Duration
}
func (f *fakeNetworkStatsSource) SetConnStatsCallback(maxPeriod time.Duration, _ int, dump func(start time.Time, end time.Time, virtual map[netlogtype.Connection]netlogtype.Counts, physical map[netlogtype.Connection]netlogtype.Counts)) {
f.Lock()
defer f.Unlock()
f.callback = dump
select {
case <-f.ctx.Done():
f.t.Error("timeout")
case f.period <- maxPeriod:
// OK
}
}
func newFakeNetworkStatsSource(ctx context.Context, t testing.TB) *fakeNetworkStatsSource {
f := &fakeNetworkStatsSource{
ctx: ctx,
t: t,
period: make(chan time.Duration, 1),
}
return f
}
type fakeCollector struct {
t testing.TB
calls chan map[netlogtype.Connection]netlogtype.Counts
stats chan *proto.Stats
}
func (f *fakeCollector) Collect(ctx context.Context, networkStats map[netlogtype.Connection]netlogtype.Counts) *proto.Stats {
select {
case <-ctx.Done():
f.t.Error("timeout on collect")
return nil
case f.calls <- networkStats:
// ok
}
select {
case <-ctx.Done():
f.t.Error("timeout on collect")
return nil
case s := <-f.stats:
return s
}
}
func newFakeCollector(t testing.TB) *fakeCollector {
return &fakeCollector{
t: t,
calls: make(chan map[netlogtype.Connection]netlogtype.Counts),
stats: make(chan *proto.Stats),
}
}
type fakeStatsDest struct {
reqs chan *proto.UpdateStatsRequest
resps chan *proto.UpdateStatsResponse
}
func (f *fakeStatsDest) UpdateStats(ctx context.Context, req *proto.UpdateStatsRequest) (*proto.UpdateStatsResponse, error) {
select {
case <-ctx.Done():
return nil, ctx.Err()
case f.reqs <- req:
// OK
}
select {
case <-ctx.Done():
return nil, ctx.Err()
case resp := <-f.resps:
return resp, nil
}
}
func newFakeStatsDest() *fakeStatsDest {
return &fakeStatsDest{
reqs: make(chan *proto.UpdateStatsRequest),
resps: make(chan *proto.UpdateStatsResponse),
}
}