Skip to content

Commit 71da81f

Browse files
erlan-zerlan-z
committed
fix for query_rejection bug
Signed-off-by: Erlan Zholdubai uulu <[email protected]>
1 parent df270ee commit 71da81f

File tree

4 files changed

+185
-58
lines changed

4 files changed

+185
-58
lines changed

‎CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
* [BUGFIX] Scheduler: Fix user queue in scheduler that was not thread-safe. #6077
4848
* [BUGFIX] Ingester: Include out-of-order head compaction when compacting TSDB head. #6108
4949
* [BUGFIX] Ingester: Fix `cortex_ingester_tsdb_mmap_chunks_total` metric. #6134
50+
* [BUGFIX] Query Frontend: Fix query rejection bug for metadata queries. #6143
5051

5152
## 1.17.1 2024-05-20
5253

‎integration/query_frontend_test.go

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -798,4 +798,44 @@ func TestQueryFrontendQueryRejection(t *testing.T) {
798798
require.Equal(t, http.StatusUnprocessableEntity, resp.StatusCode)
799799
require.Contains(t, string(body), tripperware.QueryRejectErrorMessage)
800800

801+
newRuntimeConfig = []byte(`overrides:
802+
user-1:
803+
query_rejection:
804+
enabled: true
805+
query_attributes:
806+
- query_step_limit:
807+
min: 12s
808+
- api_type: "labels"
809+
- dashboard_uid: "dash123"
810+
panel_id: "panel321"
811+
`)
812+
813+
require.NoError(t, client.Upload(context.Background(), configFileName, bytes.NewReader(newRuntimeConfig)))
814+
time.Sleep(2 * time.Second)
815+
816+
// We expect any request for speific api to be rejected if api_type is configured for that api and no other properties provided
817+
resp, body, err = c.LabelNamesRaw([]string{}, time.Time{}, time.Time{}, map[string]string{})
818+
require.NoError(t, err)
819+
require.Equal(t, http.StatusUnprocessableEntity, resp.StatusCode)
820+
require.Contains(t, string(body), tripperware.QueryRejectErrorMessage)
821+
822+
// We expect request not to be rejected if all the provided parameters are not applicable for current API type
823+
// There is no dashboardUID or panelId in metadata queries so if only those provided then metadata query shouldn't be rejected.
824+
resp, body, err = c.LabelValuesRaw("cluster", []string{}, time.Time{}, time.Time{}, map[string]string{"User-Agent": "grafana-agent/v0.19.0"})
825+
require.NoError(t, err)
826+
require.Equal(t, http.StatusOK, resp.StatusCode)
827+
require.NotContains(t, string(body), tripperware.QueryRejectErrorMessage)
828+
829+
// We expect instant request not to be rejected if only query_step_limit is provided (it's not applicable to instant queries)
830+
resp, body, err = c.QueryRaw("{instance=~\"hello.*\"}", time.Time{}, map[string]string{})
831+
require.NoError(t, err)
832+
require.Equal(t, http.StatusOK, resp.StatusCode)
833+
require.NotContains(t, string(body), tripperware.QueryRejectErrorMessage)
834+
835+
// We expect range query request to be rejected even if only query_step_limit is provided
836+
resp, body, err = c.QueryRangeRaw(`rate(test[1m])`, now.Add(-11*time.Hour), now.Add(-8*time.Hour), 20*time.Minute, map[string]string{"X-Dashboard-Uid": "dash123", "User-Agent": "grafana"})
837+
require.NoError(t, err)
838+
require.Equal(t, http.StatusUnprocessableEntity, resp.StatusCode)
839+
require.Contains(t, string(body), tripperware.QueryRejectErrorMessage)
840+
801841
}

‎pkg/querier/tripperware/query_attribute_matcher.go

Lines changed: 71 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import (
1414
"github.com/cortexproject/cortex/pkg/util/validation"
1515
)
1616

17-
const QueryRejectErrorMessage = "This query has been rejected by the service operator."
17+
const QueryRejectErrorMessage = "This query has been rejected by the service operator due to its high load on the service and impact to performance of other queries."
1818

1919
func rejectQueryOrSetPriority(r *http.Request, now time.Time, lookbackDelta time.Duration, limits Limits, userStr string, rejectedQueriesPerTenant *prometheus.CounterVec) error {
2020
if limits == nil || !(limits.QueryPriority(userStr).Enabled || limits.QueryRejection(userStr).Enabled) {
@@ -86,89 +86,120 @@ func getOperation(r *http.Request) string {
8686
}
8787

8888
func matchAttributeForExpressionQuery(attribute validation.QueryAttribute, op string, r *http.Request, query string, now time.Time, minTime, maxTime int64) bool {
89-
if attribute.ApiType != "" && attribute.ApiType != op {
90-
return false
89+
matched := false
90+
if attribute.ApiType != "" {
91+
matched = true
92+
if attribute.ApiType != op {
93+
return false
94+
}
9195
}
92-
if attribute.Regex != "" && attribute.Regex != ".*" && attribute.Regex != ".+" {
93-
if attribute.CompiledRegex != nil && !attribute.CompiledRegex.MatchString(query) {
96+
if attribute.Regex != "" {
97+
matched = true
98+
if attribute.Regex != ".*" && attribute.Regex != ".+" && attribute.CompiledRegex != nil && !attribute.CompiledRegex.MatchString(query) {
9499
return false
95100
}
96101
}
97102

98-
if !isWithinTimeAttributes(attribute.TimeWindow, now, minTime, maxTime) {
99-
return false
103+
if attribute.TimeWindow.Start != 0 || attribute.TimeWindow.End != 0 {
104+
matched = true
105+
if !isWithinTimeAttributes(attribute.TimeWindow, now, minTime, maxTime) {
106+
return false
107+
}
100108
}
101109

102-
if !isWithinTimeRangeAttribute(attribute.TimeRangeLimit, minTime, maxTime) {
103-
return false
110+
if attribute.TimeRangeLimit.Min != 0 || attribute.TimeRangeLimit.Max != 0 {
111+
matched = true
112+
if !isWithinTimeRangeAttribute(attribute.TimeRangeLimit, minTime, maxTime) {
113+
return false
114+
}
104115
}
105116

106-
if op == "query_range" && !isWithinQueryStepLimit(attribute.QueryStepLimit, r) {
107-
return false
117+
if op == "query_range" && (attribute.QueryStepLimit.Min != 0 || attribute.QueryStepLimit.Max != 0) {
118+
matched = true
119+
if !isWithinQueryStepLimit(attribute.QueryStepLimit, r) {
120+
return false
121+
}
108122
}
109123

110-
if attribute.UserAgentRegex != "" && attribute.UserAgentRegex != ".*" && attribute.CompiledUserAgentRegex != nil {
111-
if !attribute.CompiledUserAgentRegex.MatchString(r.Header.Get("User-Agent")) {
124+
if attribute.UserAgentRegex != "" {
125+
matched = true
126+
if attribute.UserAgentRegex != ".*" && attribute.CompiledUserAgentRegex != nil && !attribute.CompiledUserAgentRegex.MatchString(r.Header.Get("User-Agent")) {
112127
return false
113128
}
114129
}
115130

116-
if attribute.DashboardUID != "" && attribute.DashboardUID != r.Header.Get("X-Dashboard-Uid") {
117-
return false
131+
if attribute.DashboardUID != "" {
132+
matched = true
133+
if attribute.DashboardUID != r.Header.Get("X-Dashboard-Uid") {
134+
return false
135+
}
118136
}
119137

120-
if attribute.PanelID != "" && attribute.PanelID != r.Header.Get("X-Panel-Id") {
121-
return false
138+
if attribute.PanelID != "" {
139+
matched = true
140+
if attribute.PanelID != r.Header.Get("X-Panel-Id") {
141+
return false
142+
}
122143
}
123144

124-
return true
145+
return matched
125146
}
126147

127148
func matchAttributeForMetadataQuery(attribute validation.QueryAttribute, op string, r *http.Request, now time.Time) bool {
128-
if attribute.ApiType != "" && attribute.ApiType != op {
129-
return false
149+
matched := false
150+
if attribute.ApiType != "" {
151+
matched = true
152+
if attribute.ApiType != op {
153+
return false
154+
}
130155
}
131156
if err := r.ParseForm(); err != nil {
132157
return false
133158
}
134-
if attribute.Regex != "" && attribute.Regex != ".*" && attribute.CompiledRegex != nil {
135-
atLeastOneMatched := false
136-
for _, matcher := range r.Form["match[]"] {
137-
if attribute.CompiledRegex.MatchString(matcher) {
138-
atLeastOneMatched = true
139-
break
159+
if attribute.Regex != "" {
160+
matched = true
161+
if attribute.Regex != ".*" && attribute.CompiledRegex != nil {
162+
atLeastOneMatched := false
163+
for _, matcher := range r.Form["match[]"] {
164+
if attribute.CompiledRegex.MatchString(matcher) {
165+
atLeastOneMatched = true
166+
break
167+
}
168+
}
169+
if !atLeastOneMatched {
170+
return false
140171
}
141-
}
142-
if !atLeastOneMatched {
143-
return false
144172
}
145173
}
146174

147175
startTime, _ := util.ParseTime(r.FormValue("start"))
148176
endTime, _ := util.ParseTime(r.FormValue("end"))
149177

150-
if !isWithinTimeAttributes(attribute.TimeWindow, now, startTime, endTime) {
151-
return false
178+
if attribute.TimeWindow.Start != 0 || attribute.TimeWindow.End != 0 {
179+
matched = true
180+
if !isWithinTimeAttributes(attribute.TimeWindow, now, startTime, endTime) {
181+
return false
182+
}
152183
}
153184

154-
if !isWithinTimeRangeAttribute(attribute.TimeRangeLimit, startTime, endTime) {
155-
return false
185+
if attribute.TimeRangeLimit.Min != 0 || attribute.TimeRangeLimit.Max != 0 {
186+
matched = true
187+
if !isWithinTimeRangeAttribute(attribute.TimeRangeLimit, startTime, endTime) {
188+
return false
189+
}
156190
}
157191

158-
if attribute.UserAgentRegex != "" && attribute.UserAgentRegex != ".*" && attribute.CompiledUserAgentRegex != nil {
159-
if !attribute.CompiledUserAgentRegex.MatchString(r.Header.Get("User-Agent")) {
192+
if attribute.UserAgentRegex != "" {
193+
matched = true
194+
if attribute.UserAgentRegex != ".*" && attribute.CompiledUserAgentRegex != nil && !attribute.CompiledUserAgentRegex.MatchString(r.Header.Get("User-Agent")) {
160195
return false
161196
}
162197
}
163198

164-
return true
199+
return matched
165200
}
166201

167202
func isWithinTimeAttributes(timeWindow validation.TimeWindow, now time.Time, startTime, endTime int64) bool {
168-
if timeWindow.Start == 0 && timeWindow.End == 0 {
169-
return true
170-
}
171-
172203
if timeWindow.Start != 0 {
173204
startTimeThreshold := now.Add(-1 * time.Duration(timeWindow.Start).Abs()).Add(-1 * time.Minute).Truncate(time.Minute).UnixMilli()
174205
if startTime == 0 || startTime < startTimeThreshold {
@@ -187,9 +218,6 @@ func isWithinTimeAttributes(timeWindow validation.TimeWindow, now time.Time, sta
187218
}
188219

189220
func isWithinTimeRangeAttribute(limit validation.TimeRangeLimit, startTime, endTime int64) bool {
190-
if limit.Min == 0 && limit.Max == 0 {
191-
return true
192-
}
193221

194222
if startTime == 0 || endTime == 0 {
195223
return false
@@ -208,9 +236,6 @@ func isWithinTimeRangeAttribute(limit validation.TimeRangeLimit, startTime, endT
208236
}
209237

210238
func isWithinQueryStepLimit(queryStepLimit validation.QueryStepLimit, r *http.Request) bool {
211-
if queryStepLimit.Min == 0 && queryStepLimit.Max == 0 {
212-
return true
213-
}
214239

215240
step, err := util.ParseDurationMs(r.FormValue("step"))
216241
if err != nil {

0 commit comments

Comments
 (0)