github.com/grafana/pyroscope@v1.18.0/pkg/segmentwriter/client/distributor/placement/adaptiveplacement/adaptive_placement_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/segmentwriter/client/distributor/placement"
    10  	"github.com/grafana/pyroscope/pkg/segmentwriter/client/distributor/placement/adaptiveplacement/adaptive_placementpb"
    11  )
    12  
    13  func Test_AdaptivePlacement(t *testing.T) {
    14  	const unitSize = 512 << 10
    15  	defaults := PlacementLimits{
    16  		TenantShards:         10,
    17  		DefaultDatasetShards: 2,
    18  		MinDatasetShards:     1,
    19  		MaxDatasetShards:     10,
    20  		UnitSizeBytes:        unitSize,
    21  		BurstWindow:          17 * time.Minute,
    22  		DecayWindow:          19 * time.Minute,
    23  		LoadBalancing:        DynamicLoadBalancing,
    24  	}
    25  
    26  	withDefaults := func(fn func(*PlacementLimits)) PlacementLimits {
    27  		limits := defaults
    28  		fn(&limits)
    29  		return limits
    30  	}
    31  
    32  	m := new(mockLimits)
    33  	m.On("PlacementLimits", "tenant-a").Return(withDefaults(func(l *PlacementLimits) {}))
    34  	m.On("PlacementLimits", "tenant-b").Return(withDefaults(func(l *PlacementLimits) {
    35  		l.TenantShards = 20
    36  		l.DefaultDatasetShards = 3
    37  	}))
    38  
    39  	p := NewAdaptivePlacement(m)
    40  
    41  	policy := p.Policy(placement.Key{
    42  		TenantID:    "tenant-a",
    43  		DatasetName: "dataset-a",
    44  	})
    45  	assert.Equal(t, 10, policy.TenantShards)
    46  	assert.Equal(t, 2, policy.DatasetShards)
    47  	assert.False(t, isRoundRobin(policy.PickShard))
    48  
    49  	policy = p.Policy(placement.Key{
    50  		TenantID:    "tenant-b",
    51  		DatasetName: "dataset-b-1",
    52  	})
    53  	assert.Equal(t, 20, policy.TenantShards)
    54  	assert.Equal(t, 3, policy.DatasetShards)
    55  	assert.False(t, isRoundRobin(policy.PickShard))
    56  
    57  	// Load new rules and override limits for tenant-a dataset-a
    58  	p.Update(&adaptive_placementpb.PlacementRules{
    59  		Tenants: []*adaptive_placementpb.TenantPlacement{
    60  			{TenantId: "tenant-a"},
    61  			{TenantId: "tenant-b"},
    62  			{TenantId: "tenant-c"},
    63  		},
    64  		Datasets: []*adaptive_placementpb.DatasetPlacement{
    65  			{
    66  				// A placement rule may have a newer/different limit for the tenant.
    67  				Tenant:            0,
    68  				Name:              "dataset-a",
    69  				TenantShardLimit:  4,
    70  				DatasetShardLimit: 4,
    71  				LoadBalancing:     adaptive_placementpb.LoadBalancing_LOAD_BALANCING_ROUND_ROBIN,
    72  			},
    73  			{
    74  				Tenant:            0,
    75  				Name:              "dataset-a-2",
    76  				TenantShardLimit:  4,
    77  				DatasetShardLimit: 1,
    78  				LoadBalancing:     adaptive_placementpb.LoadBalancing_LOAD_BALANCING_FINGERPRINT,
    79  			},
    80  		},
    81  		CreatedAt: 1,
    82  	})
    83  
    84  	// Assert that the new rules impacted the placement policy for the dataset.
    85  	policy = p.Policy(placement.Key{
    86  		TenantID:    "tenant-a",
    87  		DatasetName: "dataset-a",
    88  	})
    89  	assert.Equal(t, 4, policy.TenantShards)
    90  	assert.Equal(t, 4, policy.DatasetShards)
    91  	assert.True(t, isRoundRobin(policy.PickShard))
    92  
    93  	// Other datasets of the tenant are not affected.
    94  	policy = p.Policy(placement.Key{
    95  		TenantID:    "tenant-a",
    96  		DatasetName: "dataset-b",
    97  	})
    98  	assert.Equal(t, 10, policy.TenantShards)
    99  	assert.Equal(t, 2, policy.DatasetShards)
   100  	assert.False(t, isRoundRobin(policy.PickShard))
   101  
   102  	policy = p.Policy(placement.Key{
   103  		TenantID:    "tenant-a",
   104  		DatasetName: "dataset-a-2",
   105  	})
   106  	assert.Equal(t, 4, policy.TenantShards)
   107  	assert.Equal(t, 1, policy.DatasetShards)
   108  	assert.False(t, isRoundRobin(policy.PickShard))
   109  
   110  	// Other tenants are not affected.
   111  	policy = p.Policy(placement.Key{
   112  		TenantID:    "tenant-b",
   113  		DatasetName: "dataset-b-1",
   114  	})
   115  	assert.Equal(t, 20, policy.TenantShards)
   116  	assert.Equal(t, 3, policy.DatasetShards)
   117  	assert.False(t, isRoundRobin(policy.PickShard))
   118  }
   119  
   120  // This does not test the actual round-robin behavior,
   121  // but rather that the function is not deterministic.
   122  func isRoundRobin(fn func(int) int) bool {
   123  	const N = 10
   124  	r := fn(N)
   125  	for i := 0; i < N; i++ {
   126  		if r != fn(N) {
   127  			return true
   128  		}
   129  	}
   130  	return false
   131  }