Moving code for cache functions to utils file

internal/cortex/frontend/transport/handler.go

@@ -11,17 +11,16 @@ import (
 	"io"
 	"net/http"
 	"net/url"
-	"regexp"
 	"strconv"
 	"strings"
 	"syscall"
 	"time"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/hashicorp/golang-lru"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	"github.com/thanos-io/thanos/internal/cortex/frontend/transport/utils"
 	querier_stats "github.com/thanos-io/thanos/internal/cortex/querier/stats"
 	"github.com/thanos-io/thanos/internal/cortex/tenant"
 	"github.com/thanos-io/thanos/internal/cortex/util"
@@ -40,7 +39,6 @@ var (
 	errCanceled              = httpgrpc.Errorf(StatusClientClosedRequest, context.Canceled.Error())
 	errDeadlineExceeded      = httpgrpc.Errorf(http.StatusGatewayTimeout, context.DeadlineExceeded.Error())
 	errRequestEntityTooLarge = httpgrpc.Errorf(http.StatusRequestEntityTooLarge, "http: request body too large")
-	cacheableResponseCodes   = []int{http.StatusRequestTimeout, http.StatusGatewayTimeout, http.StatusBadRequest}
 )
 
 // HandlerConfig Config for a Handler.
@@ -55,12 +53,10 @@ type HandlerConfig struct {
 // Handler accepts queries and forwards them to RoundTripper. It can log slow queries,
 // but all other logic is inside the RoundTripper.
 type Handler struct {
-	cfg          HandlerConfig
-	log          log.Logger
-	roundTripper http.RoundTripper
-	lruCache     *lru.Cache
-	regex        *regexp.Regexp
-	errorExtract *regexp.Regexp
+	cfg              HandlerConfig
+	log              log.Logger
+	roundTripper     http.RoundTripper
+	failedQueryCache *utils.FailedQueryCache
 
 	// Metrics.
 	querySeconds *prometheus.CounterVec
@@ -72,26 +68,18 @@ type Handler struct {
 
 // NewHandler creates a new frontend handler.
 func NewHandler(cfg HandlerConfig, roundTripper http.RoundTripper, log log.Logger, reg prometheus.Registerer) http.Handler {
-	var (
-		LruCache *lru.Cache
-		err      error
-	)
-
-	if cfg.FailedQueryCacheCapacity > 0 {
-		LruCache, err = lru.New(cfg.FailedQueryCacheCapacity)
-		if err != nil {
-			LruCache = nil
-			level.Warn(log).Log("msg", "Failed to create LruCache", "error", err)
-		}
-	}
-
 	h := &Handler{
 		cfg:          cfg,
 		log:          log,
 		roundTripper: roundTripper,
-		lruCache:     LruCache,
-		regex:        regexp.MustCompile(`[\s\n\t]+`),
-		errorExtract: regexp.MustCompile(`Code\((\d+)\)`),
+	}
+
+	if cfg.FailedQueryCacheCapacity > 0 {
+		FailedQueryCache, errQueryCache := utils.NewFailedQueryCache(cfg.FailedQueryCacheCapacity)
+		if errQueryCache != nil {
+			level.Warn(log).Log(errQueryCache.Error())
+		}
+		h.failedQueryCache = FailedQueryCache
 	}
 
 	if cfg.QueryStatsEnabled {
@@ -130,10 +118,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
@@ -153,22 +140,14 @@ 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.
-	if f.lruCache != nil {
-		// Store query expression.
-		queryExpressionNormalized = f.regex.ReplaceAllString(r.URL.Query().Get("query"), " ")
-
-		// Store query time range length.
-		queryExpressionRangeLength = getQueryRangeSeconds(r)
+	urlQuery = r.URL.Query()
 
-		// Check if query in cache and whether value exceeds time range length.
-		if value, ok := f.lruCache.Get(queryExpressionNormalized); ok && value.(int) >= queryExpressionRangeLength {
+	// Check if query is cached
+	if f.failedQueryCache != nil {
+		cached, message := f.failedQueryCache.QueryHitCache(urlQuery)
+		if cached {
 			w.WriteHeader(http.StatusForbidden)
-			level.Info(util_log.WithContext(r.Context(), f.log)).Log(
-				"msg", "Retrieved query from cache",
-				"normalized_query", queryExpressionNormalized,
-				"range_seconds", queryExpressionRangeLength,
-			)
+			level.Info(util_log.WithContext(r.Context(), f.log)).Log(message)
 			f.cachedHits.Inc()
 			return
 		}
@@ -182,9 +161,14 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		writeError(w, err)
 		queryString = f.parseRequestQueryString(r, buf)
 
-		// Check if caching is enabled.
-		if f.lruCache != nil {
-			f.updateFailedQueryCache(err, queryExpressionNormalized, queryExpressionRangeLength, r)
+		// Update cache for failed queries.
+		if f.failedQueryCache != nil {
+			success, message := f.failedQueryCache.UpdateFailedQueryCache(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 {
@@ -223,84 +207,6 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-func (f *Handler) updateFailedQueryCache(err error, queryExpressionNormalized string, queryExpressionRangeLength int, r *http.Request) {
-	// Extracting error code from error string.
-	codeExtract := f.errorExtract.FindStringSubmatch(err.Error())
-
-	// Checking if error code extracted successfully.
-	if codeExtract == nil || len(codeExtract) < 2 {
-		level.Error(util_log.WithContext(r.Context(), f.log)).Log(
-			"msg", "Error string regex conversion error",
-			"normalized_query", queryExpressionNormalized,
-			"range_seconds", queryExpressionRangeLength,
-			"error", err)
-		return
-	}
-
-	// Converting error code to int.
-	errCode, strConvError := strconv.Atoi(codeExtract[1])
-
-	// Checking if error code extracted properly from string.
-	if strConvError != nil {
-		level.Error(util_log.WithContext(r.Context(), f.log)).Log(
-			"msg", "String to int conversion error",
-			"normalized_query", queryExpressionNormalized,
-			"range_seconds", queryExpressionRangeLength,
-			"error", err)
-		return
-	}
-
-	// If error should be cached, store it in cache.
-	if !isCacheableError(errCode) {
-		level.Debug(util_log.WithContext(r.Context(), f.log)).Log(
-			"msg", "Query not cached due to non-cacheable error code",
-			"normalized_query", queryExpressionNormalized,
-			"range_seconds", queryExpressionRangeLength,
-			"error", err,
-		)
-		return
-	}
-
-	// Checks if queryExpression is already in cache, and updates time range length value to min of stored and new value.
-	if contains, _ := f.lruCache.ContainsOrAdd(queryExpressionNormalized, queryExpressionRangeLength); contains {
-		if oldValue, ok := f.lruCache.Get(queryExpressionNormalized); ok {
-			queryExpressionRangeLength = min(queryExpressionRangeLength, oldValue.(int))
-		}
-		f.lruCache.Add(queryExpressionNormalized, queryExpressionRangeLength)
-	}
-
-	level.Debug(util_log.WithContext(r.Context(), f.log)).Log(
-		"msg", "Cached a failed query",
-		"normalized_query", queryExpressionNormalized,
-		"range_seconds", queryExpressionRangeLength,
-		"error", err,
-	)
-
-}
-
-// isCacheableError Returns true if response code is in pre-defined cacheable errors list, else returns false.
-func isCacheableError(statusCode int) bool {
-	for _, errStatusCode := range cacheableResponseCodes {
-		if errStatusCode == statusCode {
-			return true
-		}
-	}
-	return false
-}
-
-// Time range length for queries, if either of "start" or "end" are not present, return 0.
-func getQueryRangeSeconds(r *http.Request) int {
-	start, err := strconv.Atoi(r.URL.Query().Get("start"))
-	if err != nil {
-		return 0
-	}
-	end, err := strconv.Atoi(r.URL.Query().Get("end"))
-	if err != nil {
-		return 0
-	}
-	return end - start
-}
-
 func (f *Handler) reportFailedQuery(r *http.Request, queryString url.Values, err error) {
 	// NOTE(GiedriusS): see https://github.com/grafana/grafana/pull/60301 for more info.
 	grafanaDashboardUID := "-"

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

@@ -0,0 +1,145 @@
+// Copyright (c) The Cortex Authors.
+// Licensed under the Apache License 2.0.
+
+// Package utils Monitoring platform team helper resources for frontend
+package utils
+
+import (
+	"fmt"
+	"net/http"
+	"net/url"
+	"regexp"
+	"strconv"
+
+	lru "github.com/hashicorp/golang-lru"
+)
+
+var (
+	cacheableResponseCodes = []int{http.StatusRequestTimeout, http.StatusGatewayTimeout, http.StatusBadRequest}
+)
+
+// FailedQueryCache Handler holds an instance of FailedQueryCache and calls its methods
+type FailedQueryCache struct {
+	regex        *regexp.Regexp
+	errorExtract *regexp.Regexp
+	lruCache     *lru.Cache
+}
+
+func NewFailedQueryCache(capacity int) (*FailedQueryCache, error) {
+	regex := regexp.MustCompile(`[\s\n\t]+`)
+	errorExtract := regexp.MustCompile(`Code\((\d+)\)`)
+	lruCache, err := lru.New(capacity)
+	if err != nil {
+		lruCache = nil
+		err = fmt.Errorf("Failed to create lru cache: %s", err)
+		return nil, err
+	}
+	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) {
+	// Extracting error code from error string.
+	codeExtract := f.errorExtract.FindStringSubmatch(err.Error())
+
+	// Checking if error code extracted successfully.
+	if codeExtract == nil || len(codeExtract) < 2 {
+		message := createLogMessage("String to regex conversion error", queryExpressionNormalized, -1, queryExpressionRangeLength, err)
+		return false, message
+	}
+
+	// Converting error code to int.
+	errCode, strConvError := strconv.Atoi(codeExtract[1])
+
+	// Checking if error code extracted properly from string.
+	if strConvError != nil {
+		message := createLogMessage("String to int conversion error", queryExpressionNormalized, -1, queryExpressionRangeLength, err)
+		return false, message
+	}
+
+	// If error should be cached, store it in cache.
+	if !isCacheableError(errCode) {
+		message := createLogMessage("Query not cached due to non-cacheable error code", queryExpressionNormalized, -1, queryExpressionRangeLength, err)
+		return false, message
+	}
+
+	// Checks if queryExpression is already in cache, and updates time range length value to min of stored and new value.
+	if contains, _ := lruCache.ContainsOrAdd(queryExpressionNormalized, queryExpressionRangeLength); contains {
+		if oldValue, ok := lruCache.Get(queryExpressionNormalized); ok {
+			queryExpressionRangeLength = min(queryExpressionRangeLength, oldValue.(int))
+		}
+		lruCache.Add(queryExpressionNormalized, queryExpressionRangeLength)
+	}
+
+	message := createLogMessage("Cached a failed query", queryExpressionNormalized, -1, queryExpressionRangeLength, err)
+	return true, message
+}
+
+// QueryHitCache checks if the lru cache is hit and returns whether to increment counter for cache hits along with appropriate message.
+func queryHitCache(queryExpressionNormalized string, queryExpressionRangeLength int, lruCache *lru.Cache) (bool, string) {
+	if value, ok := lruCache.Get(queryExpressionNormalized); ok && value.(int) <= queryExpressionRangeLength {
+		cachedQueryRangeSeconds := value.(int)
+		message := createLogMessage("Retrieved query from cache", queryExpressionNormalized, cachedQueryRangeSeconds, queryExpressionRangeLength, nil)
+		return true, message
+	}
+	return false, ""
+}
+
+// isCacheableError Returns true if response code is in pre-defined cacheable errors list, else returns false.
+func isCacheableError(statusCode int) bool {
+	for _, errStatusCode := range cacheableResponseCodes {
+		if errStatusCode == statusCode {
+			return true
+		}
+	}
+	return false
+}
+
+// GetQueryRangeSeconds Time range length for queries, if either of "start" or "end" are not present, return 0.
+func getQueryRangeSeconds(query url.Values) int {
+	start, err := strconv.Atoi(query.Get("start"))
+	if err != nil {
+		return 0
+	}
+	end, err := strconv.Atoi(query.Get("end"))
+	if err != nil {
+		return 0
+	}
+	return end - start
+}
+
+func (f *FailedQueryCache) normalizeQueryString(query url.Values) string {
+	return f.regex.ReplaceAllString(query.Get("query"), " ")
+}
+
+func createLogMessage(message string, queryExpressionNormalized string, cachedQueryRangeSeconds int, queryExpressionRangeLength int, err error) string {
+	if err == nil {
+		return fmt.Sprintf(
+			`%s: %s, %s: %s, %s: %d, %s: %d`, "msg", message,
+			"cached_query", queryExpressionNormalized,
+			"cached_range_seconds", cachedQueryRangeSeconds,
+			"query_range_seconds", queryExpressionRangeLength)
+	}
+	return fmt.Sprintf(
+		`%s: %s, %s: %s, %s: %d, %s: %s`, "msg", message,
+		"cached_query", queryExpressionNormalized,
+		"query_range_seconds", queryExpressionRangeLength,
+		"cached_error", err)
+}
+
+func (f *FailedQueryCache) UpdateFailedQueryCache(err error, query url.Values) (bool, string) {
+	queryExpressionNormalized := f.normalizeQueryString(query)
+	queryExpressionRangeLength := getQueryRangeSeconds(query)
+	success, message := f.updateFailedQueryCache(err, queryExpressionNormalized, queryExpressionRangeLength, f.lruCache)
+	return success, message
+}
+
+func (f *FailedQueryCache) QueryHitCache(query url.Values) (bool, string) {
+	queryExpressionNormalized := f.normalizeQueryString(query)
+	queryExpressionRangeLength := getQueryRangeSeconds(query)
+	cached, message := queryHitCache(queryExpressionNormalized, queryExpressionRangeLength, f.lruCache)
+	return cached, message
+}