Add iteration support to dynamo (#52517)

This builds on top of #52199 by updating the dynamo backend to
implement backend.BackendWithItems. Both GetRange and DeleteRange
were refactored to call Items instead of getAllRecords to unify
logic and vet the implementation of Items.

The custom pagination logic to retrieve items was also replaced with
an iterator that works in a similar fashio to the query paginator
from the aws sdk. In addition to simplifying logic this also
removed some extraneous sorting.

Author: rosstimothy

lib/backend/dynamo/dynamodbbk.go

@@ -21,9 +21,9 @@ package dynamo
 import (
 	"context"
 	"errors"
+	"iter"
 	"log/slog"
 	"net/http"
-	"sort"
 	"strconv"
 	"sync/atomic"
 	"time"
@@ -548,63 +548,132 @@ func (b *Backend) Update(ctx context.Context, item backend.Item) (*backend.Lease
 	return backend.NewLease(item), nil
 }
 
-// GetRange returns range of elements
-func (b *Backend) GetRange(ctx context.Context, startKey, endKey backend.Key, limit int) (*backend.GetResult, error) {
-	if startKey.IsZero() {
-		return nil, trace.BadParameter("missing parameter startKey")
-	}
-	if endKey.IsZero() {
-		return nil, trace.BadParameter("missing parameter endKey")
-	}
+func (b *Backend) queryOutputPages(ctx context.Context, limit int, input *dynamodb.QueryInput) iter.Seq2[*dynamodb.QueryOutput, error] {
 	if limit <= 0 {
 		limit = backend.DefaultRangeLimit
 	}
 
-	result, err := b.getAllRecords(ctx, startKey, endKey, limit)
-	if err != nil {
-		return nil, trace.Wrap(err)
-	}
-	sort.Sort(records(result.records))
-	values := make([]backend.Item, len(result.records))
-	for i, r := range result.records {
-		values[i] = backend.Item{
-			Key:      trimPrefix(r.FullPath),
-			Value:    r.Value,
-			Revision: r.Revision,
-		}
-		if r.Expires != nil {
-			values[i].Expires = time.Unix(*r.Expires, 0).UTC()
-		}
-		if values[i].Revision == "" {
-			values[i].Revision = backend.BlankRevision
+	return func(yield func(*dynamodb.QueryOutput, error) bool) {
+		const defaultPageSize = 1000
+		var nextToken map[string]types.AttributeValue
+
+		totalCount := 0
+		for {
+			input.Limit = aws.Int32(int32(min(limit-totalCount, defaultPageSize)))
+
+			result, err := b.svc.Query(ctx, input)
+			if err != nil {
+				yield(nil, trace.Wrap(err))
+				return
+			}
+
+			nextToken = result.LastEvaluatedKey
+			if !yield(result, nil) {
+				return
+			}
+
+			if nextToken == nil {
+				return
+			}
+
+			totalCount += len(result.Items)
+			if totalCount >= limit {
+				return
+			}
+			input.ExclusiveStartKey = nextToken
 		}
 	}
-	return &backend.GetResult{Items: values}, nil
 }
 
-func (b *Backend) getAllRecords(ctx context.Context, startKey, endKey backend.Key, limit int) (*getResult, error) {
-	var result getResult
+func (b *Backend) Items(ctx context.Context, params backend.IterateParams) iter.Seq2[backend.Item, error] {
+	if params.StartKey.IsZero() {
+		err := trace.BadParameter("missing parameter startKey")
+		return func(yield func(backend.Item, error) bool) { yield(backend.Item{}, err) }
+	}
+	if params.EndKey.IsZero() {
+		err := trace.BadParameter("missing parameter endKey")
+		return func(yield func(backend.Item, error) bool) { yield(backend.Item{}, err) }
+	}
 
-	// this code is being extra careful here not to introduce endless loop
-	// by some unfortunate series of events
-	for i := 0; i < backend.DefaultRangeLimit/100; i++ {
-		re, err := b.getRecords(ctx, prependPrefix(startKey), prependPrefix(endKey), limit, result.lastEvaluatedKey)
+	const (
+		query = "HashKey = :hashKey AND FullPath BETWEEN :rangeStart AND :rangeEnd"
+
+		// filter out expired items, otherwise they might show up in the query
+		// http://docs.aws.amazon.com/amazondynamodb/latest/developerguide/howitworks-ttl.html
+		filter = "attribute_not_exists(Expires) OR Expires >= :timestamp"
+	)
+
+	av := map[string]types.AttributeValue{
+		":rangeStart": &types.AttributeValueMemberS{Value: prependPrefix(params.StartKey)},
+		":rangeEnd":   &types.AttributeValueMemberS{Value: prependPrefix(params.EndKey)},
+		":timestamp":  timeToAttributeValue(b.clock.Now().UTC()),
+		":hashKey":    &types.AttributeValueMemberS{Value: hashKey},
+	}
+
+	input := dynamodb.QueryInput{
+		KeyConditionExpression:    aws.String(query),
+		TableName:                 &b.TableName,
+		ExpressionAttributeValues: av,
+		FilterExpression:          aws.String(filter),
+		ConsistentRead:            aws.Bool(true),
+		ScanIndexForward:          aws.Bool(!params.Descending),
+	}
+
+	return func(yield func(backend.Item, error) bool) {
+		count := 0
+		defer func() {
+			if count >= backend.DefaultRangeLimit {
+				b.logger.WarnContext(ctx, "Range query hit backend limit. (this is a bug!)", "start_key", params.StartKey, "limit", backend.DefaultRangeLimit)
+			}
+		}()
+
+		for page, err := range b.queryOutputPages(ctx, params.Limit, &input) {
+			if err != nil {
+				yield(backend.Item{}, convertError(err))
+				return
+			}
+
+			for _, itemAttributes := range page.Items {
+				var r record
+				if err := attributevalue.UnmarshalMap(itemAttributes, &r); err != nil {
+					yield(backend.Item{}, convertError(err))
+					return
+				}
+
+				item := backend.Item{
+					Key:      trimPrefix(r.FullPath),
+					Value:    r.Value,
+					Revision: r.Revision,
+				}
+				if r.Expires != nil {
+					item.Expires = time.Unix(*r.Expires, 0).UTC()
+				}
+				if item.Revision == "" {
+					item.Revision = backend.BlankRevision
+				}
+
+				if !yield(item, nil) {
+					return
+				}
+				count++
+				if params.Limit != backend.NoLimit && count >= params.Limit {
+					return
+				}
+			}
+		}
+	}
+}
+
+// GetRange returns range of elements
+func (b *Backend) GetRange(ctx context.Context, startKey, endKey backend.Key, limit int) (*backend.GetResult, error) {
+	var result backend.GetResult
+	for i, err := range b.Items(ctx, backend.IterateParams{StartKey: startKey, EndKey: endKey, Limit: limit}) {
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
-		result.records = append(result.records, re.records...)
-		// If the limit was exceeded or there are no more records to fetch return the current result
-		// otherwise updated lastEvaluatedKey and proceed with obtaining new records.
-		if (limit != 0 && len(result.records) >= limit) || len(re.lastEvaluatedKey) == 0 {
-			if len(result.records) == backend.DefaultRangeLimit {
-				b.logger.WarnContext(ctx, "Range query hit backend limit. (this is a bug!)", "start_key", startKey, "limit", backend.DefaultRangeLimit)
-			}
-			result.lastEvaluatedKey = nil
-			return &result, nil
-		}
-		result.lastEvaluatedKey = re.lastEvaluatedKey
+		result.Items = append(result.Items, i)
 	}
-	return nil, trace.BadParameter("backend entered endless loop")
+	return &result, nil
 }
 
 const (
@@ -617,45 +686,63 @@ const (
 
 // DeleteRange deletes range of items with keys between startKey and endKey
 func (b *Backend) DeleteRange(ctx context.Context, startKey, endKey backend.Key) error {
-	if startKey.IsZero() {
-		return trace.BadParameter("missing parameter startKey")
-	}
-	if endKey.IsZero() {
-		return trace.BadParameter("missing parameter endKey")
-	}
-	// keep fetching and deleting until no records left,
-	// keep the very large limit, just in case if someone else
-	// keeps adding records
-	for i := 0; i < backend.DefaultRangeLimit/100; i++ {
-		result, err := b.getRecords(ctx, prependPrefix(startKey), prependPrefix(endKey), batchOperationItemsLimit, nil)
+	// Attempt to pull all existing items and delete them in batches
+	// in accordance with the BatchWriteItem limits. There is a hard
+	// cap on the total number of items that can be deleted in a single
+	// DeleteRange call to avoid racing with additional records being added.
+	const maxDeletionOperations = backend.DefaultRangeLimit / 100 / batchOperationItemsLimit
+	requests := make([]types.WriteRequest, 0, batchOperationItemsLimit)
+	var deletions int
+	for item, err := range b.Items(ctx, backend.IterateParams{StartKey: startKey, EndKey: endKey}) {
 		if err != nil {
 			return trace.Wrap(err)
 		}
-		if len(result.records) == 0 {
-			return nil
+
+		if deletions >= maxDeletionOperations {
+			break
 		}
-		requests := make([]types.WriteRequest, 0, len(result.records))
-		for _, record := range result.records {
-			requests = append(requests, types.WriteRequest{
-				DeleteRequest: &types.DeleteRequest{
-					Key: map[string]types.AttributeValue{
-						hashKeyKey:  &types.AttributeValueMemberS{Value: hashKey},
-						fullPathKey: &types.AttributeValueMemberS{Value: record.FullPath},
-					},
+
+		requests = append(requests, types.WriteRequest{
+			DeleteRequest: &types.DeleteRequest{
+				Key: map[string]types.AttributeValue{
+					hashKeyKey:  &types.AttributeValueMemberS{Value: hashKey},
+					fullPathKey: &types.AttributeValueMemberS{Value: prependPrefix(item.Key)},
 				},
-			})
+			},
+		})
+
+		if len(requests) == batchOperationItemsLimit {
+			if _, err := b.svc.BatchWriteItem(ctx, &dynamodb.BatchWriteItemInput{
+				RequestItems: map[string][]types.WriteRequest{
+					b.TableName: requests,
+				},
+			}); err != nil {
+				return trace.Wrap(err)
+			}
+
+			requests = requests[:0]
+			deletions++
+			if deletions >= maxDeletionOperations {
+				break
+			}
 		}
-		input := dynamodb.BatchWriteItemInput{
+	}
+
+	if deletions >= maxDeletionOperations {
+		return trace.ConnectionProblem(nil, "not all items deleted, too many requests")
+	}
+
+	if len(requests) > 0 {
+		if _, err := b.svc.BatchWriteItem(ctx, &dynamodb.BatchWriteItemInput{
 			RequestItems: map[string][]types.WriteRequest{
 				b.TableName: requests,
 			},
-		}
-
-		if _, err = b.svc.BatchWriteItem(ctx, &input); err != nil {
+		}); err != nil {
 			return trace.Wrap(err)
 		}
 	}
-	return trace.ConnectionProblem(nil, "not all items deleted, too many requests")
+
+	return nil
 }
 
 // Get returns a single item or not found error
@@ -961,60 +1048,6 @@ func (b *Backend) createTable(ctx context.Context, tableName *string, rangeKey s
 	return trace.Wrap(err)
 }
 
-type getResult struct {
-	// lastEvaluatedKey is the primary key of the item where the operation stopped, inclusive of the
-	// previous result set. Use this value to start a new operation, excluding this
-	// value in the new request.
-	lastEvaluatedKey map[string]types.AttributeValue
-	records          []record
-}
-
-// getRecords retrieves all keys by path
-func (b *Backend) getRecords(ctx context.Context, startKey, endKey string, limit int, lastEvaluatedKey map[string]types.AttributeValue) (*getResult, error) {
-	query := "HashKey = :hashKey AND FullPath BETWEEN :fullPath AND :rangeEnd"
-	attrV := map[string]interface{}{
-		":fullPath":  startKey,
-		":hashKey":   hashKey,
-		":timestamp": b.clock.Now().UTC().Unix(),
-		":rangeEnd":  endKey,
-	}
-
-	// filter out expired items, otherwise they might show up in the query
-	// http://docs.aws.amazon.com/amazondynamodb/latest/developerguide/howitworks-ttl.html
-	filter := "attribute_not_exists(Expires) OR Expires >= :timestamp"
-	av, err := attributevalue.MarshalMap(attrV)
-	if err != nil {
-		return nil, convertError(err)
-	}
-	input := dynamodb.QueryInput{
-		KeyConditionExpression:    aws.String(query),
-		TableName:                 &b.TableName,
-		ExpressionAttributeValues: av,
-		FilterExpression:          aws.String(filter),
-		ConsistentRead:            aws.Bool(true),
-		ExclusiveStartKey:         lastEvaluatedKey,
-	}
-	if limit > 0 {
-		input.Limit = aws.Int32(int32(limit))
-	}
-	out, err := b.svc.Query(ctx, &input)
-	if err != nil {
-		return nil, trace.Wrap(err)
-	}
-	var result getResult
-	for _, item := range out.Items {
-		var r record
-		if err := attributevalue.UnmarshalMap(item, &r); err != nil {
-			return nil, trace.Wrap(err)
-		}
-		result.records = append(result.records, r)
-	}
-	sort.Sort(records(result.records))
-	result.records = removeDuplicates(result.records)
-	result.lastEvaluatedKey = out.LastEvaluatedKey
-	return &result, nil
-}
-
 // isExpired returns 'true' if the given object (record) has a TTL and
 // it's due.
 func (r *record) isExpired(now time.Time) bool {
@@ -1025,23 +1058,6 @@ func (r *record) isExpired(now time.Time) bool {
 	return now.UTC().After(expiryDateUTC)
 }
 
-func removeDuplicates(elements []record) []record {
-	// Use map to record duplicates as we find them.
-	encountered := map[string]bool{}
-	var result []record
-
-	for v := range elements {
-		if !encountered[elements[v].FullPath] {
-			// Record this element as an encountered element.
-			encountered[elements[v].FullPath] = true
-			// Append to result slice.
-			result = append(result, elements[v])
-		}
-	}
-	// Return the new slice.
-	return result
-}
-
 const (
 	modeCreate = iota
 	modePut
@@ -1235,23 +1251,6 @@ func convertError(err error) error {
 	return err
 }
 
-type records []record
-
-// Len is part of sort.Interface.
-func (r records) Len() int {
-	return len(r)
-}
-
-// Swap is part of sort.Interface.
-func (r records) Swap(i, j int) {
-	r[i], r[j] = r[j], r[i]
-}
-
-// Less is part of sort.Interface.
-func (r records) Less(i, j int) bool {
-	return r[i].FullPath < r[j].FullPath
-}
-
 func fullPathToAttributeValueMap(fullPath string) map[string]types.AttributeValue {
 	return map[string]types.AttributeValue{
 		hashKeyKey:  &types.AttributeValueMemberS{Value: hashKey},