github.com/grafana/pyroscope@v1.18.0/pkg/distributor/ingestlimits/sampler_test.go (about)

     1  package ingestlimits
     2  
     3  import (
     4  	"sync"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  // Mock ring for testing
    12  type mockRing struct {
    13  	instanceCount int
    14  }
    15  
    16  func (m *mockRing) InstancesCount() int {
    17  	return m.instanceCount
    18  }
    19  
    20  func TestAllowRequest(t *testing.T) {
    21  	sampler := NewSampler(&mockRing{instanceCount: 1})
    22  
    23  	config := SamplingConfig{
    24  		Period:      1 * time.Second,
    25  		NumRequests: 1,
    26  	}
    27  
    28  	tenantID := "test-tenant"
    29  	allowed := sampler.AllowRequest(tenantID, config)
    30  	assert.True(t, allowed, "this request is brand new and should be allowed")
    31  
    32  	allowed = sampler.AllowRequest(tenantID, config)
    33  	assert.Falsef(t, allowed, "this request should be within a second of the first and not be allowed")
    34  }
    35  
    36  func TestTenantTrackerAllowRequest(t *testing.T) {
    37  	testCases := []struct {
    38  		name           string
    39  		replicaCount   int
    40  		windowDuration time.Duration
    41  		maxRequests    int
    42  		requestCount   int
    43  		expectAllowed  int
    44  	}{
    45  		{
    46  			name:           "Within limit",
    47  			replicaCount:   2,
    48  			windowDuration: 5 * time.Second,
    49  			maxRequests:    5,
    50  			requestCount:   3,
    51  			expectAllowed:  3,
    52  		},
    53  		{
    54  			name:           "Exceed limit",
    55  			replicaCount:   2,
    56  			windowDuration: 5 * time.Second,
    57  			maxRequests:    3,
    58  			requestCount:   5,
    59  			expectAllowed:  3,
    60  		},
    61  		{
    62  			name:           "Random probability",
    63  			replicaCount:   100,
    64  			windowDuration: 5 * time.Second,
    65  			maxRequests:    5,
    66  			requestCount:   100,
    67  			expectAllowed:  5,
    68  		},
    69  	}
    70  
    71  	for _, tc := range testCases {
    72  		t.Run(tc.name, func(t *testing.T) {
    73  			tracker := &tenantTracker{
    74  				lastRequestTime:   time.Now(),
    75  				remainingRequests: tc.maxRequests,
    76  			}
    77  
    78  			allowedCnt := 0
    79  			for i := 0; i < tc.requestCount; i++ {
    80  				if tracker.AllowRequest(tc.replicaCount, tc.windowDuration, tc.maxRequests) {
    81  					allowedCnt++
    82  				}
    83  			}
    84  			assert.LessOrEqualf(t, allowedCnt, tc.expectAllowed, "request %d should match expected")
    85  		})
    86  	}
    87  }
    88  
    89  func TestConcurrentAccess(t *testing.T) {
    90  	sampler := NewSampler(&mockRing{instanceCount: 2})
    91  
    92  	config := SamplingConfig{
    93  		Period:      5 * time.Second,
    94  		NumRequests: 10,
    95  	}
    96  
    97  	var wg sync.WaitGroup
    98  	allowedCount := 0
    99  	mu := sync.Mutex{}
   100  
   101  	for i := 0; i < 100; i++ {
   102  		wg.Add(1)
   103  		go func() {
   104  			defer wg.Done()
   105  			if sampler.AllowRequest("concurrent-tenant", config) {
   106  				mu.Lock()
   107  				allowedCount++
   108  				mu.Unlock()
   109  			}
   110  		}()
   111  	}
   112  
   113  	wg.Wait()
   114  	assert.LessOrEqual(t, allowedCount, 10, "Should not exceed max requests")
   115  }