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

     1  package adaptiveplacement
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/stretchr/testify/assert"
     8  	"github.com/stretchr/testify/mock"
     9  
    10  	"github.com/grafana/pyroscope/pkg/segmentwriter/client/distributor/placement/adaptiveplacement/adaptive_placementpb"
    11  )
    12  
    13  type mockLimits struct{ mock.Mock }
    14  
    15  func (m *mockLimits) PlacementLimits(tenant string) PlacementLimits {
    16  	return m.Called(tenant).Get(0).(PlacementLimits)
    17  }
    18  
    19  func Test_Ruler(t *testing.T) {
    20  	const unitSize = 512 << 10
    21  	defaults := PlacementLimits{
    22  		TenantShards:         10,
    23  		DefaultDatasetShards: 2,
    24  		MinDatasetShards:     1,
    25  		MaxDatasetShards:     10,
    26  		UnitSizeBytes:        unitSize,
    27  		BurstWindow:          17 * time.Minute,
    28  		DecayWindow:          19 * time.Minute,
    29  		LoadBalancing:        DynamicLoadBalancing,
    30  	}
    31  
    32  	withDefaults := func(fn func(*PlacementLimits)) PlacementLimits {
    33  		limits := defaults
    34  		fn(&limits)
    35  		return limits
    36  	}
    37  
    38  	defaultLimits := withDefaults(func(l *PlacementLimits) {})
    39  
    40  	m := new(mockLimits)
    41  	m.On("PlacementLimits", "tenant-a").
    42  		Return(withDefaults(func(l *PlacementLimits) {
    43  			l.TenantShards = 20
    44  			l.DefaultDatasetShards = 3
    45  		}))
    46  
    47  	m.On("PlacementLimits", "tenant-b").Return(withDefaults(func(l *PlacementLimits) {
    48  		l.LoadBalancing = FingerprintLoadBalancing
    49  	}))
    50  
    51  	m.On("PlacementLimits", "tenant-c").Return(defaultLimits)
    52  
    53  	m.On("PlacementLimits", "tenant-d").Return(withDefaults(func(l *PlacementLimits) {
    54  		l.MinDatasetShards = 5
    55  		l.LoadBalancing = RoundRobinLoadBalancing
    56  	}))
    57  
    58  	oldRules := &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  				Tenant:            0,
    67  				Name:              "dataset-a",
    68  				TenantShardLimit:  2,
    69  				DatasetShardLimit: 5,
    70  				LoadBalancing:     adaptive_placementpb.LoadBalancing_LOAD_BALANCING_ROUND_ROBIN,
    71  			},
    72  			{
    73  				Tenant:            1,
    74  				Name:              "dataset-b",
    75  				TenantShardLimit:  2,
    76  				DatasetShardLimit: 3,
    77  				LoadBalancing:     adaptive_placementpb.LoadBalancing_LOAD_BALANCING_ROUND_ROBIN,
    78  			},
    79  			{
    80  				Tenant:            2,
    81  				Name:              "dataset-c",
    82  				TenantShardLimit:  2,
    83  				DatasetShardLimit: 3,
    84  				LoadBalancing:     adaptive_placementpb.LoadBalancing_LOAD_BALANCING_FINGERPRINT,
    85  			},
    86  		},
    87  	}
    88  
    89  	stats := &adaptive_placementpb.DistributionStats{
    90  		Tenants: []*adaptive_placementpb.TenantStats{
    91  			{TenantId: "tenant-a"},
    92  			{TenantId: "tenant-b"},
    93  			{TenantId: "tenant-c"},
    94  			{TenantId: "tenant-d"},
    95  		},
    96  		Datasets: []*adaptive_placementpb.DatasetStats{
    97  			{
    98  				Tenant: 0,
    99  				Name:   "dataset-a",
   100  				Shards: []uint32{0, 1, 2, 3, 4},
   101  				Usage:  []uint64{unitSize, unitSize, unitSize, unitSize, unitSize / 2},
   102  			},
   103  			{
   104  				Tenant: 1,
   105  				Name:   "dataset-b",
   106  				Shards: []uint32{0, 1, 2},
   107  				Usage:  []uint64{unitSize, unitSize, unitSize},
   108  			},
   109  			{
   110  				Tenant: 2,
   111  				Name:   "dataset-c",
   112  				Shards: []uint32{0, 1, 2},
   113  				Usage:  []uint64{unitSize, unitSize, unitSize / 2},
   114  			},
   115  			{
   116  				Tenant: 3,
   117  				Name:   "dataset-d",
   118  				Shards: []uint32{0},
   119  				Usage:  []uint64{unitSize},
   120  			},
   121  		},
   122  		Shards: []*adaptive_placementpb.ShardStats{
   123  			{Id: 1, Owner: "node-a"},
   124  			{Id: 2, Owner: "node-b"},
   125  			{Id: 3, Owner: "node-c"},
   126  			{Id: 4, Owner: "node-a"},
   127  			{Id: 5, Owner: "node-c"},
   128  		},
   129  		CreatedAt: 1,
   130  	}
   131  
   132  	expected := &adaptive_placementpb.PlacementRules{
   133  		CreatedAt: 1,
   134  		Tenants: []*adaptive_placementpb.TenantPlacement{
   135  			{TenantId: "tenant-a"},
   136  			{TenantId: "tenant-b"},
   137  			{TenantId: "tenant-c"},
   138  			{TenantId: "tenant-d"},
   139  		},
   140  		Datasets: []*adaptive_placementpb.DatasetPlacement{
   141  			{
   142  				Tenant:            0,
   143  				Name:              "dataset-a",
   144  				TenantShardLimit:  20,
   145  				DatasetShardLimit: 5,
   146  				LoadBalancing:     adaptive_placementpb.LoadBalancing_LOAD_BALANCING_ROUND_ROBIN,
   147  			},
   148  			{
   149  				Tenant:            1,
   150  				Name:              "dataset-b",
   151  				TenantShardLimit:  10,
   152  				DatasetShardLimit: 4,
   153  				LoadBalancing:     adaptive_placementpb.LoadBalancing_LOAD_BALANCING_FINGERPRINT,
   154  			},
   155  			{
   156  				Tenant:            2,
   157  				Name:              "dataset-c",
   158  				TenantShardLimit:  10,
   159  				DatasetShardLimit: 3,
   160  				LoadBalancing:     adaptive_placementpb.LoadBalancing_LOAD_BALANCING_FINGERPRINT,
   161  			},
   162  			{
   163  				Tenant:            3,
   164  				Name:              "dataset-d",
   165  				TenantShardLimit:  10,
   166  				DatasetShardLimit: 5,
   167  				LoadBalancing:     adaptive_placementpb.LoadBalancing_LOAD_BALANCING_ROUND_ROBIN,
   168  			},
   169  		},
   170  	}
   171  
   172  	ruler := NewRuler(m)
   173  	ruler.Load(oldRules)
   174  	assert.Equal(t, expected.String(), ruler.BuildRules(stats).String())
   175  
   176  	// Next update only includes tenant-a dataset-a.
   177  	// We expect that dataset-b and dataset-c will still be present.
   178  	update := &adaptive_placementpb.DistributionStats{
   179  		Tenants: []*adaptive_placementpb.TenantStats{
   180  			{TenantId: "tenant-a"},
   181  		},
   182  		Datasets: []*adaptive_placementpb.DatasetStats{
   183  			{
   184  				Tenant: 0,
   185  				Name:   "dataset-a",
   186  				Shards: []uint32{0, 1, 2, 3, 4},
   187  				Usage:  []uint64{unitSize, unitSize, unitSize, unitSize, unitSize / 2},
   188  			},
   189  		},
   190  		Shards: []*adaptive_placementpb.ShardStats{
   191  			{Id: 1, Owner: "node-a"},
   192  			{Id: 2, Owner: "node-b"},
   193  			{Id: 3, Owner: "node-c"},
   194  			{Id: 4, Owner: "node-a"},
   195  			{Id: 5, Owner: "node-c"},
   196  		},
   197  		CreatedAt: 2,
   198  	}
   199  
   200  	expected.CreatedAt = 2
   201  	assert.Equal(t, expected.String(), ruler.BuildRules(update).String())
   202  
   203  	ruler.Expire(time.Now())
   204  	expected = &adaptive_placementpb.PlacementRules{CreatedAt: 3}
   205  	empty := &adaptive_placementpb.DistributionStats{CreatedAt: 3}
   206  	assert.Equal(t, expected.String(), ruler.BuildRules(empty).String())
   207  }