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

     1  package adaptiveplacement
     2  
     3  import (
     4  	"fmt"
     5  	"math/rand"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  
    10  	"github.com/grafana/pyroscope/pkg/segmentwriter/client/distributor/placement/adaptiveplacement/adaptive_placementpb"
    11  )
    12  
    13  func Test_loadBalancingStrategy(t *testing.T) {
    14  	rnd := rand.New(rand.NewSource(randSeed))
    15  	const unitSize = 512 << 10
    16  
    17  	randomize := func(f float64, values ...uint64) []uint64 {
    18  		for i, v := range values {
    19  			j := uint64(float64(v) * f)
    20  			if rnd.Float64() > 0.5 {
    21  				values[i] += j
    22  			} else {
    23  				values[i] -= j
    24  			}
    25  		}
    26  		return values
    27  	}
    28  
    29  	for i, test := range []struct {
    30  		usage    []uint64
    31  		expected LoadBalancing
    32  	}{
    33  		{
    34  			expected: FingerprintLoadBalancing,
    35  		},
    36  		{
    37  			usage:    []uint64{0},
    38  			expected: FingerprintLoadBalancing,
    39  		},
    40  		{
    41  			usage:    []uint64{unitSize},
    42  			expected: FingerprintLoadBalancing,
    43  		},
    44  		{
    45  			usage:    []uint64{0, 0, 0, 0, 0},
    46  			expected: FingerprintLoadBalancing,
    47  		},
    48  		{
    49  			usage:    []uint64{unitSize, unitSize, unitSize, unitSize, unitSize},
    50  			expected: FingerprintLoadBalancing,
    51  		},
    52  		{
    53  			usage:    []uint64{2 * unitSize, unitSize, unitSize, unitSize, unitSize},
    54  			expected: FingerprintLoadBalancing,
    55  		},
    56  		{
    57  			usage:    randomize(0.1, unitSize, unitSize, unitSize, unitSize, unitSize),
    58  			expected: FingerprintLoadBalancing,
    59  		},
    60  		{
    61  			usage:    randomize(0.9, unitSize, unitSize, unitSize, unitSize, unitSize),
    62  			expected: FingerprintLoadBalancing,
    63  		},
    64  		{
    65  			usage:    randomize(0.1, 2*unitSize, 2*unitSize, 2*unitSize, 2*unitSize, 2*unitSize),
    66  			expected: FingerprintLoadBalancing,
    67  		},
    68  		{
    69  			usage:    []uint64{2 * unitSize, unitSize / 2, unitSize, unitSize, unitSize},
    70  			expected: FingerprintLoadBalancing,
    71  		},
    72  		{
    73  			usage:    []uint64{2 * unitSize, unitSize / 2, unitSize / 2, unitSize, unitSize},
    74  			expected: RoundRobinLoadBalancing,
    75  		},
    76  		{
    77  			usage:    randomize(0.9, 2*unitSize, 2*unitSize, 2*unitSize, 2*unitSize, 2*unitSize),
    78  			expected: RoundRobinLoadBalancing,
    79  		},
    80  	} {
    81  		stats := &adaptive_placementpb.DatasetStats{
    82  			Shards: make([]uint32, len(test.usage)),
    83  			Usage:  test.usage,
    84  			StdDev: stdDev(test.usage),
    85  		}
    86  		target := len(stats.Shards)
    87  		assert.Equal(t, test.expected, loadBalancingStrategy(stats, unitSize, target), fmt.Sprint(i))
    88  	}
    89  
    90  }
    91  
    92  func Test_loadBalancingStrategy_relocation(t *testing.T) {
    93  	const unitSize = 512 << 10
    94  	for i, test := range []struct {
    95  		usage    []uint64
    96  		expected LoadBalancing
    97  		target   int
    98  	}{
    99  		{
   100  			usage:    []uint64{2 * unitSize, 2 * unitSize, unitSize / 2, unitSize / 2, unitSize / 2},
   101  			expected: RoundRobinLoadBalancing,
   102  			target:   5, // 5/5
   103  		},
   104  		{
   105  			usage:    []uint64{2 * unitSize, 2 * unitSize, unitSize / 2, unitSize / 2},
   106  			expected: RoundRobinLoadBalancing,
   107  			target:   2, // 2/4
   108  		},
   109  		{
   110  			usage:    []uint64{2 * unitSize, 2 * unitSize, unitSize / 2, unitSize / 2, unitSize / 2},
   111  			expected: RoundRobinLoadBalancing,
   112  			target:   3, // 3/5
   113  		},
   114  		{
   115  			usage:    []uint64{2 * unitSize, 2 * unitSize, unitSize / 2, unitSize / 2, unitSize / 2},
   116  			expected: FingerprintLoadBalancing,
   117  			target:   2, // 2/5
   118  		},
   119  		{
   120  			usage:    []uint64{unitSize, unitSize, unitSize / 5, unitSize / 5},
   121  			expected: FingerprintLoadBalancing,
   122  			target:   2, // 2/4
   123  		},
   124  		{
   125  			usage:    []uint64{2*unitSize - 1, 0},
   126  			expected: FingerprintLoadBalancing,
   127  			target:   1, // 1/2
   128  		},
   129  	} {
   130  		stats := &adaptive_placementpb.DatasetStats{
   131  			Shards: make([]uint32, len(test.usage)),
   132  			Usage:  test.usage,
   133  			StdDev: stdDev(test.usage),
   134  		}
   135  		assert.Equal(t, test.expected, loadBalancingStrategy(stats, unitSize, test.target), fmt.Sprint(i))
   136  	}
   137  }