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

     1  package adaptiveplacement
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/go-kit/log"
    11  	"github.com/grafana/dskit/services"
    12  	"github.com/prometheus/client_golang/prometheus"
    13  	"github.com/prometheus/client_golang/prometheus/testutil"
    14  	"github.com/stretchr/testify/mock"
    15  	"github.com/stretchr/testify/suite"
    16  
    17  	"github.com/grafana/pyroscope/pkg/iter"
    18  	"github.com/grafana/pyroscope/pkg/segmentwriter/client/distributor/placement/adaptiveplacement/adaptive_placementpb"
    19  	"github.com/grafana/pyroscope/pkg/test/mocks/mockadaptiveplacement"
    20  )
    21  
    22  type managerSuite struct {
    23  	suite.Suite
    24  
    25  	logger  log.Logger
    26  	reg     *prometheus.Registry
    27  	config  Config
    28  	limits  *mockLimits
    29  	store   *mockadaptiveplacement.MockStore
    30  	manager *Manager
    31  }
    32  
    33  func (s *managerSuite) SetupTest() {
    34  	s.logger = log.NewLogfmtLogger(io.Discard)
    35  	s.reg = prometheus.NewRegistry()
    36  	s.config = Config{
    37  		PlacementUpdateInterval:  15 * time.Second,
    38  		PlacementRetentionPeriod: 15 * time.Minute,
    39  
    40  		StatsConfidencePeriod:  0,
    41  		StatsAggregationWindow: 3 * time.Minute,
    42  		StatsRetentionPeriod:   5 * time.Minute,
    43  	}
    44  	s.limits = new(mockLimits)
    45  	s.store = new(mockadaptiveplacement.MockStore)
    46  	s.manager = NewManager(
    47  		s.logger,
    48  		s.reg,
    49  		s.config,
    50  		s.limits,
    51  		s.store,
    52  	)
    53  }
    54  
    55  func (s *managerSuite) BeforeTest(_, _ string) {
    56  	ctx := context.Background()
    57  	svc := s.manager.Service()
    58  	s.Require().NoError(svc.StartAsync(ctx))
    59  	s.Require().NoError(svc.AwaitRunning(ctx))
    60  }
    61  
    62  func (s *managerSuite) AfterTest(_, _ string) {
    63  	svc := s.manager.Service()
    64  	svc.StopAsync()
    65  	s.Require().NoError(svc.AwaitTerminated(context.Background()))
    66  	s.Require().Equal(services.Terminated, svc.State())
    67  	s.limits.AssertExpectations(s.T())
    68  	s.store.AssertExpectations(s.T())
    69  }
    70  
    71  func Test_ManagerSuite(t *testing.T) { suite.Run(t, new(managerSuite)) }
    72  
    73  func (s *managerSuite) Test_Manager_only_updates_rules_if_started() {
    74  	oldRules := &adaptive_placementpb.PlacementRules{CreatedAt: 100}
    75  	s.store.On("LoadRules", mock.Anything).Return(oldRules, nil)
    76  
    77  	newRules := func(r *adaptive_placementpb.PlacementRules) bool { return r.CreatedAt > 100 }
    78  	s.store.On("StoreRules", mock.Anything, mock.MatchedBy(newRules)).Return(nil).Once()
    79  	s.store.On("StoreStats", mock.Anything, mock.Anything).Return(nil).Once()
    80  
    81  	s.manager.Start()
    82  	s.manager.updateRules(context.Background())
    83  
    84  	s.manager.Stop()
    85  	s.manager.updateRules(context.Background())
    86  }
    87  
    88  func (s *managerSuite) Test_Manager_doesnt_update_rules_until_confidence_period_expiration() {
    89  	s.manager.config.StatsConfidencePeriod = time.Minute
    90  
    91  	oldRules := &adaptive_placementpb.PlacementRules{CreatedAt: 100}
    92  	s.store.On("LoadRules", mock.Anything).Return(oldRules, nil)
    93  
    94  	s.manager.Start()
    95  	s.manager.updateRules(context.Background())
    96  }
    97  
    98  func (s *managerSuite) Test_Manager_doesnt_update_rules_if_store_fails() {
    99  	s.store.On("LoadRules", mock.Anything).
   100  		Return((*adaptive_placementpb.PlacementRules)(nil), fmt.Errorf("error"))
   101  
   102  	s.manager.Start()
   103  	s.manager.updateRules(context.Background())
   104  }
   105  
   106  func (s *managerSuite) Test_Manager_updates_rules_if_no_rules_not_found() {
   107  	s.store.On("LoadRules", mock.Anything).
   108  		Return((*adaptive_placementpb.PlacementRules)(nil), ErrRulesNotFound)
   109  
   110  	newRules := func(r *adaptive_placementpb.PlacementRules) bool { return r.CreatedAt > 0 }
   111  	s.store.On("StoreRules", mock.Anything, mock.MatchedBy(newRules)).Return(nil).Once()
   112  	s.store.On("StoreStats", mock.Anything, mock.Anything).Return(nil).Once()
   113  
   114  	s.manager.Start()
   115  	s.manager.updateRules(context.Background())
   116  }
   117  
   118  func (s *managerSuite) Test_Manager_exports_metrics() {
   119  	s.manager.config.ExportShardLimitMetrics = true
   120  	s.manager.config.ExportShardUsageMetrics = true
   121  	s.manager.config.ExportShardUsageBreakdownMetrics = true
   122  
   123  	oldRules := &adaptive_placementpb.PlacementRules{CreatedAt: 100}
   124  	s.store.On("LoadRules", mock.Anything).Return(oldRules, nil)
   125  
   126  	newRules := func(r *adaptive_placementpb.PlacementRules) bool { return r.CreatedAt > 100 }
   127  	s.store.On("StoreRules", mock.Anything, mock.MatchedBy(newRules)).Return(nil).Once()
   128  	s.store.On("StoreStats", mock.Anything, mock.Anything).Return(nil).Once()
   129  
   130  	s.limits.On("PlacementLimits", mock.Anything).Return(PlacementLimits{
   131  		DefaultDatasetShards: 2,
   132  		MinDatasetShards:     1,
   133  		UnitSizeBytes:        256 << 10,
   134  		BurstWindow:          time.Minute,
   135  		DecayWindow:          time.Minute,
   136  	})
   137  
   138  	s.manager.RecordStats(iter.NewSliceIterator([]Sample{
   139  		{TenantID: "tenant-a", DatasetName: "dataset-a", ShardID: 1, Size: 100 << 10},
   140  		{TenantID: "tenant-a", DatasetName: "dataset-a", ShardID: 2, Size: 100 << 10},
   141  		{TenantID: "tenant-b", DatasetName: "dataset-b", ShardID: 2, Size: 100 << 10},
   142  	}))
   143  
   144  	s.manager.Start()
   145  	s.manager.updateRules(context.Background())
   146  
   147  	n, err := testutil.GatherAndCount(s.reg,
   148  		"pyroscope_adaptive_placement_rules_last_update_time",
   149  		"pyroscope_adaptive_placement_rules",
   150  		"pyroscope_adaptive_placement_stats",
   151  		"pyroscope_adaptive_placement_dataset_shard_limit",
   152  		"pyroscope_adaptive_placement_dataset_shard_usage_bytes_per_second",
   153  		"pyroscope_adaptive_placement_dataset_shard_usage_per_shard_bytes_per_second",
   154  	)
   155  	s.Require().NoError(err)
   156  	s.Assert().Equal(10, n)
   157  }