github.com/grafana/pyroscope@v1.18.0/pkg/segmentwriter/client/distributor/placement/adaptiveplacement/distribution_stats_test.go (about)

     1  package adaptiveplacement
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/stretchr/testify/assert"
     8  
     9  	"github.com/grafana/pyroscope/pkg/iter"
    10  	"github.com/grafana/pyroscope/pkg/segmentwriter/client/distributor/placement/adaptiveplacement/adaptive_placementpb"
    11  )
    12  
    13  func Test_StatsTracker(t *testing.T) {
    14  	const window = time.Second * 10
    15  	stats := NewDistributionStats(window)
    16  	var now time.Duration
    17  
    18  	for ; now < window; now += time.Second {
    19  		stats.recordStats(now.Nanoseconds(), iter.NewSliceIterator([]Sample{
    20  			{TenantID: "tenant-a", DatasetName: "dataset-a", ShardID: 1, Size: 10},
    21  			{TenantID: "tenant-a", DatasetName: "dataset-b", ShardID: 1, Size: 10},
    22  		}))
    23  	}
    24  
    25  	// Note that we deal with half-life exponent decay here.
    26  	// Therefore, in 10s we only reached 50% of the target value.
    27  	expected := &adaptive_placementpb.DistributionStats{
    28  		Tenants: []*adaptive_placementpb.TenantStats{
    29  			{TenantId: "tenant-a"},
    30  		},
    31  		Datasets: []*adaptive_placementpb.DatasetStats{
    32  			{
    33  				Tenant: 0,
    34  				Name:   "dataset-a",
    35  				Shards: []uint32{0},
    36  				Usage:  []uint64{5},
    37  			},
    38  			{
    39  				Tenant: 0,
    40  				Name:   "dataset-b",
    41  				Shards: []uint32{0},
    42  				Usage:  []uint64{5},
    43  			},
    44  		},
    45  		Shards: []*adaptive_placementpb.ShardStats{
    46  			{Id: 1},
    47  		},
    48  		CreatedAt: now.Nanoseconds(),
    49  	}
    50  	assert.Equal(t, expected.String(), stats.build(now.Nanoseconds()).String())
    51  
    52  	// Reassign dataset-a to shard 2 and add dataset-c.
    53  	for ; now < time.Second*20; now += time.Second {
    54  		stats.recordStats(now.Nanoseconds(),
    55  			iter.NewSliceIterator([]Sample{
    56  				{TenantID: "tenant-a", DatasetName: "dataset-a", ShardID: 2, Size: 10}, // Moved from shard 1.
    57  				{TenantID: "tenant-a", DatasetName: "dataset-b", ShardID: 1, Size: 10}, // Not changed.
    58  				{TenantID: "tenant-b", DatasetName: "dataset-c", ShardID: 2, Size: 10}, // Added.
    59  			}))
    60  	}
    61  	expected = &adaptive_placementpb.DistributionStats{
    62  		Tenants: []*adaptive_placementpb.TenantStats{
    63  			{TenantId: "tenant-a"},
    64  			{TenantId: "tenant-b"},
    65  		},
    66  		Datasets: []*adaptive_placementpb.DatasetStats{
    67  			{
    68  				Tenant: 0,
    69  				Name:   "dataset-a",
    70  				Shards: []uint32{0, 1},
    71  				Usage:  []uint64{3, 5},
    72  			},
    73  			{
    74  				Tenant: 0,
    75  				Name:   "dataset-b",
    76  				Shards: []uint32{0},
    77  				Usage:  []uint64{8},
    78  			},
    79  			{
    80  				Tenant: 1,
    81  				Name:   "dataset-c",
    82  				Shards: []uint32{1},
    83  				Usage:  []uint64{5},
    84  			},
    85  		},
    86  		Shards: []*adaptive_placementpb.ShardStats{
    87  			{Id: 1},
    88  			{Id: 2},
    89  		},
    90  		CreatedAt: now.Nanoseconds(),
    91  	}
    92  	assert.Equal(t, expected.String(), stats.build(now.Nanoseconds()).String())
    93  
    94  	// Next 30 seconds nothing changes.
    95  	for ; now < time.Minute; now += time.Second {
    96  		stats.recordStats(now.Nanoseconds(),
    97  			iter.NewSliceIterator([]Sample{
    98  				{TenantID: "tenant-a", DatasetName: "dataset-a", ShardID: 2, Size: 10},
    99  				{TenantID: "tenant-a", DatasetName: "dataset-b", ShardID: 1, Size: 10},
   100  				{TenantID: "tenant-b", DatasetName: "dataset-c", ShardID: 2, Size: 10},
   101  			}))
   102  	}
   103  	expected = &adaptive_placementpb.DistributionStats{
   104  		Tenants: []*adaptive_placementpb.TenantStats{
   105  			{TenantId: "tenant-a"},
   106  			{TenantId: "tenant-b"},
   107  		},
   108  		Datasets: []*adaptive_placementpb.DatasetStats{
   109  			{
   110  				Tenant: 0,
   111  				Name:   "dataset-a",
   112  				Shards: []uint32{0, 1},
   113  				Usage:  []uint64{0, 10},
   114  			},
   115  			{
   116  				Tenant: 0,
   117  				Name:   "dataset-b",
   118  				Shards: []uint32{0},
   119  				Usage:  []uint64{10},
   120  			},
   121  			{
   122  				Tenant: 1,
   123  				Name:   "dataset-c",
   124  				Shards: []uint32{1},
   125  				Usage:  []uint64{10},
   126  			},
   127  		},
   128  		Shards: []*adaptive_placementpb.ShardStats{
   129  			{Id: 1},
   130  			{Id: 2},
   131  		},
   132  		CreatedAt: now.Nanoseconds(),
   133  	}
   134  	assert.Equal(t, expected.String(), stats.build(now.Nanoseconds()).String())
   135  
   136  	// See what happens when a stale counter is removed (dataset-a, shard 1).
   137  	stats.Expire(time.Unix(40, 0))
   138  	stats.build(now.Nanoseconds())
   139  	expected = &adaptive_placementpb.DistributionStats{
   140  		Tenants: []*adaptive_placementpb.TenantStats{
   141  			{TenantId: "tenant-a"},
   142  			{TenantId: "tenant-b"},
   143  		},
   144  		Datasets: []*adaptive_placementpb.DatasetStats{
   145  			{
   146  				Tenant: 0,
   147  				Name:   "dataset-a",
   148  				Shards: []uint32{0},
   149  				Usage:  []uint64{10},
   150  			},
   151  			{
   152  				Tenant: 0,
   153  				Name:   "dataset-b",
   154  				Shards: []uint32{1},
   155  				Usage:  []uint64{10},
   156  			},
   157  			{
   158  				Tenant: 1,
   159  				Name:   "dataset-c",
   160  				Shards: []uint32{0},
   161  				Usage:  []uint64{10},
   162  			},
   163  		},
   164  		Shards: []*adaptive_placementpb.ShardStats{
   165  			{Id: 2},
   166  			{Id: 1},
   167  		},
   168  		CreatedAt: now.Nanoseconds(),
   169  	}
   170  	assert.Equal(t, expected.String(), stats.build(now.Nanoseconds()).String())
   171  
   172  	// Expire all counters.
   173  	stats.Expire(time.Now())
   174  	s := stats.build(now.Nanoseconds())
   175  	assert.Empty(t, s.Tenants)
   176  	assert.Empty(t, s.Datasets)
   177  	assert.Empty(t, s.Shards)
   178  	assert.Empty(t, stats.counters)
   179  }