diff --git a/pkg/querier/blocks_store_queryable_test.go b/pkg/querier/blocks_store_queryable_test.go index 4070e0b7383..51001f1346d 100644 --- a/pkg/querier/blocks_store_queryable_test.go +++ b/pkg/querier/blocks_store_queryable_test.go @@ -44,6 +44,7 @@ import ( "github.com/cortexproject/cortex/pkg/util" "github.com/cortexproject/cortex/pkg/util/limiter" util_log "github.com/cortexproject/cortex/pkg/util/log" + "github.com/cortexproject/cortex/pkg/util/resource" "github.com/cortexproject/cortex/pkg/util/services" "github.com/cortexproject/cortex/pkg/util/validation" ) @@ -2498,6 +2499,22 @@ func TestBlocksStoreQuerier_PromQLExecution(t *testing.T) { } } +func TestBlocksStoreQuerier_ShouldRetryResourceBasedThrottlingError(t *testing.T) { + limits := map[resource.Type]float64{ + resource.CPU: 0.5, + resource.Heap: 0.5, + } + + resourceBasedLimiter, err := limiter.NewResourceBasedLimiter(&limiter.MockMonitor{ + CpuUtilization: 0.7, + HeapUtilization: 0.7, + }, limits, prometheus.DefaultRegisterer, "ingester") + require.NoError(t, err) + + err = resourceBasedLimiter.AcceptNewRequest() + require.True(t, isRetryableError(err)) +} + type blocksStoreSetMock struct { services.Service diff --git a/pkg/util/limiter/resource_based_limiter.go b/pkg/util/limiter/resource_based_limiter.go index 40415e39195..d7a1039af81 100644 --- a/pkg/util/limiter/resource_based_limiter.go +++ b/pkg/util/limiter/resource_based_limiter.go @@ -64,9 +64,22 @@ func (l *ResourceBasedLimiter) AcceptNewRequest() error { if utilization >= limit { l.limitBreachedCount.WithLabelValues(string(resType)).Inc() - return fmt.Errorf("%s utilization limit reached (limit: %.3f, utilization: %.3f)", resType, limit, utilization) + return fmt.Errorf("%s utilization limit reached (limit: %.3f, utilization: %.3f): %w", resType, limit, utilization, &ResourceLimitReachedError{}) } } return nil } + +type MockMonitor struct { + CpuUtilization float64 + HeapUtilization float64 +} + +func (m *MockMonitor) GetCPUUtilization() float64 { + return m.CpuUtilization +} + +func (m *MockMonitor) GetHeapUtilization() float64 { + return m.HeapUtilization +} diff --git a/pkg/util/limiter/resource_based_limiter_test.go b/pkg/util/limiter/resource_based_limiter_test.go index c84d59009f6..7691b9ceb48 100644 --- a/pkg/util/limiter/resource_based_limiter_test.go +++ b/pkg/util/limiter/resource_based_limiter_test.go @@ -15,16 +15,12 @@ func Test_ResourceBasedLimiter(t *testing.T) { resource.Heap: 0.5, } - _, err := NewResourceBasedLimiter(&mockMonitor{}, limits, prometheus.DefaultRegisterer, "ingester") + limiter, err := NewResourceBasedLimiter(&MockMonitor{ + CpuUtilization: 0.2, + HeapUtilization: 0.2, + }, limits, prometheus.DefaultRegisterer, "ingester") require.NoError(t, err) -} - -type mockMonitor struct{} -func (m *mockMonitor) GetCPUUtilization() float64 { - return 0 -} - -func (m *mockMonitor) GetHeapUtilization() float64 { - return 0 + err = limiter.AcceptNewRequest() + require.NoError(t, err) }