moved counter to vec to utils, better testing implemented

Author: pranavmishradatabricks

internal/cortex/frontend/transport/handler.go

@@ -62,7 +62,6 @@ type Handler struct {
 	querySeconds *prometheus.CounterVec
 	querySeries  *prometheus.CounterVec
 	queryBytes   *prometheus.CounterVec
-	cachedHits   *prometheus.CounterVec
 	activeUsers  *util.ActiveUsersCleanupService
 }
 
@@ -75,7 +74,7 @@ func NewHandler(cfg HandlerConfig, roundTripper http.RoundTripper, log log.Logge
 	}
 
 	if cfg.FailedQueryCacheCapacity > 0 {
-		FailedQueryCache, errQueryCache := utils.NewFailedQueryCache(cfg.FailedQueryCacheCapacity)
+		FailedQueryCache, errQueryCache := utils.NewFailedQueryCache(cfg.FailedQueryCacheCapacity, reg)
 		if errQueryCache != nil {
 			level.Warn(log).Log(errQueryCache.Error())
 		}
@@ -108,11 +107,6 @@ func NewHandler(cfg HandlerConfig, roundTripper http.RoundTripper, log log.Logge
 		_ = h.activeUsers.StartAsync(context.Background())
 	}
 
-	h.cachedHits = promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
-		Name: "cached_failed_queries_count",
-		Help: "Total number of queries that hit the failed query cache.",
-	}, []string{"total"})
-
 	return h
 }
 
@@ -148,7 +142,6 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		if cached {
 			w.WriteHeader(http.StatusForbidden)
 			level.Info(util_log.WithContext(r.Context(), f.log)).Log(message)
-			f.cachedHits.WithLabelValues("total").Inc()
 			return
 		}
 	}

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

@@ -12,6 +12,8 @@ import (
 	"strconv"
 
 	lru "github.com/hashicorp/golang-lru"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 )
 
 var (
@@ -23,9 +25,10 @@ type FailedQueryCache struct {
 	regex        *regexp.Regexp
 	errorExtract *regexp.Regexp
 	lruCache     *lru.Cache
+	cachedHits   *prometheus.CounterVec
 }
 
-func NewFailedQueryCache(capacity int) (*FailedQueryCache, error) {
+func NewFailedQueryCache(capacity int, reg prometheus.Registerer) (*FailedQueryCache, error) {
 	regex := regexp.MustCompile(`[\s\n\t]+`)
 	errorExtract := regexp.MustCompile(`Code\((\d+)\)`)
 	lruCache, err := lru.New(capacity)
@@ -34,10 +37,17 @@ func NewFailedQueryCache(capacity int) (*FailedQueryCache, error) {
 		err = fmt.Errorf("Failed to create lru cache: %s", err)
 		return nil, err
 	}
+	cachedHits := promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
+		Name: "cached_failed_queries_count",
+		Help: "Total number of queries that hit the failed query cache.",
+	}, []string{})
+
 	return &FailedQueryCache{
 		regex:        regex,
 		errorExtract: errorExtract,
-		lruCache:     lruCache}, err
+		lruCache:     lruCache,
+		cachedHits:   cachedHits,
+	}, err
 }
 
 // UpdateFailedQueryCache returns true if query is cached so that callsite can increase counter, returns message as a string for callsite to log outcome
@@ -79,10 +89,11 @@ 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 queryHitCache(queryExpressionNormalized string, queryExpressionRangeLength int, lruCache *lru.Cache) (bool, string) {
+func queryHitCache(queryExpressionNormalized string, queryExpressionRangeLength int, lruCache *lru.Cache, cachedHits *prometheus.CounterVec) (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)
+		cachedHits.WithLabelValues().Inc()
 		return true, message
 	}
 	return false, ""
@@ -140,6 +151,6 @@ func (f *FailedQueryCache) UpdateFailedQueryCache(err error, query url.Values) (
 func (f *FailedQueryCache) QueryHitCache(query url.Values) (bool, string) {
 	queryExpressionNormalized := f.normalizeQueryString(query)
 	queryExpressionRangeLength := getQueryRangeSeconds(query)
-	cached, message := queryHitCache(queryExpressionNormalized, queryExpressionRangeLength, f.lruCache)
+	cached, message := queryHitCache(queryExpressionNormalized, queryExpressionRangeLength, f.lruCache, f.cachedHits)
 	return cached, message
 }

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

@@ -10,12 +10,12 @@ import (
 	"testing"
 
 	"github.com/prometheus/client_golang/prometheus"
-	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/prometheus/client_golang/prometheus/testutil"
 )
 
 func TestNewFailedQueryCache(t *testing.T) {
-	cache, err := NewFailedQueryCache(2)
+	reg := prometheus.NewRegistry()
+	cache, err := NewFailedQueryCache(2, reg)
 	if cache == nil {
 		t.Fatalf("Expected cache to be created, but got nil")
 	}
@@ -25,7 +25,8 @@ func TestNewFailedQueryCache(t *testing.T) {
 }
 
 func TestUpdateFailedQueryCache(t *testing.T) {
-	cache, _ := NewFailedQueryCache(2)
+	reg := prometheus.NewRegistry()
+	cache, _ := NewFailedQueryCache(2, reg)
 
 	tests := []struct {
 		name            string
@@ -120,7 +121,8 @@ func TestUpdateFailedQueryCache(t *testing.T) {
 
 // TestQueryHitCache tests the QueryHitCache method
 func TestQueryHitCache(t *testing.T) {
-	cache, _ := NewFailedQueryCache(2)
+	reg := prometheus.NewRegistry()
+	cache, _ := NewFailedQueryCache(2, reg)
 	lruCache := cache.lruCache
 
 	lruCache.Add("test_query", 100)
@@ -202,26 +204,83 @@ func TestQueryHitCache(t *testing.T) {
 
 func TestCacheCounterVec(t *testing.T) {
 	reg := prometheus.NewRegistry()
-	cachedHits := promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
-		Name: "cached_failed_queries_count",
-		Help: "Total number of queries that hit the failed query cache.",
-	}, []string{"total"})
-
-	cachedHits.WithLabelValues("total").Inc()
-	expectedValue := 1.0
-	value := testutil.ToFloat64(cachedHits.WithLabelValues("total"))
-	if value != expectedValue {
-		t.Errorf("expected %v, got %v", expectedValue, value)
-	}
-	cachedHits.WithLabelValues("total").Inc()
-	expectedValue = 2.0
-	value = testutil.ToFloat64(cachedHits.WithLabelValues("total"))
-	if value != expectedValue {
-		t.Errorf("expected %v, got %v", expectedValue, value)
+	cache, _ := NewFailedQueryCache(2, reg)
+	lruCache := cache.lruCache
+
+	lruCache.Add("test_query", 100)
+	lruCache.Add(" tes t query ", 100)
+
+	tests := []struct {
+		name            string
+		query           url.Values
+		expectedCounter int
+	}{
+		{
+			name: "Cache miss",
+			query: url.Values{
+				"start": {"100"},
+				"end":   {"2000"},
+				"query": {"miss_query_counter_test"},
+			},
+			expectedCounter: 0,
+		}, {
+			name: "Cache hit",
+			query: url.Values{
+				"start": {"100"},
+				"end":   {"200"},
+				"query": {"test_query"},
+			},
+			expectedCounter: 1,
+		},
+		{
+			name: "Cache miss",
+			query: url.Values{
+				"start": {"100"},
+				"end":   {"200"},
+				"query": {"miss"},
+			},
+			expectedCounter: 1,
+		},
+
+		{
+			name: "Cache miss due to shorter range length",
+			query: url.Values{
+				"start": {"100"},
+				"end":   {"150"},
+				"query": {"test_query"},
+			},
+			expectedCounter: 1,
+		},
+
+		{
+			name: "Cache hit whitespace",
+			query: url.Values{
+				"start": {"100"},
+				"end":   {"200"},
+				"query": {" \n\ttes \tt \n   query \t\n  "},
+			},
+			expectedCounter: 2,
+		},
+
+		{
+			name: "Cache miss whitespace",
+			query: url.Values{
+				"start": {"100"},
+				"end":   {"200"},
+				"query": {" \n\tte s \tt \n   query \t\n  "},
+			},
+			expectedCounter: 2,
+		},
 	}
-	expectedValue = 0.0
-	value = testutil.ToFloat64(cachedHits.WithLabelValues("incorrect_label"))
-	if value != expectedValue {
-		t.Errorf("expected %v, got %v", expectedValue, value)
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cache.QueryHitCache(tt.query)
+			result := int(testutil.ToFloat64(cache.cachedHits.WithLabelValues()))
+			if result != tt.expectedCounter {
+				t.Errorf("expected counter value to be %v, got %v", tt.expectedCounter, result)
+			}
+
+		})
 	}
 }