1package validation
2
3import (
4	"encoding/json"
5	"reflect"
6	"strings"
7	"testing"
8	"time"
9
10	"github.com/grafana/dskit/flagext"
11	"github.com/prometheus/common/model"
12	"github.com/prometheus/prometheus/pkg/relabel"
13	"github.com/stretchr/testify/assert"
14	"github.com/stretchr/testify/require"
15	"golang.org/x/time/rate"
16	"gopkg.in/yaml.v2"
17)
18
19// mockTenantLimits exposes per-tenant limits based on a provided map
20type mockTenantLimits struct {
21	limits map[string]*Limits
22}
23
24// newMockTenantLimits creates a new mockTenantLimits that returns per-tenant limits based on
25// the given map
26func newMockTenantLimits(limits map[string]*Limits) *mockTenantLimits {
27	return &mockTenantLimits{
28		limits: limits,
29	}
30}
31
32func (l *mockTenantLimits) ByUserID(userID string) *Limits {
33	return l.limits[userID]
34}
35
36func (l *mockTenantLimits) AllByUserID() map[string]*Limits {
37	return l.limits
38}
39
40func TestLimits_Validate(t *testing.T) {
41	t.Parallel()
42
43	tests := map[string]struct {
44		limits           Limits
45		shardByAllLabels bool
46		expected         error
47	}{
48		"max-global-series-per-user disabled and shard-by-all-labels=false": {
49			limits:           Limits{MaxGlobalSeriesPerUser: 0},
50			shardByAllLabels: false,
51			expected:         nil,
52		},
53		"max-global-series-per-user enabled and shard-by-all-labels=false": {
54			limits:           Limits{MaxGlobalSeriesPerUser: 1000},
55			shardByAllLabels: false,
56			expected:         errMaxGlobalSeriesPerUserValidation,
57		},
58		"max-global-series-per-user disabled and shard-by-all-labels=true": {
59			limits:           Limits{MaxGlobalSeriesPerUser: 1000},
60			shardByAllLabels: true,
61			expected:         nil,
62		},
63	}
64
65	for testName, testData := range tests {
66		testData := testData
67
68		t.Run(testName, func(t *testing.T) {
69			assert.Equal(t, testData.expected, testData.limits.Validate(testData.shardByAllLabels))
70		})
71	}
72}
73
74func TestOverrides_MaxChunksPerQueryFromStore(t *testing.T) {
75	tests := map[string]struct {
76		setup    func(limits *Limits)
77		expected int
78	}{
79		"should return the default legacy setting with the default config": {
80			setup:    func(limits *Limits) {},
81			expected: 2000000,
82		},
83		"the new config option should take precedence over the deprecated one": {
84			setup: func(limits *Limits) {
85				limits.MaxChunksPerQueryFromStore = 10
86				limits.MaxChunksPerQuery = 20
87			},
88			expected: 20,
89		},
90		"the deprecated config option should be used if the new config option is unset": {
91			setup: func(limits *Limits) {
92				limits.MaxChunksPerQueryFromStore = 10
93			},
94			expected: 10,
95		},
96	}
97
98	for testName, testData := range tests {
99		t.Run(testName, func(t *testing.T) {
100			limits := Limits{}
101			flagext.DefaultValues(&limits)
102			testData.setup(&limits)
103
104			overrides, err := NewOverrides(limits, nil)
105			require.NoError(t, err)
106			assert.Equal(t, testData.expected, overrides.MaxChunksPerQueryFromStore("test"))
107		})
108	}
109}
110
111func TestOverridesManager_GetOverrides(t *testing.T) {
112	tenantLimits := map[string]*Limits{}
113
114	defaults := Limits{
115		MaxLabelNamesPerSeries: 100,
116	}
117	ov, err := NewOverrides(defaults, newMockTenantLimits(tenantLimits))
118	require.NoError(t, err)
119
120	require.Equal(t, 100, ov.MaxLabelNamesPerSeries("user1"))
121	require.Equal(t, 0, ov.MaxLabelValueLength("user1"))
122
123	// Update limits for tenant user1. We only update single field, the rest is copied from defaults.
124	// (That is how limits work when loaded from YAML)
125	l := Limits{}
126	l = defaults
127	l.MaxLabelValueLength = 150
128
129	tenantLimits["user1"] = &l
130
131	// Checking whether overrides were enforced
132	require.Equal(t, 100, ov.MaxLabelNamesPerSeries("user1"))
133	require.Equal(t, 150, ov.MaxLabelValueLength("user1"))
134
135	// Verifying user2 limits are not impacted by overrides
136	require.Equal(t, 100, ov.MaxLabelNamesPerSeries("user2"))
137	require.Equal(t, 0, ov.MaxLabelValueLength("user2"))
138}
139
140func TestLimitsLoadingFromYaml(t *testing.T) {
141	SetDefaultLimitsForYAMLUnmarshalling(Limits{
142		MaxLabelNameLength: 100,
143	})
144
145	inp := `ingestion_rate: 0.5`
146
147	l := Limits{}
148	err := yaml.UnmarshalStrict([]byte(inp), &l)
149	require.NoError(t, err)
150
151	assert.Equal(t, 0.5, l.IngestionRate, "from yaml")
152	assert.Equal(t, 100, l.MaxLabelNameLength, "from defaults")
153}
154
155func TestLimitsLoadingFromJson(t *testing.T) {
156	SetDefaultLimitsForYAMLUnmarshalling(Limits{
157		MaxLabelNameLength: 100,
158	})
159
160	inp := `{"ingestion_rate": 0.5}`
161
162	l := Limits{}
163	err := json.Unmarshal([]byte(inp), &l)
164	require.NoError(t, err)
165
166	assert.Equal(t, 0.5, l.IngestionRate, "from json")
167	assert.Equal(t, 100, l.MaxLabelNameLength, "from defaults")
168
169	// Unmarshal should fail if input contains unknown struct fields and
170	// the decoder flag `json.Decoder.DisallowUnknownFields()` is set
171	inp = `{"unknown_fields": 100}`
172	l = Limits{}
173	dec := json.NewDecoder(strings.NewReader(inp))
174	dec.DisallowUnknownFields()
175	err = dec.Decode(&l)
176	assert.Error(t, err)
177}
178
179func TestLimitsTagsYamlMatchJson(t *testing.T) {
180	limits := reflect.TypeOf(Limits{})
181	n := limits.NumField()
182	var mismatch []string
183
184	for i := 0; i < n; i++ {
185		field := limits.Field(i)
186
187		// Note that we aren't requiring YAML and JSON tags to match, just that
188		// they either both exist or both don't exist.
189		hasYAMLTag := field.Tag.Get("yaml") != ""
190		hasJSONTag := field.Tag.Get("json") != ""
191
192		if hasYAMLTag != hasJSONTag {
193			mismatch = append(mismatch, field.Name)
194		}
195	}
196
197	assert.Empty(t, mismatch, "expected no mismatched JSON and YAML tags")
198}
199
200func TestLimitsStringDurationYamlMatchJson(t *testing.T) {
201	inputYAML := `
202max_query_lookback: 1s
203max_query_length: 1s
204`
205	inputJSON := `{"max_query_lookback": "1s", "max_query_length": "1s"}`
206
207	limitsYAML := Limits{}
208	err := yaml.Unmarshal([]byte(inputYAML), &limitsYAML)
209	require.NoError(t, err, "expected to be able to unmarshal from YAML")
210
211	limitsJSON := Limits{}
212	err = json.Unmarshal([]byte(inputJSON), &limitsJSON)
213	require.NoError(t, err, "expected to be able to unmarshal from JSON")
214
215	assert.Equal(t, limitsYAML, limitsJSON)
216}
217
218func TestLimitsAlwaysUsesPromDuration(t *testing.T) {
219	stdlibDuration := reflect.TypeOf(time.Duration(0))
220	limits := reflect.TypeOf(Limits{})
221	n := limits.NumField()
222	var badDurationType []string
223
224	for i := 0; i < n; i++ {
225		field := limits.Field(i)
226		if field.Type == stdlibDuration {
227			badDurationType = append(badDurationType, field.Name)
228		}
229	}
230
231	assert.Empty(t, badDurationType, "some Limits fields are using stdlib time.Duration instead of model.Duration")
232}
233
234func TestMetricRelabelConfigLimitsLoadingFromYaml(t *testing.T) {
235	SetDefaultLimitsForYAMLUnmarshalling(Limits{})
236
237	inp := `
238metric_relabel_configs:
239- action: drop
240  source_labels: [le]
241  regex: .+
242`
243	exp := relabel.DefaultRelabelConfig
244	exp.Action = relabel.Drop
245	regex, err := relabel.NewRegexp(".+")
246	require.NoError(t, err)
247	exp.Regex = regex
248	exp.SourceLabels = model.LabelNames([]model.LabelName{"le"})
249
250	l := Limits{}
251	err = yaml.UnmarshalStrict([]byte(inp), &l)
252	require.NoError(t, err)
253
254	assert.Equal(t, []*relabel.Config{&exp}, l.MetricRelabelConfigs)
255}
256
257func TestSmallestPositiveIntPerTenant(t *testing.T) {
258	tenantLimits := map[string]*Limits{
259		"tenant-a": {
260			MaxQueryParallelism: 5,
261		},
262		"tenant-b": {
263			MaxQueryParallelism: 10,
264		},
265	}
266
267	defaults := Limits{
268		MaxQueryParallelism: 0,
269	}
270	ov, err := NewOverrides(defaults, newMockTenantLimits(tenantLimits))
271	require.NoError(t, err)
272
273	for _, tc := range []struct {
274		tenantIDs []string
275		expLimit  int
276	}{
277		{tenantIDs: []string{}, expLimit: 0},
278		{tenantIDs: []string{"tenant-a"}, expLimit: 5},
279		{tenantIDs: []string{"tenant-b"}, expLimit: 10},
280		{tenantIDs: []string{"tenant-c"}, expLimit: 0},
281		{tenantIDs: []string{"tenant-a", "tenant-b"}, expLimit: 5},
282		{tenantIDs: []string{"tenant-c", "tenant-d", "tenant-e"}, expLimit: 0},
283		{tenantIDs: []string{"tenant-a", "tenant-b", "tenant-c"}, expLimit: 0},
284	} {
285		assert.Equal(t, tc.expLimit, SmallestPositiveIntPerTenant(tc.tenantIDs, ov.MaxQueryParallelism))
286	}
287}
288
289func TestSmallestPositiveNonZeroIntPerTenant(t *testing.T) {
290	tenantLimits := map[string]*Limits{
291		"tenant-a": {
292			MaxQueriersPerTenant: 5,
293		},
294		"tenant-b": {
295			MaxQueriersPerTenant: 10,
296		},
297	}
298
299	defaults := Limits{
300		MaxQueriersPerTenant: 0,
301	}
302	ov, err := NewOverrides(defaults, newMockTenantLimits(tenantLimits))
303	require.NoError(t, err)
304
305	for _, tc := range []struct {
306		tenantIDs []string
307		expLimit  int
308	}{
309		{tenantIDs: []string{}, expLimit: 0},
310		{tenantIDs: []string{"tenant-a"}, expLimit: 5},
311		{tenantIDs: []string{"tenant-b"}, expLimit: 10},
312		{tenantIDs: []string{"tenant-c"}, expLimit: 0},
313		{tenantIDs: []string{"tenant-a", "tenant-b"}, expLimit: 5},
314		{tenantIDs: []string{"tenant-c", "tenant-d", "tenant-e"}, expLimit: 0},
315		{tenantIDs: []string{"tenant-a", "tenant-b", "tenant-c"}, expLimit: 5},
316	} {
317		assert.Equal(t, tc.expLimit, SmallestPositiveNonZeroIntPerTenant(tc.tenantIDs, ov.MaxQueriersPerUser))
318	}
319}
320
321func TestSmallestPositiveNonZeroDurationPerTenant(t *testing.T) {
322	tenantLimits := map[string]*Limits{
323		"tenant-a": {
324			MaxQueryLength: model.Duration(time.Hour),
325		},
326		"tenant-b": {
327			MaxQueryLength: model.Duration(4 * time.Hour),
328		},
329	}
330
331	defaults := Limits{
332		MaxQueryLength: 0,
333	}
334	ov, err := NewOverrides(defaults, newMockTenantLimits(tenantLimits))
335	require.NoError(t, err)
336
337	for _, tc := range []struct {
338		tenantIDs []string
339		expLimit  time.Duration
340	}{
341		{tenantIDs: []string{}, expLimit: time.Duration(0)},
342		{tenantIDs: []string{"tenant-a"}, expLimit: time.Hour},
343		{tenantIDs: []string{"tenant-b"}, expLimit: 4 * time.Hour},
344		{tenantIDs: []string{"tenant-c"}, expLimit: time.Duration(0)},
345		{tenantIDs: []string{"tenant-a", "tenant-b"}, expLimit: time.Hour},
346		{tenantIDs: []string{"tenant-c", "tenant-d", "tenant-e"}, expLimit: time.Duration(0)},
347		{tenantIDs: []string{"tenant-a", "tenant-b", "tenant-c"}, expLimit: time.Hour},
348	} {
349		assert.Equal(t, tc.expLimit, SmallestPositiveNonZeroDurationPerTenant(tc.tenantIDs, ov.MaxQueryLength))
350	}
351}
352
353func TestAlertmanagerNotificationLimits(t *testing.T) {
354	for name, tc := range map[string]struct {
355		inputYAML         string
356		expectedRateLimit rate.Limit
357		expectedBurstSize int
358	}{
359		"no email specific limit": {
360			inputYAML: `
361alertmanager_notification_rate_limit: 100
362`,
363			expectedRateLimit: 100,
364			expectedBurstSize: 100,
365		},
366		"zero limit": {
367			inputYAML: `
368alertmanager_notification_rate_limit: 100
369
370alertmanager_notification_rate_limit_per_integration:
371  email: 0
372`,
373			expectedRateLimit: rate.Inf,
374			expectedBurstSize: maxInt,
375		},
376
377		"negative limit": {
378			inputYAML: `
379alertmanager_notification_rate_limit_per_integration:
380  email: -10
381`,
382			expectedRateLimit: 0,
383			expectedBurstSize: 0,
384		},
385
386		"positive limit, negative burst": {
387			inputYAML: `
388alertmanager_notification_rate_limit_per_integration:
389  email: 222
390`,
391			expectedRateLimit: 222,
392			expectedBurstSize: 222,
393		},
394
395		"infinte limit": {
396			inputYAML: `
397alertmanager_notification_rate_limit_per_integration:
398  email: .inf
399`,
400			expectedRateLimit: rate.Inf,
401			expectedBurstSize: maxInt,
402		},
403	} {
404		t.Run(name, func(t *testing.T) {
405			limitsYAML := Limits{}
406			err := yaml.Unmarshal([]byte(tc.inputYAML), &limitsYAML)
407			require.NoError(t, err, "expected to be able to unmarshal from YAML")
408
409			ov, err := NewOverrides(limitsYAML, nil)
410			require.NoError(t, err)
411
412			require.Equal(t, tc.expectedRateLimit, ov.NotificationRateLimit("user", "email"))
413			require.Equal(t, tc.expectedBurstSize, ov.NotificationBurstSize("user", "email"))
414		})
415	}
416}
417
418func TestAlertmanagerNotificationLimitsOverrides(t *testing.T) {
419	baseYaml := `
420alertmanager_notification_rate_limit: 5
421
422alertmanager_notification_rate_limit_per_integration:
423 email: 100
424`
425
426	overrideGenericLimitsOnly := `
427testuser:
428  alertmanager_notification_rate_limit: 333
429`
430
431	overrideEmailLimits := `
432testuser:
433  alertmanager_notification_rate_limit_per_integration:
434    email: 7777
435`
436
437	overrideGenericLimitsAndEmailLimits := `
438testuser:
439  alertmanager_notification_rate_limit: 333
440
441  alertmanager_notification_rate_limit_per_integration:
442    email: 7777
443`
444
445	differentUserOverride := `
446differentuser:
447  alertmanager_notification_limits_per_integration:
448    email: 500
449`
450
451	for name, tc := range map[string]struct {
452		testedIntegration string
453		overrides         string
454		expectedRateLimit rate.Limit
455		expectedBurstSize int
456	}{
457		"no overrides, pushover": {
458			testedIntegration: "pushover",
459			expectedRateLimit: 5,
460			expectedBurstSize: 5,
461		},
462
463		"no overrides, email": {
464			testedIntegration: "email",
465			expectedRateLimit: 100,
466			expectedBurstSize: 100,
467		},
468
469		"generic override, pushover": {
470			testedIntegration: "pushover",
471			overrides:         overrideGenericLimitsOnly,
472			expectedRateLimit: 333,
473			expectedBurstSize: 333,
474		},
475
476		"generic override, email": {
477			testedIntegration: "email",
478			overrides:         overrideGenericLimitsOnly,
479			expectedRateLimit: 100, // there is email-specific override in default config.
480			expectedBurstSize: 100,
481		},
482
483		"email limit override, pushover": {
484			testedIntegration: "pushover",
485			overrides:         overrideEmailLimits,
486			expectedRateLimit: 5, // loaded from defaults when parsing YAML
487			expectedBurstSize: 5,
488		},
489
490		"email limit override, email": {
491			testedIntegration: "email",
492			overrides:         overrideEmailLimits,
493			expectedRateLimit: 7777,
494			expectedBurstSize: 7777,
495		},
496
497		"generic and email limit override, pushover": {
498			testedIntegration: "pushover",
499			overrides:         overrideGenericLimitsAndEmailLimits,
500			expectedRateLimit: 333,
501			expectedBurstSize: 333,
502		},
503
504		"generic and email limit override, email": {
505			testedIntegration: "email",
506			overrides:         overrideGenericLimitsAndEmailLimits,
507			expectedRateLimit: 7777,
508			expectedBurstSize: 7777,
509		},
510
511		"partial email limit override": {
512			testedIntegration: "email",
513			overrides: `
514testuser:
515  alertmanager_notification_rate_limit_per_integration:
516    email: 500
517`,
518			expectedRateLimit: 500, // overridden
519			expectedBurstSize: 500, // same as rate limit
520		},
521
522		"different user override, pushover": {
523			testedIntegration: "pushover",
524			overrides:         differentUserOverride,
525			expectedRateLimit: 5,
526			expectedBurstSize: 5,
527		},
528
529		"different user overridem, email": {
530			testedIntegration: "email",
531			overrides:         differentUserOverride,
532			expectedRateLimit: 100,
533			expectedBurstSize: 100,
534		},
535	} {
536		t.Run(name, func(t *testing.T) {
537			SetDefaultLimitsForYAMLUnmarshalling(Limits{})
538
539			limitsYAML := Limits{}
540			err := yaml.Unmarshal([]byte(baseYaml), &limitsYAML)
541			require.NoError(t, err, "expected to be able to unmarshal from YAML")
542
543			SetDefaultLimitsForYAMLUnmarshalling(limitsYAML)
544
545			overrides := map[string]*Limits{}
546			err = yaml.Unmarshal([]byte(tc.overrides), &overrides)
547			require.NoError(t, err, "parsing overrides")
548
549			tl := newMockTenantLimits(overrides)
550
551			ov, err := NewOverrides(limitsYAML, tl)
552			require.NoError(t, err)
553
554			require.Equal(t, tc.expectedRateLimit, ov.NotificationRateLimit("testuser", tc.testedIntegration))
555			require.Equal(t, tc.expectedBurstSize, ov.NotificationBurstSize("testuser", tc.testedIntegration))
556		})
557	}
558}
559