Handle cache-control headers also during a partial cache hit.

pkg/querier/queryrange/results_cache.go

@@ -168,25 +168,19 @@ func (s resultsCache) Do(ctx context.Context, r Request) (Response, error) {
 }
 
 // shouldCacheResponse says whether the response should be cached or not.
-func shouldCacheResponse(r Response) bool {
+func (s resultsCache) shouldCacheResponse(r Response) bool {
 	if promResp, ok := r.(*PrometheusResponse); ok {
-		shouldCache := true
-	outer:
 		for _, hv := range promResp.Headers {
-			if hv == nil {
+			if hv.GetName() != cachecontrolHeader {
 				continue
 			}
-			if hv.Name != cachecontrolHeader {
-				continue
-			}
-			for _, v := range hv.Values {
+			for _, v := range hv.GetValues() {
 				if v == noCacheValue {
-					shouldCache = false
-					break outer
+					level.Debug(s.logger).Log("msg", fmt.Sprintf("%s header in response is equal to %s, not caching the response", cachecontrolHeader, noCacheValue))
+					return false
 				}
 			}
 		}
-		return shouldCache
 	}
 	return true
 }
@@ -197,8 +191,7 @@ func (s resultsCache) handleMiss(ctx context.Context, r Request) (Response, []Ex
 		return nil, nil, err
 	}
 
-	if !shouldCacheResponse(response) {
-		level.Debug(s.logger).Log("msg", fmt.Sprintf("%s header in response is equal to %s, not caching the response", cachecontrolHeader, noCacheValue))
+	if !s.shouldCacheResponse(response) {
 		return response, []Extent{}, nil
 	}
 
@@ -238,6 +231,9 @@ func (s resultsCache) handleHit(ctx context.Context, r Request, extents []Extent
 
 	for _, reqResp := range reqResps {
 		responses = append(responses, reqResp.Response)
+		if !s.shouldCacheResponse(reqResp.Response) {
+			continue
+		}
 		extent, err := toExtent(ctx, reqResp.Request, reqResp.Response)
 		if err != nil {
 			return nil, nil, err

pkg/querier/queryrange/results_cache_test.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/cortexproject/cortex/pkg/chunk/cache"
 	"github.com/cortexproject/cortex/pkg/ingester/client"
+	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
 )
 
@@ -96,6 +97,7 @@ func mkExtent(start, end int64) Extent {
 }
 
 func TestShouldCache(t *testing.T) {
+	c := &resultsCache{logger: util.Logger}
 	for i, tc := range []struct {
 		input    Response
 		expected bool
@@ -136,10 +138,27 @@ func TestShouldCache(t *testing.T) {
 			}),
 			expected: false,
 		},
+		// some broken responses
+		{
+			input:    Response(&PrometheusResponse{}),
+			expected: true,
+		},
+		{
+			input: Response(&PrometheusResponse{
+				Headers: []*PrometheusResponseHeader{nil},
+			}),
+			expected: true,
+		},
+		{
+			input: Response(&PrometheusResponse{
+				Headers: []*PrometheusResponseHeader{{Name: cachecontrolHeader}},
+			}),
+			expected: true,
+		},
 	} {
 		{
 			t.Run(strconv.Itoa(i), func(t *testing.T) {
-				ret := shouldCacheResponse(tc.input)
+				ret := c.shouldCacheResponse(tc.input)
 				require.Equal(t, tc.expected, ret)
 			})
 		}