github.com/grafana/pyroscope@v1.18.0/pkg/segmentwriter/client/distributor/placement/adaptiveplacement/placement_agent_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/stretchr/testify/mock"
    14  	"github.com/stretchr/testify/suite"
    15  
    16  	"github.com/grafana/pyroscope/pkg/segmentwriter/client/distributor/placement"
    17  	"github.com/grafana/pyroscope/pkg/segmentwriter/client/distributor/placement/adaptiveplacement/adaptive_placementpb"
    18  	"github.com/grafana/pyroscope/pkg/test/mocks/mockadaptiveplacement"
    19  )
    20  
    21  type agentSuite struct {
    22  	suite.Suite
    23  
    24  	logger log.Logger
    25  	reg    *prometheus.Registry
    26  	config Config
    27  	limits *mockLimits
    28  	store  *mockadaptiveplacement.MockStore
    29  	agent  *Agent
    30  	error  error
    31  }
    32  
    33  func (s *agentSuite) SetupTest() {
    34  	s.logger = log.NewLogfmtLogger(io.Discard)
    35  	s.reg = prometheus.NewRegistry()
    36  	s.config.PlacementUpdateInterval = 15 * time.Second
    37  	s.limits = new(mockLimits)
    38  	s.store = new(mockadaptiveplacement.MockStore)
    39  	s.agent = NewAgent(
    40  		s.logger,
    41  		s.reg,
    42  		s.config,
    43  		s.limits,
    44  		s.store,
    45  	)
    46  }
    47  
    48  func (s *agentSuite) AfterTest(_, _ string) {
    49  	svc := s.agent.Service()
    50  	svc.StopAsync()
    51  	if s.error == nil {
    52  		s.Require().NoError(svc.AwaitTerminated(context.Background()))
    53  		s.Require().Equal(services.Terminated, svc.State())
    54  	} else {
    55  		s.Require().Equal(services.Failed, svc.State())
    56  	}
    57  	s.limits.AssertExpectations(s.T())
    58  	s.store.AssertExpectations(s.T())
    59  }
    60  
    61  func (s *agentSuite) start() error {
    62  	ctx := context.Background()
    63  	svc := s.agent.Service()
    64  	s.Require().NoError(svc.StartAsync(ctx))
    65  	s.error = svc.AwaitRunning(ctx)
    66  	return s.error
    67  }
    68  
    69  func Test_AgentSuite(t *testing.T) { suite.Run(t, new(agentSuite)) }
    70  
    71  func (s *agentSuite) Test_Agent_loads_rules_on_start() {
    72  	s.store.On("LoadRules", mock.Anything).
    73  		Return(&adaptive_placementpb.PlacementRules{}, nil)
    74  	s.Require().NoError(s.start())
    75  	s.Assert().NotNil(s.agent.rules)
    76  	s.Assert().NotNil(s.agent.Placement())
    77  }
    78  
    79  func (s *agentSuite) Test_Agent_service_doesnt_fail_if_rules_cant_be_found() {
    80  	s.store.On("LoadRules", mock.Anything).
    81  		Return((*adaptive_placementpb.PlacementRules)(nil), ErrRulesNotFound)
    82  	s.Require().NoError(s.start())
    83  	s.Assert().NotNil(s.agent.rules)
    84  	s.Assert().NotNil(s.agent.Placement())
    85  }
    86  
    87  func (s *agentSuite) Test_Agent_service_fails_if_rules_cant_be_loaded() {
    88  	s.store.On("LoadRules", mock.Anything).
    89  		Return((*adaptive_placementpb.PlacementRules)(nil), fmt.Errorf("error"))
    90  	s.Require().Error(s.start())
    91  	s.Assert().Nil(s.agent.rules)
    92  	s.Assert().NotNil(s.agent.Placement())
    93  }
    94  
    95  func (s *agentSuite) Test_Agent_updates_placement_rules() {
    96  	s.limits.On("PlacementLimits", "tenant-a").Return(PlacementLimits{
    97  		TenantShards:         1,
    98  		DefaultDatasetShards: 1,
    99  	})
   100  
   101  	s.store.On("LoadRules", mock.Anything).
   102  		Return(&adaptive_placementpb.PlacementRules{CreatedAt: 100}, nil).
   103  		Once()
   104  
   105  	s.Require().NoError(s.start())
   106  
   107  	p := s.agent.Placement()
   108  	s.Require().NotNil(p)
   109  	policy := p.Policy(placement.Key{
   110  		TenantID:    "tenant-a",
   111  		DatasetName: "dataset-a",
   112  	})
   113  	s.Assert().Equal(1, policy.TenantShards)
   114  	s.Assert().Equal(1, policy.DatasetShards)
   115  
   116  	s.store.On("LoadRules", mock.Anything).
   117  		Return(&adaptive_placementpb.PlacementRules{
   118  			CreatedAt: 150,
   119  			Tenants: []*adaptive_placementpb.TenantPlacement{
   120  				{TenantId: "tenant-a"},
   121  			},
   122  			Datasets: []*adaptive_placementpb.DatasetPlacement{
   123  				{
   124  					Name:              "dataset-a",
   125  					TenantShardLimit:  10,
   126  					DatasetShardLimit: 10,
   127  				},
   128  			},
   129  		}, nil).
   130  		Once()
   131  
   132  	s.agent.loadRules(context.Background())
   133  	policy = p.Policy(placement.Key{
   134  		TenantID:    "tenant-a",
   135  		DatasetName: "dataset-a",
   136  	})
   137  	s.Assert().Equal(10, policy.TenantShards)
   138  	s.Assert().Equal(10, policy.DatasetShards)
   139  }
   140  
   141  func (s *agentSuite) Test_Agent_ignored_outdated_rules() {
   142  	s.limits.On("PlacementLimits", "tenant-a").Return(PlacementLimits{
   143  		TenantShards:         1,
   144  		DefaultDatasetShards: 1,
   145  	})
   146  
   147  	s.store.On("LoadRules", mock.Anything).
   148  		Return(&adaptive_placementpb.PlacementRules{CreatedAt: 100}, nil).
   149  		Once()
   150  
   151  	s.Require().NoError(s.start())
   152  
   153  	p := s.agent.Placement()
   154  	s.Require().NotNil(p)
   155  	policy := p.Policy(placement.Key{
   156  		TenantID:    "tenant-a",
   157  		DatasetName: "dataset-a",
   158  	})
   159  	s.Assert().Equal(1, policy.TenantShards)
   160  	s.Assert().Equal(1, policy.DatasetShards)
   161  
   162  	s.store.On("LoadRules", mock.Anything).
   163  		Return(&adaptive_placementpb.PlacementRules{
   164  			CreatedAt: 10,
   165  			Tenants: []*adaptive_placementpb.TenantPlacement{
   166  				{TenantId: "tenant-a"},
   167  			},
   168  			Datasets: []*adaptive_placementpb.DatasetPlacement{
   169  				{
   170  					Name:              "dataset-a",
   171  					TenantShardLimit:  10,
   172  					DatasetShardLimit: 10,
   173  				},
   174  			},
   175  		}, nil).
   176  		Once()
   177  
   178  	s.agent.loadRules(context.Background())
   179  	policy = p.Policy(placement.Key{
   180  		TenantID:    "tenant-a",
   181  		DatasetName: "dataset-a",
   182  	})
   183  	s.Assert().Equal(1, policy.TenantShards)
   184  	s.Assert().Equal(1, policy.DatasetShards)
   185  }