refactored out code, changed which functions are exported

Author: pranavmishradatabricks

internal/cortex/frontend/transport/handler.go

@@ -70,12 +70,14 @@ type Handler struct {
 func NewHandler(cfg HandlerConfig, roundTripper http.RoundTripper, log log.Logger, reg prometheus.Registerer) http.Handler {
 	var (
 		FailedQueryCache *utils.FailedQueryCache
-		message          string
+		err              error
 	)
 
 	if cfg.FailedQueryCacheCapacity > 0 {
-		FailedQueryCache, message = utils.NewFailedQueryCache(cfg.FailedQueryCacheCapacity)
-		level.Warn(log).Log(message)
+		FailedQueryCache, err = utils.NewFailedQueryCache(cfg.FailedQueryCacheCapacity)
+		if err != nil {
+			level.Warn(log).Log(err.Error())
+		}
 	}
 
 	h := &Handler{
@@ -121,10 +123,9 @@ func NewHandler(cfg HandlerConfig, roundTripper http.RoundTripper, log log.Logge
 
 func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	var (
-		stats                      *querier_stats.Stats
-		queryString                url.Values
-		queryExpressionNormalized  string
-		queryExpressionRangeLength int
+		stats       *querier_stats.Stats
+		queryString url.Values
+		urlQuery    url.Values
 	)
 
 	// Initialise the stats in the context and make sure it's propagated
@@ -144,17 +145,11 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	r.Body = http.MaxBytesReader(w, r.Body, f.cfg.MaxBodySize)
 	r.Body = io.NopCloser(io.TeeReader(r.Body, &buf))
 
-	// Check if caching is enabled.
+	urlQuery = r.URL.Query()
+
+	// Check if query is cached
 	if f.failedQueryCache != nil {
-		//Store query.
-		query := r.URL.Query()
-		// Store query expression.
-		queryExpressionNormalized = f.failedQueryCache.NormalizeQueryString(query)
-		// Store query time range length.
-		queryExpressionRangeLength = utils.GetQueryRangeSeconds(query)
-
-		// Check if query in cache and whether value exceeds time range length. Log and increment counter appropriately.
-		cached, message := f.failedQueryCache.QueryHitCache(queryExpressionNormalized, queryExpressionRangeLength, f.failedQueryCache.LruCache)
+		cached, message := f.failedQueryCache.CallQueryHitCache(urlQuery)
 		if cached {
 			w.WriteHeader(http.StatusForbidden)
 			level.Info(util_log.WithContext(r.Context(), f.log)).Log(message)
@@ -171,9 +166,15 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		writeError(w, err)
 		queryString = f.parseRequestQueryString(r, buf)
 
-		// Try to update cache.
-		_, message := f.failedQueryCache.CallUpdateFailedQueryCache(err, queryExpressionNormalized, queryExpressionRangeLength)
-		level.Debug(util_log.WithContext(r.Context(), f.log)).Log(message)
+		// Update cache for failed queries.
+		if f.failedQueryCache != nil {
+			success, message := f.failedQueryCache.CallUpdateFailedQueryCache(err, urlQuery)
+			if success {
+				level.Info(util_log.WithContext(r.Context(), f.log)).Log(message)
+			} else {
+				level.Debug(util_log.WithContext(r.Context(), f.log)).Log(message)
+			}
+		}
 
 		if f.cfg.LogFailedQueries {
 			f.reportFailedQuery(r, queryString, err)

internal/cortex/frontend/transport/utils/utils.go

@@ -22,23 +22,26 @@ var (
 type FailedQueryCache struct {
 	regex        *regexp.Regexp
 	errorExtract *regexp.Regexp
-	LruCache     *lru.Cache
+	lruCache     *lru.Cache
 }
 
-func NewFailedQueryCache(capacity int) (*FailedQueryCache, string) {
+func NewFailedQueryCache(capacity int) (*FailedQueryCache, error) {
 	regex := regexp.MustCompile(`[\s\n\t]+`)
 	errorExtract := regexp.MustCompile(`Code\((\d+)\)`)
 	lruCache, err := lru.New(capacity)
-	message := ""
 	if err != nil {
 		lruCache = nil
-		message = fmt.Sprintf("Failed to create lru cache: %s", err)
+		err = fmt.Errorf("Failed to create lru cache: %s", err)
+		return nil, err
 	}
-	return &FailedQueryCache{regex, errorExtract, lruCache}, message
+	return &FailedQueryCache{
+		regex:        regex,
+		errorExtract: errorExtract,
+		lruCache:     lruCache}, err
 }
 
 // UpdateFailedQueryCache returns true if query is cached so that callsite can increase counter, returns message as a string for callsite to log outcome
-func (f *FailedQueryCache) UpdateFailedQueryCache(err error, queryExpressionNormalized string, queryExpressionRangeLength int, lruCache *lru.Cache) (bool, string) {
+func updateFailedQueryCache(err error, queryExpressionNormalized string, queryExpressionRangeLength int, lruCache *lru.Cache) (bool, string) {
 	// Extracting error code from error string.
 	codeExtract := f.errorExtract.FindStringSubmatch(err.Error())
 
@@ -100,7 +103,7 @@ func (f *FailedQueryCache) UpdateFailedQueryCache(err error, queryExpressionNorm
 }
 
 // QueryHitCache checks if the lru cache is hit and returns whether to increment counter for cache hits along with appropriate message.
-func (f *FailedQueryCache) QueryHitCache(queryExpressionNormalized string, queryExpressionRangeLength int, lruCache *lru.Cache) (bool, string) {
+func queryHitCache(queryExpressionNormalized string, queryExpressionRangeLength int, lruCache *lru.Cache) (bool, string) {
 	if value, ok := lruCache.Get(queryExpressionNormalized); ok && value.(int) >= queryExpressionRangeLength {
 		message := fmt.Sprintf(
 			`%s: %s, %s: %s, %s: %d`, "msg", "Retrieved query from cache",
@@ -124,7 +127,7 @@ func isCacheableError(statusCode int) bool {
 }
 
 // GetQueryRangeSeconds Time range length for queries, if either of "start" or "end" are not present, return 0.
-func GetQueryRangeSeconds(query url.Values) int {
+func getQueryRangeSeconds(query url.Values) int {
 	start, err := strconv.Atoi(query.Get("start"))
 	if err != nil {
 		return 0
@@ -136,14 +139,20 @@ func GetQueryRangeSeconds(query url.Values) int {
 	return end - start
 }
 
-func (f *FailedQueryCache) NormalizeQueryString(query url.Values) string {
+func normalizeQueryString(query url.Values) string {
 	return f.regex.ReplaceAllString(query.Get("query"), " ")
 }
 
-func (f *FailedQueryCache) CallUpdateFailedQueryCache(err error, queryExpressionNormalized string, queryExpressionRangeLength int) (bool, string) {
-	if f == nil {
-		return false, "Failed query cache is not enabled"
-	}
-	success, message := f.UpdateFailedQueryCache(err, queryExpressionNormalized, queryExpressionRangeLength, f.LruCache)
+func (f *FailedQueryCache) CallUpdateFailedQueryCache(err error, query url.Values) (bool, string) {
+	queryExpressionNormalized := normalizeQueryString(query)
+	queryExpressionRangeLength := getQueryRangeSeconds(query)
+	success, message := updateFailedQueryCache(err, queryExpressionNormalized, queryExpressionRangeLength, f.lruCache)
 	return success, message
 }
+
+func (f *FailedQueryCache) CallQueryHitCache(query url.Values) (bool, string) {
+	queryExpressionNormalized := normalizeQueryString(query)
+	queryExpressionRangeLength := getQueryRangeSeconds(query)
+	cached, message := queryHitCache(queryExpressionNormalized, queryExpressionRangeLength, f.lruCache)
+	return cached, message
+}

internal/cortex/frontend/transport/utils/utils_test.go

@@ -9,7 +9,6 @@ import (
 	"testing"
 )
 
-// TestNewFailedQueryCache tests the NewFailedQueryCache function
 func TestNewFailedQueryCache(t *testing.T) {
 	cache, msg := NewFailedQueryCache(2)
 	if cache == nil {
@@ -20,7 +19,6 @@ func TestNewFailedQueryCache(t *testing.T) {
 	}
 }
 
-// TestUpdateFailedQueryCache tests the UpdateFailedQueryCache method
 func TestUpdateFailedQueryCache(t *testing.T) {
 	cache, _ := NewFailedQueryCache(2)
 	lruCache := cache.LruCache