Fix AWS QueryPages backoff logic (#523)

Author: aaron7

pkg/chunk/aws_storage_client.go

@@ -272,13 +272,32 @@ func (a awsStorageClient) QueryPages(ctx context.Context, query IndexQuery, call
 	}
 
 	request := a.queryRequestFn(ctx, input)
-	backoff := minBackoff
 	pageCount := 0
 	defer func() {
 		dynamoQueryPagesCount.Observe(float64(pageCount))
 	}()
+
 	for page := request; page != nil; page = page.NextPage() {
 		pageCount++
+
+		response, err := a.queryPage(ctx, input, page)
+		if err != nil {
+			return err
+		}
+
+		if getNextPage := callback(response, !page.HasNextPage()); !getNextPage {
+			if err != nil {
+				return fmt.Errorf("QueryPages error: table=%v, err=%v", *input.TableName, page.Error())
+			}
+			return nil
+		}
+	}
+	return nil
+}
+
+func (a awsStorageClient) queryPage(ctx context.Context, input *dynamodb.QueryInput, page dynamoDBRequest) (dynamoDBReadResponse, error) {
+	backoff := minBackoff
+	for i := 0; i < maxRetries; i++ {
 		err := instrument.TimeRequestHistogram(ctx, "DynamoDB.QueryPages", dynamoRequestDuration, func(_ context.Context) error {
 			return page.Send()
 		})
@@ -290,28 +309,18 @@ func (a awsStorageClient) QueryPages(ctx context.Context, query IndexQuery, call
 
 		if err != nil {
 			recordDynamoError(*input.TableName, err, "DynamoDB.QueryPages")
-
 			if awsErr, ok := err.(awserr.Error); ok && ((awsErr.Code() == dynamodb.ErrCodeProvisionedThroughputExceededException) || page.Retryable()) {
 				time.Sleep(backoff)
 				backoff = nextBackoff(backoff)
 				continue
 			}
-
-			return fmt.Errorf("QueryPages error: table=%v, err=%v", *input.TableName, err)
+			return nil, fmt.Errorf("QueryPage error: table=%v, err=%v", *input.TableName, err)
 		}
 
 		queryOutput := page.Data().(*dynamodb.QueryOutput)
-		if getNextPage := callback(dynamoDBReadResponse(queryOutput.Items), !page.HasNextPage()); !getNextPage {
-			if err != nil {
-				return fmt.Errorf("QueryPages error: table=%v, err=%v", *input.TableName, page.Error())
-			}
-			return nil
-		}
-
-		backoff = minBackoff
+		return dynamoDBReadResponse(queryOutput.Items), nil
 	}
-
-	return nil
+	return nil, fmt.Errorf("QueryPage error: maxRetries exceeded for table %v", *input.TableName)
 }
 
 type dynamoDBRequest interface {

pkg/chunk/aws_storage_client_test.go

@@ -605,16 +605,18 @@ func TestAWSStorageClientQueryPages(t *testing.T) {
 	}
 
 	tests := []struct {
-		name  string
-		query IndexQuery
-		want  []IndexEntry
+		name           string
+		query          IndexQuery
+		provisionedErr int
+		want           []IndexEntry
 	}{
 		{
 			"check HashValue only",
 			IndexQuery{
 				TableName: "table",
 				HashValue: "flip",
 			},
+			0,
 			[]IndexEntry{entries[5], entries[6], entries[7]},
 		},
 		{
@@ -624,6 +626,7 @@ func TestAWSStorageClientQueryPages(t *testing.T) {
 				HashValue:       "foo",
 				RangeValueStart: []byte("bar:2"),
 			},
+			0,
 			[]IndexEntry{entries[1], entries[2], entries[3], entries[4]},
 		},
 		{
@@ -633,6 +636,7 @@ func TestAWSStorageClientQueryPages(t *testing.T) {
 				HashValue:        "foo",
 				RangeValuePrefix: []byte("baz:"),
 			},
+			0,
 			[]IndexEntry{entries[3], entries[4]},
 		},
 		{
@@ -643,13 +647,25 @@ func TestAWSStorageClientQueryPages(t *testing.T) {
 				RangeValuePrefix: []byte("bar"),
 				ValueEqual:       []byte("20"),
 			},
+			0,
+			[]IndexEntry{entries[1]},
+		},
+		{
+			"check retry logic",
+			IndexQuery{
+				TableName:        "table",
+				HashValue:        "foo",
+				RangeValuePrefix: []byte("bar"),
+				ValueEqual:       []byte("20"),
+			},
+			2,
 			[]IndexEntry{entries[1]},
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			dynamoDB := newMockDynamoDB(0, 0)
+			dynamoDB := newMockDynamoDB(0, tt.provisionedErr)
 			client := awsStorageClient{
 				DynamoDB:                dynamoDB,
 				queryRequestFn:          dynamoDB.queryRequest,