Change condition to check validation.LimitError

Signed-off-by: Justin Jung <[email protected]>

Author: justinjung04

pkg/ring/replication_set.go

@@ -3,13 +3,11 @@ package ring
 import (
 	"context"
 	"fmt"
-	"net/http"
 	"sort"
 	"time"
 
-	"github.com/weaveworks/common/httpgrpc"
-
 	"github.com/cortexproject/cortex/pkg/querier/partialdata"
+	"github.com/cortexproject/cortex/pkg/util/validation"
 )
 
 // ReplicationSet describes the instances to talk to for a given key, and how
@@ -83,10 +81,8 @@ func (r ReplicationSet) Do(ctx context.Context, delay time.Duration, zoneResults
 					return nil, res.err
 				}
 
-				if httpRes, ok := httpgrpc.HTTPResponseFromError(res.err); ok {
-					if httpRes.Code == http.StatusUnprocessableEntity {
-						return nil, res.err
-					}
+				if validation.IsLimitError(res.err) {
+					return nil, res.err
 				}
 
 				// force one of the delayed requests to start

pkg/ring/replication_set_test.go

@@ -3,16 +3,15 @@ package ring
 import (
 	"context"
 	"errors"
-	"net/http"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"github.com/weaveworks/common/httpgrpc"
 	"go.uber.org/atomic"
 
 	"github.com/cortexproject/cortex/pkg/querier/partialdata"
+	"github.com/cortexproject/cortex/pkg/util/validation"
 )
 
 func TestReplicationSet_GetAddresses(t *testing.T) {
@@ -221,16 +220,16 @@ func TestReplicationSet_Do(t *testing.T) {
 			queryPartialData:    true,
 		},
 		{
-			name:      "with partial data enabled, should fail on instances returning 422 in two zones",
-			instances: []InstanceDesc{{Zone: "zone1"}, {Zone: "zone2"}, {Zone: "zone3"}},
+			name:      "with partial data enabled, should fail on instances returning 422",
+			instances: []InstanceDesc{{Addr: "1", Zone: "zone1"}, {Addr: "2", Zone: "zone2"}, {Addr: "3", Zone: "zone3"}, {Addr: "4", Zone: "zone1"}, {Addr: "5", Zone: "zone2"}, {Addr: "6", Zone: "zone3"}},
 			f: func(ctx context.Context, ing *InstanceDesc) (interface{}, error) {
-				if ing.Zone == "zone1" || ing.Zone == "zone2" {
-					return nil, httpgrpc.Errorf(http.StatusUnprocessableEntity, "breached limit")
+				if ing.Addr == "1" || ing.Addr == "2" {
+					return nil, validation.LimitError("limit breached")
 				}
 				return 1, nil
 			},
 			maxUnavailableZones: 1,
-			expectedError:       httpgrpc.Errorf(http.StatusUnprocessableEntity, "breached limit"),
+			expectedError:       validation.LimitError("limit breached"),
 			queryPartialData:    true,
 		},
 		{

pkg/util/validation/limits.go

@@ -50,6 +50,11 @@ func (e LimitError) Error() string {
 	return string(e)
 }
 
+func IsLimitError(e error) bool {
+	var limitError LimitError
+	return errors.As(e, &limitError)
+}
+
 type DisabledRuleGroup struct {
 	Namespace string `yaml:"namespace" doc:"nocli|description=namespace in which the rule group belongs"`
 	Name      string `yaml:"name" doc:"nocli|description=name of the rule group"`

pkg/util/validation/limits_test.go

@@ -2,6 +2,7 @@ package validation
 
 import (
 	"encoding/json"
+	"fmt"
 	"reflect"
 	"regexp"
 	"strings"
@@ -885,3 +886,8 @@ func TestLimitsPerLabelSetsForSeries(t *testing.T) {
 		})
 	}
 }
+
+func TestIsLimitsError(t *testing.T) {
+	assert.False(t, IsLimitError(fmt.Errorf("test error")))
+	assert.True(t, IsLimitError(LimitError("test error")))
+}