Refactor: use newer AWS API for paginated queries

This is less code, and more robust when retrying requests.
We don't need an indirection on the request object for testing now.

Note the `readThrottle` is no longer used on queries: it should only
have been applied to chunk operations.

Signed-off-by: Bryan Boreham <[email protected]>

Author: bboreham

pkg/chunk/aws/dynamodb_storage_client.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/go-kit/kit/log/level"
 	ot "github.com/opentracing/opentracing-go"
+	otlog "github.com/opentracing/opentracing-go/log"
 	"golang.org/x/time/rate"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -144,7 +145,6 @@ type dynamoDBStorageClient struct {
 
 	// These functions exists for mocking, so we don't have to write a whole load
 	// of boilerplate.
-	queryRequestFn          func(ctx context.Context, input *dynamodb.QueryInput) dynamoDBRequest
 	batchGetItemRequestFn   func(ctx context.Context, input *dynamodb.BatchGetItemInput) dynamoDBRequest
 	batchWriteItemRequestFn func(ctx context.Context, input *dynamodb.BatchWriteItemInput) dynamoDBRequest
 }
@@ -172,7 +172,6 @@ func newDynamoDBStorageClient(cfg DynamoDBConfig, schemaCfg chunk.SchemaConfig)
 		DynamoDB:      dynamoDB,
 		writeThrottle: rate.NewLimiter(rate.Limit(cfg.ThrottleLimit), dynamoDBMaxWriteBatchSize),
 	}
-	client.queryRequestFn = client.queryRequest
 	client.batchGetItemRequestFn = client.batchGetItemRequest
 	client.batchWriteItemRequestFn = client.batchWriteItemRequest
 	return client, nil
@@ -327,71 +326,35 @@ func (a dynamoDBStorageClient) query(ctx context.Context, query chunk.IndexQuery
 		}
 	}
 
-	request := a.queryRequestFn(ctx, input)
 	pageCount := 0
 	defer func() {
 		dynamoQueryPagesCount.Observe(float64(pageCount))
 	}()
 
-	for page := request; page != nil; page = page.NextPage() {
-		pageCount++
-
-		response, err := a.queryPage(ctx, input, page, query.HashValue, pageCount)
-		if err != nil {
-			return err
+	retryer := newRetryer(ctx, a.cfg.backoffConfig)
+	err := instrument.CollectedRequest(ctx, "DynamoDB.QueryPages", dynamoRequestDuration, instrument.ErrorCode, func(innerCtx context.Context) error {
+		if sp := ot.SpanFromContext(innerCtx); sp != nil {
+			sp.SetTag("tableName", query.TableName)
+			sp.SetTag("hashValue", query.HashValue)
 		}
-
-		if !callback(response) {
-			if err != nil {
-				return fmt.Errorf("QueryPages error: table=%v, err=%v", *input.TableName, page.Error())
-			}
-			return nil
-		}
-		if !page.HasNextPage() {
-			return nil
-		}
-	}
-	return nil
-}
-
-func (a dynamoDBStorageClient) queryPage(ctx context.Context, input *dynamodb.QueryInput, page dynamoDBRequest, hashValue string, pageCount int) (*dynamoDBReadResponse, error) {
-	backoff := util.NewBackoff(ctx, a.cfg.backoffConfig)
-
-	var err error
-	for backoff.Ongoing() {
-		err = instrument.CollectedRequest(ctx, "DynamoDB.QueryPages", dynamoRequestDuration, instrument.ErrorCode, func(innerCtx context.Context) error {
+		return a.DynamoDB.QueryPagesWithContext(ctx, input, func(output *dynamodb.QueryOutput, _ bool) bool {
+			pageCount++
 			if sp := ot.SpanFromContext(innerCtx); sp != nil {
-				sp.SetTag("tableName", aws.StringValue(input.TableName))
-				sp.SetTag("hashValue", hashValue)
-				sp.SetTag("page", pageCount)
-				sp.SetTag("retry", backoff.NumRetries())
+				sp.LogFields(otlog.Int("page", pageCount))
 			}
-			return page.Send()
-		})
-
-		if cc := page.Data().(*dynamodb.QueryOutput).ConsumedCapacity; cc != nil {
-			dynamoConsumedCapacity.WithLabelValues("DynamoDB.QueryPages", *cc.TableName).
-				Add(float64(*cc.CapacityUnits))
-		}
 
-		if err != nil {
-			recordDynamoError(*input.TableName, err, "DynamoDB.QueryPages")
-			if awsErr, ok := err.(awserr.Error); ok && ((awsErr.Code() == dynamodb.ErrCodeProvisionedThroughputExceededException) || page.Retryable()) {
-				if awsErr.Code() != dynamodb.ErrCodeProvisionedThroughputExceededException {
-					level.Warn(util.Logger).Log("msg", "DynamoDB error", "retry", backoff.NumRetries(), "table", *input.TableName, "err", err)
-				}
-				backoff.Wait()
-				continue
+			if cc := output.ConsumedCapacity; cc != nil {
+				dynamoConsumedCapacity.WithLabelValues("DynamoDB.QueryPages", *cc.TableName).
+					Add(float64(*cc.CapacityUnits))
 			}
-			return nil, fmt.Errorf("QueryPage error: table=%v, err=%v", *input.TableName, err)
-		}
 
-		queryOutput := page.Data().(*dynamodb.QueryOutput)
-		return &dynamoDBReadResponse{
-			items: queryOutput.Items,
-		}, nil
+			return callback(&dynamoDBReadResponse{items: output.Items})
+		}, retryer.withRetrys, withErrorHandler(query.TableName, "DynamoDB.QueryPages"))
+	})
+	if err != nil {
+		return fmt.Errorf("QueryPage error: table=%v, err=%v", query.TableName, err)
 	}
-	return nil, fmt.Errorf("QueryPage error: %s for table %v, last error %v", backoff.Err(), *input.TableName, err)
+	return err
 }
 
 type dynamoDBRequest interface {
@@ -403,12 +366,6 @@ type dynamoDBRequest interface {
 	Retryable() bool
 }
 
-func (a dynamoDBStorageClient) queryRequest(ctx context.Context, input *dynamodb.QueryInput) dynamoDBRequest {
-	req, _ := a.DynamoDB.QueryRequest(input)
-	req.SetContext(ctx)
-	return dynamoDBRequestAdapter{req}
-}
-
 func (a dynamoDBStorageClient) batchGetItemRequest(ctx context.Context, input *dynamodb.BatchGetItemInput) dynamoDBRequest {
 	req, _ := a.DynamoDB.BatchGetItemRequest(input)
 	req.SetContext(ctx)
@@ -840,6 +797,16 @@ func (b dynamoDBReadRequest) TakeReqs(from dynamoDBReadRequest, max int) {
 	}
 }
 
+func withErrorHandler(tableName, operation string) func(req *request.Request) {
+	return func(req *request.Request) {
+		req.Handlers.CompleteAttempt.PushBack(func(req *request.Request) {
+			if req.Error != nil {
+				recordDynamoError(tableName, req.Error, operation)
+			}
+		})
+	}
+}
+
 func recordDynamoError(tableName string, err error, operation string) {
 	if awsErr, ok := err.(awserr.Error); ok {
 		dynamoFailures.WithLabelValues(tableName, awsErr.Code(), operation).Add(float64(1))

pkg/chunk/aws/fixtures.go

@@ -41,7 +41,6 @@ var Fixtures = []testutils.Fixture{
 			}
 			index := &dynamoDBStorageClient{
 				DynamoDB:                dynamoDB,
-				queryRequestFn:          dynamoDB.queryRequest,
 				batchGetItemRequestFn:   dynamoDB.batchGetItemRequest,
 				batchWriteItemRequestFn: dynamoDB.batchWriteItemRequest,
 				schemaCfg:               schemaConfig,
@@ -80,7 +79,6 @@ func dynamoDBFixture(provisionedErr, gangsize, maxParallelism int) testutils.Fix
 				},
 				DynamoDB:                dynamoDB,
 				writeThrottle:           rate.NewLimiter(10, dynamoDBMaxWriteBatchSize),
-				queryRequestFn:          dynamoDB.queryRequest,
 				batchGetItemRequestFn:   dynamoDB.batchGetItemRequest,
 				batchWriteItemRequestFn: dynamoDB.batchWriteItemRequest,
 				schemaCfg:               schemaCfg,

pkg/chunk/aws/mock.go

@@ -194,7 +194,7 @@ func (m *mockDynamoDBClient) batchGetItemRequest(_ context.Context, input *dynam
 	}
 }
 
-func (m *mockDynamoDBClient) queryRequest(_ context.Context, input *dynamodb.QueryInput) dynamoDBRequest {
+func (m *mockDynamoDBClient) QueryPagesWithContext(ctx aws.Context, input *dynamodb.QueryInput, fn func(*dynamodb.QueryOutput, bool) bool, opts ...request.Option) error {
 	result := &dynamodb.QueryOutput{
 		Items: []map[string]*dynamodb.AttributeValue{},
 	}
@@ -241,10 +241,8 @@ func (m *mockDynamoDBClient) queryRequest(_ context.Context, input *dynamodb.Que
 
 		result.Items = append(result.Items, item)
 	}
-
-	return &dynamoDBMockRequest{
-		result: result,
-	}
+	fn(result, true)
+	return nil
 }
 
 type dynamoDBMockRequest struct {

pkg/chunk/aws/retryer.go

@@ -0,0 +1,53 @@
+package aws
+
+import (
+	"context"
+	"time"
+
+	"github.com/aws/aws-sdk-go/aws/client"
+	"github.com/aws/aws-sdk-go/aws/request"
+	ot "github.com/opentracing/opentracing-go"
+	otlog "github.com/opentracing/opentracing-go/log"
+
+	"github.com/cortexproject/cortex/pkg/util"
+)
+
+// Map Cortex Backoff into AWS Retryer interface
+type retryer struct {
+	*util.Backoff
+	maxRetries int
+}
+
+var _ request.Retryer = &retryer{}
+
+func newRetryer(ctx context.Context, cfg util.BackoffConfig) *retryer {
+	return &retryer{
+		Backoff:    util.NewBackoff(ctx, cfg),
+		maxRetries: cfg.MaxRetries,
+	}
+}
+
+func (r *retryer) withRetrys(req *request.Request) {
+	req.Retryer = r
+}
+
+// RetryRules return the retry delay that should be used by the SDK before
+// making another request attempt for the failed request.
+func (r *retryer) RetryRules(req *request.Request) time.Duration {
+	if sp := ot.SpanFromContext(req.Context()); sp != nil {
+		sp.LogFields(otlog.Int("retry", r.NumRetries()))
+	}
+	return r.Backoff.NextDelay()
+}
+
+// ShouldRetry returns if the failed request is retryable.
+func (r *retryer) ShouldRetry(req *request.Request) bool {
+	var d client.DefaultRetryer
+	return d.ShouldRetry(req)
+}
+
+// MaxRetries is the number of times a request may be retried before
+// failing.
+func (r *retryer) MaxRetries() int {
+	return r.maxRetries
+}