Use newer AWS API for paginated queries (#2452)

* Refactor: export NextDelay() so we can call it from other packages

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

* 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.

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

* Don't need DynamoDB request wrapper to do so much now

Now we are calling QueryPagesWithContext directly we don't need the
paging interface and we never re-use request objects.

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

* Stop AWS Retryer if context is cancelled

Ingester.flushUserSeries() puts a timeout on the context, so don't
retry for longer than that.

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

Author: bboreham

CHANGELOG.md

@@ -56,6 +56,7 @@
 * [ENHANCEMENT] Redis Cache: Added `idle_timeout`, `wait_on_pool_exhaustion` and `max_conn_lifetime` options to redis cache configuration. #2550
 * [ENHANCEMENT] WAL: the experimental tag has been removed on the WAL in ingesters.
 * [BUGFIX] Ruler: Ensure temporary rule files with special characters are properly mapped and cleaned up. #2506
+* [ENHANCEMENT] Use newer AWS API for paginated queries - removes 'Deprecated' message from logfiles. #2452
 * [BUGFIX] Fixes #2411, Ensure requests are properly routed to the prometheus api embedded in the query if `-server.path-prefix` is set. #2372
 * [BUGFIX] Experimental TSDB: fixed chunk data corruption when querying back series using the experimental blocks storage. #2400
 * [BUGFIX] Cassandra Storage: Fix endpoint TLS host verification. #2109

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"
@@ -21,6 +22,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/dynamodb"
 	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
+	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	awscommon "github.com/weaveworks/common/aws"
 	"github.com/weaveworks/common/instrument"
@@ -144,7 +146,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 +173,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,88 +327,44 @@ 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
-		}
-
-		if !callback(query, response) {
-			if err != nil {
-				return fmt.Errorf("QueryPages error: table=%v, err=%v", *input.TableName, page.Error())
-			}
-			return nil
-		}
-		if !page.HasNextPage() {
-			return nil
+	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)
 		}
-	}
-	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(innerCtx, 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(query, &dynamoDBReadResponse{items: output.Items})
+		}, retryer.withRetries, withErrorHandler(query.TableName, "DynamoDB.QueryPages"))
+	})
+	if err != nil {
+		return errors.Wrapf(err, "QueryPages error: table=%v", query.TableName)
 	}
-	return nil, fmt.Errorf("QueryPage error: %s for table %v, last error %v", backoff.Err(), *input.TableName, err)
+	return err
 }
 
 type dynamoDBRequest interface {
-	NextPage() dynamoDBRequest
 	Send() error
 	Data() interface{}
 	Error() error
-	HasNextPage() bool
 	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)
@@ -425,33 +381,18 @@ type dynamoDBRequestAdapter struct {
 	request *request.Request
 }
 
-func (a dynamoDBRequestAdapter) NextPage() dynamoDBRequest {
-	next := a.request.NextPage()
-	if next == nil {
-		return nil
-	}
-	return dynamoDBRequestAdapter{next}
-}
-
 func (a dynamoDBRequestAdapter) Data() interface{} {
 	return a.request.Data
 }
 
 func (a dynamoDBRequestAdapter) Send() error {
-	// Clear error in case we are retrying the same operation - if we
-	// don't do this then the same error will come back again immediately
-	a.request.Error = nil
 	return a.request.Send()
 }
 
 func (a dynamoDBRequestAdapter) Error() error {
 	return a.request.Error
 }
 
-func (a dynamoDBRequestAdapter) HasNextPage() bool {
-	return a.request.HasNextPage()
-}
-
 func (a dynamoDBRequestAdapter) Retryable() bool {
 	return aws.BoolValue(a.request.Retryable)
 }
@@ -840,6 +781,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,20 +241,15 @@ 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 {
 	result interface{}
 	err    error
 }
 
-func (m *dynamoDBMockRequest) NextPage() dynamoDBRequest {
-	return m
-}
 func (m *dynamoDBMockRequest) Send() error {
 	return m.err
 }
@@ -264,9 +259,6 @@ func (m *dynamoDBMockRequest) Data() interface{} {
 func (m *dynamoDBMockRequest) Error() error {
 	return m.err
 }
-func (m *dynamoDBMockRequest) HasNextPage() bool {
-	return false
-}
 func (m *dynamoDBMockRequest) Retryable() bool {
 	return false
 }

pkg/chunk/aws/retryer.go

@@ -0,0 +1,52 @@
+package aws
+
+import (
+	"context"
+	"time"
+
+	"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) withRetries(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 {
+	duration := r.Backoff.NextDelay()
+	if sp := ot.SpanFromContext(req.Context()); sp != nil {
+		sp.LogFields(otlog.Int("retry", r.NumRetries()))
+	}
+	return duration
+}
+
+// ShouldRetry returns if the failed request is retryable.
+func (r *retryer) ShouldRetry(req *request.Request) bool {
+	return r.Ongoing() && (req.IsErrorRetryable() || req.IsErrorThrottle())
+}
+
+// MaxRetries is the number of times a request may be retried before
+// failing.
+func (r *retryer) MaxRetries() int {
+	return r.maxRetries
+}

pkg/util/backoff.go

@@ -74,7 +74,7 @@ func (b *Backoff) NumRetries() int {
 // Returns immediately if Context is terminated
 func (b *Backoff) Wait() {
 	// Increase the number of retries and get the next delay
-	sleepTime := b.nextDelay()
+	sleepTime := b.NextDelay()
 
 	if b.Ongoing() {
 		select {
@@ -84,7 +84,7 @@ func (b *Backoff) Wait() {
 	}
 }
 
-func (b *Backoff) nextDelay() time.Duration {
+func (b *Backoff) NextDelay() time.Duration {
 	b.numRetries++
 
 	// Handle the edge case the min and max have the same value

pkg/util/backoff_test.go

@@ -6,7 +6,7 @@ import (
 	"time"
 )
 
-func TestBackoff_nextDelay(t *testing.T) {
+func TestBackoff_NextDelay(t *testing.T) {
 	t.Parallel()
 
 	tests := map[string]struct {
@@ -92,7 +92,7 @@ func TestBackoff_nextDelay(t *testing.T) {
 			})
 
 			for _, expectedRange := range testData.expectedRanges {
-				delay := b.nextDelay()
+				delay := b.NextDelay()
 
 				if delay < expectedRange[0] || delay > expectedRange[1] {
 					t.Errorf("%d expected to be within %d and %d", delay, expectedRange[0], expectedRange[1])