iframe-proxy | Sunbelt Computer Software

223 lines (198 loc) · 6.08 KB
package agent
	"net/netip"
	"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)) {
	defer f.Unlock()
	f.callback = dump
	case <-f.ctx.Done():
		f.t.Error("timeout")
	case f.period <- maxPeriod:
func newFakeNetworkStatsSource(ctx context.Context, t testing.TB) *fakeNetworkStatsSource {
	f := &fakeNetworkStatsSource{
		ctx:    ctx,
		t:      t,
		period: make(chan time.Duration, 1),
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 {
	case <-ctx.Done():
		f.t.Error("timeout on collect")
		return nil
	case f.calls <- networkStats:
	case <-ctx.Done():
		f.t.Error("timeout on collect")
		return nil
	case s := <-f.stats:
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) {
	case <-ctx.Done():
		return nil, ctx.Err()
	case f.reqs <- req:
	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),
Sunbelt Computer Software

PL/B Language Development and Support

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Sunbelt Computer Software

PL/B Language Development and Support

Uh oh!

FilesExpand file tree

stats_internal_test.go

Latest commit

History

stats_internal_test.go

File metadata and controls
