github.com/grafana/pyroscope@v1.18.0/pkg/ingester/limiter_test.go (about)

     1  package ingester
     2  
     3  import (
     4  	"fmt"
     5  	"strconv"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/prometheus/common/model"
    10  	"github.com/stretchr/testify/assert"
    11  	"github.com/stretchr/testify/require"
    12  
    13  	phlaremodel "github.com/grafana/pyroscope/pkg/model"
    14  	"github.com/grafana/pyroscope/pkg/validation"
    15  )
    16  
    17  type fakeLimits struct {
    18  	maxLocalSeriesPerTenant  int
    19  	maxGlobalSeriesPerTenant int
    20  	ingestionTenantShardSize int
    21  }
    22  
    23  func (f *fakeLimits) MaxLocalSeriesPerTenant(userID string) int {
    24  	return f.maxLocalSeriesPerTenant
    25  }
    26  
    27  func (f *fakeLimits) MaxGlobalSeriesPerTenant(userID string) int {
    28  	return f.maxGlobalSeriesPerTenant
    29  }
    30  
    31  func (f *fakeLimits) IngestionTenantShardSize(userID string) int {
    32  	return f.ingestionTenantShardSize
    33  }
    34  
    35  func (f *fakeLimits) DistributorUsageGroups(userID string) *validation.UsageGroupConfig {
    36  	return &validation.UsageGroupConfig{}
    37  }
    38  
    39  type fakeRingCount struct {
    40  	healthyInstancesCount int
    41  }
    42  
    43  func (f *fakeRingCount) HealthyInstancesCount() int {
    44  	return f.healthyInstancesCount
    45  }
    46  
    47  func TestGlobalMaxSeries(t *testing.T) {
    48  	// 5 series per user, 2 ingesters, replication factor 3.
    49  	// We should be able to push 7.5 series. (5 / 2 * 3 = 7.5)
    50  	activeSeriesTimeout = 200 * time.Millisecond
    51  	activeSeriesCleanup = 100 * time.Millisecond
    52  
    53  	limiter := NewLimiter("foo", &fakeLimits{maxGlobalSeriesPerTenant: 5}, &fakeRingCount{2}, 3)
    54  	defer limiter.Stop()
    55  
    56  	for i := 0; i < 7; i++ {
    57  		err := limiter.AllowProfile(model.Fingerprint(i), phlaremodel.LabelsFromStrings("i", fmt.Sprintf("%d", i)), 0)
    58  		require.NoError(t, err)
    59  	}
    60  
    61  	err := limiter.AllowProfile(8, phlaremodel.LabelsFromStrings("i", "8"), 0)
    62  	require.Error(t, err)
    63  
    64  	// Wait for cleanup to happen.
    65  	time.Sleep(400 * time.Millisecond)
    66  
    67  	// Now we should be able to push 5 series.
    68  	for i := 0; i < 5; i++ {
    69  		err := limiter.AllowProfile(model.Fingerprint(i), phlaremodel.LabelsFromStrings("i", fmt.Sprintf("%d", i)), 0)
    70  		require.NoError(t, err)
    71  	}
    72  }
    73  
    74  func assertMaxSeries(t testing.TB, limiter Limiter, count int) {
    75  	var (
    76  		i   int
    77  		err error
    78  	)
    79  	for i = 1; i < count+1; i++ {
    80  		err = limiter.AllowProfile(model.Fingerprint(i), phlaremodel.LabelsFromStrings("i", strconv.Itoa(i)), 0)
    81  		require.NoError(t, err, "series %d should be allowed", i)
    82  	}
    83  
    84  	i += 1
    85  	err = limiter.AllowProfile(model.Fingerprint(i), phlaremodel.LabelsFromStrings("i", strconv.Itoa(i)), 0)
    86  	require.Error(t, err)
    87  	assert.ErrorContains(t, err, "Maximum active series limit exceeded")
    88  }
    89  
    90  func TestLocalLimit(t *testing.T) {
    91  	t.Run("local limit", func(t *testing.T) {
    92  		limiter := NewLimiter("foo", &fakeLimits{maxGlobalSeriesPerTenant: 5, maxLocalSeriesPerTenant: 1}, &fakeRingCount{5}, 3)
    93  		defer limiter.Stop()
    94  
    95  		// local limit of 1 series should take precedence over global limit of 5 series.
    96  		assertMaxSeries(t, limiter, 1)
    97  	})
    98  
    99  	t.Run("local limit enforced by diving global limit", func(t *testing.T) {
   100  		limiter := NewLimiter("foo", &fakeLimits{maxGlobalSeriesPerTenant: 3}, &fakeRingCount{9}, 3)
   101  		defer limiter.Stop()
   102  
   103  		// local limit of 1 should be per ingester (globalLimit * replicationFactor) / ingesterNum
   104  		assertMaxSeries(t, limiter, 1)
   105  	})
   106  
   107  	t.Run("ensure we do not panic with zero ingesters", func(t *testing.T) {
   108  		limiter := NewLimiter("foo", &fakeLimits{maxGlobalSeriesPerTenant: 3}, &fakeRingCount{0}, 3)
   109  		defer limiter.Stop()
   110  
   111  		// we can ingest as many series as we want
   112  		require.NoError(t, limiter.AllowProfile(1, phlaremodel.LabelsFromStrings("i", "1"), 0))
   113  	})
   114  
   115  	t.Run("ensure we handle sharding correctly", func(t *testing.T) {
   116  		limiter := NewLimiter("foo", &fakeLimits{maxGlobalSeriesPerTenant: 3, ingestionTenantShardSize: 3}, &fakeRingCount{9}, 3)
   117  		defer limiter.Stop()
   118  
   119  		// local limit of 3 should be per ingester (globalLimit * replicationFactor) / shardSize
   120  		assertMaxSeries(t, limiter, 3)
   121  	})
   122  }