Make start and end time check to be skipped if not specified

Signed-off-by: Justin Jung <jungjust@amazon.com>
cortexproject · Nov 28, 2023 · 4001fe4 · 4001fe4
1 parent 584c6f9
commit 4001fe4
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 18 deletions.
diff --git a/pkg/util/query/priority.go b/pkg/util/query/priority.go
@@ -24,25 +24,22 @@ func GetPriority(requestParams url.Values, now time.Time, queryPriority validati
 				continue
 			}
 
-			startTimeThreshold := now.Add(-1 * attribute.StartTime.Abs()).Truncate(time.Second).UTC()
-			endTimeThreshold := now.Add(-1 * attribute.EndTime.Abs()).Add(1 * time.Second).Truncate(time.Second).UTC()
-
 			if startTime, err := util.ParseTime(startParam); err == nil {
 				if endTime, err := util.ParseTime(endParam); err == nil {
-					if isBetweenThresholds(util.TimeFromMillis(startTime), util.TimeFromMillis(endTime), startTimeThreshold, endTimeThreshold) {
+					if isWithinTimeAttributes(attribute, now, startTime, endTime) {
 						return priority.Priority
 					}
 				}
 			}
 
 			if instantTime, err := util.ParseTime(timeParam); err == nil {
-				if isBetweenThresholds(util.TimeFromMillis(instantTime), util.TimeFromMillis(instantTime), startTimeThreshold, endTimeThreshold) {
+				if isWithinTimeAttributes(attribute, now, instantTime, instantTime) {
 					return priority.Priority
 				}
 			}
 
 			if timeParam == "" {
-				if isBetweenThresholds(now, now, startTimeThreshold, endTimeThreshold) {
+				if isWithinTimeAttributes(attribute, now, util.TimeToMillis(now), util.TimeToMillis(now)) {
 					return priority.Priority
 				}
 			}
@@ -52,6 +49,26 @@ func GetPriority(requestParams url.Values, now time.Time, queryPriority validati
 	return queryPriority.DefaultPriority
 }
 
-func isBetweenThresholds(start, end, startThreshold, endThreshold time.Time) bool {
-	return (start.Equal(startThreshold) || start.After(startThreshold)) && (end.Equal(endThreshold) || end.Before(endThreshold))
+func isWithinTimeAttributes(attribute validation.QueryAttribute, now time.Time, startTime, endTime int64) bool {
+	if attribute.StartTime == 0 && attribute.EndTime == 0 {
+		return true
+	}
+
+	if attribute.StartTime != 0 {
+		startTimeThreshold := now.Add(-1 * time.Duration(attribute.StartTime).Abs()).Truncate(time.Second)
+
+		if util.TimeFromMillis(startTime).Before(startTimeThreshold) {
+			return false
+		}
+	}
+
+	if attribute.EndTime != 0 {
+		endTimeThreshold := now.Add(-1 * time.Duration(attribute.EndTime).Abs()).Add(1 * time.Second).Truncate(time.Second)
+
+		if util.TimeFromMillis(endTime).After(endTimeThreshold) {
+			return false
+		}
+	}
+
+	return true
 }
diff --git a/pkg/util/query/priority_test.go b/pkg/util/query/priority_test.go
@@ -7,6 +7,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/prometheus/common/model"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/cortexproject/cortex/pkg/util/validation"
@@ -21,8 +22,6 @@ func Test_GetPriorityShouldReturnDefaultPriorityIfNotEnabledOrEmptyQueryString(t
 				{
 					Regex:         ".*",
 					CompiledRegex: regexp.MustCompile(".*"),
-					StartTime:     2 * time.Hour,
-					EndTime:       0 * time.Hour,
 				},
 			},
 		},
@@ -52,8 +51,6 @@ func Test_GetPriorityShouldConsiderRegex(t *testing.T) {
 				{
 					Regex:         "sum",
 					CompiledRegex: regexp.MustCompile("sum"),
-					StartTime:     2 * time.Hour,
-					EndTime:       0 * time.Hour,
 				},
 			},
 		},
@@ -128,8 +125,8 @@ func Test_GetPriorityShouldConsiderStartAndEndTime(t *testing.T) {
 			Priority: 1,
 			QueryAttributes: []validation.QueryAttribute{
 				{
-					StartTime: 45 * time.Minute,
-					EndTime:   15 * time.Minute,
+					StartTime: model.Duration(45 * time.Minute),
+					EndTime:   model.Duration(15 * time.Minute),
 				},
 			},
 		},
@@ -186,3 +183,42 @@ func Test_GetPriorityShouldConsiderStartAndEndTime(t *testing.T) {
 		"end":   []string{strconv.FormatInt(now.Add(-1*time.Minute).Unix(), 10)},
 	}, now, queryPriority))
 }
+
+func Test_GetPriorityShouldSKipStartAndEndTimeIfEmpty(t *testing.T) {
+	now := time.Now()
+	priorities := []validation.PriorityDef{
+		{
+			Priority: 1,
+			QueryAttributes: []validation.QueryAttribute{
+				{
+					Regex: "^test$",
+				},
+			},
+		},
+	}
+	queryPriority := validation.QueryPriority{
+		Enabled:    true,
+		Priorities: priorities,
+	}
+
+	assert.Equal(t, int64(1), GetPriority(url.Values{
+		"query": []string{"test"},
+	}, now, queryPriority))
+	assert.Equal(t, int64(1), GetPriority(url.Values{
+		"query": []string{"test"},
+		"time":  []string{strconv.FormatInt(now.Add(8760*time.Hour).Unix(), 10)},
+	}, now, queryPriority))
+	assert.Equal(t, int64(1), GetPriority(url.Values{
+		"query": []string{"test"},
+		"time":  []string{strconv.FormatInt(now.Unix(), 10)},
+	}, now, queryPriority))
+	assert.Equal(t, int64(1), GetPriority(url.Values{
+		"query": []string{"test"},
+		"time":  []string{strconv.FormatInt(now.Add(-8760*time.Hour).Unix(), 10)},
+	}, now, queryPriority))
+	assert.Equal(t, int64(1), GetPriority(url.Values{
+		"query": []string{"test"},
+		"start": []string{strconv.FormatInt(now.Add(-100000*time.Minute).Unix(), 10)},
+		"end":   []string{strconv.FormatInt(now.Add(100000*time.Minute).Unix(), 10)},
+	}, now, queryPriority))
+}
diff --git a/pkg/util/validation/limits.go b/pkg/util/validation/limits.go
@@ -62,9 +62,9 @@ type PriorityDef struct {
 }
 
 type QueryAttribute struct {
-	Regex         string        `yaml:"regex" json:"regex" doc:"nocli|description=Query string regex.|default=.*"`
-	StartTime     time.Duration `yaml:"start_time" json:"start_time" doc:"nocli|description=Query start time.|default=0s"`
-	EndTime       time.Duration `yaml:"end_time" json:"end_time" doc:"nocli|description=Query end time.|default=0s"`
+	Regex         string         `yaml:"regex" json:"regex" doc:"nocli|description=Query string regex. If set to empty string, it will not match anything.|default=\"\""`
+	StartTime     model.Duration `yaml:"start_time" json:"start_time" doc:"nocli|description=Query start time. If set to 0, the start time won't be checked.'.|default=0"`
+	EndTime       model.Duration `yaml:"end_time" json:"end_time" doc:"nocli|description=Query end time. If set to 0, the end time won't be checked.|default=0"`
 	CompiledRegex *regexp.Regexp
 }
 

diff --git a/pkg/util/validation/limits_test.go b/pkg/util/validation/limits_test.go
@@ -620,7 +620,7 @@ func TestHasQueryPriorityRegexChanged(t *testing.T) {
 
 	require.True(t, l.hasQueryPriorityRegexChanged())
 
-	l.QueryPriority.Priorities[0].QueryAttributes[0].StartTime = 2 * time.Hour
+	l.QueryPriority.Priorities[0].QueryAttributes[0].StartTime = model.Duration(2 * time.Hour)
 
 	require.False(t, l.hasQueryPriorityRegexChanged())