From c133aa4f9ce114cd05a5904f855624106a4a3cf2 Mon Sep 17 00:00:00 2001
From: Thibault Mange <22740367+thibaultmg@users.noreply.github.com>
Date: Tue, 2 May 2023 19:12:33 +0200
Subject: [PATCH] store: fix inconsistent error for series limits

Signed-off-by: Thibault Mange <22740367+thibaultmg@users.noreply.github.com>
---
 pkg/store/bucket.go          |  12 ++++
 pkg/store/bucket_e2e_test.go | 110 +++++++++++++++++++++++++++++++++++
 2 files changed, 122 insertions(+)

diff --git a/pkg/store/bucket.go b/pkg/store/bucket.go
index ac4e439ab22..15084d4791d 100644
--- a/pkg/store/bucket.go
+++ b/pkg/store/bucket.go
@@ -1574,6 +1574,12 @@ func (s *BucketStore) LabelNames(ctx context.Context, req *storepb.LabelNamesReq
 	s.mtx.RUnlock()
 
 	if err := g.Wait(); err != nil {
+		if statusErr, ok := status.FromError(err); ok {
+			if statusErr.Code() == codes.ResourceExhausted {
+				return nil, status.Error(codes.ResourceExhausted, err.Error())
+			}
+		}
+
 		return nil, status.Error(codes.Internal, err.Error())
 	}
 
@@ -1762,6 +1768,12 @@ func (s *BucketStore) LabelValues(ctx context.Context, req *storepb.LabelValuesR
 	s.mtx.RUnlock()
 
 	if err := g.Wait(); err != nil {
+		if statusErr, ok := status.FromError(err); ok {
+			if statusErr.Code() == codes.ResourceExhausted {
+				return nil, status.Error(codes.ResourceExhausted, err.Error())
+			}
+		}
+
 		return nil, status.Error(codes.Aborted, err.Error())
 	}
 
diff --git a/pkg/store/bucket_e2e_test.go b/pkg/store/bucket_e2e_test.go
index 3d30ee171ce..a229d896001 100644
--- a/pkg/store/bucket_e2e_test.go
+++ b/pkg/store/bucket_e2e_test.go
@@ -6,6 +6,7 @@ package store
 import (
 	"context"
 	"fmt"
+	"math"
 	"os"
 	"path/filepath"
 	"strings"
@@ -760,6 +761,57 @@ func TestBucketStore_LabelNames_e2e(t *testing.T) {
 	})
 }
 
+func TestBucketStore_LabelNames_SeriesLimiter_e2e(t *testing.T) {
+	cases := map[string]struct {
+		maxSeriesLimit uint64
+		expectedErr    string
+		code           codes.Code
+	}{
+		"should succeed if the max series limit is not exceeded": {
+			maxSeriesLimit: math.MaxUint64,
+		},
+		"should fail if the max series limit is exceeded - ResourceExhausted": {
+			expectedErr:    "exceeded series limit",
+			maxSeriesLimit: 1,
+			code:           codes.ResourceExhausted,
+		},
+	}
+
+	for testName, testData := range cases {
+		t.Run(testName, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			bkt := objstore.NewInMemBucket()
+			dir := t.TempDir()
+			s := prepareStoreWithTestBlocks(t, dir, bkt, false, NewChunksLimiterFactory(0), NewSeriesLimiterFactory(testData.maxSeriesLimit), NewBytesLimiterFactory(0), emptyRelabelConfig, allowAllFilterConf)
+			testutil.Ok(t, s.store.SyncBlocks(ctx))
+			req := &storepb.LabelNamesRequest{
+				Matchers: []storepb.LabelMatcher{
+					{Type: storepb.LabelMatcher_EQ, Name: "a", Value: "1"},
+				},
+				Start: minTimeDuration.PrometheusTimestamp(),
+				End:   maxTimeDuration.PrometheusTimestamp(),
+			}
+
+			s.cache.SwapWith(noopCache{})
+
+			_, err := s.store.LabelNames(context.Background(), req)
+
+			if testData.expectedErr == "" {
+				testutil.Ok(t, err)
+			} else {
+				testutil.NotOk(t, err)
+				testutil.Assert(t, strings.Contains(err.Error(), testData.expectedErr))
+
+				status, ok := status.FromError(err)
+				testutil.Equals(t, true, ok)
+				testutil.Equals(t, testData.code, status.Code())
+			}
+		})
+	}
+}
+
 func TestBucketStore_LabelValues_e2e(t *testing.T) {
 	objtesting.ForeachStore(t, func(t *testing.T, bkt objstore.Bucket) {
 		ctx, cancel := context.WithCancel(context.Background())
@@ -867,6 +919,64 @@ func TestBucketStore_LabelValues_e2e(t *testing.T) {
 	})
 }
 
+func TestBucketStore_LabelValues_ChunksLimiter_e2e(t *testing.T) {
+	cases := map[string]struct {
+		maxSeriesLimit uint64
+		expectedErr    string
+		code           codes.Code
+	}{
+		"should succeed if the max chunks limit is not exceeded": {
+			maxSeriesLimit: math.MaxUint64,
+		},
+		"should fail if the max series limit is exceeded - ResourceExhausted": {
+			expectedErr:    "exceeded series limit",
+			maxSeriesLimit: 1,
+			code:           codes.ResourceExhausted,
+		},
+	}
+
+	for testName, testData := range cases {
+		t.Run(testName, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			bkt := objstore.NewInMemBucket()
+
+			dir := t.TempDir()
+
+			s := prepareStoreWithTestBlocks(t, dir, bkt, false, NewChunksLimiterFactory(0), NewSeriesLimiterFactory(testData.maxSeriesLimit), NewBytesLimiterFactory(0), emptyRelabelConfig, allowAllFilterConf)
+			testutil.Ok(t, s.store.SyncBlocks(ctx))
+
+			req := &storepb.LabelValuesRequest{
+				Label: "a",
+				Start: minTimeDuration.PrometheusTimestamp(),
+				End:   maxTimeDuration.PrometheusTimestamp(),
+				Matchers: []storepb.LabelMatcher{
+					{
+						Type:  storepb.LabelMatcher_EQ,
+						Name:  "a",
+						Value: "1",
+					},
+				},
+			}
+
+			s.cache.SwapWith(noopCache{})
+
+			_, err := s.store.LabelValues(context.Background(), req)
+
+			if testData.expectedErr == "" {
+				testutil.Ok(t, err)
+			} else {
+				testutil.NotOk(t, err)
+				testutil.Assert(t, strings.Contains(err.Error(), testData.expectedErr))
+
+				status, ok := status.FromError(err)
+				testutil.Equals(t, true, ok)
+				testutil.Equals(t, testData.code, status.Code())
+			}
+		})
+	}
+}
+
 func emptyToNil(values []string) []string {
 	if len(values) == 0 {
 		return nil